engine_wrapper.py

"""
FALCONS.AI "Verify Anything" — engine wrapper (transport-independent).

Every function below is the FastAPI-era code, moved VERBATIM out of app.py so
the Gradio app and any future transport call exactly the same logic. The
standing rules (BUILD SPEC section 1) live here:
  1. The engine is invoked ONLY as a subprocess. Never ported/patched/inlined.
  2. No model-parsing code of any kind in this Space.
  3. Nothing retained: callers use per-request temp dirs, deleted in `finally`.
  4. A missing capability is never reported as tampering.
  5. Weights are never re-served.
"""

import base64
import fnmatch
import json
import os
import re
import subprocess
import sys
import zipfile
from pathlib import Path

# ---------------------------------------------------------------- constants

ROOT = Path(__file__).resolve().parent
ENGINE = ROOT / "engine" / "verify_attestation.py"
SAMPLES_DIR = ROOT / "samples"
MANIFEST = SAMPLES_DIR / "manifest.json"

MAX_UPLOAD_MB = int(os.environ.get("MAX_UPLOAD_MB", "1024"))
MAX_REPO_MB = int(os.environ.get("MAX_REPO_MB", "2048"))
MAX_UPLOAD_BYTES = MAX_UPLOAD_MB * 1024 * 1024
MAX_REPO_BYTES = MAX_REPO_MB * 1024 * 1024

ENGINE_TIMEOUT_S = 120
DOWNLOAD_TIMEOUT_S = 300
MAX_CONCURRENCY = 4

# SPEC 5 — appended verbatim to every error state.
LOCAL_CMD = (
    "You can always verify locally: python verify_attestation.py package.zip "
    "— github.com/Falcons-ai/surgeon-verify"
)

NOT_A_PACKAGE_COPY = (
    "This doesn't look like a Surgeon package — no signed attestation inside. "
    "Not an error, just not something this tool can vouch for. Export any "
    "model from Model Surgeon and drop it here, or try the two samples."
)

# SPEC 4.3.c — attestation filename patterns at repo root.
# TODO(SPEC 6.2): confirm the exact filename against the product fixture repo
# before enabling this path in production; patterns below follow the spec.
# Pinned per the shipped engine: it looks for exactly this name (SPEC 6.2).
ATTESTATION_PATTERNS = ("lineage.intoto.jsonl",)

# Bumped on every delivered build so a running instance is identifiable.
BUILD_STAMP = "2026-09-03.5-gradio"

# ---------------------------------------------------------------- helpers


def _result(verdict, exit_code, stdout, keyid=None, files_checked=None, detail=""):
    """Common response shape (SPEC 4). Never invent fields the engine
    didn't print."""
    return {
        "verdict": verdict,
        "exit_code": exit_code,
        "stdout": stdout,
        "keyid": keyid,
        "files_checked": files_checked,
        "detail": detail,
    }


# Defensive stdout parsing. If these miss, fields stay null and the raw
# stdout still shows — the engine's own words are always the receipt.
_KEYID_RE = re.compile(
    r"publisher\s+(\S+)"
)
# fallback for older engine builds that print "keyid: ..."
_KEYID_FALLBACK_RE = re.compile(
    r"key\s*[- ]?id\s*[:=]?\s*([0-9a-zA-Z:._-]{8,})"
)
_FILES_RE = re.compile(
    r"(\d+)\s+files?\s+(?:checked|verified|ok|passed)", re.I
)
_SUBJECT_LINE_RE = re.compile(r"^\s*[\u2714\u2716]\s+\S", re.M)


def _parse_stdout(stdout):
    """Best-effort keyid / files_checked from engine stdout (SPEC 4.5).
    Null when absent — never invented."""
    keyid = None
    m = _KEYID_RE.search(stdout)
    if not m:
        m = _KEYID_FALLBACK_RE.search(stdout)
    if m:
        keyid = m.group(1)
    files_checked = None
    m = _FILES_RE.search(stdout)
    if m:
        files_checked = int(m.group(1))
    else:
        n = len(_SUBJECT_LINE_RE.findall(stdout))
        if n:
            files_checked = n
    return keyid, files_checked


