curu-memory 0.2.0__py3-none-any.whl

curu_memory/oidc.py

@@ -0,0 +1,437 @@
+"""The check of an OIDC access token on the HTTP transport (ADR-015, ADR-017).
+
+curu-memory is a resource server. It never issues a token. It reads the keys of one
+issuer, and it accepts a token only when each rule here passes. The key sets the
+algorithm, not the token header, and each HMAC algorithm and `none` are refused.
+"""
+
+from __future__ import annotations
+
+import base64
+import hashlib
+import threading
+import time
+from collections import OrderedDict
+from typing import Any
+from urllib.parse import urlsplit
+
+import anyio
+import httpx2
+import jwt
+from mcp.server.auth.middleware.auth_context import get_access_token
+from mcp.server.auth.provider import AccessToken
+
+from curu_memory.access import ANONYMOUS, AccessStore, Principal
+from curu_memory.audit import REFUSED, AuditLog
+
+#: The algorithms of a key that curu-memory accepts. No HMAC and no `none` (ADR-015).
+ALLOWED_ALGORITHMS = frozenset({"RS256", "RS384", "RS512", "PS256", "ES256", "EdDSA"})
+#: The clock skew in seconds for `exp`, `nbf` and `iat`.
+LEEWAY = 60
+#: The longest life of a token, from `iat` (or `nbf`) to `exp`, in seconds.
+MAX_LIFETIME = 3600
+#: The shortest time between two reads of the keys for an unknown key identifier.
+REFRESH_SECONDS = 60
+#: The shortest time between two reads after a failed read, in seconds.
+FAILED_READ_SECONDS = 10
+_LOOPBACK = frozenset({"127.0.0.1", "localhost", "::1"})
+#: The token type of Authentik. With it, a token must also have `azp` and `scope` (ADR-017).
+PLAIN_JWT = "JWT"
+
+
+class TokenRefused(Exception):
+    """The token does not pass a rule. `reason` is a short text with no part of the token."""
+
+    def __init__(self, reason: str) -> None:
+        super().__init__(reason)
+        self.reason = reason
+
+
+class KeyCache:
+    """The keys of one issuer, read from its discovery document and its JWKS."""
+
+    def __init__(self, issuer: str, *, timeout: float = 5.0) -> None:
+        self._issuer = issuer
+        self._timeout = timeout
+        self._keys: dict[str, jwt.PyJWK] = {}
+        self._jwks_uri: str | None = None
+        self._read_at: float | None = None
+        self._read_ok = False
+        self._discovery: dict[str, Any] | None = None
+        self._lock = threading.Lock()
+
+    def key(self, kid: str | None) -> jwt.PyJWK:
+        """The key of `kid`. An unknown `kid` reads the keys again, at most each 60 seconds.
+
+        A known `kid` never waits for a read. Only one request reads at a time. Another
+        request with an unknown `kid` does not wait for it: it is refused at once.
+        """
+        if not isinstance(kid, str) or not kid:
+            raise TokenRefused("unknown key")
+        found = self._keys.get(kid)
+        if found is not None:
+            return found
+        # Before the first read, a request waits, so that the first requests of a new
+        # server do not fail together. After it, a request never waits for a read.
+        if self._lock.acquire(blocking=self._read_at is None):
+            try:
+                if kid not in self._keys and self._may_read():
+                    self._read()
+            finally:
+                self._lock.release()
+        found = self._keys.get(kid)
+        if found is None:
+            raise TokenRefused("unknown key")
+        return found
+
+    def _may_read(self) -> bool:
+        if self._read_at is None:
+            return True
+        wait = REFRESH_SECONDS if self._read_ok else FAILED_READ_SECONDS
+        return time.monotonic() - self._read_at >= wait
+
+    def _read(self) -> None:
+        # The time of each read counts, also of a failed one, so a down issuer gets at
+        # most one read each FAILED_READ_SECONDS (review of spec 012).
+        self._read_at = time.monotonic()
+        self._read_ok = False
+        try:
+            if self._jwks_uri is None:
+                self._jwks_uri = _safe_url(str(self.discovery()["jwks_uri"]))
+            document = self._get(self._jwks_uri)
+        except (httpx2.HTTPError, KeyError, ValueError) as error:
+            raise TokenRefused("issuer not available") from error
+        keys: dict[str, jwt.PyJWK] = {}
+        for entry in document.get("keys", []):
+            kid = entry.get("kid")
+            if not isinstance(kid, str) or not kid or entry.get("use", "sig") != "sig":
+                continue
+            try:
+                keys[kid] = jwt.PyJWK(entry)
+            except jwt.PyJWKError:
+                continue
+        self._keys = keys
+        self._read_ok = True
+
+    def discovery(self) -> dict[str, Any]:
+        """The discovery document of the issuer, read one time.
+
+        Raise `ValueError` when the document names another issuer (OIDC Discovery 1.0,
+        section 4.3), or `httpx2.HTTPError` when the issuer does not answer.
+        """
+        if self._discovery is None:
+            document = self._get(f"{self._issuer.rstrip('/')}/.well-known/openid-configuration")
+            if document.get("issuer") != self._issuer:
+                raise ValueError("the discovery document names another issuer")
+            self._discovery = document
+        return self._discovery
+
+    def _get(self, url: str) -> dict[str, Any]:
+        response = httpx2.get(url, timeout=self._timeout, follow_redirects=False)
+        response.raise_for_status()
+        return response.json()
+
+
+def _safe_url(url: str) -> str:
+    """The URL of the keys. Plain http is only for loopback, as for the issuer (R-5)."""
+    parts = urlsplit(url)
+    if parts.scheme == "https" or (parts.scheme == "http" and parts.hostname in _LOOPBACK):
+        return url
+    raise ValueError("the keys must use https")
+
+
+def verify_token(
+    token: str,
+    *,
+    issuer: str,
+    audience: str,
+    token_types: tuple[str, ...],
+    keys: KeyCache,
+) -> dict[str, Any]:
+    """The claims of a token that passes each rule. Else raise `TokenRefused`."""
+    try:
+        header = jwt.get_unverified_header(token)
+    except jwt.InvalidTokenError:
+        raise TokenRefused("malformed token") from None
+    algorithm = str(header.get("alg", ""))
+    if algorithm.lower() == "none" or algorithm.upper().startswith("HS"):
+        raise TokenRefused("algorithm not allowed")
+    token_type = header.get("typ")
+    if token_type is not None and token_type not in token_types:
+        raise TokenRefused("wrong token type")
+    key = keys.key(header.get("kid"))
+    # The key sets the algorithm. A header that names another one does not pass.
+    if key.algorithm_name not in ALLOWED_ALGORITHMS or algorithm != key.algorithm_name:
+        raise TokenRefused("algorithm not allowed")
+    try:
+        claims = jwt.decode(
+            token,
+            key.key,
+            algorithms=[key.algorithm_name],
+            audience=audience,
+            issuer=issuer,
+            leeway=LEEWAY,
+            options={"require": ["exp", "iss", "aud"]},
+        )
+    except jwt.ExpiredSignatureError:
+        raise TokenRefused("expired") from None
+    except jwt.ImmatureSignatureError:
+        raise TokenRefused("not yet valid") from None
+    except jwt.InvalidAudienceError:
+        raise TokenRefused("wrong audience") from None
+    except jwt.InvalidIssuerError:
+        raise TokenRefused("wrong issuer") from None
+    except jwt.InvalidSignatureError:
+        raise TokenRefused("bad signature") from None
+    except jwt.MissingRequiredClaimError:
+        raise TokenRefused("missing claim") from None
+    except jwt.InvalidTokenError:
+        raise TokenRefused("malformed token") from None
+    start = claims.get("iat", claims.get("nbf"))
+    if isinstance(start, int | float) and claims["exp"] - start > MAX_LIFETIME:
+        raise TokenRefused("lifetime too long")
+    # An ID token of Authentik has the same `aud`, but no `azp` and `scope` (ADR-017).
+    # A token with no typ gets the same rule, so an issuer cannot skip it.
+    plain = token_type == PLAIN_JWT or (token_type is None and PLAIN_JWT in token_types)
+    if plain and not (claims.get("azp") and claims.get("scope")):
+        raise TokenRefused("missing claim")
+    return claims
+
+
+class PrincipalToken(AccessToken):
+    """The token of one request, with the handle of its person (spec 012).
+
+    The SDK keeps it for the request. It holds no part of the bearer token.
+    """
+
+    #: The handle of the person. For an agent, the person that it acts for.
+    handle: str
+    #: The name of the agent, or None for a person (spec 013).
+    agent: str | None = None
+
+
+class OidcVerifier:
+    """The token verifier of the HTTP transport in OIDC mode (ADR-015, ADR-017).
+
+    It finds the person of a valid token by its `iss` and `sub`. Each refused token writes
+    one audit line with the reason, and the SDK answers HTTP 401.
+    """
+
+    def __init__(
+        self,
+        *,
+        issuer: str,
+        audience: str,
+        token_types: tuple[str, ...],
+        access: AccessStore,
+        audit: AuditLog,
+    ) -> None:
+        self._issuer = issuer
+        self._audience = audience
+        self._token_types = token_types
+        self._access = access
+        self._audit = audit
+        self._keys = KeyCache(issuer)
+
+    def __repr__(self) -> str:
+        return f"OidcVerifier(issuer={self._issuer!r})"
+
+    @property
+    def keys(self) -> KeyCache:
+        """The key cache, which also holds the discovery document of the issuer."""
+        return self._keys
+
+    async def verify_token(self, token: str) -> AccessToken | None:
+        return await anyio.to_thread.run_sync(self._verify, token)
+
+    def _verify(self, token: str) -> PrincipalToken | None:
+        try:
+            claims = verify_token(
+                token,
+                issuer=self._issuer,
+                audience=self._audience,
+                token_types=self._token_types,
+                keys=self._keys,
+            )
+        except TokenRefused as refused:
+            self._refuse(refused.reason)
+            return None
+        subject = claims.get("sub")
+        client = claims.get("client_id") or claims.get("azp")
+        # ADR-017: a person by sub, an agent by sub, an agent by client identifier.
+        person = self._access.person_by_subject(self._issuer, str(subject)) if subject else None
+        by_subject = self._access.agent_by_subject(self._issuer, str(subject)) if subject else None
+        # ADR-018: only a token of the client itself (no sub, or sub is the client) is the
+        # agent of that client. Else each account that logs in through the client would be it.
+        own = client and (not subject or str(subject) == str(client))
+        by_client = self._access.agent_by_client(self._issuer, str(client)) if own else None
+        found = [p for p in (person, by_subject, by_client) if p is not None]
+        if not found:
+            self._refuse("unknown subject")
+            return None
+        if len(found) > 1:
+            self._refuse("more than one principal")
+            return None
+        agent = by_subject or by_client
+        if agent is not None:
+            handle = agent.person
+        else:
+            assert person is not None
+            handle = person.handle
+        scope = claims.get("scope", "")
+        return PrincipalToken(
+            token="oidc",
+            client_id=str(claims.get("client_id") or claims.get("azp") or ""),
+            scopes=scope.split() if isinstance(scope, str) else [],
+            expires_at=int(claims["exp"]),
+            resource=self._audience,
+            # The SDK binds a session to the client, the issuer and the subject.
+            subject=str(subject or client),
+            claims={"iss": self._issuer},
+            handle=handle,
+            agent=agent.name if agent is not None else None,
+        )
+
+    def _refuse(self, reason: str) -> None:
+        self._audit.write("anonymous", "token", None, REFUSED, reason)
+
+
+def principal_of_request() -> Principal:
+    """The principal of the current HTTP request: a person, an agent, or anonymous."""
+    token = get_access_token()
+    if not isinstance(token, PrincipalToken):
+        return ANONYMOUS
+    if token.agent is not None:
+        return Principal.for_agent(token.agent, acts_for=token.handle)
+    return Principal.person(token.handle)
+
+
+#: The most exchanged tokens that the cache keeps (research R-3 of spec 013).
+BASIC_CACHE_SIZE = 1000
+#: A token leaves the cache this many seconds before it expires (ADR-018).
+BASIC_MARGIN_SECONDS = 30
+#: A refused header is kept this many seconds, so a repeat sends no request to the issuer.
+BASIC_REFUSED_SECONDS = 10
+#: The most exchanges at the same time. The exchanges have a thread limit of their own,
+#: so a flood of Basic headers cannot take the threads of the token checks.
+BASIC_THREADS = 4
+
+
+class BasicExchange:
+    """Exchange `Authorization: Basic <user name>:<app password>` for a token (ADR-018).
+
+    curu-memory sends the credentials to the token endpoint of the issuer, with the
+    client credentials grant of Authentik. The token then gets each check of a bearer
+    token. The cache key is the SHA-256 of the full header, so curu-memory keeps no
+    password, and two callers never share a token.
+    """
+
+    def __init__(
+        self, *, client_id: str, keys: KeyCache, audit: AuditLog, timeout: float = 5.0
+    ) -> None:
+        self._client_id = client_id
+        self._keys = keys
+        self._audit = audit
+        self._timeout = timeout
+        # A value of None is a refused header.
+        self._cache: OrderedDict[str, tuple[str | None, float]] = OrderedDict()
+        self._lock = threading.Lock()
+
+    def __repr__(self) -> str:
+        return f"BasicExchange(client_id={self._client_id!r})"
+
+    def token_for(self, header: str) -> str | None:
+        """The access token for a Basic header, or None after one refused audit line."""
+        key = hashlib.sha256(header.encode()).hexdigest()
+        with self._lock:
+            cached = self._cache.get(key)
+            if cached is not None and cached[1] > time.monotonic():
+                self._cache.move_to_end(key)
+                return cached[0]
+        try:
+            user, password = _credentials(header)
+            token, lifetime = self._exchange(user, password)
+        except (httpx2.HTTPError, KeyError, ValueError, TypeError):
+            # The reason holds no user name and no password (FR-016).
+            self._audit.write("anonymous", "token", None, REFUSED, "basic exchange refused")
+            self._keep(key, None, BASIC_REFUSED_SECONDS)
+            return None
+        self._keep(
+            key, token, min(lifetime, _seconds_to_exp(token), MAX_LIFETIME) - BASIC_MARGIN_SECONDS
+        )
+        return token
+
+    def _keep(self, key: str, token: str | None, seconds: float) -> None:
+        if seconds <= 0:
+            return
+        with self._lock:
+            self._cache[key] = (token, time.monotonic() + seconds)
+            self._cache.move_to_end(key)
+            while len(self._cache) > BASIC_CACHE_SIZE:
+                self._cache.popitem(last=False)
+
+    def _exchange(self, user: str, password: str) -> tuple[str, float]:
+        endpoint = _safe_url(str(self._keys.discovery()["token_endpoint"]))
+        response = httpx2.post(
+            endpoint,
+            data={
+                "grant_type": "client_credentials",
+                "client_id": self._client_id,
+                "username": user,
+                "password": password,
+                "scope": "openid profile",
+            },
+            timeout=self._timeout,
+            follow_redirects=False,
+        )
+        response.raise_for_status()
+        body = response.json()
+        return str(body["access_token"]), float(body.get("expires_in", 300))
+
+
+def _seconds_to_exp(token: str) -> float:
+    """The seconds to the `exp` of a token, read with no check: only the cache uses it."""
+    try:
+        exp = jwt.decode(token, options={"verify_signature": False}).get("exp")
+        return float(str(exp)) - time.time()
+    except (jwt.PyJWTError, TypeError, ValueError):
+        return 0.0
+
+
+def _credentials(header: str) -> tuple[str, str]:
+    scheme, _, value = header.partition(" ")
+    if scheme.lower() != "basic":
+        raise ValueError("not a Basic header")
+    user, separator, password = (
+        base64.b64decode(value.strip(), validate=True).decode().partition(":")
+    )
+    if not separator or not user or not password:
+        raise ValueError("a Basic header needs a user name and a password")
+    return user, password
+
+
+class BasicToBearer:
+    """An ASGI middleware: it replaces a Basic header with the token of the exchange.
+
+    It runs before the authentication of the SDK, which reads only Bearer tokens. A
+    refused exchange keeps the Basic header, so the SDK answers HTTP 401.
+    """
+
+    def __init__(self, app: Any, exchange: BasicExchange) -> None:
+        self._app = app
+        self._exchange = exchange
+        self._threads = anyio.CapacityLimiter(BASIC_THREADS)
+
+    async def __call__(self, scope: Any, receive: Any, send: Any) -> None:
+        if scope.get("type") == "http":
+            headers = list(scope.get("headers", []))
+            for index, (name, value) in enumerate(headers):
+                if name == b"authorization" and value[:6].lower() == b"basic ":
+                    token = await anyio.to_thread.run_sync(
+                        self._exchange.token_for, value.decode("latin-1"), limiter=self._threads
+                    )
+                    if token is not None:
+                        headers[index] = (name, b"Bearer " + token.encode())
+                        scope = {**scope, "headers": headers}
+                    break
+        await self._app(scope, receive, send)