iris_x_hermes/gateway-plugin/media.py

"""Inbound media cache + outbound chunked streaming.

Inbound: ``media.upload`` (chunked binary frames) -> ``cache_*_from_bytes``
-> a ``media_ref`` the adapter attaches to the ``MessageEvent``. Enforces
size limit + sha256 + MIME re-sniff (the client's declared kind is not
trusted). Chunks are streamed to a temp file so RAM stays bounded during
the transfer; the file is read once at the end for the hermes cache
helpers.

Outbound: ``send_*`` -> register the (already delivery-validated) file in
the media registry, mint a ``media_id``, emit ``media.offer {media_id,
kind, mime, size, filename}``; serve bytes on ``media.pull`` as chunked
binary frames. Delivery-path security via ``validate_media_delivery_path``
(re-checked at pull time, so only allowed files are servable).

Reuses hermes ``cache_image/audio/video/document_from_bytes`` + the
``_looks_like_image`` / ``sniff_container`` magic-byte sniffers. Temp files
live under ``get_hermes_home()/"android"/media/tmp``.

Milestone M4.
"""

import asyncio
import contextlib
import hashlib
import logging
import os
import re
import tempfile
import threading
import time
import uuid
from dataclasses import dataclass
from pathlib import Path

from gateway.platforms.base import (
    _looks_like_image,
    cache_audio_from_bytes,
    cache_document_from_bytes,
    cache_image_from_bytes,
    cache_video_from_bytes,
)
from tools.audio_container import sniff_container

logger = logging.getLogger(__name__)

# One WS binary frame carries at most this many media bytes (docs/07 §7.5).
DEFAULT_CHUNK_BYTES = 256 * 1024

KINDS = ("image", "audio", "video", "document", "voice")

_SHA256_RE = re.compile(r"^[0-9a-f]{64}$")

# Magic-byte containers that are unambiguously audio (vs video-in-same-box).
_AUDIO_CONTAINERS = {"m4a", "ogg", "flac", "wav", "mp3", "aac"}
_VIDEO_CONTAINERS = {"mp4", "webm"}
# Longest file extension we trust from the client (e.g. ".webm").
_MAX_EXT_LEN = 6

# Extension -> MIME for outbound offers (the app picks a player/viewer from it).
_EXT_TO_MIME: dict[str, str] = {
    ".jpg": "image/jpeg",
    ".jpeg": "image/jpeg",
    ".png": "image/png",
    ".webp": "image/webp",
    ".gif": "image/gif",
    ".bmp": "image/bmp",
    ".heic": "image/heic",
    ".mp4": "video/mp4",
    ".mov": "video/quicktime",
    ".webm": "video/webm",
    ".mkv": "video/x-matroska",
    ".avi": "video/x-msvideo",
    ".3gp": "video/3gpp",
    ".ogg": "audio/ogg",
    ".opus": "audio/ogg",
    ".m4a": "audio/mp4",
    ".mp3": "audio/mpeg",
    ".wav": "audio/wav",
    ".flac": "audio/flac",
    ".aac": "audio/aac",
    ".pdf": "application/pdf",
    ".txt": "text/plain",
    ".md": "text/markdown",
    ".csv": "text/csv",
    ".html": "text/html",
    ".json": "application/json",
    ".zip": "application/zip",
    ".doc": "application/msword",
    ".docx": "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
    ".xls": "application/vnd.ms-excel",
    ".xlsx": "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
    ".ppt": "application/vnd.ms-powerpoint",
    ".pptx": "application/vnd.openxmlformats-officedocument.presentationml.presentation",
}

# MIME -> extension for inbound caching (the cache helpers take an ext).
_MIME_TO_EXT: dict[str, str] = {
    "image/jpeg": ".jpg",
    "image/png": ".png",
    "image/webp": ".webp",
    "image/gif": ".gif",
    "image/bmp": ".bmp",
    "image/heic": ".heic",
    "video/mp4": ".mp4",
    "video/quicktime": ".mov",
    "video/webm": ".webm",
    "video/x-matroska": ".mkv",
    "video/x-msvideo": ".avi",
    "audio/ogg": ".ogg",
    "audio/mp4": ".m4a",
    "audio/mpeg": ".mp3",
    "audio/wav": ".wav",
    "audio/flac": ".flac",
    "audio/aac": ".aac",
}


class MediaError(Exception):
    """Media flow failure carrying a wire ``error`` code."""

    def __init__(self, code: str, message: str):
        super().__init__(message)
        self.code = code
        self.message = message


def mime_for_path(path: str) -> str:
    ext = os.path.splitext(path)[1].lower()
    return _EXT_TO_MIME.get(ext, "application/octet-stream")


def ext_for_mime(mime: str, filename: str, default: str) -> str:
    """Best-effort extension for the hermes cache helpers.

    Prefers the MIME map, falls back to the filename's own extension, then
    *default*.
    """
    base = (mime or "").split(";")[0].strip().lower()
    ext = _MIME_TO_EXT.get(base)
    if ext:
        return ext
    file_ext = os.path.splitext(filename or "")[1].lower()
    if file_ext and len(file_ext) <= _MAX_EXT_LEN:
        return file_ext
    return default


