stackpilot/backend/services/oidc_service.py

"""Signing in through an OpenID Connect provider.

Authorization Code flow with PKCE, which is the current recommendation even for
a confidential client: the code is useless to anyone who intercepts it without
the verifier that never left this process.

The part worth understanding is how the session is handed back to the browser.
The callback does *not* put a token in the URL. It sets the same httpOnly
refresh cookie a password login sets and redirects to the app, which already
trades that cookie for an access token on boot — so an OIDC session is exactly
a normal StackPilot session, and no token is ever written somewhere a proxy log
or a browser history could keep it.

What this trusts, and what it verifies: the ID token's signature against the
provider's published keys, its issuer, its audience, its expiry, and the nonce
minted for this particular login. An unsigned or mis-signed token is refused —
"the provider said so" is only worth anything if it really was the provider.
"""
from __future__ import annotations

import base64
import hashlib
import logging
import secrets
import time
from datetime import datetime, timedelta, timezone
from typing import Optional
from urllib.parse import urlencode

import httpx
from jose import jwt
from jose.exceptions import JWTError
from sqlmodel import Session, select

from models.oidc import OidcConfig, OidcState
from models.user import User
from services import crypto_service

logger = logging.getLogger("stackpilot.oidc")

TIMEOUT = httpx.Timeout(15.0)

#: A login has this long to come back from the provider. Long enough to type a
#: password and answer an MFA prompt, short enough that an abandoned state row
#: is not a lasting foothold.
STATE_TTL = timedelta(minutes=15)

#: Discovery documents and signing keys change rarely; refetching them on every
#: login would put the provider in the hot path of every sign-in.
_DISCOVERY_TTL = 3600.0
_cache: dict[str, tuple[float, dict]] = {}


class OidcError(Exception):
    """Anything that stops a sign-in, phrased for the person who has to fix it."""


# --------------------------------------------------------------------------- #
# Configuration
# --------------------------------------------------------------------------- #


def get_config(session: Session) -> Optional[OidcConfig]:
    return session.get(OidcConfig, 1)


def is_enabled(session: Session) -> bool:
    config = get_config(session)
    return bool(config and config.enabled and config.issuer and config.client_id)


def client_secret(config: OidcConfig) -> str:
    try:
        return crypto_service.decrypt(config.client_secret or "")
    except crypto_service.DecryptError as exc:
        raise OidcError(str(exc)) from exc


# --------------------------------------------------------------------------- #
# Provider metadata
# --------------------------------------------------------------------------- #


async def discover(issuer: str, force: bool = False) -> dict:
    """The provider's OpenID configuration document."""
    url = issuer.rstrip("/") + "/.well-known/openid-configuration"
    hit = _cache.get(url)
    if hit and not force and time.time() - hit[0] < _DISCOVERY_TTL:
        return hit[1]
    try:
        async with httpx.AsyncClient(follow_redirects=True, timeout=TIMEOUT) as client:
            response = await client.get(url)
            response.raise_for_status()
            document = response.json()
    except (httpx.HTTPError, ValueError) as exc:
        raise OidcError(f"Could not read {url}: {exc}") from exc
    for required in ("authorization_endpoint", "token_endpoint", "issuer"):
        if not document.get(required):
            raise OidcError(f"The provider's metadata is missing '{required}'")
    _cache[url] = (time.time(), document)
    return document


async def _signing_keys(document: dict, force: bool = False) -> dict:
    url = document.get("jwks_uri")
    if not url:
        raise OidcError("The provider publishes no jwks_uri, so tokens cannot be verified")
    hit = _cache.get(url)
    if hit and not force and time.time() - hit[0] < _DISCOVERY_TTL:
        return hit[1]
    try:
        async with httpx.AsyncClient(follow_redirects=True, timeout=TIMEOUT) as client:
            response = await client.get(url)
            response.raise_for_status()
            keys = response.json()
    except (httpx.HTTPError, ValueError) as exc:
        raise OidcError(f"Could not read the provider's signing keys: {exc}") from exc
    _cache[url] = (time.time(), keys)
    return keys


