ping-mee-self-service-platform/portal/adapters/idp/oidc.py

"""OpenID Connect authorization-code flow with PKCE. The only module that talks OIDC (AD-1)."""

from __future__ import annotations

import hmac
import logging
import secrets
import threading
import time
from collections.abc import Callable, Mapping
from typing import Any

import requests
from authlib.common.errors import AuthlibBaseError
from authlib.integrations.requests_client import OAuth2Session
from joserfc import jwt
from joserfc.errors import InvalidKeyIdError, JoseError
from joserfc.jwk import KeySet

from portal.config.env import OidcConfig

from .base import Flow, Identity, IdpUnavailable, LoginRejected

logger = logging.getLogger(__name__)

TIMEOUT_SECONDS = 10
CACHE_SECONDS = 60 * 60
CLOCK_LEEWAY_SECONDS = 30
SCOPE = "openid profile email"
# Asymmetric algorithms only: a token signed with a shared secret is never accepted.
ALGORITHMS = ["RS256", "RS384", "RS512", "PS256", "PS384", "PS512", "ES256", "ES384", "ES512"]

SessionFactory = Callable[[], OAuth2Session]


class OidcIdp:
    def __init__(
        self,
        config: OidcConfig,
        session_factory: SessionFactory | None = None,
        clock: Callable[[], float] = time.time,
    ) -> None:
        self._config = config
        self._clock = clock
        self._factory = session_factory or self._default_session
        self._lock = threading.Lock()
        self._metadata: dict[str, Any] | None = None
        self._metadata_at = 0.0
        self._keys: KeySet | None = None
        self._keys_at = 0.0

    def _default_session(self) -> OAuth2Session:
        return OAuth2Session(
            client_id=self._config.client_id,
            client_secret=self._config.client_secret,
            scope=SCOPE,
            redirect_uri=self._config.redirect_uri,
            code_challenge_method="S256",
        )

    # -- provider documents ----------------------------------------------------------

    def _get_json(self, session: OAuth2Session, url: str) -> dict[str, Any]:
        try:
            response = session.get(url, timeout=TIMEOUT_SECONDS, withhold_token=True)
            response.raise_for_status()
            data = response.json()
        except (requests.RequestException, ValueError) as exc:
            logger.warning("Identity provider request failed: %s", type(exc).__name__)
            raise IdpUnavailable("The identity provider is not available") from exc
        if not isinstance(data, dict):
            raise IdpUnavailable("The identity provider sent an unusable document")
        return data

    def _discovery(self, session: OAuth2Session) -> dict[str, Any]:
        with self._lock:
            if self._metadata is not None and self._clock() - self._metadata_at < CACHE_SECONDS:
                return self._metadata
        url = self._config.issuer.rstrip("/") + "/.well-known/openid-configuration"
        data = self._get_json(session, url)
        if data.get("issuer") != self._config.issuer:
            raise IdpUnavailable("The identity provider reports a different issuer")
        for name in ("authorization_endpoint", "token_endpoint", "jwks_uri"):
            if not isinstance(data.get(name), str) or not data[name].startswith("https://"):
                raise IdpUnavailable(f"The identity provider document lacks a usable {name}")
        with self._lock:
            self._metadata, self._metadata_at = data, self._clock()
        return data

    def _keyset(self, session: OAuth2Session, jwks_uri: str, *, refresh: bool) -> KeySet:
        with self._lock:
            fresh = self._clock() - self._keys_at < CACHE_SECONDS
            if self._keys is not None and not refresh and fresh:
                return self._keys
        try:
            keys = KeySet.import_key_set(self._get_json(session, jwks_uri))
        except (JoseError, ValueError, KeyError) as exc:
            raise IdpUnavailable("The identity provider sent unusable keys") from exc
        with self._lock:
            self._keys, self._keys_at = keys, self._clock()
        return keys

    # -- interface ---------------------------------------------------------------------

    def begin(self) -> tuple[str, Flow]:
        session = self._factory()
        metadata = self._discovery(session)
        flow = Flow(
            state=secrets.token_urlsafe(32),
            nonce=secrets.token_urlsafe(32),
            code_verifier=secrets.token_urlsafe(64),
        )
        url, _ = session.create_authorization_url(
            metadata["authorization_endpoint"],
            state=flow.state,
            nonce=flow.nonce,
            code_verifier=flow.code_verifier,
        )
        return url, flow

    def complete(self, params: Mapping[str, str], flow: Flow) -> Identity:
        if "error" in params:
            raise LoginRejected("The identity provider refused the sign-in")
        state, code = params.get("state", ""), params.get("code", "")
        if not code or not hmac.compare_digest(state.encode(), flow.state.encode()):
            raise LoginRejected("The sign-in response does not match the request")

        session = self._factory()
        metadata = self._discovery(session)
        try:
            token = session.fetch_token(
                metadata["token_endpoint"],
                grant_type="authorization_code",
                code=code,
                code_verifier=flow.code_verifier,
                timeout=TIMEOUT_SECONDS,
            )
        except requests.RequestException as exc:
            logger.warning("Identity provider request failed: %s", type(exc).__name__)
            raise IdpUnavailable("The identity provider is not available") from exc
        except AuthlibBaseError as exc:
            raise LoginRejected("The identity provider refused the code") from exc
        id_token = token.get("id_token") if isinstance(token, dict) else None
        if not isinstance(id_token, str):
            raise LoginRejected("The identity provider sent no ID token")
        return self._identity(session, metadata["jwks_uri"], id_token, flow)

    def _identity(self, session: OAuth2Session, jwks_uri: str, id_token: str, flow: Flow):
        decoded = None
        for refresh in (False, True):  # the second try covers a key rotation
            keys = self._keyset(session, jwks_uri, refresh=refresh)
            try:
                decoded = jwt.decode(id_token, keys, algorithms=ALGORITHMS)
                break
            except InvalidKeyIdError as exc:
                if refresh:
                    raise LoginRejected("The ID token is not valid") from exc
            except JoseError as exc:
                raise LoginRejected("The ID token is not valid") from exc
        registry = jwt.JWTClaimsRegistry(
            leeway=CLOCK_LEEWAY_SECONDS,
            now=int(self._clock()),
            iss={"essential": True, "value": self._config.issuer},
            aud={"essential": True, "value": self._config.client_id},
            exp={"essential": True},
            sub={"essential": True},
            nonce={"essential": True, "value": flow.nonce},
        )
        try:
            registry.validate(decoded.claims)
        except JoseError as exc:
            raise LoginRejected("The ID token claims are not valid") from exc
        username = decoded.claims.get("preferred_username")
        if not isinstance(username, str) or not username.strip():
            raise LoginRejected("The ID token has no usable user name")
        return Identity(username=username.strip(), subject=str(decoded.claims["sub"]))