ai-agent/backend/files/avatar.py

"""The custom avatar set (a sprite-editor "Export set" zip)."""

import io
import json
import pathlib
import shutil
import uuid as uuidlib
import zipfile

from fastapi import APIRouter, File, HTTPException, UploadFile
from fastapi.responses import FileResponse

from core.config import (
    _UNSAFE_NAME_RE,
    AVATARS_DIR,
    CUSTOM_AVATAR_DIR,
    MAX_AVATAR_ZIP_BYTES,
)

router = APIRouter()


# ── custom avatar set (a sprite-editor "Export set" zip) ──────────────────────
# The zip is one square-cell strip per clip plus a manifest.json in the Hemp
# Henry shape ({clip: {cell, fps, frames, loop, src}}). We validate it, unpack a
# clean copy into the writable data volume, and serve the strips inline so the
# roaming avatar plays a full custom emote set. Only strips referenced by the
# manifest are kept, and only as safe raster images — the zip never lands as-is.
_AVATAR_STRIP_EXTS = {".webp", ".png", ".gif"}


def _read_custom_manifest() -> dict:
    """The stored custom manifest with each src rewritten to its serve URL, or
    {} when no set is imported."""
    mf = CUSTOM_AVATAR_DIR / "manifest.json"
    if not mf.is_file():
        return {}
    try:
        raw = json.loads(mf.read_text())
    except Exception:
        return {}
    out: dict = {}
    for clip, meta in (raw or {}).items():
        if not isinstance(meta, dict):
            continue
        src = str(meta.get("src") or f"{clip}.webp")
        name = pathlib.PurePosixPath(src.replace("\\", "/")).name
        if not (CUSTOM_AVATAR_DIR / name).is_file():
            continue
        out[clip] = {**meta, "src": f"/api/avatar-asset?path={name}"}
    return out


@router.get("/api/avatar/manifest")
def avatar_manifest():
    """The imported custom avatar set's manifest (empty {} if none), with each
    clip src pointing at /api/avatar-asset so the frontend can render it."""
    return _read_custom_manifest()


@router.get("/api/avatar-asset")
def avatar_asset(path: str):
    """Serve one strip from the imported custom avatar set (inline image)."""
    name = pathlib.PurePosixPath((path or "").replace("\\", "/")).name
    ext = pathlib.Path(name).suffix.lower()
    if not name or ext not in _AVATAR_STRIP_EXTS:
        raise HTTPException(400, "bad path")
    p = (CUSTOM_AVATAR_DIR / name).resolve()
    try:
        p.relative_to(CUSTOM_AVATAR_DIR)
    except ValueError:
        raise HTTPException(400, "bad path")
    if not p.is_file():
        raise HTTPException(404, "avatar asset not found")
    # Pin the image content-type: the slim container's mimetypes DB can lack the
    # .webp mapping, and with nosniff a text/plain response won't render as an
    # image (breaking the sprite background). All _AVATAR_STRIP_EXTS are rasters.
    media = {".webp": "image/webp", ".png": "image/png", ".gif": "image/gif"}[ext]
    return FileResponse(p, media_type=media, headers={"X-Content-Type-Options": "nosniff"})


@router.post("/api/avatar/import")
async def avatar_import(file: UploadFile = File(...)):
    """Import a custom avatar set from a sprite-editor "Export set" zip.

    Validates the zip has a manifest.json plus one strip per referenced clip,
    then unpacks a clean copy into CUSTOM_AVATAR_DIR (replacing any previous
    set) and returns the served manifest. Rejects anything that isn't a small
    zip of a manifest + raster strips — no path traversal, no oversized members,
    nothing executable ever written."""
    raw = b""
    while chunk := await file.read(1 << 20):
        raw += chunk
        if len(raw) > MAX_AVATAR_ZIP_BYTES:
            raise HTTPException(413, f"zip exceeds {MAX_AVATAR_ZIP_BYTES // (1 << 20)} MB")
    if not raw:
        raise HTTPException(400, "empty upload")

    try:
        zf = zipfile.ZipFile(io.BytesIO(raw))
    except zipfile.BadZipFile:
        raise HTTPException(400, "not a valid zip")

    # Index members by basename (tolerate a wrapping top-level folder).
    members: dict[str, zipfile.ZipInfo] = {}
    total = 0
    for info in zf.infolist():
        if info.is_dir():
            continue
        name = pathlib.PurePosixPath(info.filename.replace("\\", "/")).name
        if not name or name.startswith("."):
            continue
        total += info.file_size
        if total > MAX_AVATAR_ZIP_BYTES:
            raise HTTPException(413, "zip contents too large")
        members[name] = info

    if "manifest.json" not in members:
        raise HTTPException(400, "zip is missing manifest.json")
    try:
        manifest = json.loads(zf.read(members["manifest.json"].filename))
    except Exception:
        raise HTTPException(400, "manifest.json is not valid JSON")
    if not isinstance(manifest, dict) or not manifest:
        raise HTTPException(400, "manifest.json must be a non-empty object of clips")

    # Validate every referenced strip is present and a safe raster image.
    clean: dict = {}
    strips: dict[str, bytes] = {}
    for clip, meta in manifest.items():
        if not isinstance(clip, str) or not _UNSAFE_NAME_RE.sub("", clip):
            raise HTTPException(400, f"bad clip name: {clip!r}")
        if not isinstance(meta, dict):
            raise HTTPException(400, f"clip {clip!r}: entry must be an object")
        src = str(meta.get("src") or f"{clip}.webp")
        name = pathlib.PurePosixPath(src.replace("\\", "/")).name
        if pathlib.Path(name).suffix.lower() not in _AVATAR_STRIP_EXTS:
            raise HTTPException(400, f"clip {clip!r}: src must be a webp/png/gif strip")
        if name not in members:
            raise HTTPException(400, f"clip {clip!r}: strip {name} not in zip")
        try:
            frames = int(meta.get("frames") or 0)
            cell = int(meta.get("cell") or 0)
        except (TypeError, ValueError):
            raise HTTPException(400, f"clip {clip!r}: frames/cell must be numbers")
        strips[name] = zf.read(members[name].filename)
        clean[clip] = {
            "cell": cell,
            "fps": int(meta.get("fps") or 6) or 6,
            "frames": max(1, frames),
            "loop": bool(meta.get("loop", True)),
            "src": name,
        }

    # Write a clean set atomically-ish: build in a temp dir, then swap.
    AVATARS_DIR.mkdir(parents=True, exist_ok=True)
    tmp = AVATARS_DIR / f".custom-{uuidlib.uuid4().hex[:8]}"
    tmp.mkdir(parents=True, exist_ok=True)
    try:
        for name, data in strips.items():
            (tmp / name).write_bytes(data)
        (tmp / "manifest.json").write_text(json.dumps(clean, indent=2, sort_keys=True) + "\n")
        if CUSTOM_AVATAR_DIR.exists():
            shutil.rmtree(CUSTOM_AVATAR_DIR)
        tmp.rename(CUSTOM_AVATAR_DIR)
    except Exception:
        shutil.rmtree(tmp, ignore_errors=True)
        raise
    return {"manifest": _read_custom_manifest(), "clips": sorted(clean.keys())}


@router.delete("/api/avatar/custom")
def avatar_clear():
    """Remove the imported custom avatar set."""
    if CUSTOM_AVATAR_DIR.exists():
        shutil.rmtree(CUSTOM_AVATAR_DIR, ignore_errors=True)
    return {"ok": True}