diff options
-rwxr-xr-x | fbin/fbin.py | 52 | ||||
-rw-r--r-- | fbin/file_storage/base.py | 9 | ||||
-rw-r--r-- | fbin/file_storage/filesystem.py | 15 | ||||
-rw-r--r-- | fbin/file_storage/s3.py | 32 |
4 files changed, 80 insertions, 28 deletions
diff --git a/fbin/fbin.py b/fbin/fbin.py index 983c1aa..d1dff76 100755 --- a/fbin/fbin.py +++ b/fbin/fbin.py @@ -359,33 +359,37 @@ def videos(): @app.route('/t/<hash:hash>') @app.route('/thumb/<hash:hash>') def thumb(hash): - thumbfile = os.path.join(current_app.config['THUMB_DIRECTORY'], hash + '.jpg') - if not os.access(thumbfile, os.F_OK): - f = get_file(hash, update_accessed = False) - if f.is_image(): - try: - #im = Image.open(f.get_path()) + f = get_file(hash, update_accessed=False) + response = storage.get_thumbnail(f) + if not response: + with tempfile.NamedTemporaryFile(suffix='.jpg') as ttf: # temporary thumb file + if f.is_image(): + try: + 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(ttf) + except IOError: + # We can't generate a thumbnail for this file, just say it doesn't exist. + abort(404) + elif f.is_video(): 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. + p = subprocess.run(['ffmpegthumbnailer', '-i', '-', '-o', ttf.name], stdin=tf) + if p.returncode != 0: + abort(404) + else: abort(404) - elif f.is_video(): - #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) + ttf.seek(0) + if not os.path.getsize(ttf.name): abort(404) - else: - abort(404) - return send_file(thumbfile) + storage.store_thumbnail(f, ttf) + response = storage.get_thumbnail(f) + if isinstance(response, Response): + return response + return send_file(response, attachment_filename='thumb.jpg') @app.route('/h') @app.route('/help') diff --git a/fbin/file_storage/base.py b/fbin/file_storage/base.py index 6f39665..9f09199 100644 --- a/fbin/file_storage/base.py +++ b/fbin/file_storage/base.py @@ -37,3 +37,12 @@ class BaseStorage: This is used internally for eg. thumbnails.''' raise NotImplementedError() + def get_thumbnail(self, f): + '''Return a file object for the specified file's thumbnail. + + Subclasses can also return a flask.Response instance if required.''' + raise NotImplementedError() + + def store_thumbnail(self, f, stream): + '''Store thumbnail for the specified file.''' + raise NotImplementedError() diff --git a/fbin/file_storage/filesystem.py b/fbin/file_storage/filesystem.py index 3b46e34..1259002 100644 --- a/fbin/file_storage/filesystem.py +++ b/fbin/file_storage/filesystem.py @@ -8,6 +8,7 @@ class Storage(BaseStorage): def __init__(self, app): super().__init__(app) 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): size = uploaded_file.content_length @@ -45,3 +46,17 @@ class Storage(BaseStorage): def temp_file(self, f): with open(f.get_path(), 'rb') as f: yield f + + def get_thumbnail(self, f): + path = f.get_thumb_path() + if not os.path.exists(path): + return + return path + + def store_thumbnail(self, f, stream): + path = f.get_thumb_path() + with open(path, 'wb') as f: + buf = stream.read(1024*10) + while buf: + f.write(buf) + buf = stream.read(1024*10) diff --git a/fbin/file_storage/s3.py b/fbin/file_storage/s3.py index 2f0b87b..e81f8b4 100644 --- a/fbin/file_storage/s3.py +++ b/fbin/file_storage/s3.py @@ -2,6 +2,7 @@ import contextlib import tempfile import boto3 +import botocore.exceptions from flask import request, send_file from .base import BaseStorage @@ -14,8 +15,11 @@ class Storage(BaseStorage): def _get_object_key(self, file_hash, user_id): return '{}_{}'.format(file_hash, user_id) - def get_object_key(self, f): - return self._get_object_key(f.hash, f.user_id if f.user_id else 0) + def get_object_key(self, f, thumb=False): + key = self._get_object_key(f.hash, f.user_id if f.user_id else 0) + if thumb: + key += '_thumb' + return key def store_file(self, uploaded_file, file_hash, user, ip): bucket = self.client.Bucket(self.app.config['S3_BUCKET']) @@ -27,8 +31,13 @@ class Storage(BaseStorage): size = obj.size return self.add_file(file_hash, uploaded_file.filename, size, user, ip) - def get_file(self, f): - obj = self.client.Object(self.app.config['S3_BUCKET'], self.get_object_key(f)) + def get_file(self, f, thumb=True): + key = self.get_object_key(f, thumb=thumb) + if thumb: + bucket = self.app.config['S3_THUMB_BUCKET']) + else: + bucket = self.app.config['S3_BUCKET'] + obj = self.client.Object(bucket, key) kwargs = {} if 'Range' in request.headers: kwargs['Range'] = request.headers['Range'] @@ -44,6 +53,8 @@ class Storage(BaseStorage): def delete_file(self, f): obj = self.client.Object(self.app.config['S3_BUCKET'], self.get_object_key(f)) obj.delete() + obj = self.client.Object(self.app.config['S3_BUCKET'], self.get_object_key(f, thumb=True)) + obj.delete() @contextlib.contextmanager def temp_file(self, f): @@ -53,3 +64,16 @@ class Storage(BaseStorage): f.seek(0) yield f + def get_thumbnail(self, f): + try: + return self.get_file(f, thumb=True) + except botocore.exceptions.ClientError as e: + if e.response['Error']['Code'] == 'NoSuchKey': + # If thumbnail does not exist, just return None. + return + raise + + def store_thumbnail(self, f, stream): + bucket = self.client.Bucket(self.app.config['S3_THUMB_BUCKET'])) + key = self.get_object_key(f, thumb=True) + obj = bucket.upload_fileobj(Fileobj=stream, Key=key) |