acdl/tests/test_kms_roundtrip.py

"""CAP-037 KMS round-trip test (REQ-350).

Sign a test JWT via ``core.kms_signing.sign_jwt()`` (mock KMS with a
test keypair) → fetch JWKS via ``nova_idp_jwks.lambda_handler()`` (mock
KMS) → verify the JWT with ``pyjwt`` using the JWKS key. Round-trip
succeeds — proves the DER→raw conversion + JWK export are mutually
consistent (the #1 gotcha from RESEARCH §5).
"""

from __future__ import annotations

import importlib.util
import json
import os
import sys
from pathlib import Path

import pytest

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

os.environ.setdefault("NOVA_LAMBDA_LOCAL_BYPASS", "1")

_JWKS_PATH = Path(__file__).resolve().parent.parent / "core" / "lambda" / "nova_idp_jwks.py"
_spec = importlib.util.spec_from_file_location("nova_idp_jwks_rt", _JWKS_PATH)
jwks_mod = importlib.util.module_from_spec(_spec)
_spec.loader.exec_module(jwks_mod)

import core.kms_signing as kms_signing
import jwt as pyjwt
from cryptography.hazmat.primitives.asymmetric import ec
from cryptography.hazmat.primitives import hashes, serialization


class _MockKms:
    def __init__(self, priv, pub_der):
        self._priv = priv
        self._pub_der = pub_der

    def sign(self, KeyId, Message, MessageType, SigningAlgorithm):
        return {"Signature": self._priv.sign(Message, ec.ECDSA(hashes.SHA256()))}

    def get_public_key(self, KeyId):
        return {"PublicKey": self._pub_der}


@pytest.fixture(autouse=True)
def _reset():
    yield
    kms_signing.set_kms_client_for_testing(None)


def test_cap037_kms_roundtrip():
    """Sign JWT → JWKS → pyjwt verify. The full KMS round-trip (REQ-350)."""
    priv = ec.generate_private_key(ec.SECP256R1())
    pub = priv.public_key()
    pub_der = pub.public_bytes(
        encoding=serialization.Encoding.DER,
        format=serialization.PublicFormat.SubjectPublicKeyInfo,
    )
    kms_signing.set_kms_client_for_testing(_MockKms(priv, pub_der))

    # 1. Sign a JWT via kms_signing.sign_jwt (uses DER→raw conversion).
    claims = {
        "sub": "roundtrip-user", "aud": "nova-cli", "iss": "nova-idp",
        "exp": 9999999999, "iat": 1700000000, "jti": "rt-jti",
        "roles": ["developer"], "typ": "nova_oidc_token",
    }
    token = kms_signing.sign_jwt(claims, key_id="alias/nova-oidc-signing")

    # 2. Fetch the JWKS via the JWKS Lambda (mock KMS get_public_key).
    resp = jwks_mod.lambda_handler({}, None)
    assert resp["statusCode"] == 200
    jwks_body = json.loads(resp["body"])
    jwk = jwks_body["keys"][0]
    assert jwk["kty"] == "EC" and jwk["crv"] == "P-256"

    # 3. Verify the JWT with pyjwt using the JWKS key.
    key = pyjwt.PyJWK(jwk).key
    decoded = pyjwt.decode(token, key, algorithms=["ES256"], audience="nova-cli")
    assert decoded["sub"] == "roundtrip-user"
    assert decoded["jti"] == "rt-jti"
    assert decoded["roles"] == ["developer"]
    assert decoded["typ"] == "nova_oidc_token"


# ---------------------------------------------------------------------------
# Live-KMS round-trip (REQ-362, Edge 5 item 6).
#
# This test is marked ``@pytest.mark.live_aws`` and is SKIPPED in acdl CI
# (the live KMS key ``alias/nova-oidc-signing`` is not provisioned here).
# It runs in nova-platform-ops CI against the real KMS key, REQ-362
# (covered-reference — verification surface is the nova-platform-ops
# pipeline, not acdl's). It exercises the same sign → JWKS → verify path
# against the production key/alias so the DER→raw conversion + JWK export
# are verified end-to-end against real AWS KMS.
# ---------------------------------------------------------------------------


def _live_kms_available() -> bool:
    """Return True iff a live ``alias/nova-oidc-signing`` KMS key is
    reachable (best-effort probe; any error → False)."""
    try:
        import boto3
        client = boto3.client("kms")
        client.describe_key(KeyId="alias/nova-oidc-signing")
        return True
    except Exception:
        return False


@pytest.mark.live_aws
def test_cap037_kms_roundtrip_live():
    """Sign → JWKS → pyjwt verify against the LIVE KMS key
    (``alias/nova-oidc-signing``). Edge 5 item 6, REQ-362.

    Skipped unless a live KMS key is reachable (acdl CI has none; this
    runs in nova-platform-ops CI). The mock-based ``test_cap037_kms_roundtrip``
    above is the acdl-CI-runnable covered-path.
    """
    if not _live_kms_available():
        pytest.skip(
            "live KMS key alias/nova-oidc-signing not reachable "
            "(acdl CI; runs in nova-platform-ops CI, REQ-362)"
        )
    # Use the real KMS client (reset any test-injected mock client).
    kms_signing.set_kms_client_for_testing(None)

    claims = {
        "sub": "live-roundtrip-user", "aud": "nova-cli", "iss": "nova-idp",
        "exp": 9999999999, "iat": 1700000000, "jti": "live-rt-jti",
        "roles": ["developer"], "typ": "nova_oidc_token",
    }
    token = kms_signing.sign_jwt(claims, key_id="alias/nova-oidc-signing")

    resp = jwks_mod.lambda_handler({}, None)
    assert resp["statusCode"] == 200, resp
    jwk = json.loads(resp["body"])["keys"][0]
    assert jwk["kty"] == "EC" and jwk["crv"] == "P-256"

    key = pyjwt.PyJWK(jwk).key
    decoded = pyjwt.decode(token, key, algorithms=["ES256"], audience="nova-cli")
    assert decoded["sub"] == "live-roundtrip-user"
    assert decoded["jti"] == "live-rt-jti"
    assert decoded["typ"] == "nova_oidc_token"