Add file storage modules

Allows for storing files other places than the local file system.
Currently the local filesystem and S3 are supported.

1 files changed, 30 insertions, 39 deletions

diff --git a/fbin/fbin.py b/fbin/fbin.py
index 19f82ed..d3065e2 100755
--- a/fbin/fbin.py
+++ b/fbin/fbin.py
@@ -4,6 +4,7 @@ import base64
 import cgi
 import datetime
 import hashlib
+import importlib
 import io
 import json
 import mimetypes
@@ -14,7 +15,7 @@ import tempfile
 import urllib
 from urllib.parse import urlencode, urljoin
 
-from flask import Blueprint, redirect, current_app, url_for, request, render_template, session, flash, send_file, abort, jsonify, Markup
+from flask import Blueprint, redirect, current_app, url_for, request, render_template, session, flash, send_file, abort, jsonify, Markup, Response
 from flask_login import login_user, logout_user, current_user, login_required
 import jwt
 from PIL import Image
@@ -25,13 +26,12 @@ from . import db
 from .monkey import patch as monkey_patch
 from .login import login_manager, load_user
 
+storage = importlib.import_module(current_app.config.get('STORAGE_MODULE', '.file_storage.filesystem'), package='fbin').Storage(current_app)
+
 monkey_patch()
 
 base62_alphabet = '0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ'
 
-if not os.path.isdir(current_app.config['FILE_DIRECTORY']):
-    os.mkdir(current_app.config['FILE_DIRECTORY'])
-
 if not os.path.isdir(current_app.config['THUMB_DIRECTORY']):
     os.mkdir(current_app.config['THUMB_DIRECTORY'])
 
@@ -57,19 +57,6 @@ def get_or_create_user(username, jab_id):
             except db.IntegrityError:
                 return None
 
-def add_file(path, filename, user = None, ip = None):
-    file_hash = ''.join(random.choice(base62_alphabet) for x in range(5))
-    new_path = os.path.join(current_app.config['FILE_DIRECTORY'], file_hash + os.path.splitext(filename)[1])
-    os.rename(path, new_path)
-    if current_app.config.get('DESTINATION_MODE'):
-        os.chmod(new_path, current_app.config.get('DESTINATION_MODE'))
-    with db.session_scope() as sess:
-        f = db.File(file_hash, filename, datetime.datetime.utcnow(), user.id if user else None, ip)
-        sess.add(f)
-        sess.commit()
-        sess.refresh(f)
-    return f
-
 def get_file(file_hash, user_id=None, update_accessed=False):
     with db.session_scope() as sess:
         try:
@@ -101,8 +88,7 @@ def delete_file(file):
         sess.delete(file)
         sess.commit()
         filename = file.get_path()
-        if os.path.exists(filename):
-            os.unlink(filename)
+        storage.delete_file(file)
         thumbfile = file.get_thumb_path()
         if os.path.exists(thumbfile):
             os.unlink(thumbfile)
@@ -149,14 +135,8 @@ def upload(api=False, user=None):
     uploaded_file = request.files.get('file')
     if not uploaded_file or not uploaded_file.filename:
         return error('No valid file or filename was provided.')
-    if hasattr(uploaded_file.stream, 'file'):
-        temp = None
-        temp_path = uploaded_file.stream.name
-    else:
-        temp = tempfile.NamedTemporaryFile(prefix = 'upload_', dir = current_app.config['FILE_DIRECTORY'], delete = False)
-        uploaded_file.save(temp.file)
-        temp_path = temp.name
-    new_file = add_file(temp_path, uploaded_file.filename, user, request.remote_addr)
+    file_hash = ''.join(random.choice(base62_alphabet) for x in range(5))
+    new_file = storage.store_file(uploaded_file, file_hash, user, request.remote_addr)
 
     mime = new_file.get_mime_type()
     # TODO: Apparently TIFF also supports EXIF, test this.
@@ -207,9 +187,14 @@ def uploaded(hash):
 @app.route('/file/<hash:hash>/<path:filename>', endpoint = 'file')
 def _file(hash, ext=None, filename=None):
     f = get_file(hash)
-    if not f or not f.exists or f.blocked_reason:
+    if not f or f.blocked_reason:
+        abort(404)
+    path = storage.get_file(f)
+    if isinstance(path, Response):
+        return path
+    if not path or not os.path.exists(path):
         abort(404)
-    return send_file(f.get_path())
+    return send_file(path)
 
 @app.route('/l')
 @app.route('/login')
@@ -302,7 +287,7 @@ def files():
     context = {
         'title': 'Files',
         'files': files,
-        'total_size': db.File.pretty_size(sum(f.get_size() for f in files if f.exists)),
+        'total_size': db.File.pretty_size(sum(f.size for f in files if f.size)),
     }
     return render_template('files.html', **context)
 
@@ -354,7 +339,7 @@ def images():
         'title': 'Images',
         'fullwidth': True,
         'files': files,
-        'total_size': db.File.pretty_size(sum(f.get_size() for f in files if f.exists)),
+        'total_size': db.File.pretty_size(sum(f.size for f in files if f.size)),
     }
     return render_template('images.html', **context)
 
@@ -367,7 +352,7 @@ def videos():
         'title': 'Videos',
         'fullwidth': True,
         'files': files,
-        'total_size': db.File.pretty_size(sum(f.get_size() for f in files if f.exists)),
+        'total_size': db.File.pretty_size(sum(f.size for f in files if f.size)),
     }
     return render_template('images.html', **context)
 
@@ -379,18 +364,24 @@ def thumb(hash):
         f = get_file(hash, update_accessed = False)
         if f.is_image():
             try:
-                im = Image.open(f.get_path())
+                #im = Image.open(f.get_path())
+                with storage.temp_file(f) as tf:
+                    im = Image.open(tf)
+                    # Check for valid JPEG modes.
+                    if im.mode not in ('1', 'L', 'RGB', 'RGBX', 'CMYK', 'YCbCr'):
+                        im = im.convert('RGB')
+                    im.thumbnail(current_app.config.get('THUMB_SIZE', (128, 128)), Image.ANTIALIAS)
+                    im.save(thumbfile)
             except IOError:
                 # We can't generate a thumbnail for this file, just say it doesn't exist.
                 abort(404)
-            # Check for valid JPEG modes.
-            if im.mode not in ('1', 'L', 'RGB', 'RGBX', 'CMYK', 'YCbCr'):
-                im = im.convert('RGB')
-            im.thumbnail(current_app.config.get('THUMB_SIZE', (128, 128)), Image.ANTIALIAS)
-            im.save(thumbfile)
         elif f.is_video():
-            p = subprocess.run(['ffmpegthumbnailer', '-i', f.get_path(), '-o', thumbfile])
+            #p = subprocess.run(['ffmpegthumbnailer', '-i', f.get_path(), '-o', thumbfile])
+            with storage.temp_file(f) as tf:
+                p = subprocess.run(['ffmpegthumbnailer', '-i', '-', '-o', thumbfile], stdin=tf)
             if p.returncode != 0:
+                if os.path.exists(thumbfile):
+                    os.unlink(thumbfile)
                 abort(404)
         else:
             abort(404)