Attestation Workflow for Spatial Enclaves

from __future__ import annotations

import hashlib
import hmac
import secrets
from dataclasses import dataclass
from typing import Callable


class AttestationError(RuntimeError):
    """Any failed check. Callers must not catch this and continue."""


@dataclass(frozen=True)
class Quote:
    measurement: str
    ephemeral_pubkey: bytes
    tcb_version: int
    nonce: bytes
    issued_at: float
    signature: bytes


@dataclass(frozen=True)
class VerifierPolicy:
    approved_measurements: frozenset[str]
    min_tcb_version: int
    vendor_key: bytes
    quote_max_age_s: float = 30.0

    def _expected_signature(self, quote: Quote) -> bytes:
        payload = b"|".join([
            quote.measurement.encode(),
            quote.ephemeral_pubkey,
            str(quote.tcb_version).encode(),
            quote.nonce,
            f"{quote.issued_at:.3f}".encode(),
        ])
        return hmac.new(self.vendor_key, payload, hashlib.sha256).digest()

    def verify(self, quote: Quote, *, expected_nonce: bytes, now: float) -> None:
        """Run every check, in the order that fails fastest and loudest.

        Ordering matters for diagnostics, not for security: all checks must pass.
        The measurement check is listed last deliberately, because it is the one
        engineers are tempted to make advisory, and putting it after the cheap
        checks makes its absence obvious in a diff.
        """
        if not hmac.compare_digest(self._expected_signature(quote), quote.signature):
            raise AttestationError("vendor signature invalid")
        if not hmac.compare_digest(quote.nonce, expected_nonce):
            raise AttestationError("nonce mismatch — this quote answers a different request")
        age = now - quote.issued_at
        if age < -1.0 or age > self.quote_max_age_s:
            raise AttestationError(f"quote age {age:.1f}s outside the accepted window")
        if quote.tcb_version < self.min_tcb_version:
            raise AttestationError(
                f"TCB {quote.tcb_version} below floor {self.min_tcb_version}"
            )
        if len(quote.ephemeral_pubkey) < 32:
            raise AttestationError("ephemeral key absent or too short")
        if quote.measurement not in self.approved_measurements:
            raise AttestationError(
                f"measurement {quote.measurement[:12]}… not approved — a valid "
                "signature says an enclave exists, not that it runs your code"
            )


@dataclass
class AttestedChannel:
    """A channel that cannot transmit before a successful verification."""

    policy: VerifierPolicy
    _key: bytes | None = None

    def open(self, request_quote: Callable[[bytes], Quote], *, now: float) -> None:
        nonce = secrets.token_bytes(32)
        quote = request_quote(nonce)
        self.policy.verify(quote, expected_nonce=nonce, now=now)
        self._key = quote.ephemeral_pubkey

    def send(self, payload: bytes) -> bytes:
        if self._key is None:
            raise AttestationError("channel not attested — refusing to transmit")
        stream = hashlib.blake2b(self._key, digest_size=len(payload)).digest()
        return bytes(a ^ b for a, b in zip(payload, stream))

def _validate() -> None:
    vendor_key = b"vendor-signing-key"
    good_measurement = hashlib.sha256(b"enclave-v4.2.1").hexdigest()
    policy = VerifierPolicy(
        approved_measurements=frozenset({good_measurement}),
        min_tcb_version=7,
        vendor_key=vendor_key,
    )

    def make_quote(nonce: bytes, *, measurement=good_measurement, tcb=9,
                   issued_at=1_000.0, key=b"k" * 32) -> Quote:
        q = Quote(measurement, key, tcb, nonce, issued_at, b"")
        return Quote(measurement, key, tcb, nonce, issued_at,
                     policy._expected_signature(q))

    nonce = secrets.token_bytes(32)

    # 1. A well-formed, approved quote passes.
    policy.verify(make_quote(nonce), expected_nonce=nonce, now=1_005.0)

    # 2. A signature from the wrong key fails.
    bad = make_quote(nonce)
    forged = Quote(bad.measurement, bad.ephemeral_pubkey, bad.tcb_version,
                   bad.nonce, bad.issued_at, b"x" * 32)
    for case, quote, kwargs in [
        ("forged signature", forged, {}),
        ("unapproved build", make_quote(nonce, measurement="deadbeef" * 8), {}),
        ("downlevel TCB", make_quote(nonce, tcb=3), {}),
        ("stale quote", make_quote(nonce, issued_at=900.0), {}),
        ("short key", make_quote(nonce, key=b"short"), {}),
    ]:
        try:
            policy.verify(quote, expected_nonce=nonce, now=1_005.0, **kwargs)
        except AttestationError:
            pass
        else:
            raise AssertionError(f"{case} must be rejected")

    # 3. A replayed quote answering a different nonce fails.
    other = secrets.token_bytes(32)
    try:
        policy.verify(make_quote(other), expected_nonce=nonce, now=1_005.0)
    except AttestationError as exc:
        assert "nonce" in str(exc)
    else:
        raise AssertionError("nonce mismatch must be rejected")

    # 4. The channel refuses to transmit before attestation.
    channel = AttestedChannel(policy)
    try:
        channel.send(b"coordinates")
    except AttestationError:
        pass
    else:
        raise AssertionError("an unattested channel must refuse to send")

    # 5. After a successful open, transmission is bound to the attested key.
    channel.open(lambda n: make_quote(n), now=1_002.0)
    assert channel.send(b"coordinates") != b"coordinates"

    print("attestation workflow: all assertions passed")


_validate()