def _run_engine_sync(target):
    """Invoke the pinned engine as a subprocess (SPEC 4, all paths)."""
    try:
        proc = subprocess.run(
            [sys.executable, str(ENGINE), str(target)],
            capture_output=True,
            text=True,
            encoding="utf-8",
            errors="replace",
            timeout=ENGINE_TIMEOUT_S,
        )
    except subprocess.TimeoutExpired:
        return _result(
            "error",
            None,
            "",
            detail=(
                "Verification timed out after "
                + str(ENGINE_TIMEOUT_S)
                + " s. "
                + LOCAL_CMD
            ),
        )

    stdout = (proc.stdout or "") + (("\n" + proc.stderr) if proc.stderr else "")
    keyid, files_checked = _parse_stdout(proc.stdout or "")

    if proc.returncode == 0:
        return _result("verified", 0, stdout, keyid, files_checked)
    if proc.returncode == 1:
        if "VERDICT:" not in (proc.stdout or ""):
            # SCAN F2: the engine crashed before reaching a verdict (e.g. the
            # legacy-no-key path raises TypeError at the signature compare).
            # A crash is OUR bug — never evidence about the user's package,
            # and never rendered as TAMPERED (SPEC 1.4).
            if "legacy shared-key" in (proc.stdout or ""):
                _d = ("This is an older shared-key Surgeon package — the demo "
                      "can't check its signature without the key. Not evidence "
                      "of tampering. " + LOCAL_CMD +
                      " (add --key or FALCONSAI_HMAC_KEY)")
            else:
                _d = ("The verifier itself hit a bug on this input — that's "
                      "on our side, not evidence about your package. "
                      + LOCAL_CMD)
            return _result("error", 1, stdout, keyid, files_checked, _d)
        if "VERDICT: INCOMPLETE" in (proc.stdout or ""):
            return _result(
                "error", 1, stdout, keyid, files_checked,
                detail=(
                    "The engine could not fully establish the signature in "
                    "this environment (verdict INCOMPLETE). That is a missing "
                    "capability, not evidence of tampering. " + LOCAL_CMD
                ),
            )
        return _result("tampered", 1, stdout, keyid, files_checked)

    # Any other exit: classify by the engine's own words (SPEC 4). A missing
    # capability is never reported as tampering (SPEC 1.4).
    lower = (proc.stdout or "").lower() + (proc.stderr or "").lower()
    package_signals = (
        "no signed attestation",
        "no attestation",
        "attestation missing",
        "missing attestation",
        "not a zip",
        "bad zip",
        "not a valid",
        "not a package",
        "not a surgeon package",
        "no lineage.intoto.jsonl",
        "cannot open",
    )
    if any(sig in lower for sig in package_signals):
        return _result(
            "not_a_package", proc.returncode, stdout, keyid, files_checked,
            NOT_A_PACKAGE_COPY,
        )
    return _result(
        "error", proc.returncode, stdout, keyid, files_checked,
        "The verifier couldn't process this input. " + LOCAL_CMD,
    )


def _junk_entry(name):
    """Archiver metadata that no OS's archive tool should make load-bearing:
    macOS Finder resource forks and .DS_Store, Windows Thumbs.db /
    desktop.ini. Matched per path segment, case-insensitively."""
    parts = [seg.lower() for seg in name.replace("\\", "/").split("/") if seg]
    if not parts:
        return True
    if parts[0] == "__macosx":
        return True
    leaf = parts[-1]
    return leaf in (".ds_store", "thumbs.db", "desktop.ini") or leaf.startswith("._")


