core-hub/tests/test_iam_profile.py

from __future__ import annotations

import base64
import hashlib
import json
import time
from typing import Any
from urllib.parse import parse_qs, urlencode, urlparse

from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.asymmetric import padding, rsa
from fastapi import Depends, FastAPI, Request
from fastapi.responses import JSONResponse, RedirectResponse
from fastapi.testclient import TestClient

from core_hub.iam_profile import IamProfileClaims, IamProfileVerifier, iam_profile_dependency


def b64url(data: bytes) -> str:
    return base64.urlsafe_b64encode(data).rstrip(b"=").decode("ascii")


def jwk_from_key(private_key: rsa.RSAPrivateKey, kid: str) -> dict[str, str]:
    numbers = private_key.public_key().public_numbers()
    n = b64url(numbers.n.to_bytes((numbers.n.bit_length() + 7) // 8, "big"))
    e = b64url(numbers.e.to_bytes((numbers.e.bit_length() + 7) // 8, "big"))
    return {"kty": "RSA", "use": "sig", "alg": "RS256", "kid": kid, "n": n, "e": e}


def sign_jwt(private_key: rsa.RSAPrivateKey, kid: str, payload: dict[str, Any]) -> str:
    header = {"alg": "RS256", "typ": "JWT", "kid": kid}
    header_b64 = b64url(json.dumps(header, separators=(",", ":")).encode())
    payload_b64 = b64url(json.dumps(payload, separators=(",", ":")).encode())
    signing_input = f"{header_b64}.{payload_b64}".encode("ascii")
    signature = private_key.sign(signing_input, padding.PKCS1v15(), hashes.SHA256())
    return f"{header_b64}.{payload_b64}.{b64url(signature)}"


def pkce_challenge(verifier: str) -> str:
    return b64url(hashlib.sha256(verifier.encode("ascii")).digest())


class FixtureIssuer:
    def __init__(self, issuer: str, audience: str) -> None:
        self.issuer = issuer.rstrip("/")
        self.audience = audience
        self.private_key = rsa.generate_private_key(public_exponent=65537, key_size=2048)
        self.kid = "iam-profile-fixture"
        self.codes: dict[str, str] = {}

    @property
    def jwks(self) -> dict[str, list[dict[str, str]]]:
        return {"keys": [jwk_from_key(self.private_key, self.kid)]}

    def token(self, *, subject: str = "user:alice", audience: str | None = None) -> str:
        now = int(time.time())
        return sign_jwt(
            self.private_key,
            self.kid,
            {
                "iss": self.issuer,
                "sub": subject,
                "aud": [audience or self.audience, "profile-consumer"],
                "exp": now + 600,
                "iat": now,
                "nbf": now - 5,
                "jti": "fixture-token",
                "tenant": "tenant:platform",
                "principal_type": "human",
                "preferred_username": "alice",
                "email": "alice@example.test",
                "groups": ["netkingdom-admins"],
                "roles": ["admin"],
                "scope": "openid profile email hub:read",
                "assurance": {
                    "level": "aal2",
                    "methods": ["pwd", "otp"],
                    "mfa": True,
                    "source": "local-identity",
                    "at": now,
                },
            },
        )


def make_issuer_app(fixture: FixtureIssuer) -> FastAPI:
    app = FastAPI()

    @app.get("/.well-known/openid-configuration")
    def discovery() -> dict[str, Any]:
        return {
            "issuer": fixture.issuer,
            "authorization_endpoint": f"{fixture.issuer}/authorize",
            "token_endpoint": f"{fixture.issuer}/token",
            "userinfo_endpoint": f"{fixture.issuer}/userinfo",
            "jwks_uri": f"{fixture.issuer}/jwks",
            "response_types_supported": ["code"],
            "grant_types_supported": ["authorization_code", "client_credentials"],
            "id_token_signing_alg_values_supported": ["RS256"],
            "code_challenge_methods_supported": ["S256"],
            "scopes_supported": ["openid", "profile", "email", "hub:read"],
        }

    @app.get("/jwks")
    def jwks() -> dict[str, list[dict[str, str]]]:
        return fixture.jwks

    @app.get("/authorize")
    def authorize(request: Request) -> RedirectResponse:
        query = request.query_params
        redirect_uri = query["redirect_uri"]
        if query.get("response_type") != "code" or query.get("code_challenge_method") != "S256":
            return RedirectResponse(f"{redirect_uri}?error=invalid_request")
        code_challenge = query.get("code_challenge")
        if not code_challenge:
            return RedirectResponse(f"{redirect_uri}?error=invalid_request&error_description=pkce")
        code = "fixture-code"
        fixture.codes[code] = code_challenge
        params = urlencode({"code": code, "state": query.get("state", "")})
        return RedirectResponse(f"{redirect_uri}?{params}")

    @app.post("/token")
    async def token(request: Request) -> JSONResponse:
        body = (await request.body()).decode("utf-8")
        form = {name: values[0] for name, values in parse_qs(body).items()}
        code = form.get("code", "")
        verifier = form.get("code_verifier", "")
        if fixture.codes.get(code) != pkce_challenge(verifier):
            return JSONResponse({"error": "invalid_grant"}, status_code=400)
        return JSONResponse(
            {
                "access_token": fixture.token(),
                "id_token": fixture.token(),
                "token_type": "Bearer",
                "expires_in": 600,
            }
        )

    return app


def make_protected_app(verifier: IamProfileVerifier) -> FastAPI:
    app = FastAPI()
    require_claims = iam_profile_dependency(verifier)
    claims_dependency = Depends(require_claims)

    @app.get("/protected")
    def protected(claims: IamProfileClaims = claims_dependency) -> dict[str, Any]:
        return {
            "sub": claims.subject,
            "tenant": claims.tenant,
            "roles": list(claims.roles),
            "groups": list(claims.groups),
            "scopes": list(claims.scopes),
            "assuranceLevel": claims.assurance["level"],
        }

    return app


def obtain_access_token(client: TestClient) -> str:
    verifier = "0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"
    auth = client.get(
        "/authorize",
        params={
            "response_type": "code",
            "client_id": "core-hub-test",
            "redirect_uri": "http://127.0.0.1/callback",
            "scope": "openid profile email hub:read",
            "state": "state-1",
            "nonce": "nonce-1",
            "code_challenge": pkce_challenge(verifier),
            "code_challenge_method": "S256",
        },
        follow_redirects=False,
    )
    assert auth.status_code in {302, 307}
    code = parse_qs(urlparse(auth.headers["location"]).query)["code"][0]
    token = client.post(
        "/token",
        data={
            "grant_type": "authorization_code",
            "client_id": "core-hub-test",
            "code": code,
            "redirect_uri": "http://127.0.0.1/callback",
            "code_verifier": verifier,
        },
    )
    assert token.status_code == 200
    return token.json()["access_token"]


def test_fastapi_service_accepts_iam_profile_pkce_token() -> None:
    fixture = FixtureIssuer("http://issuer.local", "core-hub")
    issuer_client = TestClient(make_issuer_app(fixture), base_url=fixture.issuer)
    token = obtain_access_token(issuer_client)
    verifier = IamProfileVerifier(
        issuer=fixture.issuer,
        audience="core-hub",
        jwks=fixture.jwks,
        environment="local",
    )
    protected_client = TestClient(make_protected_app(verifier))

    response = protected_client.get("/protected", headers={"Authorization": f"Bearer {token}"})

    assert response.status_code == 200
    assert response.json() == {
        "sub": "user:alice",
        "tenant": "tenant:platform",
        "roles": ["admin"],
        "groups": ["netkingdom-admins"],
        "scopes": ["email", "hub:read", "openid", "profile"],
        "assuranceLevel": "aal2",
    }


def test_fastapi_service_rejects_missing_bearer_token() -> None:
    fixture = FixtureIssuer("http://issuer.local", "core-hub")
    verifier = IamProfileVerifier(
        issuer=fixture.issuer,
        audience="core-hub",
        jwks=fixture.jwks,
        environment="local",
    )
    protected_client = TestClient(make_protected_app(verifier))

    response = protected_client.get("/protected")

    assert response.status_code == 401
    assert response.headers["www-authenticate"] == "Bearer"


def test_fastapi_service_rejects_wrong_audience() -> None:
    fixture = FixtureIssuer("http://issuer.local", "core-hub")
    token = fixture.token(audience="another-service")
    verifier = IamProfileVerifier(
        issuer=fixture.issuer,
        audience="core-hub",
        jwks=fixture.jwks,
        environment="local",
    )
    protected_client = TestClient(make_protected_app(verifier))

    response = protected_client.get("/protected", headers={"Authorization": f"Bearer {token}"})

    assert response.status_code == 403
    assert response.json()["detail"] == "token audience does not include this service"


def test_production_verifier_rejects_local_development_issuer() -> None:
    fixture = FixtureIssuer("http://issuer.local", "core-hub")
    token = fixture.token()
    verifier = IamProfileVerifier(issuer=fixture.issuer, audience="core-hub", jwks=fixture.jwks)
    protected_client = TestClient(make_protected_app(verifier))

    response = protected_client.get("/protected", headers={"Authorization": f"Bearer {token}"})

    assert response.status_code == 403
    assert response.json()["detail"] == "production mode must reject local-development issuers"