"""Intel TDX quote parsing and verification for Clairwave cloud workspaces (2026-10-23).

Used by the workspace controller (before a VM gets its TLS certificate) and by the standalone customer verifier
verify_workspace.py. Python 3.9+ and `cryptography` only; no network access here.

What a successful verify() proves:
  1. the quote was signed by an attestation key that an Intel Quoting Enclave vouched for, whose PCK certificate chains
     to the Intel SGX Root CA (pinned below) — i.e. a genuine Intel TDX CPU produced it;
  2. the TD is not a debug TD (a debug TD's memory is readable by the host);
  3. MRTD / RTMR0..3 / REPORTDATA are what the CPU measured — the caller compares them with the release manifest
     (expected_rtmrs) and with the report data it expects (nonce + TLS key + controller key, see report_data()).
Not checked here (prototype): the TCB level of the platform (Intel PCS collateral: TCB info, QE identity, CRLs). A
revoked or out-of-date platform still passes; verify_workspace.py prints the TCB SVNs so they can be checked by hand.

Report data binding (both sides MUST compute it the same way):
  SHA-512( b"clairwave-ws-attest-v1\0" + nonce(utf-8) + b"\0" + SHA-256(TLS SubjectPublicKeyInfo DER)
           + SHA-256(controller Ed25519 public key, raw 32 bytes) )
"""
import hashlib
import struct

from cryptography import x509
from cryptography.exceptions import InvalidSignature
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.asymmetric import ec
from cryptography.hazmat.primitives.asymmetric.utils import encode_dss_signature
from cryptography.hazmat.primitives.serialization import Encoding, PublicFormat

# Intel SGX Root CA (CN=Intel SGX Root CA, valid to 2049-12-31), from
# https://certificates.trustedservices.intel.com/Intel_SGX_Provisioning_Certification_RootCA.cer (fetched 2026-10-23).
INTEL_ROOT_SPKI_SHA256 = "a0af031289f5d5d4132f9186068a7fc13628633ba235777472e29b6b6c67a49e"

TEE_TDX = 0x81
_BODY_V4 = 584      # TDX 1.0 TD report body
_BODY_15 = 648      # TDX 1.5 (quote v5 body type 3): + TEE_TCB_SVN2 (16) + MRSERVICETD (48)


class QuoteError(Exception):
    pass


def report_data(nonce, tls_spki_der, ctl_pub_raw=None, ctl_pub_sha256=None):
    """ctl_pub_sha256 (32 bytes) instead of the raw key: for verifiers that only know the key's hash."""
    h = hashlib.sha512()
    h.update(b"clairwave-ws-attest-v1\0" + nonce.encode() + b"\0")
    h.update(hashlib.sha256(tls_spki_der).digest())
    h.update(ctl_pub_sha256 if ctl_pub_sha256 is not None else hashlib.sha256(ctl_pub_raw).digest())
    return h.digest()


def _sig(raw64):
    return encode_dss_signature(int.from_bytes(raw64[:32], "big"), int.from_bytes(raw64[32:], "big"))


def _p256(raw64):
    return ec.EllipticCurvePublicKey.from_encoded_point(ec.SECP256R1(), b"\x04" + raw64)


