#!/usr/bin/env python3
"""Rebuild and verify the registered UML L0 RCA with its existing private key."""

from __future__ import annotations

import getpass
import hashlib
import os
from pathlib import Path
import shutil
import subprocess
import sys
import tarfile
import zipfile

UPSTREAM_COMMIT = "e3bb3b82480d6df4237e2a8c35ea0dd7eade25b4"
REGISTERED_SHA256 = "7168E341764188E140EF591634A9D2488D7FDACDA667EE53AFD566A8034ED5DB"
REGISTERED_ID8 = "AFD566A8034ED5DB"

HERE = Path(__file__).resolve().parent
ROOT = HERE.parent
REGISTERED_CERT = ROOT / "rca" / f"{REGISTERED_ID8}.oer"
REGISTERED_PUBLIC_KEY = ROOT / "rca" / "UML_L0_RCA_public.pem"
REGISTERED_PRIVATE_KEY = ROOT / "private" / "UML_L0_RCA_private_encrypted.pem"
PATCH_SOURCE = HERE / "src" / "generate_root.rs"
BUNDLED_SOURCE = HERE / "vendor" / "c-its-source-e3bb3b8.zip"
VENDORED_CRATES = HERE / "vendor" / "cargo-crates.tar.gz"
UPSTREAM_DIR = HERE / "_build_cits"
CARGO_VENDOR_DIR = HERE / "_cargo_vendor"
OUTPUT_DIR = HERE / "rebuild-output"
TEMPLATE_NAME = "registered_rca.oer"


def die(message: str) -> None:
    raise SystemExit(f"\nERROR: {message}")


def run(command: list[str], *, cwd: Path, input_bytes: bytes | None = None,
        capture: bool = False, env: dict[str, str] | None = None):
    print("+", " ".join(command))
    return subprocess.run(
        command, cwd=cwd, input=input_bytes, check=True, capture_output=capture, env=env
    )


def require_program(name: str) -> None:
    if shutil.which(name) is None:
        die(f"'{name}' is required and was not found on PATH.")


def verify_inputs() -> None:
    for path in (REGISTERED_CERT, REGISTERED_PUBLIC_KEY, REGISTERED_PRIVATE_KEY, PATCH_SOURCE):
        if not path.is_file():
            die(f"Required file is missing: {path}")
    actual = hashlib.sha256(REGISTERED_CERT.read_bytes()).hexdigest().upper()
    if actual != REGISTERED_SHA256:
        die(f"Registered certificate hash mismatch: {actual}")


def prepare_upstream() -> None:
    if not BUNDLED_SOURCE.is_file() or not VENDORED_CRATES.is_file():
        die("Bundled c-its source or Cargo crate archive is missing.")
    if UPSTREAM_DIR.exists():
        shutil.rmtree(UPSTREAM_DIR)
    if CARGO_VENDOR_DIR.exists():
        shutil.rmtree(CARGO_VENDOR_DIR)
    UPSTREAM_DIR.mkdir()
    CARGO_VENDOR_DIR.mkdir()
    with zipfile.ZipFile(BUNDLED_SOURCE) as archive:
        archive.extractall(UPSTREAM_DIR)
    with tarfile.open(VENDORED_CRATES, "r:gz") as archive:
        archive.extractall(CARGO_VENDOR_DIR)

    shutil.copy2(PATCH_SOURCE, UPSTREAM_DIR / "src" / "util" / "generate_root.rs")
    shutil.copy2(REGISTERED_CERT, UPSTREAM_DIR / TEMPLATE_NAME)

    cargo_toml = UPSTREAM_DIR / "Cargo.toml"
    text = cargo_toml.read_text(encoding="utf-8")
    old = 'ieee80211 = "0.5.9"'
    new = 'ieee80211 = { version = "0.5.9", default-features = false }'
    if old in text:
        cargo_toml.write_text(text.replace(old, new, 1), encoding="utf-8")
    elif new not in text:
        die("Expected ieee80211 dependency was not found in the pinned source.")

    cargo_config = UPSTREAM_DIR / ".cargo" / "config.toml"
    cargo_config.parent.mkdir()
    cargo_config.write_text(
        '[source.crates-io]\nreplace-with = "vendored-sources"\n\n'
        '[source.vendored-sources]\ndirectory = "../_cargo_vendor"\n\n'
        '[net]\noffline = true\n',
        encoding="utf-8",
    )


def build_generator() -> Path:
    run(
        ["cargo", "build", "--offline", "--release", "--features", "build-binary", "--bin", "c-its-generate-root"],
        cwd=UPSTREAM_DIR,
    )
    executable = UPSTREAM_DIR / "target" / "release" / (
        "c-its-generate-root.exe" if os.name == "nt" else "c-its-generate-root"
    )
    if not executable.is_file():
        die(f"Generator executable is missing: {executable}")
    return executable


def load_registered_scalar() -> bytearray:
    try:
        from cryptography.hazmat.primitives import serialization
        from cryptography.hazmat.primitives.asymmetric import ec
    except ImportError:
        die("Install the Python package 'cryptography'.")

    password = getpass.getpass("Registered RCA private-key passphrase: ").encode("utf-8")
    try:
        key = serialization.load_pem_private_key(REGISTERED_PRIVATE_KEY.read_bytes(), password=password)
    except Exception as exc:
        die(f"Could not unlock the registered private key: {exc}")

    if not isinstance(key, ec.EllipticCurvePrivateKey) or not isinstance(key.curve, ec.SECP256R1):
        die("Registered key is not NIST P-256.")
    expected_public = serialization.load_pem_public_key(REGISTERED_PUBLIC_KEY.read_bytes())
    if key.public_key().public_numbers() != expected_public.public_numbers():
        die("Private key does not match the registered public key.")
    return bytearray(key.private_numbers().private_value.to_bytes(32, "big"))


def rebuild(executable: Path) -> None:
    scalar = load_registered_scalar()
    if OUTPUT_DIR.exists():
        shutil.rmtree(OUTPUT_DIR)
    OUTPUT_DIR.mkdir()

    env = os.environ.copy()
    env["UML_TEMPLATE"] = str(UPSTREAM_DIR / TEMPLATE_NAME)
    env["UML_OUTPUT_DIR"] = str(OUTPUT_DIR)
    try:
        result = run(
            [str(executable)], cwd=UPSTREAM_DIR,
            input_bytes=(bytes(scalar).hex().upper() + "\n").encode("ascii"),
            capture=True, env=env,
        )
    finally:
        for index in range(len(scalar)):
            scalar[index] = 0

    sys.stdout.write(result.stdout.decode("utf-8", errors="replace"))
    sys.stderr.write(result.stderr.decode("utf-8", errors="replace"))
    rebuilt = OUTPUT_DIR / f"{REGISTERED_ID8}.oer"
    if not rebuilt.is_file():
        die(f"Expected rebuilt certificate is missing: {rebuilt}")
    if rebuilt.read_bytes() != REGISTERED_CERT.read_bytes():
        die("Rebuilt certificate is not byte-for-byte identical to the EU-registered certificate.")
    print("\nPASS: rebuilt certificate exactly matches the registered certificate.")
    print(f"Validation output: {OUTPUT_DIR}")


def main() -> None:
    require_program("cargo")
    verify_inputs()
    prepare_upstream()
    rebuild(build_generator())


if __name__ == "__main__":
    main()