226 lines
6.8 KiB
Python
226 lines
6.8 KiB
Python
import boto
|
|
import os
|
|
import logging
|
|
import hashlib
|
|
import magic
|
|
|
|
from boto.s3.key import Key
|
|
from uuid import uuid4
|
|
from flask import url_for, request, send_file, make_response, abort
|
|
from flask.views import View
|
|
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class FakeUserfiles(object):
|
|
def prepare_for_drop(self, mime_type):
|
|
return ('http://fake/url', uuid4())
|
|
|
|
def store_file(self, file_like_obj, content_type):
|
|
raise NotImplementedError()
|
|
|
|
def get_file_url(self, file_id, expires_in=300):
|
|
return ('http://fake/url')
|
|
|
|
def get_file_checksum(self, file_id):
|
|
return 'abcdefg'
|
|
|
|
|
|
class S3FileWriteException(Exception):
|
|
pass
|
|
|
|
|
|
class S3Userfiles(object):
|
|
def __init__(self, path, s3_access_key, s3_secret_key, bucket_name):
|
|
self._initialized = False
|
|
self._bucket_name = bucket_name
|
|
self._access_key = s3_access_key
|
|
self._secret_key = s3_secret_key
|
|
self._prefix = path
|
|
self._s3_conn = None
|
|
self._bucket = None
|
|
|
|
def _initialize_s3(self):
|
|
if not self._initialized:
|
|
self._s3_conn = boto.connect_s3(self._access_key, self._secret_key)
|
|
self._bucket = self._s3_conn.get_bucket(self._bucket_name)
|
|
self._initialized = True
|
|
|
|
def prepare_for_drop(self, mime_type):
|
|
""" Returns a signed URL to upload a file to our bucket. """
|
|
self._initialize_s3()
|
|
logger.debug('Requested upload url with content type: %s' % mime_type)
|
|
file_id = str(uuid4())
|
|
full_key = os.path.join(self._prefix, file_id)
|
|
k = Key(self._bucket, full_key)
|
|
url = k.generate_url(300, 'PUT', headers={'Content-Type': mime_type},
|
|
encrypt_key=True)
|
|
return (url, file_id)
|
|
|
|
def store_file(self, file_like_obj, content_type):
|
|
self._initialize_s3()
|
|
file_id = str(uuid4())
|
|
full_key = os.path.join(self._prefix, file_id)
|
|
k = Key(self._bucket, full_key)
|
|
logger.debug('Setting s3 content type to: %s' % content_type)
|
|
k.set_metadata('Content-Type', content_type)
|
|
bytes_written = k.set_contents_from_file(file_like_obj, encrypt_key=True,
|
|
rewind=True)
|
|
|
|
if bytes_written == 0:
|
|
raise S3FileWriteException('Unable to write file to S3')
|
|
|
|
return file_id
|
|
|
|
def get_file_url(self, file_id, expires_in=300, mime_type=None):
|
|
self._initialize_s3()
|
|
full_key = os.path.join(self._prefix, file_id)
|
|
k = Key(self._bucket, full_key)
|
|
headers = None
|
|
if mime_type:
|
|
headers={'Content-Type': mime_type}
|
|
|
|
return k.generate_url(expires_in, headers=headers)
|
|
|
|
def get_file_checksum(self, file_id):
|
|
self._initialize_s3()
|
|
full_key = os.path.join(self._prefix, file_id)
|
|
k = self._bucket.lookup(full_key)
|
|
return k.etag[1:-1][:7]
|
|
|
|
|
|
class UserfilesHandlers(View):
|
|
methods = ['GET', 'PUT']
|
|
|
|
def __init__(self, local_userfiles):
|
|
self._userfiles = local_userfiles
|
|
self._magic = magic.Magic(mime=True)
|
|
|
|
def get(self, file_id):
|
|
path = self._userfiles.file_path(file_id)
|
|
if not os.path.exists(path):
|
|
abort(404)
|
|
|
|
logger.debug('Sending path: %s' % path)
|
|
return send_file(path, mimetype=self._magic.from_file(path))
|
|
|
|
def put(self, file_id):
|
|
input_stream = request.stream
|
|
if request.headers.get('transfer-encoding') == 'chunked':
|
|
# Careful, might work only with WSGI servers supporting chunked
|
|
# encoding (Gunicorn)
|
|
input_stream = request.environ['wsgi.input']
|
|
|
|
self._userfiles.store_stream(input_stream, file_id)
|
|
|
|
return make_response('Okay')
|
|
|
|
def dispatch_request(self, file_id):
|
|
if request.method == 'GET':
|
|
return self.get(file_id)
|
|
elif request.method == 'PUT':
|
|
return self.put(file_id)
|
|
|
|
|
|
class LocalUserfiles(object):
|
|
def __init__(self, app, path):
|
|
self._root_path = path
|
|
self._buffer_size = 64 * 1024 # 64 KB
|
|
self._app = app
|
|
|
|
def _build_url_adapter(self):
|
|
return self._app.url_map.bind(self._app.config['SERVER_HOSTNAME'],
|
|
script_name=self._app.config['APPLICATION_ROOT'] or '/',
|
|
url_scheme=self._app.config['PREFERRED_URL_SCHEME'])
|
|
|
|
def prepare_for_drop(self, mime_type):
|
|
file_id = str(uuid4())
|
|
with self._app.app_context() as ctx:
|
|
ctx.url_adapter = self._build_url_adapter()
|
|
return (url_for('userfiles_handlers', file_id=file_id, _external=True), file_id)
|
|
|
|
def file_path(self, file_id):
|
|
if '..' in file_id or file_id.startswith('/'):
|
|
raise RuntimeError('Invalid Filename')
|
|
return os.path.join(self._root_path, file_id)
|
|
|
|
def store_stream(self, stream, file_id):
|
|
path = self.file_path(file_id)
|
|
dirname = os.path.dirname(path)
|
|
if not os.path.exists(dirname):
|
|
os.makedirs(dirname)
|
|
|
|
with open(path, 'w') as to_write:
|
|
while True:
|
|
try:
|
|
buf = stream.read(self._buffer_size)
|
|
if not buf:
|
|
break
|
|
to_write.write(buf)
|
|
except IOError:
|
|
break
|
|
|
|
def store_file(self, file_like_obj, content_type):
|
|
file_id = str(uuid4())
|
|
|
|
# Rewind the file to match what s3 does
|
|
file_like_obj.seek(0, os.SEEK_SET)
|
|
|
|
self.store_stream(file_like_obj, file_id)
|
|
return file_id
|
|
|
|
def get_file_url(self, file_id, expires_in=300):
|
|
with self._app.app_context() as ctx:
|
|
ctx.url_adapter = self._build_url_adapter()
|
|
return url_for('userfiles_handlers', file_id=file_id, _external=True)
|
|
|
|
def get_file_checksum(self, file_id):
|
|
path = self.file_path(file_id)
|
|
sha_hash = hashlib.sha256()
|
|
with open(path, 'r') as to_hash:
|
|
while True:
|
|
buf = to_hash.read(self._buffer_size)
|
|
if not buf:
|
|
break
|
|
sha_hash.update(buf)
|
|
return sha_hash.hexdigest()[:7]
|
|
|
|
|
|
class Userfiles(object):
|
|
def __init__(self, app=None):
|
|
self.app = app
|
|
if app is not None:
|
|
self.state = self.init_app(app)
|
|
else:
|
|
self.state = None
|
|
|
|
def init_app(self, app):
|
|
storage_type = app.config.get('USERFILES_TYPE', 'LocalUserfiles')
|
|
path = app.config.get('USERFILES_PATH', '')
|
|
|
|
if storage_type == 'LocalUserfiles':
|
|
userfiles = LocalUserfiles(app, path)
|
|
app.add_url_rule('/userfiles/<file_id>',
|
|
view_func=UserfilesHandlers.as_view('userfiles_handlers',
|
|
local_userfiles=userfiles))
|
|
|
|
elif storage_type == 'S3Userfiles':
|
|
access_key = app.config.get('USERFILES_AWS_ACCESS_KEY', '')
|
|
secret_key = app.config.get('USERFILES_AWS_SECRET_KEY', '')
|
|
bucket = app.config.get('USERFILES_S3_BUCKET', '')
|
|
userfiles = S3Userfiles(path, access_key, secret_key, bucket)
|
|
|
|
elif storage_type == 'FakeUserfiles':
|
|
userfiles = FakeUserfiles()
|
|
|
|
else:
|
|
raise RuntimeError('Unknown userfiles type: %s' % storage_type)
|
|
|
|
# register extension with app
|
|
app.extensions = getattr(app, 'extensions', {})
|
|
app.extensions['userfiles'] = userfiles
|
|
return userfiles
|
|
|
|
def __getattr__(self, name):
|
|
return getattr(self.state, name, None)
|