validate_sideload_zip.py

"""Validate the canonical Nemotron-3-Diarization debug sideload archive."""

from __future__ import annotations

import argparse
import hashlib
import json
import posixpath
import zipfile
from pathlib import PurePosixPath
from typing import Any


ADAPTER_ID = "nemotron3-diarization-low-latency-v1"
FUSED_ADAPTER_ID = "nemotron3-diarization-low-latency-v2"
SUPPORTED_ADAPTER_IDS = {ADAPTER_ID, FUSED_ADAPTER_ID}
MANIFEST_PATH = "model-manifest.json"
MAX_ENTRIES = 64
MAX_ENTRY_BYTES = 128 * 1024 * 1024
MAX_TOTAL_BYTES = 256 * 1024 * 1024


def parse_args() -> argparse.Namespace:
    """Parse archive validation arguments."""

    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("archive", type=str)
    return parser.parse_args()


def _validate_path(path: str) -> None:
    """Reject non-canonical, absolute, or traversing ZIP paths."""

    pure = PurePosixPath(path)
    if (
        not path
        or path.startswith("/")
        or "\\" in path
        or any(part in ("", ".", "..") for part in pure.parts)
        or posixpath.normpath(path) != path
    ):
        raise ValueError(f"Unsafe ZIP path: {path!r}")


def validate_archive(path: str) -> dict[str, Any]:
    """Verify structure, limits, manifest identity, and every payload hash."""

    with zipfile.ZipFile(path) as archive:
        entries = archive.infolist()
        names = [entry.filename for entry in entries]
        if len(entries) > MAX_ENTRIES:
            raise ValueError(f"Too many ZIP entries: {len(entries)}")
        if len(names) != len(set(names)):
            raise ValueError("Duplicate ZIP entry")
        if MANIFEST_PATH not in names:
            raise ValueError(f"Missing {MANIFEST_PATH}")
        total = 0
        for entry in entries:
            _validate_path(entry.filename)
            if entry.is_dir():
                raise ValueError(f"Directory entries are not allowed: {entry.filename}")
            if entry.file_size > MAX_ENTRY_BYTES:
                raise ValueError(f"ZIP entry is too large: {entry.filename}")
            total += entry.file_size
            if total > MAX_TOTAL_BYTES:
                raise ValueError("ZIP uncompressed size exceeds limit")
            unix_type = (entry.external_attr >> 16) & 0o170000
            if unix_type not in (0, 0o100000):
                raise ValueError(f"Non-regular ZIP entry: {entry.filename}")

        manifest = json.loads(archive.read(MANIFEST_PATH))
        if (
            manifest.get("schema_version") != 1
            or manifest.get("adapter_id") not in SUPPORTED_ADAPTER_IDS
        ):
            raise ValueError("Unsupported sideload manifest")
        inventory = manifest.get("files")
        if not isinstance(inventory, list):
            raise ValueError("Manifest files must be a list")
        expected = {MANIFEST_PATH}
        inventory_paths: set[str] = set()
        for item in inventory:
            item_path = item.get("path")
            if not isinstance(item_path, str):
                raise ValueError("Manifest file path must be a string")
            _validate_path(item_path)
            if item_path in inventory_paths:
                raise ValueError(f"Duplicate manifest path: {item_path}")
            inventory_paths.add(item_path)
            expected.add(item_path)
            try:
                entry = archive.getinfo(item_path)
            except KeyError as error:
                raise ValueError(f"Missing inventoried file: {item_path}") from error
            content = archive.read(entry)
            if entry.file_size != item.get("bytes"):
                raise ValueError(f"Size mismatch: {item_path}")
            if hashlib.sha256(content).hexdigest() != item.get("sha256"):
                raise ValueError(f"SHA-256 mismatch: {item_path}")
            if item_path.endswith(".tflite") and content[4:8] != b"TFL3":
                raise ValueError(f"Invalid LiteRT flatbuffer: {item_path}")
        if set(names) != expected:
            raise ValueError(f"Unexpected ZIP entries: {sorted(set(names) ^ expected)}")

        graphs = manifest.get("ordered_graphs", [])
        graph_paths = [graph.get("path") for graph in graphs]
        if manifest["adapter_id"] == FUSED_ADAPTER_ID:
            expected_graphs = [
                "graphs/feature_stacker.tflite",
                "graphs/encoder_head.tflite",
            ]
        else:
            expected_graphs = ["graphs/feature_stacker.tflite"]
            expected_graphs.extend(
                f"graphs/encoder_layer_{index:02d}.tflite" for index in range(31)
            )
            expected_graphs.append("graphs/classification_head.tflite")
        if graph_paths != expected_graphs:
            raise ValueError("Unexpected graph count or invocation order")
        if manifest.get("total_uncompressed_bytes") != sum(
            int(item["bytes"]) for item in inventory
        ):
            raise ValueError("Manifest total_uncompressed_bytes mismatch")
        return manifest


def main() -> int:
    """Validate one archive and print its release identity."""

    args = parse_args()
    manifest = validate_archive(args.archive)
    print(
        json.dumps(
            {
                "adapter_id": manifest["adapter_id"],
                "evaluation_only": manifest["evaluation_only"],
                "graph_count": len(manifest["ordered_graphs"]),
                "source_revision": manifest["source_revision"],
                "total_uncompressed_bytes": manifest["total_uncompressed_bytes"],
            },
            indent=2,
            sort_keys=True,
        )
    )
    return 0


if __name__ == "__main__":
    raise SystemExit(main())