53 lines
1.8 KiB
Python
53 lines
1.8 KiB
Python
from tempfile import NamedTemporaryFile
|
|
from uuid import uuid4
|
|
|
|
from libcloud.storage.types import Provider
|
|
from libcloud.storage.providers import get_driver
|
|
|
|
|
|
class UploadError(Exception):
|
|
pass
|
|
|
|
|
|
class Uploader:
|
|
_PERMITTED_MIMETYPES = ["application/pdf"]
|
|
|
|
def __init__(self, provider, container=None, key=None, secret=None):
|
|
self.container = self._get_container(provider, container, key, secret)
|
|
|
|
def upload(self, fyle):
|
|
# TODO: for hardening, we should probably use a better library for
|
|
# determining mimetype and not rely on FileUpload's determination
|
|
# TODO: we should set MAX_CONTENT_LENGTH in the config to prevent large
|
|
# uploads
|
|
if not fyle.mimetype in self._PERMITTED_MIMETYPES:
|
|
raise UploadError(
|
|
"could not upload {} with mimetype {}".format(
|
|
fyle.filename, fyle.mimetype
|
|
)
|
|
)
|
|
|
|
object_name = uuid4().hex
|
|
with NamedTemporaryFile() as tempfile:
|
|
tempfile.write(fyle.stream.read())
|
|
self.container.upload_object(
|
|
file_path=tempfile.name,
|
|
object_name=object_name,
|
|
extra={"acl": "private"},
|
|
)
|
|
return (fyle.filename, object_name)
|
|
|
|
def download_stream(self, object_name):
|
|
obj = self.container.get_object(object_name=object_name)
|
|
with NamedTemporaryFile() as tempfile:
|
|
obj.download(tempfile.name, overwrite_existing=True)
|
|
return open(tempfile.name, "rb")
|
|
|
|
def _get_container(self, provider, container, key, secret):
|
|
if provider == "LOCAL":
|
|
key = container
|
|
container = ""
|
|
|
|
driver = get_driver(getattr(Provider, provider))(key=key, secret=secret)
|
|
return driver.get_container(container)
|