stackpilot/backend/tests/test_oidc.py

"""Signing in through an identity provider.

No provider is contacted: a throwaway RSA key stands in for one, so ID tokens
can be minted here and the verification exercised for real rather than mocked
away. That is the point — the whole feature rests on "the provider said so"
being worth something, which it only is if the signature is actually checked.

The cases that matter are the ones where a token looks fine and must still be
refused: signed by the wrong key, issued to somebody else, from another issuer,
replayed with an old nonce, or arriving on a state that was already used.
"""
from __future__ import annotations

import time
from datetime import datetime, timedelta, timezone

import pytest
from jose import jwt
from sqlmodel import Session, delete, select

ISSUER = "https://idp.test/realms/homelab"
CLIENT_ID = "stackpilot"

# A 2048-bit RSA key, generated once for these tests and used nowhere else.
KEY = None
WRONG_KEY = None


def _keypair():
    """Build a JWK pair lazily — generating RSA keys is slow enough to matter."""
    global KEY, WRONG_KEY
    if KEY is None:
        from cryptography.hazmat.primitives import serialization
        from cryptography.hazmat.primitives.asymmetric import rsa

        def make(kid):
            private = rsa.generate_private_key(public_exponent=65537, key_size=2048)
            pem = private.private_bytes(
                encoding=serialization.Encoding.PEM,
                format=serialization.PrivateFormat.PKCS8,
                encryption_algorithm=serialization.NoEncryption(),
            ).decode()
            from jose import jwk

            public = jwk.construct(
                private.public_key()
                .public_bytes(
                    encoding=serialization.Encoding.PEM,
                    format=serialization.PublicFormat.SubjectPublicKeyInfo,
                )
                .decode(),
                algorithm="RS256",
            ).to_dict()
            public = {k: (v.decode() if isinstance(v, bytes) else v) for k, v in public.items()}
            public["kid"] = kid
            public["alg"] = "RS256"
            public["use"] = "sig"
            return pem, public

        KEY = make("test-key-1")
        WRONG_KEY = make("test-key-1")  # same kid, different key: the nasty case
    return KEY, WRONG_KEY


def _id_token(nonce: str, *, key=None, claims=None, kid="test-key-1") -> str:
    (good_pem, _), (bad_pem, _) = _keypair()
    payload = {
        "iss": ISSUER,
        "aud": CLIENT_ID,
        "sub": "idp-subject-1",
        "preferred_username": "alice",
        "email": "alice@example.test",
        "nonce": nonce,
        "exp": int(time.time()) + 300,
        "iat": int(time.time()),
    }
    payload.update(claims or {})
    return jwt.encode(payload, key or good_pem, algorithm="RS256", headers={"kid": kid})


@pytest.fixture
def svc(db):
    from services import oidc_service

    oidc_service._cache.clear()
    return oidc_service


@pytest.fixture(autouse=True)
def clean(db):
    from database import engine
    from models.oidc import OidcConfig, OidcState
    from models.user import User

    def wipe():
        with Session(engine) as session:
            session.exec(delete(OidcState))
            session.exec(delete(OidcConfig))
            for user in session.exec(select(User)).all():
                if user.oidc_subject or user.username in ("alice", "bob"):
                    session.delete(user)
            session.commit()

    wipe()
    yield
    wipe()


@pytest.fixture
def config(db, svc, monkeypatch):
    """A configured provider, with discovery and JWKS served from memory."""
    from database import engine
    from models.oidc import OidcConfig
    from services import crypto_service

    (_pem, public), _ = _keypair()
    document = {
        "issuer": ISSUER,
        "authorization_endpoint": f"{ISSUER}/protocol/openid-connect/auth",
        "token_endpoint": f"{ISSUER}/protocol/openid-connect/token",
        "jwks_uri": f"{ISSUER}/protocol/openid-connect/certs",
    }

    async def fake_discover(issuer, force=False):
        return document

    async def fake_keys(doc, force=False):
        return {"keys": [public]}

    monkeypatch.setattr(svc, "discover", fake_discover)
    monkeypatch.setattr(svc, "_signing_keys", fake_keys)

    with Session(engine) as session:
        row = OidcConfig(
            id=1,
            enabled=True,
            issuer=ISSUER,
            client_id=CLIENT_ID,
            client_secret=crypto_service.encrypt("shhh"),
            auto_create=True,
            default_role="user",
        )
        session.add(row)
        session.commit()
        session.refresh(row)
        yield session, row


