auth-ai 0.1.0__py3-none-any.whl

auth_ai/__init__.py

@@ -0,0 +1,127 @@
+"""Python SDK for applications relying on a CIRCLE auth-service identity provider."""
+
+from __future__ import annotations
+
+from importlib.metadata import PackageNotFoundError, version
+
+from .ai import AiClient, UserAi
+from .claims import (
+    AccessTokenClaims,
+    FeatureValue,
+    IdTokenClaims,
+    get_feature,
+    has_feature,
+    has_permission,
+    has_role,
+)
+from .client import AuthClient
+from .config import DEFAULT_SCOPES, WALLET_SELF_SCOPE, AuthClientConfig
+from .errors import (
+    ERR,
+    AuthAiError,
+    AuthAiErrorCode,
+    AuthAiErrorDetails,
+    is_auth_ai_error_code,
+)
+from .gateway import (
+    GatewayAgent,
+    GatewayAuth,
+    GatewayClient,
+    GatewayMcpServer,
+    GatewayModel,
+)
+from .gateway_errors import (
+    is_insufficient_credit,
+    is_insufficient_credit_status,
+    to_gateway_error,
+)
+from .grants import (
+    AccessToken,
+    GrantLock,
+    GrantManager,
+    GrantStore,
+    InProcessGrantLock,
+    MemoryGrantStore,
+)
+from .introspection import IntrospectionResult
+from .session_user import SessionUser, to_session_user
+from .types import (
+    LoginHandshake,
+    LoginRequest,
+    LoginResult,
+    Prompt,
+    TokenSet,
+    TokenTypeHint,
+)
+from .wallet import (
+    MICRO_PER_EUR,
+    TopupOptions,
+    WalletAccount,
+    WalletClient,
+    WalletSessionUsage,
+    WalletStats,
+    WalletTransaction,
+    WalletTransactionPage,
+    WalletTransactionType,
+    format_euros,
+)
+
+try:
+    __version__ = version("auth-ai")
+except PackageNotFoundError:  # pragma: no cover - running from a source tree
+    __version__ = "0.0.0"
+
+__all__ = [
+    "DEFAULT_SCOPES",
+    "ERR",
+    "MICRO_PER_EUR",
+    "WALLET_SELF_SCOPE",
+    "AccessToken",
+    "AccessTokenClaims",
+    "AiClient",
+    "AuthAiError",
+    "AuthAiErrorCode",
+    "AuthAiErrorDetails",
+    "AuthClient",
+    "AuthClientConfig",
+    "FeatureValue",
+    "GatewayAgent",
+    "GatewayAuth",
+    "GatewayClient",
+    "GatewayMcpServer",
+    "GatewayModel",
+    "GrantLock",
+    "GrantManager",
+    "GrantStore",
+    "IdTokenClaims",
+    "InProcessGrantLock",
+    "IntrospectionResult",
+    "LoginHandshake",
+    "LoginRequest",
+    "LoginResult",
+    "MemoryGrantStore",
+    "Prompt",
+    "SessionUser",
+    "TokenSet",
+    "TokenTypeHint",
+    "TopupOptions",
+    "UserAi",
+    "WalletAccount",
+    "WalletClient",
+    "WalletSessionUsage",
+    "WalletStats",
+    "WalletTransaction",
+    "WalletTransactionPage",
+    "WalletTransactionType",
+    "__version__",
+    "format_euros",
+    "get_feature",
+    "has_feature",
+    "has_permission",
+    "has_role",
+    "is_auth_ai_error_code",
+    "is_insufficient_credit",
+    "is_insufficient_credit_status",
+    "to_gateway_error",
+    "to_session_user",
+]

auth_ai/ai.py

