import os import tempfile import uuid from auth import login_required import boto3 from botocore.exceptions import BotoCoreError, ClientError from starlette.applications import Starlette from starlette.responses import JSONResponse from starlette.routing import Route from starlette.requests import Request STORJ_ACCESS_KEY = os.environ.get('STORJ_ACCESS_KEY') STORJ_SECRET_KEY = os.environ.get('STORJ_SECRET_KEY') STORJ_END_POINT = os.environ.get('STORJ_END_POINT') STORJ_BUCKET_NAME = os.environ.get('STORJ_BUCKET_NAME') CDN_DOMAIN = os.environ.get('CDN_DOMAIN') @login_required async def upload_handler(request: Request): form = await request.form() file = form.get('file') if file is None: return JSONResponse({'error': 'No file uploaded'}, status_code=400) file_name, file_extension = os.path.splitext(file.filename) key = str(uuid.uuid4()) + file_extension s3 = boto3.client('s3', aws_access_key_id=STORJ_ACCESS_KEY, aws_secret_access_key=STORJ_SECRET_KEY, endpoint_url=STORJ_END_POINT) try: with tempfile.NamedTemporaryFile() as tmp_file: while True: chunk = await file.read(8192) # 8192 bytes by default. if not chunk: break tmp_file.write(chunk) s3.upload_file( Filename=tmp_file.name, Bucket=STORJ_BUCKET_NAME, Key=key, ExtraArgs={ "ContentType": file.content_type } ) url = 'http://' + CDN_DOMAIN + '/' + key return JSONResponse({'url': url, 'originalFilename': file.filename}) except (BotoCoreError, ClientError) as e: print(e) return JSONResponse({'error': 'Failed to upload file'}, status_code=500) routes = [ Route('/', upload_handler, methods=['POST']), ] app = Starlette(debug=True, routes=routes) if __name__ == "__main__": import uvicorn uvicorn.run(app, host='0.0.0.0', port=8080)