test_oidc_callback.py

"""Tests for the generic-OIDC ``/auth/callback`` email resolution gate.

These drive the *real* generic-OIDC validation path with genuinely signed
``id_token`` values (real signature + ``iss`` / ``aud`` / ``exp``
verification). The callback-route tests vary the ``email_verified`` claim;
the algorithm regression test uses an ES384 key to cover providers that do
not advertise ES256.

This is the regression coverage for "OIDC login accepts
unverified email claim (account takeover)". Before the fix the
callback minted a session for any signature-valid ``id_token``,
ignoring ``email_verified``; an IdP that lets a user assert an
arbitrary unverified email could be used to sign in as a victim in an
allowed domain.

The token endpoint (``httpx``) and the JWKS signing-key lookup are the
only mocked boundaries — everything between the HTTP request and the
minted session cookie is the production code path.
"""

from __future__ import annotations

import json
import time
from collections.abc import Iterator
from pathlib import Path

import httpx
import jwt
import pytest
from cryptography.hazmat.primitives.asymmetric import ec, rsa
from fastapi import FastAPI
from fastapi.testclient import TestClient
from jwt.algorithms import ECAlgorithm, RSAAlgorithm

from omnigent.server.admin_list import AdminList
from omnigent.server.auth import UnifiedAuthProvider
from omnigent.server.oidc import OIDCConfig
from omnigent.server.routes.auth import (
    _AUTH_STATE_COOKIE_PLAIN,
    _resolve_oidc_email,
    create_auth_router,
)
from omnigent.stores.permission_store.sqlalchemy_store import SqlAlchemyPermissionStore

_TEST_SECRET = bytes.fromhex("aa" * 32)
_ISSUER = "https://accounts.google.com"
_CLIENT_ID = "cid"


def _oidc_config(
    skip_email_verification: bool = False,
    email_claim: str = "email",
) -> OIDCConfig:
    """Build a generic-OIDC config over plain HTTP (so TestClient cookies stick).

    ``allowed_domains=None`` means admit-all, so the test isolates the
    ``email_verified`` gate from the domain-allowlist check.

    :param skip_email_verification: Waive the ``email_verified`` gate,
        as ``OMNIGENT_OIDC_SKIP_EMAIL_VERIFICATION`` would.
    :param email_claim: Claim carrying the email identity, as
        ``OMNIGENT_OIDC_EMAIL_CLAIM`` would set it.
    """
    return OIDCConfig(
        issuer=_ISSUER,
        client_id=_CLIENT_ID,
        client_secret="secret",
        redirect_uri="http://localhost:8000/auth/callback",
        cookie_secret=_TEST_SECRET,
        scopes="openid email profile",
        session_ttl_hours=8,
        logout_redirect_uri=None,
        allowed_domains=None,
        provider_type="oidc",
        authorization_endpoint=f"{_ISSUER}/o/oauth2/v2/auth",
        token_endpoint="https://oauth2.googleapis.com/token",
        jwks_uri="https://www.googleapis.com/oauth2/v3/certs",
        userinfo_endpoint=None,
        allow_invites=False,
        skip_email_verification=skip_email_verification,
        email_claim=email_claim,
    )


class _IdpKeys:
    """An RSA keypair plus the JWKS signing key derived from its public half.

    :param private_key: The RSA private key used to sign test
        ``id_token`` JWTs.
    :param signing_key: A :class:`jwt.PyJWK` wrapping the public key,
        shaped exactly like what ``PyJWKClient.get_signing_key_from_jwt``
        returns — its ``.key`` is consumed by the production
        ``jwt.decode`` call for real signature verification.
    """

    def __init__(self) -> None:
        self.private_key = rsa.generate_private_key(public_exponent=65537, key_size=2048)
        jwk_dict = json.loads(RSAAlgorithm.to_jwk(self.private_key.public_key()))
        jwk_dict["alg"] = "RS256"
        self.signing_key = jwt.PyJWK.from_dict(jwk_dict)

    def sign_id_token(self, claims: dict[str, object]) -> str:
        """Sign ``claims`` into an RS256 ``id_token``, filling iss/aud/exp.

        :param claims: Claims to embed, e.g.
            ``{"email": "alice@example.com", "email_verified": True}``.
            ``iss``/``aud``/``exp``/``iat`` are added if absent.
        :returns: A compact-serialized signed JWT string.
        """
        now = int(time.time())
        payload: dict[str, object] = {
            "iss": _ISSUER,
            "aud": _CLIENT_ID,
            "iat": now,
            "exp": now + 300,
            "sub": "idp-subject-123",
            **claims,
        }
        return jwt.encode(payload, self.private_key, algorithm="RS256")


