80 lines
3.1 KiB
Python
80 lines
3.1 KiB
Python
from __future__ import annotations
|
|
|
|
import io
|
|
import uuid
|
|
from functools import lru_cache
|
|
|
|
import boto3
|
|
from botocore.client import BaseClient
|
|
from botocore.exceptions import ClientError
|
|
from PIL import Image, UnidentifiedImageError
|
|
|
|
from .config import settings
|
|
|
|
|
|
class ScreenshotError(ValueError):
|
|
pass
|
|
|
|
|
|
ALLOWED_FORMATS = {"JPEG": ("jpg", "image/jpeg"), "PNG": ("png", "image/png"), "WEBP": ("webp", "image/webp")}
|
|
ALLOWED_UPLOADS = {"image/jpeg": {".jpg", ".jpeg"}, "image/png": {".png"}, "image/webp": {".webp"}}
|
|
|
|
|
|
def validate_upload_metadata(filename: str | None, content_type: str | None) -> None:
|
|
mime = (content_type or "").lower()
|
|
suffix = f".{(filename or '').rsplit('.', 1)[-1].lower()}" if "." in (filename or "") else ""
|
|
if mime not in ALLOWED_UPLOADS or suffix not in ALLOWED_UPLOADS[mime]:
|
|
raise ScreenshotError("filename extension and MIME type must match JPEG, PNG or WebP")
|
|
|
|
|
|
@lru_cache
|
|
def client() -> BaseClient:
|
|
return boto3.client("s3", endpoint_url=settings.s3_endpoint_url, aws_access_key_id=settings.s3_access_key, aws_secret_access_key=settings.s3_secret_key, region_name="us-east-1")
|
|
|
|
|
|
@lru_cache
|
|
def public_client() -> BaseClient:
|
|
return boto3.client("s3", endpoint_url=settings.s3_public_endpoint_url, aws_access_key_id=settings.s3_access_key, aws_secret_access_key=settings.s3_secret_key, region_name="us-east-1")
|
|
|
|
|
|
def prepare_image(raw: bytes) -> tuple[bytes, str, str]:
|
|
if not raw or len(raw) > settings.screenshot_max_bytes:
|
|
raise ScreenshotError("screenshot must be between 1 byte and 8 MB")
|
|
try:
|
|
with Image.open(io.BytesIO(raw)) as source:
|
|
source.verify()
|
|
with Image.open(io.BytesIO(raw)) as source:
|
|
fmt = (source.format or "").upper()
|
|
if fmt not in ALLOWED_FORMATS:
|
|
raise ScreenshotError("only JPEG, PNG and WebP images are accepted")
|
|
if source.width * source.height > 40_000_000:
|
|
raise ScreenshotError("image dimensions are too large")
|
|
extension, mime = ALLOWED_FORMATS[fmt]
|
|
clean = source.convert("RGB") if fmt in {"JPEG", "WEBP"} else source.convert("RGBA")
|
|
output = io.BytesIO()
|
|
clean.save(output, format=fmt, optimize=True)
|
|
return output.getvalue(), extension, mime
|
|
except (UnidentifiedImageError, OSError) as exc:
|
|
raise ScreenshotError("file is not a valid image") from exc
|
|
|
|
|
|
def upload_screenshot(raw: bytes, *, filename: str | None = None, content_type: str | None = None) -> str:
|
|
validate_upload_metadata(filename, content_type)
|
|
body, extension, mime = prepare_image(raw)
|
|
key = f"reports/{uuid.uuid4()}.{extension}"
|
|
s3 = client()
|
|
try:
|
|
s3.head_bucket(Bucket=settings.s3_bucket)
|
|
except ClientError:
|
|
s3.create_bucket(Bucket=settings.s3_bucket)
|
|
s3.put_object(Bucket=settings.s3_bucket, Key=key, Body=body, ContentType=mime)
|
|
return key
|
|
|
|
|
|
def signed_screenshot_url(key: str) -> str:
|
|
return public_client().generate_presigned_url("get_object", Params={"Bucket": settings.s3_bucket, "Key": key}, ExpiresIn=900)
|
|
|
|
|
|
def delete_screenshot(key: str) -> None:
|
|
client().delete_object(Bucket=settings.s3_bucket, Key=key)
|