@@ -0,0 +1,101 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Protocol
+
+from .errors import ERR
+from .gateway import GatewayClient
+from .http import JsonHttp
+from .wallet import WalletClient
+
+if TYPE_CHECKING:
+    import httpx
+
+    from .config import AuthClientConfig
+    from .grants import AccessToken
+
+
+class AiAuth(Protocol):
+    @property
+    def config(self) -> AuthClientConfig: ...
+
+    @property
+    def http(self) -> httpx.AsyncClient: ...
+
+
+class AiGrants(Protocol):
+    async def get_access_token(
+        self, grant_id: str, *, resource: str | None = None
+    ) -> AccessToken: ...
+
+    def invalidate(self, grant_id: str, resource: str | None = None) -> None: ...
+
+
+class _AudienceTokens:
+    """Mints the access token of one user for one audience, from the grant manager.
+
+    Serves both the wallet (a single audience, so the key is ignored) and the gateway.
+    """
+
+    def __init__(self, grants: AiGrants, grant_id: str, resource: str) -> None:
+        self._grants = grants
+        self._grant_id = grant_id
+        self._resource = resource
+
+    async def get_token(self, key: str | None = None) -> str:
+        token = await self._grants.get_access_token(self._grant_id, resource=self._resource)
+        return token.access_token
+
+    def invalidate(self, key: str | None = None) -> None:
+        self._grants.invalidate(self._grant_id, self._resource)
+
+
+class UserAi:
+    """What an application does for one signed-in user: read the wallet, call the gateway."""
+
+    def __init__(self, auth: AiAuth, grants: AiGrants, grant_id: str) -> None:
+        self._auth = auth
+        config = auth.config
+        self.wallet = WalletClient(
+            JsonHttp(
+                base_url=config.issuer,
+                http=auth.http,
+                timeout_seconds=config.timeout_seconds,
+                tokens=_AudienceTokens(grants, grant_id, config.issuer),
+            ),
+            config.issuer,
+        )
+        self._gateway_tokens = (
+            _AudienceTokens(grants, grant_id, config.gateway_url) if config.gateway_url else None
+        )
+        self._gateway: GatewayClient | None = None
+
+    @property
+    def gateway(self) -> GatewayClient:
+        """Needs `gateway_url` in the client configuration."""
+        config = self._auth.config
+        if not config.gateway_url or self._gateway_tokens is None:
+            raise ERR.CFG_001("The gateway requires `gateway_url` in the configuration")
+        if self._gateway is None:
+            self._gateway = GatewayClient(
+                gateway_url=config.gateway_url,
+                tokens=self._gateway_tokens,
+                http=self._auth.http,
+                timeout_seconds=config.timeout_seconds,
+            )
+        return self._gateway
+
+
+class AiClient:
+    """Wallet and gateway clients bound to a user's grant.
+
+    Each call mints (or reuses) a short-lived access token for the right audience:
+    the wallet is served by auth-service, the gateway by its own URL.
+    """
+
+    def __init__(self, *, auth: AiAuth, grants: AiGrants) -> None:
+        self._auth = auth
+        self._grants = grants
+
+    def for_user(self, grant_id: str) -> UserAi:
+        """`grant_id` is the key the user's refresh token was saved under (typically `sub`)."""
+        return UserAi(self._auth, self._grants, grant_id)

auth_ai/claims.py

@@ -0,0 +1,88 @@
+from __future__ import annotations
+
+from typing import Literal, Protocol
+
+from pydantic import BaseModel, ConfigDict
+
+type FeatureValue = str | int | float | bool
+type PermissionAction = Literal["read", "write"]
+
+
+class _Claims(BaseModel):
+    # Unknown claims are kept as is.
+    model_config = ConfigDict(extra="allow", frozen=True)
+
+    iss: str
+    sub: str
+    aud: str | list[str]
+    exp: int | float
+    iat: int | float
+    # Present only when the matching scope was granted.
+    roles: list[str] | None = None
+    permissions: list[str] | None = None
+    features: dict[str, FeatureValue] | None = None
+    plan: str | None = None
+    org_id: str | None = None
+    sid: str | None = None
+    azp: str | None = None
+
+
+class IdTokenClaims(_Claims):
+    """Claims of an ID token issued by auth-service."""
+
+    name: str | None = None
+    picture: str | None = None
+    given_name: str | None = None
+    family_name: str | None = None
+    company: str | None = None
+    updated_at: int | float | None = None
+    email: str | None = None
+    email_verified: bool | None = None
+    phone_number: str | None = None
+
+
+class AccessTokenClaims(_Claims):
+    """Claims of a JWT access token.
+
+    `azp` is the OAuth client that obtained the token and `scope` is a
+    space-separated list.
+    """
+
+    azp: str
+    scope: str | None = None
+
+
+class EntitlementClaims(Protocol):
+    @property
+    def roles(self) -> list[str] | None: ...
+    @property
+    def permissions(self) -> list[str] | None: ...
+    @property
+    def features(self) -> dict[str, FeatureValue] | None: ...
+
+
+def has_role(claims: EntitlementClaims, role: str) -> bool:
+    return role in (claims.roles or ())
+
+
+def has_permission(
+    claims: EntitlementClaims, resource: str, action: PermissionAction = "read"
+) -> bool:
+    """`resource` grants read access and `resource.write` grants write access,
+    which also implies read: auth-service keeps only the highest action."""
+    permissions = claims.permissions
+    if not permissions:
+        return False
+    if f"{resource}.write" in permissions:
+        return True
+    return action == "read" and resource in permissions
+
+
+def has_feature(claims: EntitlementClaims, key: str) -> bool:
+    """A feature is enabled when present and neither `false` nor `0` nor `""`."""
+    value = get_feature(claims, key)
+    return value is not None and value not in {"", 0}
+
+
+def get_feature(claims: EntitlementClaims, key: str) -> FeatureValue | None:
+    return (claims.features or {}).get(key)