@pytest.fixture
def callback_client(
    tmp_path: Path,
    db_uri: str,
    monkeypatch: pytest.MonkeyPatch,
    request: pytest.FixtureRequest,
) -> Iterator[tuple[TestClient, _IdpKeys]]:
    """Mount the OIDC router and stub the IdP token endpoint + JWKS lookup.

    The token endpoint is driven per-test by mutating the mutable
    one-element ``pending_id_token`` list captured by the monkeypatched
    ``post`` — exposed on ``app.state.pending_id_token`` so ``_do_callback``
    can set the signed token the IdP should return.

    Indirect parametrization (``request.param``, default ``False``)
    sets the config's ``skip_email_verification`` flag; a dict param is
    passed to :func:`_oidc_config` as keyword arguments instead.
    """
    keys = _IdpKeys()
    perm_store = SqlAlchemyPermissionStore(db_uri)
    admins = tmp_path / "admins"
    admins.write_text("")

    param = getattr(request, "param", False)
    if isinstance(param, dict):
        config = _oidc_config(**param)
    else:
        config = _oidc_config(skip_email_verification=param)
    provider = UnifiedAuthProvider(source="oidc", oidc_config=config)

    # The signed id_token the mocked token endpoint will return. Each
    # test sets this before calling /auth/callback.
    pending_id_token: list[str] = [""]

    async def _fake_post(
        self: httpx.AsyncClient,
        url: str,
        *,
        data: dict[str, str] | None = None,
        headers: dict[str, str] | None = None,
        timeout: float | None = None,
    ) -> httpx.Response:
        """Stand in for the IdP token endpoint, returning the test's id_token."""
        return httpx.Response(200, json={"id_token": pending_id_token[0]})

    monkeypatch.setattr(httpx.AsyncClient, "post", _fake_post)
    # The production code constructs PyJWKClient(jwks_uri) then calls
    # get_signing_key_from_jwt; return our public key so the real
    # jwt.decode performs genuine signature verification offline.
    monkeypatch.setattr(
        jwt.PyJWKClient,
        "get_signing_key_from_jwt",
        lambda self, token: keys.signing_key,
    )

    app = FastAPI()
    app.include_router(
        create_auth_router(provider, perm_store, AdminList(admins)),
        prefix="/auth",
    )
    app.state.pending_id_token = pending_id_token

    with TestClient(app) as client:
        yield client, keys


def _do_callback(client: TestClient, id_token: str) -> httpx.Response:
    """Drive a full ``/auth/callback`` with a valid state cookie.

    Crafts the signed state cookie the way ``/auth/login`` would, sets
    the matching ``state`` query param, and stashes ``id_token`` for the
    mocked token endpoint to return. Redirects are not followed so the
    302 (and its ``Set-Cookie``) is observable.

    :param client: The TestClient mounting the OIDC router.
    :param id_token: The signed ``id_token`` the IdP should return.
    :returns: The raw callback response.
    """
    client.app.state.pending_id_token[0] = id_token
    state = "state-token-xyz"
    state_jwt = jwt.encode(
        {
            "state": state,
            "code_verifier": "verifier",
            "return_to": "/",
            "exp": int(time.time()) + 300,
        },
        _TEST_SECRET,
        algorithm="HS256",
    )
    client.cookies.set(_AUTH_STATE_COOKIE_PLAIN, state_jwt)
    return client.get(
        f"/auth/callback?code=auth-code&state={state}",
        follow_redirects=False,
    )


def test_oidc_accepts_es384_id_token(monkeypatch: pytest.MonkeyPatch) -> None:
    """Accept a valid token from an IdP that only advertises ES384."""
    private_key = ec.generate_private_key(ec.SECP384R1())
    jwk_dict = json.loads(ECAlgorithm.to_jwk(private_key.public_key()))
    jwk_dict["alg"] = "ES384"
    signing_key = jwt.PyJWK.from_dict(jwk_dict)
    monkeypatch.setattr(
        jwt.PyJWKClient,
        "get_signing_key_from_jwt",
        lambda self, token: signing_key,
    )

    now = int(time.time())
    token = jwt.encode(
        {
            "iss": _ISSUER,
            "aud": _CLIENT_ID,
            "iat": now,
            "exp": now + 300,
            "email": "Alice@Example.com",
            "email_verified": True,
        },
        private_key,
        algorithm="ES384",
    )

    assert _resolve_oidc_email({"id_token": token}, _oidc_config()) == "Alice@Example.com"


