coulomb-social/coulomb_social/apps/identity/oidc.py

"""OIDC/PKCE helpers for NetKingdom IAM Profile issuers."""

from __future__ import annotations

import logging
import secrets
from typing import Any

import httpx
from authlib.integrations.httpx_client import OAuth2Client
from authlib.jose import JsonWebKey, jwt
from django.conf import settings

logger = logging.getLogger(__name__)


class OIDCConfigurationError(RuntimeError):
    pass


def oidc_configured() -> bool:
    return bool(
        settings.OIDC_ENABLED
        and settings.OIDC_ISSUER
        and settings.OIDC_CLIENT_ID
        and settings.OIDC_REDIRECT_URI
    )


def discovery_document() -> dict[str, Any]:
    if not settings.OIDC_ISSUER and not settings.OIDC_DISCOVERY_URL:
        raise OIDCConfigurationError("OIDC_ISSUER or OIDC_DISCOVERY_URL required")
    url = settings.OIDC_DISCOVERY_URL or (
        settings.OIDC_ISSUER.rstrip("/") + "/.well-known/openid-configuration"
    )
    resp = httpx.get(url, timeout=15.0)
    resp.raise_for_status()
    return resp.json()


def _oauth_client() -> OAuth2Client:
    """Public clients (KeyCape default for browser apps) use PKCE without a secret."""
    secret = (settings.OIDC_CLIENT_SECRET or "").strip() or None
    kwargs: dict[str, Any] = {
        "client_id": settings.OIDC_CLIENT_ID,
        "redirect_uri": settings.OIDC_REDIRECT_URI,
        "scope": settings.OIDC_SCOPES,
        "code_challenge_method": "S256",
    }
    if secret:
        kwargs["client_secret"] = secret
    else:
        # Authlib: omit secret for public clients
        kwargs["token_endpoint_auth_method"] = "none"
    return OAuth2Client(**kwargs)


def build_authorization_url(
    *, state: str, code_verifier: str, acr_values: str | None = None
) -> str:
    if not oidc_configured():
        raise OIDCConfigurationError("OIDC is not enabled/configured")
    doc = discovery_document()
    auth_endpoint = doc["authorization_endpoint"]
    client = _oauth_client()
    parameters = {"state": state, "code_verifier": code_verifier}
    if acr_values:
        parameters["acr_values"] = acr_values
    uri, _ = client.create_authorization_url(auth_endpoint, **parameters)
    return uri


def exchange_code(code: str, *, code_verifier: str) -> dict[str, Any]:
    doc = discovery_document()
    token_endpoint = doc["token_endpoint"]
    client = _oauth_client()
    token = client.fetch_token(
        token_endpoint,
        code=code,
        code_verifier=code_verifier,
        grant_type="authorization_code",
    )
    return token


def fetch_userinfo(access_token: str) -> dict[str, Any]:
    """Best-effort userinfo. KeyCape may 401 for some subjects; id_token is enough."""
    if not access_token:
        return {}
    doc = discovery_document()
    userinfo_endpoint = doc.get("userinfo_endpoint")
    if not userinfo_endpoint:
        return {}
    try:
        resp = httpx.get(
            userinfo_endpoint,
            headers={"Authorization": f"Bearer {access_token}"},
            timeout=15.0,
        )
        if resp.status_code >= 400:
            logger.warning(
                "OIDC userinfo returned %s; continuing with id_token claims",
                resp.status_code,
            )
            return {}
        return resp.json()
    except Exception:
        logger.exception("OIDC userinfo request failed; continuing with id_token claims")
        return {}


def decode_id_token(id_token: str) -> dict[str, Any]:
    """Verify id_token with issuer JWKS and return claims."""
    if not id_token or id_token.count(".") != 2:
        raise OIDCConfigurationError("token response missing a usable id_token")
    doc = discovery_document()
    jwks_uri = doc.get("jwks_uri")
    if not jwks_uri:
        raise OIDCConfigurationError("issuer discovery missing jwks_uri")
    jwks = httpx.get(jwks_uri, timeout=15.0)
    jwks.raise_for_status()
    key_set = JsonWebKey.import_key_set(jwks.json())
    claims = jwt.decode(
        id_token,
        key_set,
        claims_options={
            "iss": {"essential": True, "value": settings.OIDC_ISSUER.rstrip("/")},
            "aud": {"essential": True, "value": settings.OIDC_CLIENT_ID},
            "exp": {"essential": True},
            "sub": {"essential": True},
        },
    )
    claims.validate()
    return dict(claims)


def new_pkce_pair() -> tuple[str, str]:
    """Return (state, code_verifier). Challenge is computed by Authlib client."""
    state = secrets.token_urlsafe(24)
    code_verifier = secrets.token_urlsafe(48)
    return state, code_verifier


def claims_from_token_response(token: dict[str, Any], userinfo: dict[str, Any]) -> dict[str, Any]:
    """Prefer verified id_token claims; overlay optional userinfo."""
    claims: dict[str, Any] = {}
    id_token = token.get("id_token")
    if isinstance(id_token, str) and id_token:
        claims.update(decode_id_token(id_token))
    elif isinstance(token.get("userinfo"), dict):
        claims.update(token["userinfo"])
    # userinfo is optional enrichment (KeyCape may 401 for some subjects)
    if userinfo:
        claims.update(userinfo)
    if not claims.get("sub"):
        raise OIDCConfigurationError("OIDC response has no subject claim")
    return claims