# --------------------------------------------------------------------------- #
# Starting a login
# --------------------------------------------------------------------------- #


def _b64url(raw: bytes) -> str:
    return base64.urlsafe_b64encode(raw).decode("ascii").rstrip("=")


def begin(session: Session, redirect_uri: str) -> OidcState:
    """Mint and store the state, PKCE verifier and nonce for one login."""
    prune_states(session)
    row = OidcState(
        state=_b64url(secrets.token_bytes(24)),
        verifier=_b64url(secrets.token_bytes(48)),
        nonce=_b64url(secrets.token_bytes(24)),
        redirect_uri=redirect_uri,
    )
    session.add(row)
    session.commit()
    session.refresh(row)
    return row


def take_state(session: Session, state: str) -> Optional[OidcState]:
    """Consume a state row. Returns None if unknown, used already, or expired.

    Deleting it here is what makes a login single-use: a replayed callback finds
    nothing and is refused.
    """
    row = session.get(OidcState, state or "")
    if not row:
        return None
    session.delete(row)
    session.commit()
    created = row.created_at
    if created.tzinfo is None:
        created = created.replace(tzinfo=timezone.utc)
    if datetime.now(timezone.utc) - created > STATE_TTL:
        return None
    return row


def prune_states(session: Session) -> int:
    cutoff = datetime.now(timezone.utc) - STATE_TTL
    gone = 0
    for row in session.exec(select(OidcState)).all():
        created = row.created_at
        if created.tzinfo is None:
            created = created.replace(tzinfo=timezone.utc)
        if created < cutoff:
            session.delete(row)
            gone += 1
    if gone:
        session.commit()
    return gone


async def authorize_url(config: OidcConfig, state: OidcState) -> str:
    document = await discover(config.issuer)
    challenge = _b64url(hashlib.sha256(state.verifier.encode("ascii")).digest())
    query = {
        "response_type": "code",
        "client_id": config.client_id,
        "redirect_uri": state.redirect_uri,
        "scope": config.scopes or "openid profile email",
        "state": state.state,
        "nonce": state.nonce,
        "code_challenge": challenge,
        "code_challenge_method": "S256",
    }
    return f"{document['authorization_endpoint']}?{urlencode(query)}"


# --------------------------------------------------------------------------- #
# Finishing a login
# --------------------------------------------------------------------------- #


async def exchange(config: OidcConfig, code: str, state: OidcState) -> dict:
    """Trade the authorization code for tokens."""
    document = await discover(config.issuer)
    data = {
        "grant_type": "authorization_code",
        "code": code,
        "redirect_uri": state.redirect_uri,
        "client_id": config.client_id,
        "code_verifier": state.verifier,
    }
    secret = client_secret(config)
    auth = (config.client_id, secret) if secret else None
    try:
        async with httpx.AsyncClient(follow_redirects=True, timeout=TIMEOUT) as client:
            response = await client.post(document["token_endpoint"], data=data, auth=auth)
            if response.status_code >= 400:
                # The provider's own error is far more useful than a generic one.
                raise OidcError(_token_error(response))
            payload = response.json()
    except httpx.HTTPError as exc:
        raise OidcError(f"Could not reach the provider's token endpoint: {exc}") from exc
    except ValueError as exc:
        raise OidcError("The provider's token response was not JSON") from exc
    if not payload.get("id_token"):
        raise OidcError("The provider returned no ID token — is 'openid' in the scopes?")
    return payload


def _token_error(response: httpx.Response) -> str:
    try:
        body = response.json()
        detail = body.get("error_description") or body.get("error") or ""
    except ValueError:
        detail = response.text[:200]
    return f"The provider rejected the login ({response.status_code}): {detail}".strip()


