s3_storage.py

import os
import re
import uuid
from typing import Any, Mapping, Optional

import boto3
from botocore.client import BaseClient


def _client() -> BaseClient:
    region = os.getenv("AWS_REGION", "ap-southeast-1")
    kwargs = {"region_name": region}
    key = os.getenv("AWS_ACCESS_KEY_ID", "").strip()
    secret = os.getenv("AWS_SECRET_ACCESS_KEY", "").strip()
    if key and secret:
        kwargs["aws_access_key_id"] = key
        kwargs["aws_secret_access_key"] = secret
    return boto3.client("s3", **kwargs)


def s3_enabled() -> bool:
    return bool(os.getenv("AWS_S3_BUCKET", "").strip())


def bucket_name() -> str:
    return os.getenv("AWS_S3_BUCKET", "").strip()


def upload_bytes(
    *,
    prefix: str,
    ext: str,
    body: bytes,
    content_type: str,
) -> dict:
    """
    Upload to s3://{bucket}/{prefix}/cat_{uuid}.{ext}
    prefix is 'upload' or 'AI_PHOTO'.
    """
    b = bucket_name()
    if not b:
        raise RuntimeError("AWS_S3_BUCKET is not set")

    safe_prefix = prefix.strip("/").replace("\\", "")
    e = _sanitize_ext(ext)
    uid = uuid.uuid4().hex
    key = f"{safe_prefix}/cat_{uid}.{e}"

    client = _client()
    client.put_object(Bucket=b, Key=key, Body=body, ContentType=content_type)

    uri = f"s3://{b}/{key}"
    return {"bucket": b, "key": key, "s3_uri": uri, "object_id": uid}


def _sanitize_ext(raw: str) -> str:
    e = (raw or "bin").lower().split(";")[0].strip()
    e = re.sub(r"[^a-z0-9]", "", e)[:8] or "bin"
    return e


def ext_to_content_type(ext: str) -> str:
    e = _sanitize_ext(ext)
    return {
        "png": "image/png",
        "jpg": "image/jpeg",
        "jpeg": "image/jpeg",
        "gif": "image/gif",
        "webp": "image/webp",
        "bmp": "image/bmp",
        "bin": "application/octet-stream",
    }.get(e, "application/octet-stream")


def presigned_get_url(s3_uri: str, expires: int = 604_800) -> Optional[str]:
    if not s3_uri or not s3_uri.startswith("s3://"):
        return None
    rest = s3_uri[5:]
    parts = rest.split("/", 1)
    if len(parts) != 2 or not parts[0] or not parts[1]:
        return None
    bucket, key = parts[0], parts[1]
    client = _client()
    return client.generate_presigned_url(
        "get_object",
        Params={"Bucket": bucket, "Key": key},
        ExpiresIn=expires,
    )


def resolve_public_or_presigned(config: Mapping[str, Any], stored: Optional[str]) -> Optional[str]:
    """
    - https?://  -> unchanged
    - s3://      -> presigned GET (or override via config)
    - /...       -> relative (caller / frontend adds API origin)
    """
    if not stored:
        return None
    s = stored.strip()
    if s.startswith("http://") or s.startswith("https://"):
        return s
    if s.startswith("s3://"):
        ttl = int(config.get("AWS_S3_PRESIGN_TTL", 604_800) or 604_800)
        return presigned_get_url(s, expires=ttl)
    return s