# --------------------------------------------------------------------------- #
# State, PKCE and single use
# --------------------------------------------------------------------------- #


def test_a_login_gets_its_own_state_verifier_and_nonce(svc, config):
    session, _row = config
    first = svc.begin(session, "https://sp.test/cb")
    second = svc.begin(session, "https://sp.test/cb")
    assert first.state != second.state
    assert first.verifier != second.verifier
    assert first.nonce != second.nonce
    # PKCE verifiers must be long enough to be worth anything.
    assert len(first.verifier) >= 43


def test_the_authorize_url_carries_pkce_and_the_nonce(svc, config):
    import asyncio
    from urllib.parse import parse_qs, urlparse

    session, row = config
    state = svc.begin(session, "https://sp.test/cb")
    url = asyncio.run(svc.authorize_url(row, state))
    query = parse_qs(urlparse(url).query)

    assert query["code_challenge_method"] == ["S256"]
    # The challenge is the hash, never the verifier itself.
    assert query["code_challenge"][0] != state.verifier
    assert query["state"] == [state.state]
    assert query["nonce"] == [state.nonce]
    assert query["redirect_uri"] == ["https://sp.test/cb"]


def test_a_state_can_only_be_used_once(svc, config):
    session, _row = config
    state = svc.begin(session, "https://sp.test/cb")
    assert svc.take_state(session, state.state) is not None
    # A replayed callback finds nothing.
    assert svc.take_state(session, state.state) is None


def test_an_expired_state_is_refused(svc, config):
    from models.oidc import OidcState

    session, _row = config
    state = svc.begin(session, "https://sp.test/cb")
    row = session.get(OidcState, state.state)
    row.created_at = datetime.now(timezone.utc) - svc.STATE_TTL - timedelta(minutes=1)
    session.add(row)
    session.commit()
    assert svc.take_state(session, state.state) is None


def test_an_unknown_state_is_refused(svc, config):
    session, _row = config
    assert svc.take_state(session, "never-issued") is None


# --------------------------------------------------------------------------- #
# ID token verification
# --------------------------------------------------------------------------- #


def _verify(svc, row, state, token, access_token=None):
    import asyncio

    tokens = {"id_token": token}
    if access_token:
        tokens["access_token"] = access_token
    return asyncio.run(svc.verify_id_token(row, tokens, state))


def test_a_properly_signed_token_is_accepted(svc, config):
    session, row = config
    state = svc.begin(session, "https://sp.test/cb")
    claims = _verify(svc, row, state, _id_token(state.nonce))
    assert claims["preferred_username"] == "alice"


def test_a_token_signed_with_the_wrong_key_is_refused(svc, config):
    """The one that matters: same kid, different key."""
    session, row = config
    state = svc.begin(session, "https://sp.test/cb")
    (_good, _), (bad_pem, _) = _keypair()
    with pytest.raises(svc.OidcError):
        _verify(svc, row, state, _id_token(state.nonce, key=bad_pem))


def test_a_token_for_another_audience_is_refused(svc, config):
    session, row = config
    state = svc.begin(session, "https://sp.test/cb")
    with pytest.raises(svc.OidcError):
        _verify(svc, row, state, _id_token(state.nonce, claims={"aud": "some-other-app"}))


def test_a_token_from_another_issuer_is_refused(svc, config):
    session, row = config
    state = svc.begin(session, "https://sp.test/cb")
    with pytest.raises(svc.OidcError):
        _verify(svc, row, state, _id_token(state.nonce, claims={"iss": "https://evil.test"}))