async def verify_id_token(config: OidcConfig, tokens: dict, state: OidcState) -> dict:
    """Validate the ID token and return its claims."""
    document = await discover(config.issuer)
    id_token = tokens["id_token"]

    try:
        header = jwt.get_unverified_header(id_token)
    except JWTError as exc:
        raise OidcError("The ID token is malformed") from exc

    keys = await _signing_keys(document)
    key = _pick_key(keys, header.get("kid"))
    if key is None:
        # A provider that has rotated its keys since we cached them.
        keys = await _signing_keys(document, force=True)
        key = _pick_key(keys, header.get("kid"))
    if key is None:
        raise OidcError("The ID token was signed with a key the provider does not publish")

    try:
        claims = jwt.decode(
            id_token,
            key,
            algorithms=[header.get("alg", "RS256")],
            audience=config.client_id,
            issuer=document["issuer"],
            access_token=tokens.get("access_token"),
            options={"leeway": 60},
        )
    except JWTError as exc:
        raise OidcError(f"The ID token failed verification: {exc}") from exc

    if claims.get("nonce") != state.nonce:
        # Without this an attacker could replay an ID token obtained elsewhere.
        raise OidcError("The ID token's nonce does not match this login")
    return claims


def _pick_key(keys: dict, kid: Optional[str]) -> Optional[dict]:
    candidates = keys.get("keys") or []
    if kid:
        for key in candidates:
            if key.get("kid") == kid:
                return key
        return None
    return candidates[0] if len(candidates) == 1 else None


# --------------------------------------------------------------------------- #
# Turning claims into a user
# --------------------------------------------------------------------------- #


def _username_from(config: OidcConfig, claims: dict) -> str:
    for claim in (config.username_claim or "preferred_username", "email", "sub"):
        value = claims.get(claim)
        if isinstance(value, str) and value.strip():
            return value.strip()
    raise OidcError("The ID token carries no usable username claim")


def _is_admin(config: OidcConfig, claims: dict) -> Optional[bool]:
    """True/False from the configured claim, or None when no mapping is set."""
    if not config.admin_claim or not config.admin_value:
        return None
    value = claims.get(config.admin_claim)
    if isinstance(value, str):
        return value == config.admin_value
    if isinstance(value, (list, tuple)):
        return config.admin_value in value
    return False


def resolve_user(session: Session, config: OidcConfig, claims: dict) -> User:
    """Find, link or create the account this login belongs to."""
    subject = str(claims.get("sub") or "")
    if not subject:
        raise OidcError("The ID token has no subject")
    username = _username_from(config, claims)
    admin = _is_admin(config, claims)

    # The subject is the only identifier the provider promises is stable, so it
    # wins over the username — somebody renamed upstream stays the same account.
    user = session.exec(select(User).where(User.oidc_subject == subject)).first()
    if user is None:
        existing = session.exec(select(User).where(User.username == username)).first()
        if existing is not None:
            # First OIDC sign-in for somebody who already had a local account:
            # link them rather than creating a duplicate.
            existing.oidc_subject = subject
            user = existing
        elif config.auto_create:
            user = User(
                username=username,
                # No usable password: this account signs in through the provider.
                # A random unusable hash, never a blank one that might verify.
                hashed_password="!oidc:" + secrets.token_urlsafe(16),
                role=config.default_role if config.default_role in ("admin", "user") else "user",
                oidc_subject=subject,
            )
            session.add(user)
        else:
            raise OidcError(
                f"No StackPilot account for '{username}', and automatic creation is off"
            )

    if not user.is_active:
        # Disabling an account here has to outrank the provider's opinion.
        raise OidcError(f"The account '{user.username}' is disabled")

    if admin is not None:
        role = "admin" if admin else "user"
        if user.role != role:
            logger.info("OIDC role mapping: %s -> %s", user.username, role)
            user.role = role
    session.add(user)
    session.commit()
    session.refresh(user)
    return user