user-engine/src/user_engine/oidc.py

"""OIDC Authorization Code + PKCE relying-party support."""

from __future__ import annotations

from dataclasses import dataclass
import base64
import hashlib
import json
import secrets
import time
from typing import Any, Mapping
from urllib.parse import urlencode
from urllib.request import Request, urlopen


@dataclass
class PendingLogin:
    verifier: str
    created_at: float
    recovery: bool = False
    return_path: str = "/"


@dataclass
class BrowserSession:
    claims: Mapping[str, Any]
    expires_at: float
    csrf_token: str = ""
    id_token: str = ""
    recovery: bool = False
    return_path: str = "/"


class OIDCClient:
    """Minimal confidential-state/public-client OIDC adapter."""

    def __init__(
        self,
        *,
        issuer: str,
        client_id: str,
        redirect_uri: str,
        audience: str,
        backend_url: str | None = None,
        session_ttl: int = 3600,
    ) -> None:
        self.issuer = issuer.rstrip("/")
        self.client_id = client_id
        self.redirect_uri = redirect_uri
        self.audience = audience
        self.backend_url = (backend_url or issuer).rstrip("/")
        self.session_ttl = session_ttl
        self.pending: dict[str, PendingLogin] = {}
        self.sessions: dict[str, BrowserSession] = {}

    def begin(
        self,
        *,
        tenant_hint: str | None = None,
        recovery: bool = False,
        fresh: bool = False,
        return_path: str = "/",
    ) -> str:
        if return_path not in {"/", "/platform", "/platform/factor-recovery", "/platform/authentication-policy"}:
            raise ValueError("unsupported login return path")
        state = secrets.token_urlsafe(32)
        verifier = secrets.token_urlsafe(64)
        challenge = _b64(hashlib.sha256(verifier.encode("ascii")).digest())
        self.pending[state] = PendingLogin(verifier=verifier, created_at=time.time(), recovery=recovery, return_path=return_path)
        self._prune()
        parameters = {
            'response_type': 'code',
            'client_id': self.client_id,
            'redirect_uri': self.redirect_uri,
            'scope': 'openid profile email groups',
            'state': state,
            'code_challenge': challenge,
            'code_challenge_method': 'S256',
        }
        if recovery:
            parameters.update(prompt="login", max_age="0", acr_values="aal2")
        elif fresh:
            # KeyCape asks Authelia for a session inside its short fresh window
            # and then requires that authentication to be newer than this request.
            parameters.update(prompt="login", max_age="0")
        if tenant_hint:
            parameters["tenant_hint"] = tenant_hint
        return f"{self.issuer}/authorize?{urlencode(parameters)}"

    def complete(self, *, code: str, state: str) -> str:
        pending = self.pending.pop(state, None)
        if pending is None or time.time() - pending.created_at > 600:
            raise ValueError("invalid or expired OIDC state")
        form = urlencode(
            {
                "grant_type": "authorization_code",
                "client_id": self.client_id,
                "redirect_uri": self.redirect_uri,
                "code": code,
                "code_verifier": pending.verifier,
            }
        ).encode()
        request = Request(
            f"{self.backend_url}/token",
            data=form,
            headers={"Content-Type": "application/x-www-form-urlencoded"},
            method="POST",
        )
        with urlopen(request, timeout=10) as response:
            tokens = json.loads(response.read())
        token = str(tokens.get("id_token") or tokens.get("access_token") or "")
        import jwt
        try:
            claims = self._verify(token)
        except jwt.PyJWTError:
            raise ValueError("OIDC token verification failed") from None
        session_id = secrets.token_urlsafe(32)
        expiry = min(float(claims.get("exp", time.time() + self.session_ttl)), time.time() + self.session_ttl)
        self.sessions[session_id] = BrowserSession(
            claims=claims,
            expires_at=expiry,
            csrf_token=secrets.token_urlsafe(32),
            id_token=token,
            recovery=pending.recovery,
            return_path=pending.return_path,
        )
        self._prune()
        return session_id

    def claims(self, session_id: str) -> Mapping[str, Any] | None:
        session = self.sessions.get(session_id)
        if session is None or session.expires_at <= time.time():
            self.sessions.pop(session_id, None)
            return None
        return session.claims

    def logout(self, session_id: str) -> None:
        self.sessions.pop(session_id, None)

    def csrf_token(self, session_id: str) -> str | None:
        session = self.sessions.get(session_id)
        if session is None or session.expires_at <= time.time():
            self.sessions.pop(session_id, None)
            return None
        return session.csrf_token

    def _verify(self, token: str) -> Mapping[str, Any]:
        if not token:
            raise ValueError("OIDC token response is missing a token")
        try:
            import jwt
        except ImportError as exc:  # pragma: no cover
            raise RuntimeError("install user-engine[oidc] for OIDC login") from exc
        jwks = jwt.PyJWKClient(f"{self.backend_url}/jwks")
        key = jwks.get_signing_key_from_jwt(token)
        claims = jwt.decode(
            token,
            key.key,
            algorithms=["RS256"],
            audience=self.audience,
            issuer=self.issuer,
            options={"require": ["exp", "iat", "iss", "sub", "aud"]},
        )
        return dict(claims)

    def _prune(self) -> None:
        now = time.time()
        self.pending = {
            key: value for key, value in self.pending.items() if now - value.created_at <= 600
        }
        self.sessions = {
            key: value for key, value in self.sessions.items() if value.expires_at > now
        }


def cookie_value(header: str, name: str) -> str | None:
    for part in header.split(";"):
        key, separator, value = part.strip().partition("=")
        if separator and key == name:
            return value
    return None


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