File size: 3,169 Bytes
2c57606 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 | #!/usr/bin/env python3
"""Validate external model or third-party asset manifests without loading them."""
from __future__ import annotations
import argparse
import hashlib
import json
import sys
from pathlib import Path
from typing import Any, Iterable
def _digest(path: Path) -> str:
hasher = hashlib.sha256()
with path.open("rb") as source:
for chunk in iter(lambda: source.read(1024 * 1024), b""):
hasher.update(chunk)
return hasher.hexdigest()
def _model_entries(manifest: dict[str, Any], backend: str | None) -> Iterable[dict[str, Any]]:
yield from manifest.get("common", [])
if backend:
try:
yield from manifest["backends"][backend]
except KeyError as error:
available = ", ".join(sorted(manifest.get("backends", {})))
raise ValueError(f"Unknown backend {backend!r}; expected one of: {available}") from error
else:
for entries in manifest.get("backends", {}).values():
yield from entries
def _third_party_entries(manifest: dict[str, Any]) -> Iterable[dict[str, Any]]:
for entry in manifest.get("assets", []):
yield entry
def main() -> int:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--manifest", required=True, type=Path)
parser.add_argument("--root", required=True, type=Path)
parser.add_argument("--backend", help="validate one model backend instead of every model asset")
parser.add_argument(
"--verify-sha256",
action="store_true",
help="also calculate and compare SHA-256 values (optional, potentially slow)",
)
parser.add_argument("--verify-size", action="store_true", help="also compare recorded byte sizes")
args = parser.parse_args()
manifest = json.loads(args.manifest.read_text(encoding="utf-8"))
if "backends" in manifest:
entries = list(_model_entries(manifest, args.backend))
else:
if args.backend:
parser.error("--backend is only valid for a model manifest")
entries = list(_third_party_entries(manifest))
errors: list[str] = []
for entry in entries:
path = args.root / entry["path" if "path" in entry else "target"]
label = entry.get("role", entry.get("id", path.name))
if not path.is_file():
errors.append(f"missing {label}: {path}")
continue
expected_size = entry.get("size_bytes")
if args.verify_size and expected_size is not None and path.stat().st_size != expected_size:
errors.append(f"wrong size for {label}: {path} (expected {expected_size}, got {path.stat().st_size})")
continue
if args.verify_sha256 and entry.get("sha256") and _digest(path) != entry["sha256"]:
errors.append(f"SHA-256 mismatch for {label}: {path}")
if errors:
print("Asset verification failed:", file=sys.stderr)
print("\n".join(f"- {error}" for error in errors), file=sys.stderr)
return 1
print(f"Verified {len(entries)} asset(s) from {args.manifest} against {args.root}")
return 0
if __name__ == "__main__":
raise SystemExit(main())
|