meal-order-manager/tests/test_cognito.py

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

import json
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_2 = "https://cognito-idp.us-east-1.amazonaws.com/us-east-1_PROD"
AUDIENCE_2 = "portal-prod-client"
ISSUER_PARAM = "/meal-order-manager/portal-cognito-issuer"
AUDIENCE_PARAM = "/meal-order-manager/portal-cognito-audience"
TRUST_PARAM = "/meal-order-manager/portal-cognito-trust"


@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": False,
        "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,
        TRUST_PARAM: json.dumps(
            [
                {"issuer": ISSUER, "audience": AUDIENCE},
                {"issuer": ISSUER_2, "audience": AUDIENCE_2},
            ]
        ),
    }[name]


def _verify(token, public_key, trust_param=""):
    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, trust_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",
    }


def test_unverified_federated_email_is_accepted(signing_key):
    private_key, public_key = signing_key

    result = _verify(_token(private_key, email_verified=False), public_key)

    assert result["email"] == "employee@seahavenind.com"


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

    assert _verify(token, public_key) == {
        "name": "Ada",
        "email": "employee@seahavenind.com",
    }


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

    assert _verify(token, public_key) == {
        "name": "employee",
        "email": "employee@seahavenind.com",
    }


def test_trust_list_accepts_second_issuer(signing_key):
    private_key, public_key = signing_key
    token = _token(private_key, iss=ISSUER_2, aud=AUDIENCE_2)

    result = _verify(token, public_key, trust_param=TRUST_PARAM)

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


def test_untrusted_issuer_is_rejected(signing_key):
    private_key, public_key = signing_key
    token = _token(
        private_key,
        iss="https://cognito-idp.us-east-1.amazonaws.com/us-east-1_OTHER",
        aud=AUDIENCE,
    )

    assert _verify(token, public_key, trust_param=TRUST_PARAM) is None


@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"},
    ],
)
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
    )