def reclassify_kind(head: bytes, declared_kind: str) -> str:
    """Re-sniff the real kind from magic bytes (don't trust the client).

    Images are detected first (RIFF/WEBP would otherwise look like a
    container); then AV containers. Video declared by the client is kept
    when no magic matches (some codecs have no leading magic we check);
    anything else falls back to ``document``.
    """
    if _looks_like_image(head):
        return "image"
    container = sniff_container(head)
    if container in _AUDIO_CONTAINERS:
        return "voice" if declared_kind == "voice" else "audio"
    if container in _VIDEO_CONTAINERS:
        return "video"
    return declared_kind if declared_kind == "video" else "document"


@dataclass
class MediaEntry:
    """One servable/cached media object (inbound ref or outbound offer)."""

    media_id: str  # mu_* (inbound ref, app-chosen) or md_* (outbound, minted)
    kind: str
    mime: str
    size: int
    filename: str
    path: str  # absolute local path
    ts: float = 0.0


class UploadSession:
    """One in-flight ``media.upload`` (per connection; one at a time).

    Chunks are appended to a temp file (bounded RAM); the sha256 is
    accumulated incrementally. The declared size is enforced as bytes
    arrive so an over-limit transfer is rejected early.
    """

    def __init__(
        self,
        media_ref: str,
        kind: str,
        mime: str,
        filename: str,
        declared_size: int,
        request_id: int | None,
        max_bytes: int,
        tmp_dir: Path,
    ):
        self.media_ref = media_ref
        self.kind = kind
        self.mime = mime
        self.filename = filename
        self.declared_size = declared_size
        self.request_id = request_id
        self.max_bytes = max_bytes
        self.received = 0
        self.failed = False
        self.error_code = ""
        self.error_message = ""
        self._sha = hashlib.sha256()
        fd, self.tmp_path = tempfile.mkstemp(prefix="upl_", dir=str(tmp_dir))
        self._fh = os.fdopen(fd, "wb")

    def feed(self, chunk: bytes) -> None:
        if self.failed:
            return
        self.received += len(chunk)
        if self.max_bytes and self.received > self.max_bytes:
            self._fail(
                "media_too_large",
                f"upload exceeds limit ({self.received} > {self.max_bytes} bytes)",
            )
            return
        try:
            self._fh.write(chunk)
            self._fh.flush()
            self._sha.update(chunk)
        except OSError as e:
            self._fail("internal", f"temp file write failed: {e}")

    def _fail(self, code: str, message: str) -> None:
        self.failed = True
        self.error_code = code
        self.error_message = message
        logger.warning("android: upload %s failed: %s", self.media_ref, message)

    def digest(self) -> str:
        return self._sha.hexdigest()

    def read_bytes(self) -> bytes:
        self._fh.close()
        with open(self.tmp_path, "rb") as f:
            return f.read()

    def close(self) -> None:
        """Discard the session and remove the temp file."""
        # Best-effort cleanup: a file that is already gone (or a handle that
        # is already closed) needs no further handling.
        with contextlib.suppress(Exception):
            self._fh.close()
        with contextlib.suppress(OSError):
            os.unlink(self.tmp_path)