def test_an_expired_token_is_refused(svc, config):
    session, row = config
    state = svc.begin(session, "https://sp.test/cb")
    with pytest.raises(svc.OidcError):
        _verify(svc, row, state, _id_token(state.nonce, claims={"exp": int(time.time()) - 600}))


def test_a_token_with_the_wrong_nonce_is_refused(svc, config):
    """Replaying an ID token obtained during a different login."""
    session, row = config
    state = svc.begin(session, "https://sp.test/cb")
    with pytest.raises(svc.OidcError) as caught:
        _verify(svc, row, state, _id_token("a-nonce-from-somewhere-else"))
    assert "nonce" in str(caught.value)


def test_a_token_signed_with_an_unknown_key_is_refused(svc, config):
    session, row = config
    state = svc.begin(session, "https://sp.test/cb")
    with pytest.raises(svc.OidcError):
        _verify(svc, row, state, _id_token(state.nonce, kid="some-other-kid"))


def test_garbage_is_refused(svc, config):
    session, row = config
    state = svc.begin(session, "https://sp.test/cb")
    with pytest.raises(svc.OidcError):
        _verify(svc, row, state, "not.a.token")


# --------------------------------------------------------------------------- #
# Claims to accounts
# --------------------------------------------------------------------------- #


def test_a_first_sign_in_creates_the_account(svc, config):
    session, row = config
    user = svc.resolve_user(session, row, {"sub": "s1", "preferred_username": "alice"})
    assert user.username == "alice"
    assert user.role == "user"
    assert user.oidc_subject == "s1"


def test_the_created_account_cannot_be_signed_into_with_a_password(svc, config):
    import auth as auth_mod

    session, row = config
    user = svc.resolve_user(session, row, {"sub": "s1", "preferred_username": "alice"})
    # Whatever is stored must not verify against anything, empty string included
    # — and must fail cleanly rather than raising on an unparseable hash.
    assert not auth_mod.verify_password("", user.hashed_password)
    assert not auth_mod.verify_password("hunter2", user.hashed_password)
    assert auth_mod.authenticate(session, "alice", "") is None


def test_a_password_login_against_an_oidc_account_is_a_clean_401(client, svc, config):
    """Not a 500: an unparseable stored hash used to raise out of passlib."""
    session, row = config
    svc.resolve_user(session, row, {"sub": "s1", "preferred_username": "alice"})
    response = client.post("/api/auth/login", json={"username": "alice", "password": "x"})
    assert response.status_code == 401


def test_an_existing_local_account_is_linked_not_duplicated(svc, config):
    import auth as auth_mod
    from models.user import User

    session, row = config
    session.add(
        User(username="alice", hashed_password=auth_mod.hash_password("pw"), role="admin")
    )
    session.commit()

    user = svc.resolve_user(session, row, {"sub": "s1", "preferred_username": "alice"})
    assert user.oidc_subject == "s1"
    # The local role survives: linking must not quietly demote an admin.
    assert user.role == "admin"
    assert len(session.exec(select(User).where(User.username == "alice")).all()) == 1


def test_the_subject_wins_over_a_renamed_username(svc, config):
    session, row = config
    first = svc.resolve_user(session, row, {"sub": "s1", "preferred_username": "alice"})
    again = svc.resolve_user(session, row, {"sub": "s1", "preferred_username": "alice-renamed"})
    assert again.id == first.id


def test_auto_create_can_be_turned_off(svc, config):
    session, row = config
    row.auto_create = False
    with pytest.raises(svc.OidcError) as caught:
        svc.resolve_user(session, row, {"sub": "s1", "preferred_username": "nobody"})
    assert "automatic creation is off" in str(caught.value)


def test_a_disabled_account_cannot_sign_in_through_the_provider(svc, config):
    session, row = config
    user = svc.resolve_user(session, row, {"sub": "s1", "preferred_username": "alice"})
    user.is_active = False
    session.add(user)
    session.commit()
    with pytest.raises(svc.OidcError) as caught:
        svc.resolve_user(session, row, {"sub": "s1", "preferred_username": "alice"})
    assert "disabled" in str(caught.value)