def test_callback_verified_email_mints_session(
    callback_client: tuple[TestClient, _IdpKeys],
) -> None:
    """A signed id_token with ``email_verified=true`` logs the user in.

    Proves the golden path still works after the fix: 302 back to the
    app and a session cookie whose ``sub`` is the verified email
    (normalized to lowercase).
    """
    client, keys = callback_client
    token = keys.sign_id_token({"email": "Alice@Example.com", "email_verified": True})

    resp = _do_callback(client, token)

    # 302 redirect (not 400) means the email was accepted as identity.
    assert resp.status_code == 302, resp.text
    session_cookie = resp.cookies.get("ap_session")
    # The session cookie must be set on success; absence would mean the
    # callback bailed before minting (the bug we're guarding against,
    # inverted).
    assert session_cookie is not None
    decoded = jwt.decode(session_cookie, _TEST_SECRET, algorithms=["HS256"])
    # sub is the normalized (lowercased) verified email — proves the
    # decoded claim flowed all the way into the minted session.
    assert decoded["sub"] == "alice@example.com"


@pytest.mark.parametrize(
    "claims",
    [
        pytest.param({"email": "victim@example.com", "email_verified": False}, id="false"),
        pytest.param({"email": "victim@example.com", "email_verified": "false"}, id="str-false"),
        pytest.param({"email": "victim@example.com"}, id="absent"),
        pytest.param({"email": "victim@example.com", "email_verified": None}, id="null"),
        pytest.param({"email": "victim@example.com", "email_verified": 1}, id="int-one"),
    ],
)
def test_callback_unverified_email_rejected(
    callback_client: tuple[TestClient, _IdpKeys],
    claims: dict[str, object],
) -> None:
    """An unverified/absent ``email_verified`` claim is rejected.

    The id_token is genuinely signature-valid (same RSA key as the
    happy path), so a failure here is *exclusively* the missing
    verification gate, not a signature/iss/aud rejection. Before the
    fix every one of these minted a session for ``victim@example.com``.
    """
    client, keys = callback_client
    token = keys.sign_id_token(claims)

    resp = _do_callback(client, token)

    # 400 (not 302): the callback refused to treat an unverified email
    # as identity. Anything else means the gate let it through.
    assert resp.status_code == 400, resp.text
    assert "Could not determine user email" in resp.json()["error"]
    # No session was minted for the spoofable email.
    assert resp.cookies.get("ap_session") is None


@pytest.mark.parametrize("callback_client", [True], indirect=True)
@pytest.mark.parametrize(
    "claims",
    [
        pytest.param({"email": "carol@example.com"}, id="absent"),
        pytest.param({"email": "carol@example.com", "email_verified": False}, id="false"),
    ],
)
def test_callback_skip_verification_flag_admits_unverified(
    callback_client: tuple[TestClient, _IdpKeys],
    claims: dict[str, object],
) -> None:
    """With ``skip_email_verification`` on, the gate is waived.

    Models Okta tiers that drop ``email_verified`` for
    directory-provisioned users: the same absent-claim token rejected
    by default (covered above) mints a session when the operator has
    opted out via ``OMNIGENT_OIDC_SKIP_EMAIL_VERIFICATION``.
    """
    client, keys = callback_client
    token = keys.sign_id_token(claims)

    resp = _do_callback(client, token)

    assert resp.status_code == 302, resp.text
    session_cookie = resp.cookies.get("ap_session")
    assert session_cookie is not None
    decoded = jwt.decode(session_cookie, _TEST_SECRET, algorithms=["HS256"])
    assert decoded["sub"] == "carol@example.com"


