import os from io import BytesIO import boto3 import requests from app.config import ( AWS_REGION, BUCKET, AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY, LOCAL_FILE_UPLOAD, UPLOAD_DIR, URL, ) if not LOCAL_FILE_UPLOAD: _session = boto3.Session( aws_access_key_id=AWS_ACCESS_KEY_ID, aws_secret_access_key=AWS_SECRET_ACCESS_KEY, region_name=AWS_REGION, ) def upload_from_bytesio(key: str, bs: BytesIO, content_type="string"): bs.seek(0) if LOCAL_FILE_UPLOAD: file_path = os.path.join(UPLOAD_DIR, key) file_dir = os.path.dirname(file_path) os.makedirs(file_dir, exist_ok=True) with open(file_path, "wb") as f: f.write(bs.read()) else: _session.resource("s3").Bucket(BUCKET).put_object( Key=key, Body=bs, ContentType=content_type, ) def upload_email_from_bytesio(path: str, bs: BytesIO, filename): bs.seek(0) if LOCAL_FILE_UPLOAD: file_path = os.path.join(UPLOAD_DIR, path) file_dir = os.path.dirname(file_path) os.makedirs(file_dir, exist_ok=True) with open(file_path, "wb") as f: f.write(bs.read()) else: _session.resource("s3").Bucket(BUCKET).put_object( Key=path, Body=bs, # Support saving a remote file using Http header # Also supports Safari. More info at # https://github.com/eligrey/FileSaver.js/wiki/Saving-a-remote-file#using-http-header ContentDisposition=f'attachment; filename="{filename}.eml";', ) def upload_from_url(url: str, upload_path): r = requests.get(url) upload_from_bytesio(upload_path, BytesIO(r.content)) def get_url(key: str, expires_in=3600) -> str: if LOCAL_FILE_UPLOAD: return URL + "/static/upload/" + key else: s3_client = _session.client("s3") return s3_client.generate_presigned_url( ExpiresIn=expires_in, ClientMethod="get_object", Params={"Bucket": BUCKET, "Key": key}, ) def delete(path: str): if LOCAL_FILE_UPLOAD: os.remove(os.path.join(UPLOAD_DIR, path)) else: o = _session.resource("s3").Bucket(BUCKET).Object(path) o.delete()