def test_the_username_claim_is_configurable_with_fallbacks(svc, config):
    session, row = config
    row.username_claim = "email"
    user = svc.resolve_user(
        session, row, {"sub": "s1", "email": "bob@example.test", "preferred_username": "x"}
    )
    assert user.username == "bob@example.test"


def test_a_group_claim_can_grant_and_remove_admin(svc, config):
    session, row = config
    row.admin_claim = "groups"
    row.admin_value = "stackpilot-admins"

    promoted = svc.resolve_user(
        session, row, {"sub": "s1", "preferred_username": "alice", "groups": ["stackpilot-admins"]}
    )
    assert promoted.role == "admin"

    # Removed from the group upstream: the next sign-in takes it away again.
    demoted = svc.resolve_user(
        session, row, {"sub": "s1", "preferred_username": "alice", "groups": ["other"]}
    )
    assert demoted.role == "user"


def test_without_a_mapping_the_local_role_is_left_alone(svc, config):
    session, row = config
    user = svc.resolve_user(session, row, {"sub": "s1", "preferred_username": "alice"})
    user.role = "admin"
    session.add(user)
    session.commit()

    again = svc.resolve_user(session, row, {"sub": "s1", "preferred_username": "alice"})
    assert again.role == "admin"


# --------------------------------------------------------------------------- #
# Through the API
# --------------------------------------------------------------------------- #


def test_the_status_endpoint_is_public_and_says_little(client, config):
    response = client.get("/api/auth/oidc/status")
    assert response.status_code == 200
    body = response.json()
    assert body["enabled"] is True
    # No issuer, no client id: a stranger learns only that SSO exists.
    assert set(body) == {"enabled", "button_label"}


def test_status_is_false_when_nothing_is_configured(client, db):
    assert client.get("/api/auth/oidc/status").json()["enabled"] is False


def test_the_secret_never_comes_back_out(as_admin, config):
    body = as_admin.get("/api/auth/oidc/config").text
    assert "shhh" not in body
    assert '"has_client_secret":true' in body.replace(" ", "")


def test_saving_without_a_secret_keeps_the_stored_one(as_admin, config, svc):
    session, _row = config
    saved = as_admin.put(
        "/api/auth/oidc/config",
        json={
            "enabled": True,
            "issuer": ISSUER,
            "client_id": CLIENT_ID,
            "button_label": "Sign in with Authentik",
            "default_role": "user",
        },
    )
    assert saved.status_code == 200, saved.text
    session.expire_all()
    assert svc.client_secret(svc.get_config(session)) == "shhh"


def test_enabling_without_an_issuer_is_refused(as_admin, db):
    response = as_admin.put(
        "/api/auth/oidc/config",
        json={"enabled": True, "issuer": "", "client_id": "", "default_role": "user"},
    )
    assert response.status_code == 400


def test_a_callback_with_an_unknown_state_lands_back_on_the_login_page(client, config):
    response = client.get(
        "/api/auth/oidc/callback", params={"code": "x", "state": "made-up"},
        follow_redirects=False,
    )
    assert response.status_code == 302
    assert "/login?sso_error=" in response.headers["location"]


def test_a_provider_error_is_passed_through_to_the_login_page(client, config):
    response = client.get(
        "/api/auth/oidc/callback",
        params={"error": "access_denied", "error_description": "User said no"},
        follow_redirects=False,
    )
    assert response.status_code == 302
    assert "User%20said%20no" in response.headers["location"]


def test_the_read_only_role_cannot_read_or_change_the_config(as_user, config):
    assert as_user.get("/api/auth/oidc/config").status_code == 403
    assert (
        as_user.put(
            "/api/auth/oidc/config",
            json={"enabled": False, "issuer": "", "client_id": "", "default_role": "user"},
        ).status_code
        == 403
    )