@pytest.mark.parametrize(
    "callback_client",
    [{"email_claim": "preferred_username", "skip_email_verification": True}],
    indirect=True,
)
def test_callback_custom_email_claim_admits_upn(
    callback_client: tuple[TestClient, _IdpKeys],
) -> None:
    """``email_claim`` reads the identity from an alternate claim.

    Models Microsoft Entra ID id_tokens that carry only
    ``preferred_username`` (the UPN) and no ``email`` claim: with the
    claim configured via ``OMNIGENT_OIDC_EMAIL_CLAIM`` and the
    verification opt-out set (a custom claim has no ``email_verified``
    marker), the UPN mints the session. Before the fix this token was
    rejected outright. Surrounding whitespace is removed before the
    identity is normalized.
    """
    client, keys = callback_client
    token = keys.sign_id_token({"preferred_username": " Dana@Example.com "})

    resp = _do_callback(client, token)

    assert resp.status_code == 302, resp.text
    session_cookie = resp.cookies.get("ap_session")
    assert session_cookie is not None
    decoded = jwt.decode(session_cookie, _TEST_SECRET, algorithms=["HS256"])
    # The UPN flowed into the session sub, normalized like an email.
    assert decoded["sub"] == "dana@example.com"


@pytest.mark.parametrize(
    "callback_client",
    [{"email_claim": "preferred_username", "skip_email_verification": True}],
    indirect=True,
)
@pytest.mark.parametrize(
    "claim_value",
    [
        pytest.param(["dana@example.com"], id="list"),
        pytest.param({"value": "dana@example.com"}, id="object"),
        pytest.param("   ", id="blank"),
    ],
)
def test_callback_custom_email_claim_rejects_invalid_value(
    callback_client: tuple[TestClient, _IdpKeys],
    claim_value: object,
) -> None:
    """A malformed configured identity claim is rejected cleanly."""
    client, keys = callback_client
    token = keys.sign_id_token({"preferred_username": claim_value})

    resp = _do_callback(client, token)

    assert resp.status_code == 400, resp.text
    assert resp.cookies.get("ap_session") is None


@pytest.mark.parametrize(
    "callback_client",
    [{"email_claim": "preferred_username"}],
    indirect=True,
)
@pytest.mark.parametrize(
    "claims",
    [
        pytest.param({"preferred_username": "dana@example.com"}, id="no-verified-marker"),
        pytest.param(
            {
                "preferred_username": "dana@example.com",
                "email": "attacker@evil.example",
                "email_verified": True,
            },
            id="verified-refers-to-a-different-claim",
        ),
    ],
)
def test_callback_custom_email_claim_still_requires_verification_optout(
    callback_client: tuple[TestClient, _IdpKeys],
    claims: dict[str, object],
) -> None:
    """A custom claim always requires the verification opt-out.

    ``email_verified`` refers to the ``email`` claim, so it vouches
    nothing about a custom identity claim — a token carrying
    ``email_verified: true`` for a *different* address must not smuggle
    the custom claim past the gate. Without
    ``OMNIGENT_OIDC_SKIP_EMAIL_VERIFICATION`` both shapes are rejected.
    """
    client, keys = callback_client
    token = keys.sign_id_token(claims)

    resp = _do_callback(client, token)

    assert resp.status_code == 400, resp.text
    assert resp.cookies.get("ap_session") is None


@pytest.mark.parametrize(
    "callback_client",
    [{"email_claim": "preferred_username", "skip_email_verification": True}],
    indirect=True,
)
def test_callback_custom_email_claim_absent_rejected(
    callback_client: tuple[TestClient, _IdpKeys],
) -> None:
    """When the configured claim is absent, the login is rejected.

    A verified ``email`` claim is not a silent fallback: the operator
    configured ``preferred_username`` as the identity claim, so a token
    without it must not mint a session from a different claim.
    """
    client, keys = callback_client
    token = keys.sign_id_token({"email": "dana@example.com", "email_verified": True})

    resp = _do_callback(client, token)

    assert resp.status_code == 400, resp.text
    assert resp.cookies.get("ap_session") is None


@pytest.mark.parametrize("verified_value", [True, "true", "True", "TRUE"])
def test_callback_accepts_boolean_and_string_true(
    callback_client: tuple[TestClient, _IdpKeys],
    verified_value: object,
) -> None:
    """Both boolean ``true`` and the string ``"true"`` are accepted.

    OIDC Core §5.1 notes ``email_verified`` may arrive as a string;
    accepting ``"true"`` keeps spec-compliant-but-string IdPs working
    while still rejecting ``"false"`` / absent (covered above).
    """
    client, keys = callback_client
    token = keys.sign_id_token({"email": "bob@example.com", "email_verified": verified_value})

    resp = _do_callback(client, token)

    # Accepted as a verified identity → redirect + session.
    assert resp.status_code == 302, resp.text
    assert resp.cookies.get("ap_session") is not None