ai-video-fullstack/backend/services/input_assets.py

"""Short-lived uploaded images consumed by the realtime user-input protocol."""

from __future__ import annotations

import base64
import hashlib
import hmac
import json
import time
from dataclasses import dataclass
from uuid import uuid4

import settings
from services.image_processing import normalize_image
from services.object_storage import delete_object, get_object, put_object


INPUT_ASSET_PREFIX = "conversation-inputs/"


@dataclass(frozen=True)
class StoredInputImage:
    token: str
    width: int
    height: int
    size_bytes: int


def _b64encode(data: bytes) -> str:
    return base64.urlsafe_b64encode(data).rstrip(b"=").decode("ascii")


def _b64decode(data: str) -> bytes:
    padding = "=" * (-len(data) % 4)
    return base64.urlsafe_b64decode(f"{data}{padding}".encode("ascii"))


def _signature(payload: str) -> str:
    digest = hmac.new(
        settings.AUTH_SECRET_KEY.encode("utf-8"),
        payload.encode("ascii"),
        hashlib.sha256,
    ).digest()
    return _b64encode(digest)


def store_input_image(data: bytes) -> StoredInputImage:
    normalized, width, height = normalize_image(data)
    key = f"{INPUT_ASSET_PREFIX}{uuid4().hex}.jpg"
    put_object(key, normalized, "image/jpeg")
    expires_at = int(time.time()) + settings.INPUT_IMAGE_TOKEN_TTL_SECONDS
    encoded_payload = _b64encode(
        json.dumps(
            {"key": key, "exp": expires_at},
            separators=(",", ":"),
        ).encode("utf-8")
    )
    token = f"{encoded_payload}.{_signature(encoded_payload)}"
    return StoredInputImage(
        token=token,
        width=width,
        height=height,
        size_bytes=len(normalized),
    )


def _key_from_token(token: str) -> str:
    if not token or "." not in token:
        raise ValueError("图片附件 token 无效")
    encoded_payload, signature = token.rsplit(".", 1)
    if not hmac.compare_digest(_signature(encoded_payload), signature):
        raise ValueError("图片附件 token 签名无效")
    try:
        payload = json.loads(_b64decode(encoded_payload))
        key = str(payload.get("key") or "")
        expires_at = int(payload.get("exp") or 0)
    except (ValueError, TypeError, json.JSONDecodeError) as exc:
        raise ValueError("图片附件 token 内容无效") from exc
    if expires_at < int(time.time()):
        raise ValueError("图片附件已过期，请重新添加")
    if not key.startswith(INPUT_ASSET_PREFIX):
        raise ValueError("图片附件存储位置无效")
    return key


def consume_input_image(token: str) -> bytes:
    """Read a signed image once and remove its temporary object."""

    key = _key_from_token(token)
    data = get_object(key)
    try:
        delete_object(key)
    except Exception:
        # Consumption succeeded. A stale temporary object must not break the turn.
        pass
    return data


def discard_input_image(token: str) -> None:
    delete_object(_key_from_token(token))