113 lines
4.1 KiB
Python
Executable File
113 lines
4.1 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
|
|
import argparse
|
|
import hashlib
|
|
import json
|
|
import re
|
|
import sys
|
|
from pathlib import Path
|
|
from typing import Dict, Iterable, Tuple
|
|
|
|
|
|
REQUIRED_ROLES = ("model", "lex", "srcvocab", "trgvocab")
|
|
|
|
|
|
class RuntimeDataError(RuntimeError):
|
|
pass
|
|
|
|
|
|
def _version_key(version: object) -> Tuple[Tuple[int, object], ...]:
|
|
parts = re.split(r"(\d+)", str(version))
|
|
return tuple((0, int(part)) if part.isdigit() else (1, part) for part in parts)
|
|
|
|
|
|
def _sha256(path: Path) -> str:
|
|
digest = hashlib.sha256()
|
|
with path.open("rb") as source:
|
|
for chunk in iter(lambda: source.read(1024 * 1024), b""):
|
|
digest.update(chunk)
|
|
return digest.hexdigest()
|
|
|
|
|
|
def _load_records(records_path: Path) -> Iterable[dict]:
|
|
try:
|
|
document = json.loads(records_path.read_text(encoding="utf-8"))
|
|
except (OSError, UnicodeError, json.JSONDecodeError) as error:
|
|
raise RuntimeDataError(f"invalid records.json: {error}") from error
|
|
|
|
records = document.get("data") if isinstance(document, dict) else None
|
|
if not isinstance(records, list):
|
|
raise RuntimeDataError("invalid records.json: data must be a list")
|
|
return records
|
|
|
|
|
|
def verify_runtime_data(records_path: Path, model_dir: Path) -> Dict[str, Path]:
|
|
candidates = [
|
|
record
|
|
for record in _load_records(records_path)
|
|
if isinstance(record, dict)
|
|
and record.get("sourceLanguage") == "en"
|
|
and record.get("targetLanguage") == "zh-Hans"
|
|
]
|
|
|
|
selected = {}
|
|
for role in REQUIRED_ROLES:
|
|
role_records = [record for record in candidates if record.get("fileType") == role]
|
|
if not role_records:
|
|
raise RuntimeDataError(f"missing required record role: {role}")
|
|
selected[role] = max(role_records, key=lambda record: _version_key(record.get("version", "")))
|
|
|
|
verified = {}
|
|
for role, record in selected.items():
|
|
attachment = record.get("attachment")
|
|
filename = attachment.get("filename") if isinstance(attachment, dict) else None
|
|
if not isinstance(filename, str) or not filename:
|
|
raise RuntimeDataError(f"invalid attachment filename for role: {role}")
|
|
decompressed_name = filename[:-4] if filename.endswith(".zst") else filename
|
|
if Path(decompressed_name).name != decompressed_name:
|
|
raise RuntimeDataError(f"unsafe attachment filename for role: {role}")
|
|
|
|
model_path = model_dir / decompressed_name
|
|
if not model_path.is_file():
|
|
raise RuntimeDataError(f"missing model file for role {role}: {decompressed_name}")
|
|
|
|
expected_size = record.get("decompressedSize")
|
|
if expected_size is not None and model_path.stat().st_size != expected_size:
|
|
raise RuntimeDataError(f"size mismatch for role {role}: {decompressed_name}")
|
|
|
|
expected_hash = record.get("decompressedHash")
|
|
if not isinstance(expected_hash, str) or not expected_hash:
|
|
raise RuntimeDataError(f"missing decompressedHash for role: {role}")
|
|
if _sha256(model_path) != expected_hash.lower():
|
|
raise RuntimeDataError(f"hash mismatch for role {role}: {decompressed_name}")
|
|
verified[role] = model_path
|
|
|
|
expected_names = {path.name for path in verified.values()}
|
|
actual_names = {entry.name for entry in model_dir.iterdir()}
|
|
extras = sorted(actual_names - expected_names)
|
|
if extras:
|
|
raise RuntimeDataError(f"unexpected model files: {', '.join(extras)}")
|
|
|
|
return verified
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser(description="Verify the fixed en to zh-Hans runtime data")
|
|
parser.add_argument("--records", required=True, type=Path)
|
|
parser.add_argument("--model-dir", required=True, type=Path)
|
|
args = parser.parse_args()
|
|
|
|
try:
|
|
verified = verify_runtime_data(args.records, args.model_dir)
|
|
except RuntimeDataError as error:
|
|
print(f"runtime data verification failed: {error}", file=sys.stderr)
|
|
return 1
|
|
|
|
for role in REQUIRED_ROLES:
|
|
print(f"verified {role}: {verified[role]}")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|