scripts/references.py

#!/usr/bin/env python3
"""Verify, safely unpack and index the original Microbrush archives; never execute them."""
import argparse
from contextlib import contextmanager
import hashlib
import json
import os
from pathlib import Path, PurePosixPath
import re
import shutil
import stat
import sys
import tempfile
import unicodedata
import zipfile

ROOT = Path(__file__).resolve().parents[1]
MANIFEST = ROOT / "config/references.json"
REFERENCES = ROOT / "assets/references"
MAX_FILES = 10000
MAX_BYTES = 512 * 1024 * 1024
MAX_RATIO = 1000
RESERVED = re.compile(r"^(CON|PRN|AUX|NUL|COM[1-9]|LPT[1-9])(?:\..*)?$", re.I)


def sha256(path):
    digest = hashlib.sha256()
    with path.open("rb") as stream:
        for chunk in iter(lambda: stream.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def safe_directory(path):
    # Resolve neither symlinks nor an existing payload through a symlinked ancestor.
    for part in [path, *path.parents]:
        if part.is_symlink():
            raise ValueError(f"Symlinked destination/source is not allowed: {part}")
    if path.exists() and not path.is_dir():
        raise ValueError(f"Not a directory: {path}")


def checked_members(archive):
    """Validate every name/type and total expansion before reading/extracting payloads."""
    members = archive.infolist()
    if len(members) > MAX_FILES or sum(i.file_size for i in members) > MAX_BYTES:
        raise ValueError("Archive exceeds extraction limits")
    seen = {}
    spellings = {}
    for info in members:
        name = info.filename
        pure = PurePosixPath(name)
        parts = name.rstrip("/").split("/")
        if (info.orig_filename != name or not name or pure.is_absolute() or "\\" in name or ":" in name
                or any(c < " " or c == "\x7f" for c in name)
                or any(p in ("", ".", "..") or p.endswith((" ", ".")) or RESERVED.match(p) for p in parts)):
            raise ValueError(f"Unsafe archive path: {name!r}")
        mode = info.external_attr >> 16
        kind = stat.S_IFMT(mode)
        if kind not in (0, stat.S_IFDIR if info.is_dir() else stat.S_IFREG):
            raise ValueError(f"Unsupported archive entry type: {name}")
        if info.flag_bits & 1:
            raise ValueError(f"Encrypted member is not supported: {name}")
        if info.file_size / max(1, info.compress_size) > MAX_RATIO:
            raise ValueError(f"Member exceeds compression ratio limit: {name}")
        key = unicodedata.normalize("NFC", name.rstrip("/")).casefold()
        # Include implicit ancestors: A/file and a/other collide on macOS too.
        for count in range(1, len(parts) + 1):
            spelling = "/".join(parts[:count])
            folded = unicodedata.normalize("NFC", spelling).casefold()
            if folded in spellings and spellings[folded] != spelling:
                raise ValueError(f"Case/Unicode path collision: {name}")
            spellings[folded] = spelling
        if key in seen:
            raise ValueError(f"Duplicate/case-fold collision: {name}")
        seen[key] = info.is_dir()
    for key in seen:
        for parent in PurePosixPath(key).parents:
            if str(parent) in seen and not seen[str(parent)]:
                raise ValueError(f"File-directory collision: {key}")
    return members


def inspect_archive(path, package):
    if path.is_symlink() or not path.is_file():
        raise ValueError(f"Missing archive (hydrate assets first): {path}")
    safe_directory(path.parent)
    if path.stat().st_size != package["bytes"] or sha256(path) != package["sha256"]:
        raise ValueError(f"Archive size/SHA-256 mismatch (possibly an LFS pointer): {path}")
    with zipfile.ZipFile(path) as archive:
        members = checked_members(archive)
        if len(members) != package["archive_entries"]:
            raise ValueError(f"Archive entry count mismatch: {path}")
        # Reading every member checks its ZIP CRC in addition to the outer hash.
        files = [{"path": i.filename, "bytes": i.file_size,
                  "sha256": hashlib.sha256(archive.read(i)).hexdigest(),
                  "crc32": f"{i.CRC:08x}", "zip_attributes": i.external_attr,
                  "zip_timestamp": list(i.date_time)} for i in members if not i.is_dir()]
    return {"schema": 1, "build": package["build"], "archive": package["filename"],
            "archive_sha256": package["sha256"], "files": files,
            "directories": [i.filename.rstrip("/") for i in members if i.is_dir()]}


def check_extracted(destination, inventory, require_directories=True):
    safe_directory(destination)
    if not destination.is_dir():
        raise ValueError(f"Missing extraction: {destination}")
    expected = {f["path"] for f in inventory["files"]}
    actual = set()
    expected_dirs = set(inventory["directories"])
    for name in expected | set(inventory["directories"]):
        expected_dirs.update(str(p) for p in PurePosixPath(name).parents if str(p) != ".")
    for path in destination.rglob("*"):
        relative = path.relative_to(destination).as_posix()
        if path.is_symlink() or not (path.is_file() or path.is_dir()):
            raise ValueError(f"Unexpected extracted entry: {path}")
        if path.is_file():
            actual.add(relative)
        elif relative not in expected_dirs:
            raise ValueError(f"Unexpected directory: {path}")
    if actual != expected:
        raise ValueError(f"Extraction member set differs: {destination}")
    for info in inventory["files"]:
        path = destination / info["path"]
        if path.stat().st_size != info["bytes"] or sha256(path) != info["sha256"]:
            raise ValueError(f"Modified extracted member: {path}")
    for directory in inventory["directories"]:
        if require_directories and not (destination / directory).is_dir():
            raise ValueError(f"Missing original empty directory: {directory}")


def extract(path, destination, inventory):
    safe_directory(destination.parent)
    if destination.exists() or destination.is_symlink():
        # Git stores files only. A hydrated clone is missing the ZIP's empty
        # directories; verify every existing file before restoring only those.
        check_extracted(destination, inventory, require_directories=False)
        missing = [d for d in inventory["directories"] if not (destination / d).is_dir()]
        for directory in missing:
            (destination / directory).mkdir(parents=True, exist_ok=True)
        check_extracted(destination, inventory)
        return "restored directories" if missing else "unchanged"
    destination.parent.mkdir(parents=True, exist_ok=True)
    if shutil.disk_usage(destination.parent).free < sum(f["bytes"] for f in inventory["files"]) + 1024 * 1024:
        raise ValueError("Insufficient disk space for extraction")
    with tempfile.TemporaryDirectory(prefix=".unpack-", dir=destination.parent) as staging:
        temporary = Path(staging) / "tree"
        temporary.mkdir()
        with zipfile.ZipFile(path) as archive:
            for info in checked_members(archive):
                target = temporary / info.filename
                if info.is_dir():
                    target.mkdir(parents=True, exist_ok=True)
                else:
                    target.parent.mkdir(parents=True, exist_ok=True)
                    with archive.open(info) as source, target.open("xb") as output:
                        shutil.copyfileobj(source, output)
                    # Preserve bytes, not executable permissions; originals remain in ZIP metadata.
                    target.chmod(0o644)
        check_extracted(temporary, inventory)
        temporary.rename(destination)
    return "extracted"


def write_json(path, data):
    safe_directory(path.parent)
    path.parent.mkdir(parents=True, exist_ok=True)
    if path.is_symlink():
        raise ValueError(f"Refusing symlink: {path}")
    content = json.dumps(data, indent=2, ensure_ascii=False) + "\n"
    if path.exists() and path.read_text() == content:
        return
    fd, temporary = tempfile.mkstemp(prefix=".inventory-", dir=path.parent)
    try:
        with os.fdopen(fd, "w") as output:
            output.write(content)
        os.replace(temporary, path)
    finally:
        if os.path.exists(temporary):
            os.unlink(temporary)


@contextmanager
def exclusive(root):
    safe_directory(root)
    root.mkdir(parents=True, exist_ok=True)
    lock = root / ".unpack.lock"
    try:
        lock.mkdir()
    except FileExistsError:
        raise ValueError(f"Another operation or interrupted lock exists: {lock}; inspect before removing") from None
    try:
        yield
    finally:
        lock.rmdir()


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("command", choices=("list", "verify", "unpack", "index"))
    parser.add_argument("--build", action="append", help="Select a build; repeat for multiple builds")
    parser.add_argument("--dry-run", action="store_true")
    parser.add_argument("--check-extracted", action="store_true", help="Also verify unpacked members")
    args = parser.parse_args()
    packages = json.loads(MANIFEST.read_text())["packages"]
    known = {p["build"] for p in packages}
    if args.build and set(args.build) - known:
        parser.error("Unknown builds: " + ", ".join(sorted(set(args.build) - known)))
    packages = [p for p in packages if not args.build or p["build"] in args.build]
    if args.command == "list":
        for p in packages:
            print(f"{p['build']}: {p['filename']} -> assets/references/extracted/{p['build']}")
        return
    # Verify all selected archives before making any output changes.
    inventories = [(p, inspect_archive(REFERENCES / "archives" / p["filename"], p)) for p in packages]
    def process():
        for p, inventory in inventories:
            destination = REFERENCES / "extracted" / p["build"]
            status = "verified"
            if args.command == "unpack" and not args.dry_run:
                status = extract(REFERENCES / "archives" / p["filename"], destination, inventory)
            if args.check_extracted:
                check_extracted(destination, inventory)
            if args.command in ("unpack", "index") and not args.dry_run:
                write_json(REFERENCES / "inventories" / (p["build"] + ".json"), inventory)
            print(f"{p['build']}: {status}, {len(inventory['files'])} files" + (" (dry run)" if args.dry_run else ""))
    if args.command in ("unpack", "index") and not args.dry_run:
        with exclusive(REFERENCES):
            process()
    else:
        process()


if __name__ == "__main__":
    try:
        main()
    except (ValueError, OSError, zipfile.BadZipFile, RuntimeError) as error:
        print(f"References: {error}", file=sys.stderr)
        sys.exit(1)