FabledCurator/backend/app/utils/safe_probe.py

"""Subprocess-isolated media probes (Layer 3 of import resilience).

A malformed video or archive can hard-crash the worker process — a
decoder OOM, a native-lib segfault, or a decompression bomb. A hard
crash leaves no terminal flip, so the recovery sweep re-queues the row
and it crashes again: a poison-pill loop (the Layer-1 cap is the
backstop, but isolating the crash is better — the file gets a clean
terminal failure and the worker never dies).

These probes run the risky read in a way that contains the blast:

- Video: `ffprobe` is a separate binary, so a crash decoding the
  container kills only ffprobe (non-zero exit), never the worker. Also
  returns width/height, which the importer didn't previously capture
  for videos.
- Archive: an uncompressed-size guard (catches decompression bombs
  before they OOM anything) plus an integrity test in a spawned child
  (catches native-lib crashes on a malformed archive). A child segfault
  / OOM shows up as a non-zero exit code, not a dead worker.

Images are intentionally NOT probed here: Pillow raises (it doesn't
segfault) on the realistic corrupt-image cases, the importer already
catches that as an invalid_image skip, and a subprocess per image would
wreck deep-scan throughput on a large library. Add an image branch only
if a real image-induced worker crash is ever observed.

Operator-requested 2026-05-28 (Layer 3).
"""

import json
import multiprocessing as mp
import subprocess
from dataclasses import dataclass
from pathlib import Path

VIDEO_PROBE_TIMEOUT_SECONDS = 60
ARCHIVE_PROBE_TIMEOUT_SECONDS = 120
# Refuse archives whose total UNCOMPRESSED size exceeds this — the
# classic decompression-bomb guard (a 4 GB cap comfortably clears real
# art-pack archives while stopping a few-KB zip that expands to TB).
MAX_ARCHIVE_UNCOMPRESSED_BYTES = 4 * 1024 * 1024 * 1024


@dataclass(frozen=True)
class ProbeResult:
    ok: bool
    # crashed=True means the probe HARD-FAILED (subprocess killed by a
    # signal, OOM, or timeout) — the poison-pill signature. crashed=False
    # with ok=False means a clean rejection (corrupt-but-handled,
    # bomb-size-exceeded, integrity mismatch). Callers map crashed → a
    # terminal 'failed', clean → a 'skipped'/'failed' of their choosing.
    crashed: bool = False
    reason: str | None = None
    width: int | None = None
    height: int | None = None


def probe_video(path: Path, *, timeout: float = VIDEO_PROBE_TIMEOUT_SECONDS) -> ProbeResult:
    """Validate a video container + first video stream via ffprobe."""
    try:
        out = subprocess.run(
            [
                "ffprobe", "-v", "error",
                "-select_streams", "v:0",
                "-show_entries", "stream=width,height",
                "-of", "json", str(path),
            ],
            capture_output=True, text=True, timeout=timeout,
        )
    except subprocess.TimeoutExpired:
        return ProbeResult(ok=False, crashed=True, reason="ffprobe timed out")
    except OSError as exc:
        # ffprobe missing / not executable — environmental, not the
        # file's fault. Treat as a clean non-crash failure so the import
        # path can decide (it currently proceeds without dims).
        return ProbeResult(ok=False, crashed=False, reason=f"ffprobe unavailable: {exc}")
    if out.returncode != 0:
        return ProbeResult(
            ok=False, crashed=False,
            reason=f"ffprobe rejected the file: {out.stderr.strip()[:200]}",
        )
    try:
        streams = (json.loads(out.stdout) or {}).get("streams") or []
    except json.JSONDecodeError as exc:
        return ProbeResult(ok=False, crashed=False, reason=f"ffprobe output parse failed: {exc}")
    if not streams:
        return ProbeResult(ok=False, crashed=False, reason="no decodable video stream")
    return ProbeResult(
        ok=True, width=streams[0].get("width"), height=streams[0].get("height"),
    )


def probe_archive(path: Path, *, timeout: float = ARCHIVE_PROBE_TIMEOUT_SECONDS) -> ProbeResult:
    """Bomb-size guard + isolated integrity test for an archive."""
    ctx = mp.get_context("spawn")
    q = ctx.Queue()
    proc = ctx.Process(target=_archive_probe_target, args=(str(path), q))
    proc.start()
    proc.join(timeout)
    if proc.is_alive():
        proc.terminate()
        proc.join(5)
        return ProbeResult(ok=False, crashed=True, reason="archive probe timed out")
    if proc.exitcode != 0:
        # Negative exitcode = killed by signal (segfault); positive =
        # the child os._exit'd or was OOM-killed. Either way the file
        # hard-crashed the probe — the poison-pill signature.
        return ProbeResult(
            ok=False, crashed=True,
            reason=f"archive probe crashed (exit {proc.exitcode})",
        )
    try:
        outcome = q.get(timeout=5)
    except Exception:  # noqa: BLE001 — empty queue / broken pipe
        return ProbeResult(ok=False, crashed=True, reason="archive probe produced no result")
    status, detail = outcome
    if status == "ok":
        return ProbeResult(ok=True)
    return ProbeResult(ok=False, crashed=False, reason=detail)


def _archive_probe_target(path_str: str, q) -> None:
    """Runs in the spawned child. Reads member sizes (bomb guard) then
    runs the format's integrity test. Puts ('ok', None) or
    ('error', reason). A crash/OOM here never reaches the queue — the
    parent reads the non-zero exit code instead."""
    path = Path(path_str)
    ext = path.suffix.lower()
    try:
        total, test_bad = _inspect_archive(path, ext)
    except Exception as exc:  # noqa: BLE001 — clean rejection
        q.put(("error", f"{type(exc).__name__}: {exc}"))
        return
    if total is not None and total > MAX_ARCHIVE_UNCOMPRESSED_BYTES:
        gib = total / (1024 ** 3)
        q.put(("error", f"uncompressed size {gib:.1f} GiB exceeds the bomb-guard cap"))
        return
    if test_bad is not None:
        q.put(("error", f"integrity test failed at member {test_bad!r}"))
        return
    q.put(("ok", None))


def _inspect_archive(path: Path, ext: str):
    """Return (total_uncompressed_bytes | None, first_bad_member | None)
    for the archive. Format-specific; raises on a structurally-broken
    container (caught by the child as a clean rejection)."""
    if ext in (".zip", ".cbz"):
        import zipfile

        with zipfile.ZipFile(path) as zf:
            total = sum(zi.file_size for zi in zf.infolist())
            return total, zf.testzip()
    if ext == ".rar":
        import rarfile

        with rarfile.RarFile(path) as rf:
            total = sum(getattr(ri, "file_size", 0) for ri in rf.infolist())
            rf.testrar()
            return total, None
    if ext == ".7z":
        import py7zr

        with py7zr.SevenZipFile(path, "r") as zf:
            info = zf.archiveinfo()
            total = getattr(info, "uncompressed", None)
            ok = zf.test()  # True / None when all members pass
            return total, (None if ok in (True, None) else "7z test reported corruption")
    # Unknown extension — nothing to test; treat as clean.
    return None, None