#!/usr/bin/env python3
"""Verify and restore a downloaded runtime. Offline only; never executes its source."""
import argparse
import hashlib
import json
from pathlib import Path, PurePosixPath
import stat
import tempfile
from zipfile import ZipFile

MAX_BYTES = 80 * 1024 * 1024


def safe_name(name):
    path = PurePosixPath(name)
    if (not name or path.is_absolute() or ".." in path.parts or "\\" in name or ":" in name
            or str(path) != name or str(path) == "."):
        raise ValueError("Unsafe runtime path")
    return name


def restore(source, destination, expected_sha256):
    source, destination = Path(source), Path(destination)
    if source.stat().st_size > MAX_BYTES:
        raise ValueError("Package too large")
    raw = source.read_bytes()
    if hashlib.sha256(raw).hexdigest() != expected_sha256:
        raise ValueError("Download hash mismatch; do not run")
    if destination.exists() or destination.is_symlink():
        raise ValueError("Choose a fresh destination")
    files = {}
    if source.suffix == ".zip":
        import io
        with ZipFile(io.BytesIO(raw)) as archive:
            total = 0
            for entry in archive.infolist():
                safe_name(entry.filename)
                if entry.is_dir():
                    continue
                if stat.S_ISLNK(entry.external_attr >> 16) or not entry.filename.startswith("researchstudio/"):
                    raise ValueError("Unexpected archive member")
                name = safe_name(entry.filename.removeprefix("researchstudio/"))
                total += entry.file_size
                if total > MAX_BYTES or len(files) >= 3000 or name in files:
                    raise ValueError("Duplicate or oversized runtime")
                files[name] = archive.read(entry)
    else:
        packet = json.loads(raw)
        if packet.get("format") != "researchstudio-text-runtime-v1":
            raise ValueError("Unknown runtime packet")
        for entry in packet["files"]:
            name = safe_name(entry["path"])
            data = entry["text"].encode("utf-8")
            if name in files or hashlib.sha256(data).hexdigest() != entry["sha256"]:
                raise ValueError("Duplicate or corrupt file")
            files[name] = data
        files["INTEGRITY.json"] = json.dumps({"format": "researchstudio-integrity-v1",
            "upstream_commit": packet["upstream_commit"], "files": {
                n: {"sha256": hashlib.sha256(d).hexdigest()} for n, d in files.items()}}).encode()
    if sum(map(len, files.values())) > MAX_BYTES or len(files) > 3000:
        raise ValueError("Runtime size limit")
    manifest = json.loads(files["INTEGRITY.json"])
    if manifest.get("format") != "researchstudio-integrity-v1":
        raise ValueError("Invalid integrity manifest")
    if set(manifest["files"]) != set(files) - {"INTEGRITY.json"}:
        raise ValueError("Manifest must cover every runtime file")
    for name, record in manifest["files"].items():
        safe_name(name)
        if hashlib.sha256(files[name]).hexdigest() != record["sha256"]:
            raise ValueError("Runtime file hash mismatch")
    destination.parent.mkdir(parents=True, exist_ok=True)
    with tempfile.TemporaryDirectory(prefix="rs-restore-", dir=destination.parent) as tmp:
        stage = Path(tmp) / "runtime"
        stage.mkdir()
        for name, data in files.items():
            path = stage / name
            path.parent.mkdir(parents=True, exist_ok=True)
            path.write_bytes(data)
        stage.rename(destination)
    return {"directory": str(destination.resolve()), "files": len(files),
            "upstream_commit": manifest["upstream_commit"]}


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("source", type=Path)
    parser.add_argument("destination", type=Path)
    parser.add_argument("--sha256", required=True)
    args = parser.parse_args()
    print(json.dumps(restore(args.source, args.destination, args.sha256), indent=2))