traceforge-pcb 1.0.0__py3-none-any.whl

traceforge/remote/oidc.py

@@ -0,0 +1,240 @@
+"""Single sign-on for the hosted server through OpenID Connect (authorization code + PKCE).
+
+Flow: ``start`` builds the IdP login URL (state, nonce and PKCE verifier are kept in
+memory for ten minutes); the IdP redirects back with a code; ``finish`` exchanges it at the
+token endpoint, validates the ID token with PyJWT against the IdP's JWKS, maps groups to a
+role and returns the claims. The caller turns that into a seat and an API key.
+
+ID token checks: asymmetric signature only (``none`` and ``HS*`` are refused, so a
+token cannot be forged with the public key or without a key), key chosen by ``kid`` from
+the JWKS (refetched once when unknown), ``iss`` equals the configured issuer (and the
+discovery document's issuer must match too), ``aud`` contains the client id (``azp`` must
+equal it when there are several), ``exp``/``iat``/``sub`` required with 60 s clock leeway,
+and ``nonce`` equals the one issued for that login.
+"""
+
+from __future__ import annotations
+
+import base64
+import hashlib
+import secrets
+import time
+from dataclasses import dataclass, field
+from typing import Any
+from urllib.parse import urlencode
+
+import httpx
+import jwt
+
+ALLOWED_ALGS = ("RS256", "RS384", "RS512", "ES256", "ES384", "PS256", "EdDSA")
+LEEWAY = 60
+LOGIN_TTL = 600
+
+
+class OIDCError(Exception):
+    def __init__(self, code: str, message: str):
+        super().__init__(message)
+        self.code = code
+        self.message = message
+
+
+@dataclass
+class OIDCConfig:
+    issuer: str
+    client_id: str
+    redirect_uri: str
+    org: str  # organisation that SSO users join
+    client_secret: str = ""  # from the environment, never from a file in a project
+    groups_claim: str = "groups"
+    group_roles: dict[str, str] = field(default_factory=dict)  # IdP group -> admin | member
+    default_role: str = "member"
+    allowed_domains: list[str] = field(default_factory=list)  # empty: any email domain
+    auto_provision: bool = True  # create a seat on first login (else invite first)
+    scopes: str = "openid email profile"
+    require_https: bool = True
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any], client_secret: str = "") -> OIDCConfig:
+        known = {k: v for k, v in data.items() if k in cls.__dataclass_fields__}
+        return cls(**{**known, "client_secret": client_secret or known.get("client_secret", "")})
+
+
+def role_for(config: OIDCConfig, claims: dict[str, Any]) -> str:
+    """admin when any group maps to admin, else the first mapped role, else the default."""
+    groups = claims.get(config.groups_claim) or []
+    if isinstance(groups, str):
+        groups = [groups]
+    roles = {config.group_roles[g] for g in groups if g in config.group_roles}
+    if "admin" in roles:
+        return "admin"
+    return "member" if roles else config.default_role
+
+
+class OIDCClient:
+    def __init__(self, config: OIDCConfig, http: httpx.Client | None = None):
+        self.config = config
+        self.http = http or httpx.Client(timeout=10, follow_redirects=False)
+        self._discovery: dict[str, Any] | None = None
+        self._keys: dict[str, Any] = {}
+        self._keys_fetched = 0.0
+        self._logins: dict[str, dict[str, str]] = {}
+
+    def _check_url(self, url: str) -> str:
+        if self.config.require_https and not url.startswith("https://"):
+            local = url.startswith(("http://127.0.0.1", "http://localhost"))
+            if not local:
+                raise OIDCError("insecure_idp", f"refusing non-HTTPS identity provider URL {url}")
+        return url
+
+    def discover(self) -> dict[str, Any]:
+        if self._discovery is None:
+            base = self._check_url(self.config.issuer.rstrip("/"))
+            try:
+                doc = (
+                    self.http.get(f"{base}/.well-known/openid-configuration")
+                    .raise_for_status()
+                    .json()
+                )
+            except (httpx.HTTPError, ValueError) as exc:
+                raise OIDCError("discovery_failed", f"OIDC discovery failed: {exc}") from exc
+            if str(doc.get("issuer", "")).rstrip("/") != self.config.issuer.rstrip("/"):
+                raise OIDCError(
+                    "issuer_mismatch",
+                    "discovery document issuer differs from the configured issuer",
+                )
+            for key in ("authorization_endpoint", "token_endpoint", "jwks_uri"):
+                if not doc.get(key):
+                    raise OIDCError("discovery_failed", f"discovery document lacks {key}")
+                self._check_url(doc[key])
+            self._discovery = doc
+        return self._discovery
+
+    def _jwks(self, force: bool = False) -> dict[str, Any]:
+        if force or not self._keys or time.time() - self._keys_fetched > 3600:
+            if force and time.time() - self._keys_fetched < 10:
+                return self._keys  # do not let a stream of unknown kids hammer the IdP
+            try:
+                data = self.http.get(self.discover()["jwks_uri"]).raise_for_status().json()
+            except (httpx.HTTPError, ValueError) as exc:
+                raise OIDCError(
+                    "jwks_failed", f"could not fetch the IdP signing keys: {exc}"
+                ) from exc
+            self._keys = {k["kid"]: k for k in data.get("keys", []) if k.get("kid")}
+            self._keys_fetched = time.time()
+        return self._keys
+
+    # ------------------------------------------------------------------ tokens
+
+    def validate_id_token(self, token: str, nonce: str | None) -> dict[str, Any]:
+        try:
+            header = jwt.get_unverified_header(token)
+        except jwt.PyJWTError as exc:
+            raise OIDCError("bad_token", f"malformed ID token: {exc}") from exc
+        alg = header.get("alg")
+        if alg not in ALLOWED_ALGS:
+            raise OIDCError("bad_algorithm", f"ID token algorithm {alg!r} is not accepted")
+        kid = header.get("kid")
+        jwk = self._jwks().get(kid) or self._jwks(force=True).get(kid)
+        if jwk is None:
+            raise OIDCError("unknown_key", "ID token is signed with a key the IdP does not publish")
+        try:
+            key = jwt.PyJWK(jwk).key
+            claims = jwt.decode(
+                token,
+                key,
+                algorithms=[alg],
+                audience=self.config.client_id,
+                issuer=self.config.issuer,
+                leeway=LEEWAY,
+                options={"require": ["exp", "iat", "iss", "aud", "sub"]},
+            )
+        except jwt.ExpiredSignatureError as exc:
+            raise OIDCError("expired", "ID token has expired") from exc
+        except jwt.InvalidAudienceError as exc:
+            raise OIDCError("bad_audience", "ID token was issued for a different client") from exc
+        except jwt.InvalidIssuerError as exc:
+            raise OIDCError("bad_issuer", "ID token has the wrong issuer") from exc
+        except jwt.InvalidSignatureError as exc:
+            raise OIDCError("bad_signature", "ID token signature does not verify") from exc
+        except jwt.PyJWTError as exc:
+            raise OIDCError("bad_token", f"invalid ID token: {exc}") from exc
+        aud = claims.get("aud")
+        if isinstance(aud, list) and len(aud) > 1 and claims.get("azp") != self.config.client_id:
+            raise OIDCError("bad_audience", "ID token azp does not match this client")
+        if nonce is not None and not secrets.compare_digest(str(claims.get("nonce", "")), nonce):
+            raise OIDCError("bad_nonce", "ID token nonce does not match this login")
+        return claims
+
+    # ------------------------------------------------------------------- login
+
+    def start(self) -> dict[str, str]:
+        doc = self.discover()
+        state, nonce = secrets.token_urlsafe(24), secrets.token_urlsafe(24)
+        verifier = secrets.token_urlsafe(48)
+        challenge = base64.urlsafe_b64encode(hashlib.sha256(verifier.encode()).digest()).rstrip(
+            b"="
+        )
+        now = time.time()
+        self._logins = {s: v for s, v in self._logins.items() if now - float(v["t"]) < LOGIN_TTL}
+        self._logins[state] = {"nonce": nonce, "verifier": verifier, "t": str(now)}
+        query = urlencode(
+            {
+                "response_type": "code",
+                "client_id": self.config.client_id,
+                "redirect_uri": self.config.redirect_uri,
+                "scope": self.config.scopes,
+                "state": state,
+                "nonce": nonce,
+                "code_challenge": challenge.decode(),
+                "code_challenge_method": "S256",
+            }
+        )
+        return {"url": f"{doc['authorization_endpoint']}?{query}", "state": state}
+
+    def finish(self, code: str, state: str) -> dict[str, Any]:
+        """Exchange the code, validate the ID token; returns the claims plus the mapped role."""
+        login = self._logins.pop(state, None)
+        if login is None or time.time() - float(login["t"]) > LOGIN_TTL:
+            raise OIDCError("bad_state", "unknown or expired login; start again")
+        data = {
+            "grant_type": "authorization_code",
+            "code": code,
+            "redirect_uri": self.config.redirect_uri,
+            "client_id": self.config.client_id,
+            "code_verifier": login["verifier"],
+        }
+        if self.config.client_secret:
+            data["client_secret"] = self.config.client_secret
+        try:
+            response = self.http.post(self.discover()["token_endpoint"], data=data)
+            response.raise_for_status()
+            token = response.json().get("id_token")
+        except (httpx.HTTPError, ValueError) as exc:
+            raise OIDCError("token_exchange_failed", f"token exchange failed: {exc}") from exc
+        if not token:
+            raise OIDCError("token_exchange_failed", "the IdP returned no id_token")
+        claims = self.validate_id_token(token, login["nonce"])
+        email = str(claims.get("email") or "").lower()
+        if not email:
+            raise OIDCError("no_email", "the ID token has no email claim (request the email scope)")
+        if claims.get("email_verified") is False:
+            raise OIDCError("email_unverified", "the IdP says this email address is not verified")
+        domains = [d.lower().lstrip("@") for d in self.config.allowed_domains]
+        if domains and email.rsplit("@", 1)[-1] not in domains:
+            raise OIDCError("domain_not_allowed", f"{email} is not in an allowed email domain")
+        return {**claims, "email": email, "role": role_for(self.config, claims)}
+
+
+def load_oidc_from_env() -> OIDCClient | None:
+    """``TRACEFORGE_OIDC_CONFIG`` (a YAML file) plus ``TRACEFORGE_OIDC_CLIENT_SECRET``."""
+    import os
+    from pathlib import Path
+
+    import yaml
+
+    path = os.environ.get("TRACEFORGE_OIDC_CONFIG")
+    if not path:
+        return None
+    data = yaml.safe_load(Path(path).read_text(encoding="utf-8")) or {}
+    secret = os.environ.get("TRACEFORGE_OIDC_CLIENT_SECRET", "")
+    return OIDCClient(OIDCConfig.from_dict(data, secret))