class MediaStore:
    """In-memory media registry + upload session tracking.

    Inbound refs (``mu_*``, app-chosen) are short-lived: created on
    ``media.upload.end``, consumed by the next ``message.send``. Outbound
    offers (``md_*``, minted here) live until pulled (or process restart —
    the outbox replays the offer frame, but the id->path map is in-memory
    for M4).
    """

    def __init__(self, hermes_home: Path):
        self._tmp_dir = hermes_home / "android" / "media" / "tmp"
        self._tmp_dir.mkdir(parents=True, exist_ok=True)
        self._lock = threading.Lock()
        # (device_id, media_ref) -> UploadSession (one active per device)
        self._uploads: dict[tuple[str, str], UploadSession] = {}
        self._inbound: dict[str, MediaEntry] = {}
        self._outbound: dict[str, MediaEntry] = {}

    # ── Inbound uploads ───────────────────────────────────────────────────

    def create_upload(
        self,
        device_id: str,
        media_ref: str,
        kind: str,
        mime: str,
        filename: str,
        declared_size: int,
        request_id: int | None,
        max_bytes: int,
    ) -> UploadSession:
        with self._lock:
            for dev, _ref in self._uploads:
                if dev == device_id:
                    raise MediaError(
                        "unsupported", "an upload is already in progress on this connection"
                    )
            if media_ref in self._inbound:
                raise MediaError("unsupported", f"media_ref {media_ref} already used")
            sess = UploadSession(
                media_ref,
                kind,
                mime,
                filename,
                declared_size,
                request_id,
                max_bytes,
                self._tmp_dir,
            )
            self._uploads[(device_id, media_ref)] = sess
            return sess

    def get_upload(self, device_id: str, media_ref: str | None = None) -> UploadSession | None:
        with self._lock:
            if media_ref is not None:
                return self._uploads.get((device_id, media_ref))
            for (dev, _ref), sess in self._uploads.items():
                if dev == device_id:
                    return sess
            return None

    def discard_upload(self, device_id: str, media_ref: str) -> None:
        with self._lock:
            sess = self._uploads.pop((device_id, media_ref), None)
        if sess is not None:
            sess.close()

    def complete_upload(self, device_id: str, media_ref: str, sha256_hex: str) -> MediaEntry:
        """Verify + cache a finished upload. Raises ``MediaError`` on failure."""
        with self._lock:
            sess = self._uploads.pop((device_id, media_ref), None)
        if sess is None:
            raise MediaError("not_found", f"no active upload for media_ref {media_ref}")
        try:
            if sess.failed:
                raise MediaError(sess.error_code, sess.error_message)
            if sess.received != sess.declared_size:
                raise MediaError(
                    "internal",
                    f"size mismatch (declared {sess.declared_size}, received {sess.received})",
                )
            if not _SHA256_RE.match(sha256_hex or ""):
                raise MediaError("unsupported", "sha256 must be 64 hex chars")
            if sess.digest() != sha256_hex:
                raise MediaError("internal", "sha256 mismatch")

            data = sess.read_bytes()
            kind = reclassify_kind(data[:65536], sess.kind)
            try:
                if kind == "image":
                    path = cache_image_from_bytes(
                        data, ext_for_mime(sess.mime, sess.filename, ".jpg")
                    )
                elif kind in ("audio", "voice"):
                    path = cache_audio_from_bytes(
                        data, ext_for_mime(sess.mime, sess.filename, ".ogg")
                    )
                elif kind == "video":
                    path = cache_video_from_bytes(
                        data, ext_for_mime(sess.mime, sess.filename, ".mp4")
                    )
                else:
                    path = cache_document_from_bytes(data, sess.filename or "document")
            except ValueError as e:
                # hermes cap (gateway.max_inbound_media_bytes) or a
                # non-image payload masquerading as an image.
                if "too large" in str(e):
                    raise MediaError("media_too_large", str(e)) from e
                raise MediaError("unsupported", str(e)) from e

            entry = MediaEntry(
                media_id=media_ref,
                kind=kind,
                mime=mime_for_path(path),
                size=len(data),
                filename=os.path.basename(path),
                path=path,
                ts=time.time(),
            )
            with self._lock:
                self._inbound[media_ref] = entry
            logger.info(
                "android: upload %s cached as %s (%s, %d bytes)",
                media_ref,
                kind,
                path,
                len(data),
            )
            return entry
        finally:
            sess.close()

    def get_inbound(self, media_ref: str) -> MediaEntry | None:
        with self._lock:
            return self._inbound.get(media_ref)

    def pop_inbound(self, media_ref: str) -> MediaEntry | None:
        with self._lock:
            return self._inbound.pop(media_ref, None)

    # ── Outbound offers ───────────────────────────────────────────────────

    def register_outbound(
        self, path: str, kind: str, mime: str, filename: str, size: int
    ) -> MediaEntry:
        entry = MediaEntry(
            media_id=f"md_{uuid.uuid4().hex[:12]}",
            kind=kind,
            mime=mime,
            size=size,
            filename=filename or os.path.basename(path),
            path=path,
            ts=time.time(),
        )
        with self._lock:
            self._outbound[entry.media_id] = entry
        return entry

    def get_outbound(self, media_id: str) -> MediaEntry | None:
        with self._lock:
            return self._outbound.get(media_id)

    # ── Cleanup ───────────────────────────────────────────────────────────

    def discard_device(self, device_id: str) -> None:
        """Drop all in-flight uploads for a disconnected device."""
        with self._lock:
            stale = [k for k in self._uploads if k[0] == device_id]
            sessions = [self._uploads.pop(k) for k in stale]
        for sess in sessions:
            sess.close()

    def prune_outbound(self, max_age_s: float = 24 * 3600) -> int:
        """Drop offers older than *max_age_s* (their files stay in the hermes
        cache; only the id->path map is pruned)."""
        now = time.time()
        with self._lock:
            stale = [k for k, e in self._outbound.items() if now - e.ts > max_age_s]
            for k in stale:
                del self._outbound[k]
        return len(stale)


async def stream_file(
    ws, path: str, chunk_bytes: int = DEFAULT_CHUNK_BYTES, timeout: float = 10.0
) -> int:
    """Stream *path* to *ws* as binary frames. Returns bytes sent.

    Ordering is guaranteed by the WebSocket; the caller sends the terminal
    ``media.pull.end`` frame afterwards. Each chunk send is bounded by
    *timeout* so a stalled puller can't wedge the handler forever (the
    caller treats the raised error as an aborted pull).
    """
    sent = 0
    # Safe: ``path`` is produced by hermes ``cache_*_from_bytes`` (a path inside
    # hermes's own media cache dir), never derived from raw user input.
    # pi-lens-ignore: python-path-traversal
    with open(path, "rb") as f:
        while True:
            chunk = f.read(chunk_bytes)
            if not chunk:
                break
            await asyncio.wait_for(ws.send(chunk), timeout=timeout)
            sent += len(chunk)
    return sent