3 files changed, 41 insertions, 9 deletions
diff --git a/fbin/file_storage/base.py b/fbin/file_storage/base.py
index abdf580..aa2c510 100644
--- a/fbin/file_storage/base.py
+++ b/fbin/file_storage/base.py
@@ -13,6 +13,10 @@ class BaseStorage:
         if size_limit is not None and file.size > size_limit:
             raise FileSizeError('The file size is too large (max {})'.format(File.pretty_size(size_limit)))
 
+    def upload_file(self, uploaded_file, file_hash, user):
+        '''Upload data from uploaded_file.'''
+        raise NotImplementedError()
+
     def add_file(self, file_hash, filename, size, user=None, ip=None, verify=True):
         '''Adds the file to the database.
 
@@ -29,6 +33,10 @@ class BaseStorage:
         '''Store uploaded_file.'''
         raise NotImplementedError()
 
+    def file_exists(self, f):
+        '''Return True if the specified file exists. '''
+        raise NotImplementedError()
+
     def get_file(self, f):
         '''Return a file object for the specified file.
 
diff --git a/fbin/file_storage/filesystem.py b/fbin/file_storage/filesystem.py
index 3a640bb..7951d88 100644
--- a/fbin/file_storage/filesystem.py
+++ b/fbin/file_storage/filesystem.py
@@ -10,7 +10,7 @@ class Storage(BaseStorage):
         os.makedirs(self.app.config['FILE_DIRECTORY'], exist_ok=True)
         os.makedirs(self.app.config['THUMB_DIRECTORY'], exist_ok=True)
 
-    def store_file(self, uploaded_file, file_hash, user, ip):
+    def upload_file(self, uploaded_file, file_hash, user):
         size = uploaded_file.content_length
         if hasattr(uploaded_file.stream, 'file'):
             temp = None
@@ -20,17 +20,25 @@ class Storage(BaseStorage):
             uploaded_file.save(temp.file)
             temp_path = temp.name
             size = os.path.getsize(temp_path)
+        new_path = os.path.join(self.app.config['FILE_DIRECTORY'], file_hash + os.path.splitext(uploaded_file.filename)[1])
+        os.rename(temp_path, new_path)
+        if self.app.config.get('DESTINATION_MODE'):
+            os.chmod(new_path, self.app.config.get('DESTINATION_MODE'))
+        return new_path, size
+
+    def store_file(self, uploaded_file, file_hash, user, ip):
+        file_path, size = self.upload_file(uploaded_file, file_hash, user)
         try:
-            new_file = self.add_file(file_hash, uploaded_file.filename, size, user, ip)
-            if new_file:
-                os.rename(temp_path, new_file.get_path())
-                if self.app.config.get('DESTINATION_MODE'):
-                    os.chmod(new_file.get_path(), self.app.config.get('DESTINATION_MODE'))
-            return new_file
+            return self.add_file(file_hash, uploaded_file.filename, size, user, ip)
         except:
-            os.unlink(temp.name)
+            if os.path.exists(file_path):
+                os.unlink(file_path)
             raise
 
+    def file_exists(self, f):
+        path = f.get_path()
+        return os.path.exists(path)
+
     def get_file(self, f):
         path = f.get_path()
         if not os.path.exists(path):
diff --git a/fbin/file_storage/s3.py b/fbin/file_storage/s3.py
index a11488f..e2dd1ea 100644
--- a/fbin/file_storage/s3.py
+++ b/fbin/file_storage/s3.py
@@ -21,7 +21,7 @@ class Storage(BaseStorage):
             key += '_thumb'
         return key
 
-    def store_file(self, uploaded_file, file_hash, user, ip):
+    def upload_file(self, uploaded_file, file_hash, user):
         bucket = self.client.Bucket(self.app.config['S3_BUCKET'])
         key = self._get_object_key(file_hash, user.id if user else 0)
         obj = bucket.upload_fileobj(Fileobj=uploaded_file.stream, Key=key)
@@ -29,12 +29,28 @@ class Storage(BaseStorage):
         if not size:
             obj = self.client.ObjectSummary(self.app.config['S3_BUCKET'], key)
             size = obj.size
+        return size
+
+    def store_file(self, uploaded_file, file_hash, user, ip):
+        size = self.upload_file(uploaded_file, file_hash, user)
         try:
             return self.add_file(file_hash, uploaded_file.filename, size, user, ip)
         except:
             obj.delete()
             raise
 
+    def file_exists(self, f):
+        key = self.get_object_key(f)
+        bucket = self.app.config['S3_BUCKET']
+        obj = self.client.Object(bucket, key)
+        try:
+            meta = obj.load()
+            return True
+        except botocore.exceptions.ClientError as e:
+            if e.response['Error']['Code'] == '404':
+                return False
+            raise
+
     def get_file(self, f, thumb=False):
         key = self.get_object_key(f, thumb=thumb)
         if thumb: