acdl/tests/test_token_vend.py

"""E2E token-vend Lambda test (REQ-336, C-6.1) with moto + mock KMS.

Valid PAT (active) + ABAC allow → KMS-signed OIDC token returned.
Revoked PAT → 403. Unknown PAT → 403. ABAC deny → 403. The returned
JWT verifies with pyjwt + the mock public key.
"""

from __future__ import annotations

import importlib.util
import json
import os
import sys
import time
from pathlib import Path
from unittest import mock

import pytest

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

os.environ.setdefault("AWS_DEFAULT_REGION", "us-east-1")
os.environ.setdefault("AWS_ACCESS_KEY_ID", "test")
os.environ.setdefault("AWS_SECRET_ACCESS_KEY", "test")
os.environ.setdefault("NOVA_LAMBDA_LOCAL_BYPASS", "1")

_SOURCE_PATH = (
    Path(__file__).resolve().parent.parent / "core" / "lambda" / "nova_idp_token_vend.py"
)
_spec = importlib.util.spec_from_file_location("nova_idp_token_vend_e2e", _SOURCE_PATH)
tv = importlib.util.module_from_spec(_spec)
_spec.loader.exec_module(tv)

import boto3
from moto import mock_aws
import jwt as pyjwt
from cryptography.hazmat.primitives.asymmetric import ec
from cryptography.hazmat.primitives import hashes, serialization

import core.kms_signing as kms_signing


def _create_pats_table(ddb):
    ddb.create_table(
        TableName="nova-pats",
        KeySchema=[{"AttributeName": "jti", "KeyType": "HASH"}],
        AttributeDefinitions=[
            {"AttributeName": "jti", "AttributeType": "S"},
            {"AttributeName": "sub", "AttributeType": "S"},
            {"AttributeName": "pat_hash", "AttributeType": "S"},
        ],
        GlobalSecondaryIndexes=[
            {"IndexName": "sub-index", "KeySchema": [{"AttributeName": "sub", "KeyType": "HASH"}], "Projection": {"ProjectionType": "ALL"}},
            {"IndexName": "pat_hash-index", "KeySchema": [{"AttributeName": "pat_hash", "KeyType": "HASH"}], "Projection": {"ProjectionType": "ALL"}},
        ],
        BillingMode="PAY_PER_REQUEST",
    )


def _put_active_pat(ddb, jti="pat-active", sub="user-1", owner="t1", role="developer"):
    ddb.put_item(
        TableName="nova-pats",
        Item={
            "jti": {"S": jti},
            "sub": {"S": sub},
            "pat_hash": {"S": "hash-" + jti},
            "status": {"S": "active"},
            "issued_at": {"S": "2026-01-01T00:00:00Z"},
            "expires_at": {"N": str(int(time.time()) + 3600)},
            "claims": {"S": json.dumps({"sub": sub, "roles": [role], "owner": owner})},
        },
    )


def _make_pat_jwt(jti="pat-active", sub="user-1", role="developer", owner="t1"):
    import base64
    header = base64.urlsafe_b64encode(json.dumps({"alg": "none", "typ": "JWT"}).encode()).rstrip(b"=").decode()
    payload = base64.urlsafe_b64encode(json.dumps({
        "jti": jti, "sub": sub, "exp": int(time.time()) + 3600,
        "iat": int(time.time()), "roles": [role], "owner": owner,
        "typ": "developer_pat",
    }).encode()).rstrip(b"=").decode()
    return f"{header}.{payload}.sig"


def _vend_event(pat_jwt, env="dev", owner="t1"):
    return {"body": json.dumps({
        "token": pat_jwt, "environment": env,
        "target_resource": {"type": "contract", "id": "c1", "owner": owner, "environment": env},
        "requested_claims": ["sub", "roles"],
    })}


class _MockKmsSign:
    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
def test_keypair():
    priv = ec.generate_private_key(ec.SECP256R1())
    pub = priv.public_key()
    pub_der = pub.public_bytes(
        encoding=serialization.Encoding.DER,
        format=serialization.PublicFormat.SubjectPublicKeyInfo,
    )
    return priv, pub, pub_der


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


@mock_aws
def test_valid_pat_abac_allow_vends_token(test_keypair):
    priv, pub, pub_der = test_keypair
    ddb = boto3.client("dynamodb", region_name="us-east-1")
    _create_pats_table(ddb)
    _put_active_pat(ddb)
    kms_signing.set_kms_client_for_testing(_MockKmsSign(priv, pub_der))
    pat = _make_pat_jwt()
    # ABAC allow: developer + dev
    resp = tv.lambda_handler(_vend_event(pat, env="dev", owner="t1"), None)
    assert resp["statusCode"] == 200, resp["body"]
    body = json.loads(resp["body"])
    assert "token" in body
    assert "expires_at" in body
    # Verify the JWT with pyjwt + the test public key.
    pem = pub.public_bytes(
        encoding=serialization.Encoding.PEM,
        format=serialization.PublicFormat.SubjectPublicKeyInfo,
    ).decode("ascii")
    decoded = pyjwt.decode(body["token"], pem, algorithms=["ES256"], options={"verify_aud": False})
    assert decoded["sub"] == "user-1"
    assert decoded["typ"] == "nova_oidc_token"
    assert decoded["roles"] == ["developer"]
    assert "jti" in decoded and "iat" in decoded and "exp" in decoded and "iss" in decoded


@mock_aws
def test_revoked_pat_denied():
    ddb = boto3.client("dynamodb", region_name="us-east-1")
    _create_pats_table(ddb)
    _put_active_pat(ddb, jti="pat-r")
    ddb.update_item(
        TableName="nova-pats", Key={"jti": {"S": "pat-r"}},
        UpdateExpression="SET #s = :v",
        ExpressionAttributeNames={"#s": "status"},
        ExpressionAttributeValues={":v": {"S": "revoked"}},
    )
    pat = _make_pat_jwt(jti="pat-r")
    resp = tv.lambda_handler(_vend_event(pat), None)
    assert resp["statusCode"] == 403
    assert json.loads(resp["body"])["reason"] == "pat_revoked"


@mock_aws
def test_unknown_pat_denied():
    ddb = boto3.client("dynamodb", region_name="us-east-1")
    _create_pats_table(ddb)
    pat = _make_pat_jwt(jti="pat-missing")
    resp = tv.lambda_handler(_vend_event(pat), None)
    assert resp["statusCode"] == 403
    assert json.loads(resp["body"])["reason"] == "pat_unknown"


@mock_aws
def test_abac_deny_denied(test_keypair):
    priv, _pub, pub_der = test_keypair
    ddb = boto3.client("dynamodb", region_name="us-east-1")
    _create_pats_table(ddb)
    _put_active_pat(ddb)
    kms_signing.set_kms_client_for_testing(_MockKmsSign(priv, pub_der))
    pat = _make_pat_jwt()
    # ABAC deny: developer + prod (developer not allowed in prod)
    resp = tv.lambda_handler(_vend_event(pat, env="prod", owner="t1"), None)
    assert resp["statusCode"] == 403
    assert json.loads(resp["body"])["reason"] == "abac_denied"


@mock_aws
def test_missing_token_field():
    resp = tv.lambda_handler({"body": json.dumps({"environment": "dev"})}, None)
    assert resp["statusCode"] == 400