diff options
Diffstat (limited to 'fbin')
| -rwxr-xr-x | fbin/fbin.py | 6 | ||||
| -rw-r--r-- | fbin/file_storage/base.py | 11 | ||||
| -rw-r--r-- | fbin/file_storage/exceptions.py | 2 | ||||
| -rw-r--r-- | fbin/file_storage/filesystem.py | 3 |
4 files changed, 19 insertions, 3 deletions
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) |