#!/usr/bin/env python3
"""Strict Ritual/GCP measurement adapter for the pinned cvm-measure release."""

from __future__ import annotations

import argparse
import hashlib
import json
import sys
from pathlib import Path

import cvm_measure
from cvm_measure.disk import compute_gpt_digest, extract_uki
from cvm_measure.tdx.baseline import load as load_baseline
from cvm_measure.tdx.ccel import parse_event_log
from cvm_measure.tdx.registers import _compute_mrtd, _compute_rtmr1
from cvm_measure.tdx.rtmr import replay_event_log


VERSION = "ritual-gcp-measure 5"
PINNED_CVM_MEASURE_VERSION = "0.1.0"
MEASUREMENT_MODE = "ritual-gcp-two-disk-complete-ccel-v3"
MAX_CCEL_BYTES = 1024 * 1024


def _require_file(path: Path, name: str, max_bytes: int | None = None) -> bytes:
    if not path.is_file():
        raise ValueError(f"{name} must be a regular file: {path}")
    if max_bytes is None:
        return path.read_bytes()
    with path.open("rb") as handle:
        data = handle.read(max_bytes + 1)
    if not data or len(data) > max_bytes:
        raise ValueError(f"{name} size is outside the allowed bounds")
    return data


def _parse_candidate_ccel(path: Path) -> object:
    return parse_event_log(
        _require_file(path, "candidate CCEL", max_bytes=MAX_CCEL_BYTES)
    )


def measure(args: argparse.Namespace) -> dict[str, str]:
    if cvm_measure.__version__ != PINNED_CVM_MEASURE_VERSION:
        raise ValueError(
            f"cvm-measure version must be {PINNED_CVM_MEASURE_VERSION}, "
            f"got {cvm_measure.__version__}"
        )
    firmware = _require_file(args.firmware, "firmware")
    uki = _require_file(args.uki, "UKI")
    baseline = load_baseline(args.baseline)
    candidate_log = _parse_candidate_ccel(args.candidate_ccel)
    replayed = replay_event_log(candidate_log)
    firmware_sha384 = hashlib.sha384(firmware).hexdigest()
    if baseline.firmware_sha384.lower() != firmware_sha384:
        raise ValueError("baseline firmware SHA-384 does not match firmware bytes")
    if baseline.machine_type != args.machine_type:
        raise ValueError(
            f"baseline machine type {baseline.machine_type!r} does not match "
            f"{args.machine_type!r}"
        )
    gpt_digest = compute_gpt_digest(args.disk, args.max_extract_bytes).hex()
    independent_rtmr1 = _compute_rtmr1(uki, baseline, gpt_digest)
    if replayed[1].hex() != independent_rtmr1:
        raise ValueError(
            "candidate CCEL replay rtmr1 differs from the independently "
            "reconstructed UKI/GPT measurement"
        )
    result = {
        "mrtd": _compute_mrtd(
            firmware,
            args.ram,
            args.numa_nodes,
            args.max_per_node,
        ),
        "rtmr0": replayed[0].hex(),
        "rtmr1": independent_rtmr1,
        "rtmr2": replayed[2].hex(),
        "rtmr3": replayed[3].hex(),
    }
    return result


def parse_args(argv: list[str]) -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--version", action="version", version=VERSION)
    parser.add_argument("--firmware", required=True, type=Path)
    parser.add_argument("--uki", required=True, type=Path)
    parser.add_argument("--disk", required=True, type=Path)
    parser.add_argument("--baseline", required=True, type=Path)
    parser.add_argument("--baseline-ccel", required=True, type=Path)
    parser.add_argument("--candidate-ccel", required=True, type=Path)
    parser.add_argument("--measurement-mode", required=True, choices=(MEASUREMENT_MODE,))
    parser.add_argument("--machine-type", required=True)
    parser.add_argument("--ram", required=True, type=int)
    parser.add_argument("--numa-nodes", required=True, type=int)
    parser.add_argument("--max-per-node", required=True, type=int)
    parser.add_argument("--max-extract-bytes", type=int, default=8 * 1024**3)
    parser.add_argument("--output-format", choices=("json",), default="json")
    args = parser.parse_args(argv)
    if args.ram <= 0 or args.numa_nodes <= 0 or args.max_per_node <= 0:
        parser.error("RAM and NUMA topology values must be positive")
    if args.numa_nodes * args.max_per_node < args.ram:
        parser.error("NUMA topology cannot hold requested RAM")
    return args


def main(argv: list[str] | None = None) -> int:
    try:
        result = measure(parse_args(sys.argv[1:] if argv is None else argv))
    except (OSError, ValueError) as error:
        print(f"error: {error}", file=sys.stderr)
        return 1
    print(json.dumps(result, sort_keys=True))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
