meal-order-manager/tests/test_cognito.py

"""Unit tests for portal Cognito ID-token verification."""

import time
from types import SimpleNamespace
from unittest.mock import MagicMock, patch

import jwt
import pytest
from cryptography.hazmat.primitives.asymmetric import rsa

from shared import cognito

ISSUER = "https://cognito-idp.us-east-1.amazonaws.com/us-east-1_TEST"
AUDIENCE = "portal-client-id"
ISSUER_PARAM = "/meal-order-manager/portal-cognito-issuer"
AUDIENCE_PARAM = "/meal-order-manager/portal-cognito-audience"


@pytest.fixture(scope="module")
def signing_key():
    private_key = rsa.generate_private_key(public_exponent=65537, key_size=2048)
    return private_key, private_key.public_key()


def _claims(**overrides):
    now = int(time.time())
    claims = {
        "iss": ISSUER,
        "aud": AUDIENCE,
        "token_use": "id",
        "exp": now + 300,
        "iat": now,
        "email": "employee@seahavenind.com",
        "email_verified": True,
        "name": "Test Employee",
    }
    claims.update(overrides)
    return claims


def _token(private_key, **overrides):
    return jwt.encode(
        _claims(**overrides),
        private_key,
        algorithm="RS256",
        headers={"kid": "test-key"},
    )


def _parameter(name, decrypt=False):
    assert decrypt is False
    return {ISSUER_PARAM: ISSUER, AUDIENCE_PARAM: AUDIENCE}[name]


def _verify(token, public_key):
    jwks = MagicMock()
    jwks.get_signing_key_from_jwt.return_value = SimpleNamespace(key=public_key)
    with (
        patch("shared.cognito.get_parameter", side_effect=_parameter),
        patch("shared.cognito._jwk_client", return_value=jwks),
    ):
        return cognito.verify_cognito_id_token(token, ISSUER_PARAM, AUDIENCE_PARAM)


def test_valid_id_token_returns_trusted_identity(signing_key):
    private_key, public_key = signing_key

    result = _verify(_token(private_key), public_key)

    assert result == {
        "name": "Test Employee",
        "email": "employee@seahavenind.com",
    }


@pytest.mark.parametrize(
    "overrides",
    [
        {"iss": "https://cognito-idp.us-east-1.amazonaws.com/us-east-1_OTHER"},
        {"aud": "other-client"},
        {"token_use": "access"},
        {"exp": int(time.time()) - 60},
        {"email": "employee@example.com"},
        {"email_verified": False},
        {"name": ""},
    ],
)
def test_invalid_claims_are_rejected(signing_key, overrides):
    private_key, public_key = signing_key

    assert _verify(_token(private_key, **overrides), public_key) is None


def test_missing_required_claim_is_rejected(signing_key):
    private_key, public_key = signing_key
    claims = _claims()
    del claims["email"]
    token = jwt.encode(
        claims, private_key, algorithm="RS256", headers={"kid": "test-key"}
    )

    assert _verify(token, public_key) is None


def test_ssm_failure_is_rejected(signing_key):
    private_key, _ = signing_key
    with patch(
        "shared.cognito.get_parameter", side_effect=RuntimeError("SSM unavailable")
    ):
        with pytest.raises(cognito.CognitoVerificationUnavailable):
            cognito.verify_cognito_id_token(
                _token(private_key), ISSUER_PARAM, AUDIENCE_PARAM
            )


def test_jwks_failure_is_rejected(signing_key):
    private_key, _ = signing_key
    jwks = MagicMock()
    jwks.get_signing_key_from_jwt.side_effect = OSError("JWKS unavailable")
    with (
        patch("shared.cognito.get_parameter", side_effect=_parameter),
        patch("shared.cognito._jwk_client", return_value=jwks),
    ):
        with pytest.raises(cognito.CognitoVerificationUnavailable):
            cognito.verify_cognito_id_token(
                _token(private_key), ISSUER_PARAM, AUDIENCE_PARAM
            )


def test_cognito_token_detection_is_untrusted_routing_hint(signing_key):
    private_key, _ = signing_key

    assert cognito.looks_like_cognito_token(_token(private_key)) is True
    assert cognito.looks_like_cognito_token("not-a-jwt") is False
    assert (
        cognito.looks_like_cognito_token(
            _token(private_key, iss="https://accounts.google.com")
        )
        is False
    )