auth_ai/client.py

@@ -0,0 +1,274 @@
+from __future__ import annotations
+
+import base64
+import hashlib
+import hmac
+import secrets
+from typing import TYPE_CHECKING, Any, Self
+from urllib.parse import parse_qs, urlencode, urlsplit, urlunsplit
+
+import httpx
+from pydantic import ValidationError
+
+from .claims import IdTokenClaims
+from .config import AuthClientConfig
+from .errors import ERR
+from .id_token import check_id_token_binding
+from .introspection import IntrospectionResult
+from .provider import Provider, ProviderLoader
+from .token_endpoint import ProviderEndpoints
+from .types import (
+    LoginHandshake,
+    LoginRequest,
+    LoginResult,
+    Prompt,
+    TokenSet,
+    TokenTypeHint,
+)
+
+if TYPE_CHECKING:
+    from collections.abc import Mapping, Sequence
+    from types import TracebackType
+
+_PKCE_VERIFIER_BYTES = 48
+_RANDOM_BYTES = 32
+
+
+def _s256_challenge(code_verifier: str) -> str:
+    digest = hashlib.sha256(code_verifier.encode("ascii")).digest()
+    return base64.urlsafe_b64encode(digest).rstrip(b"=").decode("ascii")
+
+
+class AuthClient:
+    """OIDC relying party for auth-service.
+
+    Authorization-code grant with PKCE, refresh with RFC 8707 resource narrowing,
+    logout, revocation and introspection. It keeps no session: callers persist the
+    handshake, the tokens and the claims where it suits them.
+
+    `http_client` replaces the default `httpx.AsyncClient`, e.g. to add a proxy, a
+    transport or an observability layer. A client given by the caller is never
+    closed by the SDK; otherwise call `aclose()` (or use `async with`).
+    """
+
+    def __init__(
+        self,
+        config: AuthClientConfig | Mapping[str, Any],
+        *,
+        http_client: httpx.AsyncClient | None = None,
+    ) -> None:
+        self.config = (
+            config
+            if isinstance(config, AuthClientConfig)
+            else AuthClientConfig.model_validate(config)
+        )
+        self._owns_http = http_client is None
+        #: The client used for every request, shared with the REST API clients built on this one.
+        self.http = http_client or httpx.AsyncClient(
+            timeout=self.config.timeout_seconds, follow_redirects=False
+        )
+        self._provider = ProviderLoader(self.config, self.http)
+        self._endpoints = ProviderEndpoints(self.config, self.http)
+
+    async def __aenter__(self) -> Self:
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> None:
+        await self.aclose()
+
+    async def aclose(self) -> None:
+        if self._owns_http:
+            await self.http.aclose()
+
+    async def start_login(
+        self, *, prompt: Prompt | None = None, login_hint: str | None = None
+    ) -> LoginRequest:
+        provider = await self._provider.get()
+        handshake = LoginHandshake(
+            state=secrets.token_urlsafe(_RANDOM_BYTES),
+            nonce=secrets.token_urlsafe(_RANDOM_BYTES),
+            code_verifier=secrets.token_urlsafe(_PKCE_VERIFIER_BYTES),
+        )
+        parameters: list[tuple[str, str]] = [
+            ("response_type", "code"),
+            ("client_id", self.config.client_id),
+            ("redirect_uri", self.config.redirect_uri),
+            ("scope", " ".join(self.config.scopes)),
+            ("state", handshake.state),
+            ("nonce", handshake.nonce),
+            ("code_challenge", _s256_challenge(handshake.code_verifier)),
+            ("code_challenge_method", "S256"),
+        ]
+        if prompt:
+            parameters.append(("prompt", prompt))
+        if login_hint:
+            parameters.append(("login_hint", login_hint))
+        parameters.extend(("resource", resource) for resource in self.config.resources)
+        url = _with_query(provider.metadata.authorization_endpoint, parameters)
+        return LoginRequest(url=url, handshake=handshake)
+
+    async def handle_callback(self, callback_url: str, handshake: LoginHandshake) -> LoginResult:
+        """Completes the sign-in from the callback URL.
+
+        `resource` is deliberately not sent on the code exchange: the token is then
+        issued for every audience requested at sign-in and each refresh narrows it
+        to one.
+        """
+        provider = await self._provider.get()
+        code = _authorization_code(callback_url, handshake, provider)
+        tokens = await self._endpoints.token_request(
+            provider.metadata.token_endpoint,
+            {
+                "grant_type": "authorization_code",
+                "code": code,
+                "redirect_uri": self.config.redirect_uri,
+                "code_verifier": handshake.code_verifier,
+            },
+        )
+        if tokens.id_token is None:
+            raise ERR.OIDC_002("The token response carries no ID token")
+        claims = await self._verify_id_token(provider, tokens.id_token, nonce=handshake.nonce)
+        return LoginResult(tokens=tokens, claims=claims, session_id=claims.sid)
+
+    async def verify_id_token(self, id_token: str) -> IdTokenClaims:
+        return await self._verify_id_token(await self._provider.get(), id_token)
+
+    async def refresh(self, refresh_token: str, *, resource: str | None = None) -> TokenSet:
+        """Exchanges a refresh token for fresh tokens.
+
+        The provider rotates the refresh token on every use: the returned one
+        replaces the presented one, and a replayed stale token revokes the whole
+        family. Concurrent callers should go through a grant manager. `scope` is
+        never sent: the rotated token would be left with the scopes of that
+        request only.
+        """
+        provider = await self._provider.get()
+        form = {"grant_type": "refresh_token", "refresh_token": refresh_token}
+        if resource:
+            form["resource"] = resource
+        return await self._endpoints.token_request(provider.metadata.token_endpoint, form)
+
+    async def logout_url(
+        self, *, id_token_hint: str | None = None, state: str | None = None
+    ) -> str:
+        """URL ending the provider session (RP-initiated logout)."""
+        endpoint = (await self._provider.get()).metadata.end_session_endpoint
+        if not endpoint:
+            raise ERR.OIDC_004("The provider advertises no end_session_endpoint")
+        parameters: list[tuple[str, str]] = [("client_id", self.config.client_id)]
+        if self.config.post_logout_redirect_uri:
+            parameters.append(("post_logout_redirect_uri", self.config.post_logout_redirect_uri))
+        if id_token_hint:
+            parameters.append(("id_token_hint", id_token_hint))
+        if state:
+            parameters.append(("state", state))
+        return _with_query(endpoint, parameters)
+
+    async def revoke(self, token: str, hint: TokenTypeHint | None = None) -> None:
+        endpoint = (await self._provider.get()).metadata.revocation_endpoint
+        if not endpoint:
+            raise ERR.OIDC_004("The provider advertises no revocation_endpoint")
+        response = await self._endpoints.post(endpoint, _token_form(token, hint))
+        if not response.is_success:
+            raise ERR.OIDC_003(
+                f"Revocation endpoint answered HTTP {response.status_code}",
+                details={"status": response.status_code},
+            )
+
+    async def introspect(
+        self, token: str, hint: TokenTypeHint | None = None
+    ) -> IntrospectionResult:
+        endpoint = (await self._provider.get()).metadata.introspection_endpoint
+        if not endpoint:
+            raise ERR.OIDC_004("The provider advertises no introspection_endpoint")
+        response = await self._endpoints.post(endpoint, _token_form(token, hint))
+        if not response.is_success:
+            raise ERR.OIDC_003(
+                f"Introspection endpoint answered HTTP {response.status_code}",
+                details={"status": response.status_code},
+            )
+        try:
+            return IntrospectionResult.model_validate(response.json())
+        except (ValueError, ValidationError) as err:
+            raise ERR.OIDC_003(
+                "The introspection response has an unexpected shape", cause=err
+            ) from err
+
+    async def verify_jwt(
+        self, token: str, *, audience: str, typ: str | None = None
+    ) -> dict[str, Any]:
+        """Verifies the signature, issuer, audience and expiry of a JWT signed by the provider.
+
+        Low-level: prefer the bearer and logout token verifiers. Returns the claims.
+        """
+        provider = await self._provider.get()
+        return (await provider.jwt.verify(token, audience=audience, typ=typ)).claims
+
+    async def client_credentials(
+        self, *, scope: str | None = None, resource: str | None = None
+    ) -> TokenSet:
+        """`client_credentials` grant for the application itself.
+
+        No user is bound to the token. Needs a confidential client.
+        """
+        if self.config.client_secret is None:
+            raise ERR.CFG_001("client_credentials requires a client_secret")
+        provider = await self._provider.get()
+        form = {"grant_type": "client_credentials"}
+        if scope:
+            form["scope"] = scope
+        if resource:
+            form["resource"] = resource
+        return await self._endpoints.token_request(provider.metadata.token_endpoint, form)
+
+    async def _verify_id_token(
+        self, provider: Provider, id_token: str, *, nonce: str | None = None
+    ) -> IdTokenClaims:
+        verified = await provider.jwt.verify(id_token, audience=self.config.client_id)
+        check_id_token_binding(verified.claims, client_id=self.config.client_id, nonce=nonce)
+        try:
+            return IdTokenClaims.model_validate(verified.claims)
+        except ValidationError as err:
+            raise ERR.TOK_001("The ID token claims have an unexpected shape", cause=err) from err
+
+
+def _with_query(url: str, parameters: Sequence[tuple[str, str]]) -> str:
+    """`url` with `parameters` appended to its own query, repeated names kept."""
+    parts = urlsplit(url)
+    query = "&".join(part for part in (parts.query, urlencode(parameters)) if part)
+    return urlunsplit(parts._replace(query=query))
+
+
+def _token_form(token: str, hint: TokenTypeHint | None) -> dict[str, str]:
+    return {"token": token, "token_type_hint": hint} if hint else {"token": token}
+
+
+def _authorization_code(callback_url: str, handshake: LoginHandshake, provider: Provider) -> str:
+    """Checks the authorization response (RFC 6749 §4.1.2, RFC 9207) and returns its code."""
+    query = parse_qs(urlsplit(callback_url).query, keep_blank_values=True)
+    if any(len(values) != 1 for values in query.values()):
+        raise ERR.OIDC_002(details={"reason": "repeated_parameter"})
+    params = {name: values[0] for name, values in query.items()}
+
+    error = params.get("error")
+    if error is not None:
+        raise ERR.OIDC_002(f"Authorization failed: {error}", details={"oauthError": error})
+
+    issuer = params.get("iss")
+    if issuer is None and provider.metadata.authorization_response_iss_parameter_supported:
+        raise ERR.OIDC_002(details={"reason": "missing_iss"})
+    if issuer is not None and issuer != provider.metadata.issuer:
+        raise ERR.OIDC_002(details={"reason": "iss_mismatch"})
+
+    state = params.get("state")
+    if state is None or not hmac.compare_digest(state.encode(), handshake.state.encode()):
+        raise ERR.OIDC_002(details={"reason": "state_mismatch"})
+    code = params.get("code")
+    if not code:
+        raise ERR.OIDC_002(details={"reason": "missing_code"})
+    return code