1 files changed, 55 insertions, 0 deletions
diff --git a/fbin/file_storage/s3.py b/fbin/file_storage/s3.py
new file mode 100644
index 0000000..2f0b87b
--- /dev/null
+++ b/fbin/file_storage/s3.py
@@ -0,0 +1,55 @@
+import contextlib
+import tempfile
+
+import boto3
+from flask import request, send_file
+
+from .base import BaseStorage
+
+class Storage(BaseStorage):
+    def __init__(self, app):
+        super().__init__(app)
+        self.client = boto3.resource('s3', **self.app.config['S3_CONFIG'])
+
+    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 store_file(self, uploaded_file, file_hash, user, ip):
+        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)
+        size = uploaded_file.content_length
+        if not size:
+            obj = self.client.ObjectSummary(self.app.config['S3_BUCKET'], key)
+            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))
+        kwargs = {}
+        if 'Range' in request.headers:
+            kwargs['Range'] = request.headers['Range']
+        data = obj.get(**kwargs)
+        rv = send_file(data['Body'], attachment_filename=f.filename)
+        rv.headers['Content-Length'] = data['ContentLength']
+        rv.headers['Accept-Ranges'] = data['AcceptRanges']
+        if 'ContentRange' in data:
+            rv.headers['Content-Range'] = data['ContentRange']
+            rv.status_code = 206
+        return rv
+
+    def delete_file(self, f):
+        obj = self.client.Object(self.app.config['S3_BUCKET'], self.get_object_key(f))
+        obj.delete()
+
+    @contextlib.contextmanager
+    def temp_file(self, f):
+        obj = self.client.Object(self.app.config['S3_BUCKET'], self.get_object_key(f))
+        with tempfile.NamedTemporaryFile() as f:
+            obj.download_fileobj(f)
+            f.seek(0)
+            yield f
+