def _normalize_container(path, workdir):
    r"""Canonicalize a zip CONTAINER before handing it to the engine, so the
    verdict is OS-agnostic (SPEC intent: the signature covers file digests,
    not the archiver's packaging quirks).

    Repairs exactly two classes of archiver quirk:
      * Windows (Explorer / Compress-Archive): backslash entry names,
        which the ZIP APPNOTE forbids;
      * macOS (Finder "Compress"): __MACOSX/ resource-fork twins and
        .DS_Store entries, which break single-root detection.

    Entry NAMES are rewritten (\ -> /) and junk entries dropped; every
    remaining content byte is copied exactly as stored, CRC checks bypassed
    so corrupted (tampered) bytes still reach the engine's own digest
    comparison. On any doubt (name collision, unreadable container) the
    original file is returned untouched and the engine speaks for itself.
    Never touches package content (SPEC 1.1 / 1.2).
    """
    try:
        if not zipfile.is_zipfile(path):
            return path, False
        src = zipfile.ZipFile(path)
        names = src.namelist()
        needs = any("\\" in n for n in names) or any(_junk_entry(n) for n in names)
        if not needs:
            return path, False
        mapped = {}
        for info in src.infolist():
            new = info.filename.replace("\\", "/")
            if new.endswith("/") or _junk_entry(new):
                continue
            if new in mapped:
                return path, False  # collision -- refuse to guess
            mapped[new] = info
        if not mapped:
            return path, False
        fixed = Path(workdir) / "normalized.zip"
        with zipfile.ZipFile(fixed, "w", zipfile.ZIP_DEFLATED) as out:
            for new, info in mapped.items():
                f = src.open(info)
                f._expected_crc = None  # deliver bytes even if the CRC is broken
                out.writestr(new, f.read())
        return fixed, True
    except Exception:
        return path, False


NORMALIZED_NOTE = (
    "Note: the demo repaired this archive's container before verification "
    "(archiver quirks such as Windows path separators or macOS metadata "
    "entries) -- package contents untouched. "
)


def _zip_declared_size(path):
    """Sum of declared uncompressed sizes (zip-bomb guard, SPEC 4.2).
    Reads only the central directory — no member is extracted here."""
    with zipfile.ZipFile(path) as zf:
        return sum(i.file_size for i in zf.infolist())


_REPO_ID_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._\-]*/[A-Za-z0-9._\-]+$")


