"""AWS Nitro COSE verification anchored exclusively in the bundled AWS root."""
import base64
import datetime
import hashlib
import hmac
import io
import json
import re
from importlib.resources import files

import cbor2
from OpenSSL import crypto
from cryptography import x509
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import ec, utils

ROOT_SHA256 = "641a0321a3e244efe456463195d606317ed7cdcc3c1756e09893f3c68f79bb5b"
RANDOMX_VERSION = "2.0.1"
RANDOMX_COMMIT = "aaafe71322df6602c21a5c72937ac284724ae561"
MAX_ITERATIONS = 1_000_000_000_000


class AttestationError(ValueError):
    pass


def _cbor(data):
    source = io.BytesIO(data)
    result = cbor2.CBORDecoder(source).decode()
    if source.read(1):
        raise AttestationError("trailing CBOR bytes")
    return result


def _json(raw):
    def unique(pairs):
        result = {}
        for key, value in pairs:
            if key in result:
                raise AttestationError("duplicate JSON field")
            result[key] = value
        return result
    return json.loads(raw, object_pairs_hook=unique)


def _hex(value, length):
    if not isinstance(value, str) or not re.fullmatch("[0-9a-f]{%d}" % (length*2), value):
        raise AttestationError("noncanonical hexadecimal value")
    return bytes.fromhex(value)


def _authenticated_document(raw, expected_pcr0, *, now=None, historical=False):
    """Shared signature checks. Only the separate archive API uses signed time."""
    if not isinstance(raw, bytes) or len(raw) > 32768:
        raise AttestationError("invalid attestation size")
    expected = _hex(expected_pcr0, 48)
    if expected == bytes(48):
        raise AttestationError("pin an exact nonzero 48-byte PCR0")
    envelope = _cbor(raw)
    if isinstance(envelope, cbor2.CBORTag):
        if envelope.tag != 18:
            raise AttestationError("unexpected COSE tag")
        envelope = envelope.value
    if not isinstance(envelope, (tuple, list)) or len(envelope) != 4:
        raise AttestationError("invalid COSE Sign1 structure")
    protected, unprotected, payload, signature = envelope
    if _cbor(protected) != {1: -35} or unprotected != {} or len(signature) != 96:
        raise AttestationError("expected ES384 COSE signature")
    doc = _cbor(payload)
    if not isinstance(doc, dict) or doc.get("digest") != "SHA384":
        raise AttestationError("attestation digest mismatch")
    pcrs = doc.get("pcrs", {})
    for index in (0, 1, 2, 4):
        if not isinstance(pcrs.get(index), bytes) or len(pcrs[index]) != 48 or pcrs[index] == bytes(48):
            raise AttestationError("missing/zero Nitro PCR")
    if not hmac.compare_digest(pcrs[0], expected):
        raise AttestationError("PCR0 does not match the client's pinned build")
    now = now or datetime.datetime.now(datetime.timezone.utc)
    timestamp = doc.get("timestamp")
    if type(timestamp) is not int or timestamp <= 0:
        raise AttestationError("invalid signed Nitro timestamp")
    if historical:
        if timestamp > now.timestamp()*1000 + 300000:
            raise AttestationError("archived attestation claims a future timestamp")
        verification_time = datetime.datetime.fromtimestamp(timestamp/1000, datetime.timezone.utc)
    else:
        if abs(now.timestamp()*1000 - timestamp) > 300000:
            raise AttestationError("attestation timestamp outside five-minute freshness window")
        verification_time = now
    root = x509.load_pem_x509_certificate(files(__package__).joinpath("nitro-root.pem").read_bytes())
    if root.fingerprint(hashes.SHA256()).hex() != ROOT_SHA256:
        raise AttestationError("bundled AWS root fingerprint mismatch")
    leaf = x509.load_der_x509_certificate(doc["certificate"])
    bundle = doc["cabundle"]
    if not isinstance(bundle, (list, tuple)) or not 1 <= len(bundle) <= 8:
        raise AttestationError("invalid certificate chain size")
    store = crypto.X509Store()
    store.add_cert(crypto.X509.from_cryptography(root))
    store.set_time(verification_time)
    intermediates = [crypto.load_certificate(crypto.FILETYPE_ASN1, cert) for cert in bundle]
    crypto.X509StoreContext(store, crypto.X509.from_cryptography(leaf), intermediates).verify_certificate()
    public = leaf.public_key()
    if not isinstance(public, ec.EllipticCurvePublicKey) or not isinstance(public.curve, ec.SECP384R1):
        raise AttestationError("Nitro signing key is not P-384")
    r, s = int.from_bytes(signature[:48], "big"), int.from_bytes(signature[48:], "big")
    public.verify(utils.encode_dss_signature(r, s), cbor2.dumps(["Signature1", protected, b"", payload]), ec.ECDSA(hashes.SHA384()))
    return doc


def _bound_policy(doc, peer_der, *, require_graviton5, historical=False):
    peer = x509.load_der_x509_certificate(peer_der)
    spki = peer.public_key().public_bytes(serialization.Encoding.DER, serialization.PublicFormat.SubjectPublicKeyInfo)
    if doc.get("public_key") != spki:
        raise AttestationError("attested TLS key differs from actual inner TLS peer")
    policy = _json(doc["user_data"])
    if not isinstance(policy, dict) or type(policy.get("protocol_version")) is not int or policy["protocol_version"] != 2:
        raise AttestationError("unsupported measured relay protocol")
    if (policy.get("randomx_version") != RANDOMX_VERSION or policy.get("randomx_commit") != RANDOMX_COMMIT
            or policy.get("randomx_algorithm") != "v2" or type(policy.get("segments")) is not int or policy["segments"] not in (7, 84, 96)):
        raise AttestationError("unsupported measured relay protocol")
    if type(policy.get("iterations")) is not int or not 1 <= policy["iterations"] <= MAX_ITERATIONS:
        raise AttestationError("invalid RandomX policy")
    _hex(policy.get("service_signing_public_key"), 32)
    if require_graviton5 and policy.get("graviton5_verified") is not True:
        raise AttestationError("enclave has not verified Graviton5 hardware")
    if policy.get("state") not in (("warming", "ready") if historical else ("ready",)):
        raise AttestationError("enclave is not ready" if not historical else "invalid archived enclave state")
    return policy, spki


def verify_document(raw, *, nonce, peer_der, expected_pcr0, require_graviton5=True, now=None):
    """LIVE verification: fresh client nonce, current chain validity and ready gate."""
    try:
        if not isinstance(nonce, bytes) or len(nonce) != 32:
            raise AttestationError("invalid nonce size")
        doc = _authenticated_document(raw, expected_pcr0, now=now)
        if not hmac.compare_digest(doc.get("nonce", b""), nonce):
            raise AttestationError("attestation nonce mismatch")
        policy, spki = _bound_policy(doc, peer_der, require_graviton5=require_graviton5)
        return {"verified": True, "mode": "nitro", "policy": policy, "pcr0": expected_pcr0,
                "timestamp_ms": doc["timestamp"], "document_b64": base64.b64encode(raw).decode(),
                "tls_spki_sha256": hashlib.sha256(spki).hexdigest()}
    except AttestationError:
        raise
    except Exception as exc:
        raise AttestationError("invalid Nitro attestation") from exc