#!/usr/bin/env python3
"""Check the native libraries a release ships against a checked-in allow-list.

The inventory is keyed by PACKAGED path — `lib/<abi>/<name>` — so it covers every
native library in the artifact, not only the ones committed under jniLibs. A
library that arrives through a dependency (AAR/prefab) never appears in the source
tree, which is why the packaged artifact is checked directly.

Controls:
  1. Allow-list. Every native library in the source set and in the packaged APK
     must have an inventory record, and every record for a shipped ABI must be
     present. A newly added .so is refused until it is recorded and reviewed.
  2. Integrity. Each record's digest must match the bytes: the committed source
     file where one exists, otherwise the packaged entry.
  3. Provenance. The Noise core records must match the authoritative cohort in
     docs/NOISE-CORE-PROVENANCE.md.

Usage: check-native-provenance.py [release.apk]
"""

from hashlib import sha256
from pathlib import Path
import re
import sys
import zipfile

ROOT = Path(__file__).resolve().parents[1]
INVENTORY = ROOT / "docs" / "android-native-provenance-inventory.txt"
PROVENANCE = ROOT / "docs" / "NOISE-CORE-PROVENANCE.md"
JNI_LIBS = ROOT / "app" / "src" / "main" / "jniLibs"
NOISE_CORE = "libpeach_noise_core.so"


def digest_bytes(data: bytes) -> str:
    return sha256(data).hexdigest()


def digest_file(path: Path) -> str:
    hasher = sha256()
    with path.open("rb") as stream:
        for chunk in iter(lambda: stream.read(1024 * 1024), b""):
            hasher.update(chunk)
    return hasher.hexdigest()


def inventory_records() -> dict[str, str]:
    if not INVENTORY.is_file():
        raise SystemExit(f"checked-in native inventory not found: {INVENTORY}")
    records = {}
    for line in INVENTORY.read_text().splitlines():
        if not line.strip() or line.startswith("#"):
            continue
        try:
            path, expected = line.split("|", 1)
        except ValueError:
            raise SystemExit(f"invalid native inventory line: {line}")
        records[path.strip()] = expected.strip()
    return records


def abi_of(packaged_path: str) -> str:
    return Path(packaged_path).parent.name


records = inventory_records()
if not records:
    raise SystemExit("native inventory is empty")

# ── 1a. Source set: every committed library must be recorded, and match. ─────
if not JNI_LIBS.is_dir():
    raise SystemExit(f"native source set not found: {JNI_LIBS}")
source_set = {}
for path in sorted(JNI_LIBS.rglob("*.so")):
    if not path.is_file():
        continue
    packaged = f"lib/{path.parent.name}/{path.name}"
    source_set[packaged] = digest_file(path)

unrecorded_source = sorted(set(source_set) - set(records))
if unrecorded_source:
    raise SystemExit(
        "native source set contains unrecorded libraries: " + ", ".join(unrecorded_source)
    )
for packaged, actual in source_set.items():
    if records[packaged] != actual:
        raise SystemExit(
            f"committed native library digest mismatch for {packaged}: "
            f"recorded {records[packaged]}, actual {actual}"
        )

# ── 1b/2b. Packaged artifact: the shipped set for each ABI must equal the
# recorded set for that ABI, and every packaged entry must match its record. ──
if len(sys.argv) > 1:
    apk = Path(sys.argv[1])
    if not apk.is_file():
        raise SystemExit(f"release APK not found: {apk}")
    with zipfile.ZipFile(apk) as archive:
        packaged = {
            name: digest_bytes(archive.read(name))
            for name in archive.namelist()
            if name.startswith("lib/") and name.endswith(".so")
        }
    if not packaged:
        raise SystemExit("release APK contains no native libraries")
    shipped_abis = {abi_of(name) for name in packaged}
    expected = {name for name in records if abi_of(name) in shipped_abis}
    unrecorded = sorted(set(packaged) - set(records))
    missing = sorted(expected - set(packaged))
    if unrecorded or missing:
        raise SystemExit(
            "packaged native libraries do not match the inventory: "
            f"packaged but unrecorded {unrecorded}; recorded but absent {missing}"
        )
    for name, actual in packaged.items():
        if records[name] != actual:
            raise SystemExit(
                f"packaged native library digest mismatch for {name}: "
                f"recorded {records[name]}, actual {actual}"
            )

# ── 3. Provenance document agrees with the Noise core records. ──────────────
provenance = PROVENANCE.read_text()
for packaged, expected in records.items():
    if Path(packaged).name != NOISE_CORE:
        continue
    abi = abi_of(packaged)
    match = re.search(rf"^- `{re.escape(abi)}`: `([0-9a-f]+)`$", provenance, re.MULTILINE)
    if match is None:
        raise SystemExit(f"provenance document has no authoritative hash for {abi}")
    if match.group(1) != expected:
        raise SystemExit(
            f"provenance digest mismatch for {packaged}: "
            f"recorded {expected}, documented {match.group(1)}"
        )

print(f"PASS: native provenance verified ({len(records)} packaged libraries)")