def _resolve_and_verify_repo_sync(repo_id, workdir):
    """SPEC 4.3 resolver. Runs in a worker thread; wall-clock capped by
    the caller."""
    from huggingface_hub import HfApi, hf_hub_download
    from huggingface_hub.errors import GatedRepoError, RepositoryNotFoundError

    token = os.environ.get("HF_TOKEN") or None
    api = HfApi(token=token)

    tree = None
    repo_type = None
    last_err = None
    for rt in ("model", "dataset"):
        try:
            tree = list(
                api.list_repo_tree(repo_id, repo_type=rt, recursive=True)
            )
            repo_type = rt
            break
        except GatedRepoError:
            return _result(
                "error",
                None,
                "",
                detail=(
                    "this demo verifies public repos — run the verifier "
                    "locally on files you have access to. " + LOCAL_CMD
                ),
            )
        except RepositoryNotFoundError as e:
            last_err = e
            continue
        except Exception as e:  # network etc.
            last_err = e
            tree = None
            break

    if tree is None:
        if isinstance(last_err, RepositoryNotFoundError):
            # Private repos also surface as not-found when anonymous.
            return _result(
                "error",
                None,
                "",
                detail=(
                    "Couldn't find that repo on the Hub. If it's private or "
                    "gated: this demo verifies public repos — run the "
                    "verifier locally on files you have access to. "
                    + LOCAL_CMD
                ),
            )
        return _result(
            "error",
            None,
            "",
            detail=(
                "Couldn't reach the Hugging Face Hub for this repo. "
                + LOCAL_CMD
            ),
        )

    files = [t for t in tree if getattr(t, "size", None) is not None]
    names = [t.path for t in files]
    sizes = {t.path: t.size for t in files}

    def _download(fname):
        p = hf_hub_download(
            repo_id=repo_id,
            filename=fname,
            repo_type=repo_type,
            local_dir=workdir,
            token=token,
        )
        return Path(p)

    # ---- SPEC 4.3.b — exactly one *.zip in the repo → verify as a package
    zips = [n for n in names if n.lower().endswith(".zip")]
    if len(zips) == 1:
        zname = zips[0]
        if sizes.get(zname, 0) > MAX_UPLOAD_BYTES:
            return _result(
                "error",
                None,
                "",
                detail=(
                    "The package in this repo is over the "
                    + str(MAX_UPLOAD_MB)
                    + " MB demo cap. "
                    + LOCAL_CMD
                ),
            )
        try:
            zpath = _download(zname)
        except Exception:
            return _result(
                "error",
                None,
                "",
                detail=(
                    "Downloading the package from the Hub failed. "
                    + LOCAL_CMD
                ),
            )
        zpath, normalized = _normalize_container(zpath, workdir)
        result = _run_engine_sync(zpath)
        if normalized and isinstance(result, dict):
            result["detail"] = NORMALIZED_NOTE + (result.get("detail") or "")
        return result

    # ---- SPEC 4.3.c — DSSE attestation at repo root → subjects-only verify
    root_files = [n for n in names if "/" not in n]
    att_name = None
    for n in root_files:
        if any(fnmatch.fnmatch(n, pat) for pat in ATTESTATION_PATTERNS):
            att_name = n
            break
    if att_name:
        try:
            att_path = _download(att_name)
            subjects = _dsse_subject_names(att_path)
        except Exception:
            return _result(
                "error",
                None,
                "",
                detail=(
                    "Found an attestation file in this repo but couldn't "
                    "fetch or read it. " + LOCAL_CMD
                ),
            )

        # Download only the files the attestation names as subjects
        # (SPEC 4.3.c).
        total = 0
        for s in subjects:
            if s not in sizes:
                # Missing subject: nothing to download — the engine will
                # flag it, and a missing subject IS tampered (SPEC 4.3.c).
                continue
            if sizes[s] > MAX_UPLOAD_BYTES:
                return _result(
                    "error",
                    None,
                    "",
                    detail=(
                        "A file named by the attestation is over the "
                        + str(MAX_UPLOAD_MB)
                        + " MB per-file demo cap. "
                        + LOCAL_CMD
                    ),
                )
            total += sizes[s]
            if total > MAX_REPO_BYTES:
                return _result(
                    "error",
                    None,
                    "",
                    detail=(
                        "The attested files in this repo total more than "
                        "the "
                        + str(MAX_REPO_MB)
                        + " MB demo cap. "
                        + LOCAL_CMD
                    ),
                )
            try:
                _download(s)
            except Exception:
                # Leave it absent; the engine decides what that means.
                pass

        # Hub-added files not named by the attestation: informational only,
        # never TAMPERED (SPEC 4.3.c).
        uncovered = [n for n in names if n not in subjects and n != att_name]

        # SCAN F1: the engine verifies ZIP packages only (a directory input
        # exits 2 "cannot open"). Package the attestation + downloaded
        # subjects into a canonical zip; a subject missing from the repo is
        # simply absent, and the engine's own "<missing from zip>" makes
        # that TAMPERED — exactly SPEC 4.3.c.
        pkg = Path(workdir) / "__attested_package.zip"
        with zipfile.ZipFile(pkg, "w", zipfile.ZIP_DEFLATED) as zf:
            zf.write(att_path, att_name)
            for s in subjects:
                sp = Path(workdir) / s
                if sp.is_file():
                    zf.write(sp, s)
        res = _run_engine_sync(pkg)
        if uncovered:
            note = (
                "present but not covered by the attestation: "
                + ", ".join(sorted(uncovered))
            )
            res["detail"] = (
                (res["detail"] + " · ") if res["detail"] else ""
            ) + note
        return res

    # ---- SPEC 4.3.d — nothing to verify against
    return _result(
        "repo_not_attested",
        None,
        "",
        detail=(
            "This repo doesn't carry a Surgeon attestation, so there's "
            "nothing to verify against — that's a statement about "
            "provenance, not quality. Repos published through Model "
            "Surgeon's attested push verify here."
        ),
    )


def _dsse_subject_names(att_path):
    """Read subject filenames out of a DSSE envelope (or a bare in-toto
    statement). This is JSON metadata only — no model content is parsed."""
    data = json.loads(att_path.read_text())
    if isinstance(data, dict) and "payload" in data:
        payload = base64.b64decode(data["payload"])
        data = json.loads(payload)
    subjects = data.get("subject", []) if isinstance(data, dict) else []
    return [
        s.get("name")
        for s in subjects
        if isinstance(s, dict) and s.get("name")
    ]