def parse(q):
    """Split a TDX quote (v4 or v5) into fields. Raises QuoteError on anything malformed."""
    if len(q) < 48 + _BODY_V4 + 4:
        raise QuoteError("quote too short")
    ver, key_type, tee = struct.unpack_from("<HHI", q, 0)
    if tee != TEE_TDX:
        raise QuoteError(f"not a TDX quote (tee type {tee:#x})")
    if key_type != 2:
        raise QuoteError(f"unsupported attestation key type {key_type}")
    off = 48
    if ver == 4:
        body_len = _BODY_V4
    elif ver == 5:
        btype, body_len = struct.unpack_from("<HI", q, off)
        off += 6
        if body_len not in (_BODY_V4, _BODY_15):
            raise QuoteError(f"unexpected v5 body type {btype} size {body_len}")
    else:
        raise QuoteError(f"unsupported quote version {ver}")
    body = q[off:off + body_len]
    signed = q[:off + body_len]
    off += body_len
    (sig_len,) = struct.unpack_from("<I", q, off)
    off += 4
    sd = q[off:off + sig_len]
    if len(sd) != sig_len or sig_len < 128 + 6:
        raise QuoteError("signature data truncated")
    f = {"version": ver, "signed": signed, "sig": sd[:64], "attest_key": sd[64:128]}
    o = 0
    f["tee_tcb_svn"] = body[o:o + 16]; o += 16
    f["mrseam"] = body[o:o + 48]; o += 48
    f["mrsignerseam"] = body[o:o + 48]; o += 48
    f["seam_attributes"] = body[o:o + 8]; o += 8
    f["td_attributes"] = body[o:o + 8]; o += 8
    f["xfam"] = body[o:o + 8]; o += 8
    f["mrtd"] = body[o:o + 48]; o += 48
    f["mrconfigid"] = body[o:o + 48]; o += 48
    f["mrowner"] = body[o:o + 48]; o += 48
    f["mrownerconfig"] = body[o:o + 48]; o += 48
    f["rtmr"] = [body[o + 48 * i:o + 48 * (i + 1)] for i in range(4)]; o += 192
    f["report_data"] = body[o:o + 64]
    # certification data: type 6 = QE report certification data (QE report, its signature, QE auth data, PCK chain)
    ctype, csize = struct.unpack_from("<HI", sd, 128)
    cd = sd[134:134 + csize]
    if ctype != 6 or len(cd) < 384 + 64 + 2:
        raise QuoteError(f"unexpected certification data type {ctype}")
    f["qe_report"] = cd[:384]
    f["qe_sig"] = cd[384:448]
    (alen,) = struct.unpack_from("<H", cd, 448)
    f["qe_auth"] = cd[450:450 + alen]
    p = 450 + alen
    itype, isize = struct.unpack_from("<HI", cd, p)
    if itype != 5:
        raise QuoteError(f"PCK chain missing (inner certification data type {itype})")
    f["pck_pem"] = cd[p + 6:p + 6 + isize]
    return f


def verify(q):
    """Check the signature chain of a TDX quote. Returns the parsed fields; raises QuoteError if anything fails."""
    f = parse(q)
    try:
        _p256(f["attest_key"]).verify(_sig(f["sig"]), f["signed"], ec.ECDSA(hashes.SHA256()))
    except (InvalidSignature, ValueError) as e:
        raise QuoteError(f"quote signature invalid ({type(e).__name__})")
    # the QE report binds the attestation key: report_data[0:32] = SHA-256(attest_key || qe_auth_data)
    qrd = f["qe_report"][320:384]
    if qrd[:32] != hashlib.sha256(f["attest_key"] + f["qe_auth"]).digest() or any(qrd[32:]):
        raise QuoteError("QE report does not bind the attestation key")
    certs = x509.load_pem_x509_certificates(f["pck_pem"])
    if len(certs) < 3:
        raise QuoteError("PCK chain incomplete")
    leaf, inter, root = certs[0], certs[1], certs[-1]
    if hashlib.sha256(root.public_key().public_bytes(Encoding.DER, PublicFormat.SubjectPublicKeyInfo)).hexdigest() \
            != INTEL_ROOT_SPKI_SHA256:
        raise QuoteError("PCK chain does not end at the Intel SGX Root CA")
    try:
        root.verify_directly_issued_by(root)
        inter.verify_directly_issued_by(root)
        leaf.verify_directly_issued_by(inter)
        leaf.public_key().verify(_sig(f["qe_sig"]), f["qe_report"], ec.ECDSA(hashes.SHA256()))
    except (InvalidSignature, ValueError, TypeError) as e:
        raise QuoteError(f"PCK chain / QE report signature invalid ({type(e).__name__})")
    if f["td_attributes"][0] & 0x01:
        raise QuoteError("debug TD — memory readable by the host")
    f["pck_subject"] = leaf.subject.rfc4514_string()
    return f


def summary(f):
    return {"mrtd": f["mrtd"].hex(), "rtmr": [r.hex() for r in f["rtmr"]], "report_data": f["report_data"].hex(),
            "td_attributes": f["td_attributes"].hex(), "tee_tcb_svn": f["tee_tcb_svn"].hex(),
            "mrseam": f["mrseam"].hex(), "quote_version": f["version"]}


def check(q, nonce, tls_spki_der, ctl_pub_raw, manifest, ctl_pub_sha256=None):
    """Full check against a release manifest {"rtmr1": hex, "rtmr2": hex, ...}. Returns (ok, reasons, summary)."""
    reasons = []
    try:
        f = verify(q)
    except QuoteError as e:
        return False, [str(e)], {}
    if f["report_data"] != report_data(nonce, tls_spki_der, ctl_pub_raw, ctl_pub_sha256):
        reasons.append("report data does not match this nonce / TLS key / controller key")
    for i in (1, 2):
        want = (manifest.get(f"rtmr{i}") or "").lower()
        if not want or f["rtmr"][i].hex() != want:
            reasons.append(f"RTMR{i} is not the measurement of image {manifest.get('version', '?')}")
    return not reasons, reasons, summary(f)