From f1cfaef5b0f358c94fe7f5e7c2eb5a6d261a85ae Mon Sep 17 00:00:00 2001
From: Jon Bergli Heier <snakebite@jvnv.net>
Date: Sat, 7 Dec 2019 11:14:01 +0100
Subject: Add max file size configuration

This allows configuring max file sizes for both registered and anonymous
users. For registered users the USER_FILE_SIZE_LIMIT is used, and
ANONYMOUS_FILE_SIZE_LIMIT for anonymous users. If the size is not
specified or None, the limit is not enforced. Setting the limit to 0
effectively disables uploads.
---
 fbin/fbin.py                    |  6 +++++-
 fbin/file_storage/base.py       | 11 ++++++++++-
 fbin/file_storage/exceptions.py |  2 ++
 fbin/file_storage/filesystem.py |  3 ++-
 4 files changed, 19 insertions(+), 3 deletions(-)
 create mode 100644 fbin/file_storage/exceptions.py

diff --git a/fbin/fbin.py b/fbin/fbin.py
index 35fd1fc..cf7de02 100755
--- a/fbin/fbin.py
+++ b/fbin/fbin.py
@@ -25,6 +25,7 @@ from werkzeug.utils import secure_filename
 from . import db
 from .monkey import patch as monkey_patch
 from .login import login_manager, load_user
+from .file_storage.exceptions import StorageError
 
 storage = importlib.import_module(current_app.config.get('STORAGE_MODULE', '.file_storage.filesystem'), package='fbin').Storage(current_app)
 
@@ -136,7 +137,10 @@ def upload(api=False, user=None):
     if not uploaded_file or not uploaded_file.filename:
         return error('No valid file or filename was provided.')
     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)
+    try:
+        new_file = storage.store_file(uploaded_file, file_hash, user, request.remote_addr)
+    except StorageError as e:
+        return error(str(e))
 
     mime = new_file.get_mime_type()
     # TODO: Apparently TIFF also supports EXIF, test this.
diff --git a/fbin/file_storage/base.py b/fbin/file_storage/base.py
index 9f09199..e2ca1a6 100644
--- a/fbin/file_storage/base.py
+++ b/fbin/file_storage/base.py
@@ -1,17 +1,26 @@
 import datetime
 
 from .. import db
+from .exceptions import *
 
 class BaseStorage:
     def __init__(self, app):
         self.app = app
 
-    def add_file(self, file_hash, filename, size, user=None, ip=None):
+    def verify_file(self, file):
+        user = file.user_id is not None
+        size_limit = self.app.config.get('USER_FILE_SIZE_LIMIT' if user else 'ANONYMOUS_FILE_SIZE_LIMIT')
+        if size_limit is not None and file.size > size_limit:
+            raise FileSizeError('The file size is too large (max {})'.format(db.File.pretty_size(size_limit)))
+
+    def add_file(self, file_hash, filename, size, user=None, ip=None, verify=True):
         '''Adds the file to the database.
 
         Call from store_file after the file is successfully stored.'''
         with db.session_scope() as sess:
             f = db.File(file_hash, filename, size, datetime.datetime.utcnow(), user.id if user else None, ip)
+            # Raises on invalid files
+            self.verify_file(f)
             sess.add(f)
             sess.commit()
             sess.refresh(f)
diff --git a/fbin/file_storage/exceptions.py b/fbin/file_storage/exceptions.py
new file mode 100644
index 0000000..140221b
--- /dev/null
+++ b/fbin/file_storage/exceptions.py
@@ -0,0 +1,2 @@
+class StorageError(Exception): pass
+class FileSizeError(StorageError): pass
diff --git a/fbin/file_storage/filesystem.py b/fbin/file_storage/filesystem.py
index 1259002..07d29db 100644
--- a/fbin/file_storage/filesystem.py
+++ b/fbin/file_storage/filesystem.py
@@ -22,7 +22,8 @@ class Storage(BaseStorage):
             size = os.path.getsize(temp_path)
         try:
             new_file = self.add_file(file_hash, uploaded_file.filename, size, user, ip)
-            os.rename(temp_path, new_file.get_path())
+            if new_file:
+                os.rename(temp_path, new_file.get_path())
             return new_file
         except:
             os.unlink(temp.name)
-- 
cgit v1.2.3