byoi 0.1.0a1__py3-none-any.whl

byoi/providers.py

@@ -0,0 +1,434 @@
+"""OIDC Provider management and discovery.
+
+Handles OIDC discovery, JWKS fetching, and provider configuration.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Self
+from urllib.parse import urljoin
+
+import httpx
+
+from byoi.cache import CacheProtocol, NullCache
+from byoi.errors import JWKSFetchError, ProviderDiscoveryError, ProviderNotFoundError
+from byoi.models import OIDCProviderConfig, OIDCProviderInfo
+
+if TYPE_CHECKING:
+    from types import TracebackType
+
+logger = logging.getLogger("byoi.providers")
+
+__all__ = (
+    "OIDCDiscoveryDocument",
+    "ProviderManager",
+)
+
+
+@dataclass(slots=True)
+class OIDCDiscoveryDocument:
+    """OIDC Discovery document (OpenID Connect Discovery 1.0)."""
+
+    issuer: str
+    authorization_endpoint: str
+    token_endpoint: str
+    userinfo_endpoint: str | None
+    jwks_uri: str
+    scopes_supported: list[str]
+    response_types_supported: list[str]
+    subject_types_supported: list[str]
+    id_token_signing_alg_values_supported: list[str]
+    claims_supported: list[str]
+    raw: dict[str, Any]
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> OIDCDiscoveryDocument:
+        """Create a discovery document from a dictionary."""
+        return cls(
+            issuer=data["issuer"],
+            authorization_endpoint=data["authorization_endpoint"],
+            token_endpoint=data["token_endpoint"],
+            userinfo_endpoint=data.get("userinfo_endpoint"),
+            jwks_uri=data["jwks_uri"],
+            scopes_supported=data.get("scopes_supported", ["openid"]),
+            response_types_supported=data.get("response_types_supported", ["code"]),
+            subject_types_supported=data.get("subject_types_supported", ["public"]),
+            id_token_signing_alg_values_supported=data.get(
+                "id_token_signing_alg_values_supported", ["RS256"]
+            ),
+            claims_supported=data.get("claims_supported", []),
+            raw=data,
+        )
+
+
+class ProviderManager:
+    """Manages OIDC providers, discovery, and JWKS caching."""
+
+    # Cache key prefixes
+    DISCOVERY_CACHE_PREFIX = "discovery:"
+    JWKS_CACHE_PREFIX = "jwks:"
+
+    # Default TTLs
+    DEFAULT_DISCOVERY_TTL = 3600  # 1 hour
+    DEFAULT_JWKS_TTL = 3600  # 1 hour
+
+    def __init__(
+        self,
+        cache: CacheProtocol | None = None,
+        http_client: httpx.AsyncClient | None = None,
+        discovery_ttl_seconds: int = DEFAULT_DISCOVERY_TTL,
+        jwks_ttl_seconds: int = DEFAULT_JWKS_TTL,
+        http_timeout_seconds: float = 30.0,
+    ) -> None:
+        """Initialize the provider manager.
+
+        Args:
+            cache: Cache implementation for discovery docs and JWKS.
+            http_client: HTTP client for making requests.
+            discovery_ttl_seconds: TTL for cached discovery documents.
+            jwks_ttl_seconds: TTL for cached JWKS.
+            http_timeout_seconds: Timeout for HTTP requests in seconds.
+        """
+        self._providers: dict[str, OIDCProviderConfig] = {}
+        self._cache = cache or NullCache()
+        self._http_client = http_client
+        self._owns_http_client = http_client is None
+        self._discovery_ttl = discovery_ttl_seconds
+        self._jwks_ttl = jwks_ttl_seconds
+        self._http_timeout_seconds = http_timeout_seconds
+        # Locks for request deduplication to prevent concurrent fetches for the same provider
+        self._discovery_locks: dict[str, asyncio.Lock] = {}
+        self._jwks_locks: dict[str, asyncio.Lock] = {}
+        self._locks_lock = asyncio.Lock()
+
+    async def __aenter__(self) -> Self:
+        """Enter the async context manager."""
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
+        """Exit the async context manager and close the provider manager."""
+        await self.close()
+
+    async def _get_lock(self, locks_dict: dict[str, asyncio.Lock], key: str) -> asyncio.Lock:
+        """Get or create a lock for request deduplication."""
+        async with self._locks_lock:
+            if key not in locks_dict:
+                locks_dict[key] = asyncio.Lock()
+            return locks_dict[key]
+
+    async def _get_http_client(self) -> httpx.AsyncClient:
+        """Get or create the HTTP client."""
+        if self._http_client is None:
+            self._http_client = httpx.AsyncClient(
+                timeout=httpx.Timeout(self._http_timeout_seconds),
+                follow_redirects=True,
+            )
+        return self._http_client
+
+    def register_provider(self, config: OIDCProviderConfig) -> None:
+        """Register an OIDC provider.
+
+        Args:
+            config: The provider configuration.
+
+        Raises:
+            ValueError: If a provider with the same name is already registered.
+        """
+        if config.name in self._providers:
+            raise ValueError(f"Provider '{config.name}' is already registered")
+        self._providers[config.name] = config
+
+    def unregister_provider(self, name: str) -> bool:
+        """Unregister an OIDC provider.
+
+        Args:
+            name: The provider name.
+
+        Returns:
+            True if the provider was unregistered, False if not found.
+        """
+        if name in self._providers:
+            del self._providers[name]
+            return True
+        return False
+
+    def get_provider(self, name: str) -> OIDCProviderConfig:
+        """Get a provider configuration by name.
+
+        Args:
+            name: The provider name.
+
+        Returns:
+            The provider configuration.
+
+        Raises:
+            ProviderNotFoundError: If the provider is not registered.
+        """
+        if name not in self._providers:
+            raise ProviderNotFoundError(name)
+        return self._providers[name]
+
+    def get_provider_info(self, name: str) -> OIDCProviderInfo:
+        """Get public information about a provider.
+
+        Args:
+            name: The provider name.
+
+        Returns:
+            Public provider information.
+
+        Raises:
+            ProviderNotFoundError: If the provider is not registered.
+        """
+        config = self.get_provider(name)
+        return OIDCProviderInfo(
+            name=config.name,
+            display_name=config.display_name,
+            issuer=config.issuer,
+        )
+
+    def list_providers(self) -> list[OIDCProviderInfo]:
+        """List all registered providers.
+
+        Returns:
+            List of public provider information.
+        """
+        return [
+            OIDCProviderInfo(
+                name=config.name,
+                display_name=config.display_name,
+                issuer=config.issuer,
+            )
+            for config in self._providers.values()
+        ]
+
+    async def discover(self, provider_name: str) -> OIDCDiscoveryDocument:
+        """Discover OIDC endpoints for a provider.
+
+        Args:
+            provider_name: The provider name.
+
+        Returns:
+            The OIDC discovery document.
+
+        Raises:
+            ProviderNotFoundError: If the provider is not registered.
+            ProviderDiscoveryError: If discovery fails.
+        """
+        config = self.get_provider(provider_name)
+
+        # Check cache first (before acquiring lock for better performance)
+        cache_key = f"{self.DISCOVERY_CACHE_PREFIX}{provider_name}"
+        cached = await self._cache.get(cache_key)
+        if cached is not None:
+            logger.debug("Using cached discovery document for provider=%s", provider_name)
+            return OIDCDiscoveryDocument.from_dict(cached)
+
+        # Acquire per-provider lock to deduplicate concurrent requests
+        lock = await self._get_lock(self._discovery_locks, provider_name)
+        async with lock:
+            # Check cache again after acquiring lock (another request may have populated it)
+            cached = await self._cache.get(cache_key)
+            if cached is not None:
+                logger.debug("Using cached discovery document for provider=%s", provider_name)
+                return OIDCDiscoveryDocument.from_dict(cached)
+
+            # Fetch discovery document
+            discovery_url = urljoin(
+                config.issuer.rstrip("/") + "/",
+                ".well-known/openid-configuration",
+            )
+
+            logger.debug("Fetching discovery document from %s", discovery_url)
+
+            try:
+                client = await self._get_http_client()
+                response = await client.get(discovery_url)
+                response.raise_for_status()
+                data = response.json()
+            except httpx.HTTPError as e:
+                logger.error(
+                    "Failed to fetch discovery document for provider=%s: %s",
+                    provider_name,
+                    str(e),
+                )
+                raise ProviderDiscoveryError(provider_name, str(e)) from e
+            except Exception as e:
+                logger.error(
+                    "Unexpected error fetching discovery document for provider=%s: %s",
+                    provider_name,
+                    str(e),
+                )
+                raise ProviderDiscoveryError(provider_name, str(e)) from e
+
+            # Validate required fields
+            required_fields = [
+                "issuer",
+                "authorization_endpoint",
+                "token_endpoint",
+                "jwks_uri",
+            ]
+            for field in required_fields:
+                if field not in data:
+                    logger.error(
+                        "Discovery document missing required field '%s' for provider=%s",
+                        field,
+                        provider_name,
+                    )
+                    raise ProviderDiscoveryError(
+                        provider_name, f"Missing required field: {field}"
+                    )
+
+            # Cache the discovery document
+            await self._cache.set(cache_key, data, self._discovery_ttl)
+
+            logger.info("Discovery document fetched and cached for provider=%s", provider_name)
+
+            return OIDCDiscoveryDocument.from_dict(data)
+
+    async def get_jwks(self, provider_name: str, *, force_refresh: bool = False) -> dict[str, Any]:
+        """Get the JWKS for a provider.
+
+        Args:
+            provider_name: The provider name.
+            force_refresh: If True, bypass the cache and fetch fresh JWKS.
+                Useful for handling key rotation when a kid is not found.
+
+        Returns:
+            The JWKS as a dictionary.
+
+        Raises:
+            ProviderNotFoundError: If the provider is not registered.
+            JWKSFetchError: If fetching JWKS fails.
+        """
+        config = self.get_provider(provider_name)
+
+        cache_key = f"{self.JWKS_CACHE_PREFIX}{provider_name}"
+
+        # Check cache first (unless force_refresh is requested)
+        if not force_refresh:
+            cached = await self._cache.get(cache_key)
+            if cached is not None:
+                logger.debug("Using cached JWKS for provider=%s", provider_name)
+                return cached
+
+        # Acquire per-provider lock to deduplicate concurrent requests
+        lock = await self._get_lock(self._jwks_locks, provider_name)
+        async with lock:
+            # Check cache again after acquiring lock (unless force_refresh)
+            if not force_refresh:
+                cached = await self._cache.get(cache_key)
+                if cached is not None:
+                    logger.debug("Using cached JWKS for provider=%s", provider_name)
+                    return cached
+
+            # Get JWKS URI from config override or discovery
+            if config.jwks_uri_override:
+                jwks_uri = config.jwks_uri_override
+            else:
+                discovery = await self.discover(provider_name)
+                jwks_uri = discovery.jwks_uri
+
+            logger.debug("Fetching JWKS from %s for provider=%s", jwks_uri, provider_name)
+
+            # Fetch JWKS
+            try:
+                client = await self._get_http_client()
+                response = await client.get(jwks_uri)
+                response.raise_for_status()
+                jwks = response.json()
+            except httpx.HTTPError as e:
+                logger.error("Failed to fetch JWKS for provider=%s: %s", provider_name, str(e))
+                raise JWKSFetchError(provider_name, str(e)) from e
+            except Exception as e:
+                logger.error("Unexpected error fetching JWKS for provider=%s: %s", provider_name, str(e))
+                raise JWKSFetchError(provider_name, str(e)) from e
+
+            # Validate JWKS structure
+            if "keys" not in jwks:
+                logger.error("Invalid JWKS structure for provider=%s: missing 'keys' field", provider_name)
+                raise JWKSFetchError(provider_name, "Invalid JWKS: missing 'keys' field")
+
+            # Cache the JWKS
+            await self._cache.set(cache_key, jwks, self._jwks_ttl)
+
+            logger.info("JWKS fetched and cached for provider=%s", provider_name)
+
+            return jwks
+
+    async def get_authorization_endpoint(self, provider_name: str) -> str:
+        """Get the authorization endpoint for a provider.
+
+        Args:
+            provider_name: The provider name.
+
+        Returns:
+            The authorization endpoint URL.
+        """
+        config = self.get_provider(provider_name)
+        if config.authorization_endpoint_override:
+            return config.authorization_endpoint_override
+
+        discovery = await self.discover(provider_name)
+        return discovery.authorization_endpoint
+
+    async def get_token_endpoint(self, provider_name: str) -> str:
+        """Get the token endpoint for a provider.
+
+        Args:
+            provider_name: The provider name.
+
+        Returns:
+            The token endpoint URL.
+        """
+        config = self.get_provider(provider_name)
+        if config.token_endpoint_override:
+            return config.token_endpoint_override
+
+        discovery = await self.discover(provider_name)
+        return discovery.token_endpoint
+
+    async def get_userinfo_endpoint(self, provider_name: str) -> str | None:
+        """Get the userinfo endpoint for a provider.
+
+        Args:
+            provider_name: The provider name.
+
+        Returns:
+            The userinfo endpoint URL, or None if not available.
+        """
+        config = self.get_provider(provider_name)
+        if config.userinfo_endpoint_override:
+            return config.userinfo_endpoint_override
+
+        discovery = await self.discover(provider_name)
+        return discovery.userinfo_endpoint
+
+    async def invalidate_cache(self, provider_name: str | None = None) -> None:
+        """Invalidate cached data for a provider or all providers.
+
+        Args:
+            provider_name: The provider name, or None to invalidate all.
+        """
+        if provider_name:
+            await self._cache.delete(f"{self.DISCOVERY_CACHE_PREFIX}{provider_name}")
+            await self._cache.delete(f"{self.JWKS_CACHE_PREFIX}{provider_name}")
+        else:
+            # Invalidate all providers
+            for name in self._providers:
+                await self._cache.delete(f"{self.DISCOVERY_CACHE_PREFIX}{name}")
+                await self._cache.delete(f"{self.JWKS_CACHE_PREFIX}{name}")
+
+    async def close(self) -> None:
+        """Close the provider manager and release resources."""
+        if self._owns_http_client and self._http_client is not None:
+            await self._http_client.aclose()
+            self._http_client = None

byoi/repositories.py

@@ -0,0 +1,252 @@
+"""Repository protocols for the BYOI library.
+
+These protocols define the interfaces that implementing applications must satisfy
+for data persistence. The implementing application is responsible for providing
+concrete implementations backed by their database of choice.
+"""
+
+from datetime import datetime
+from typing import Any, Protocol, runtime_checkable
+
+from byoi.types import AuthStateProtocol, LinkedIdentityProtocol, UserProtocol
+
+__all__ = (
+    "UserRepositoryProtocol",
+    "LinkedIdentityRepositoryProtocol",
+    "AuthStateRepositoryProtocol",
+)
+
+
+@runtime_checkable
+class UserRepositoryProtocol(Protocol):
+    """Protocol for user data persistence.
+
+    Implementing applications must provide a repository that satisfies this protocol
+    for managing user entities.
+    """
+
+    async def get_by_id(self, user_id: str) -> UserProtocol | None:
+        """Retrieve a user by their unique identifier.
+
+        Args:
+            user_id: The unique identifier of the user.
+
+        Returns:
+            The user if found, None otherwise.
+        """
+        ...
+
+    async def get_by_email(self, email: str) -> UserProtocol | None:
+        """Retrieve a user by their email address.
+
+        Args:
+            email: The email address to search for.
+
+        Returns:
+            The user if found, None otherwise.
+        """
+        ...
+
+    async def create(
+        self,
+        email: str | None = None,
+        is_active: bool = True,
+        **extra_data: Any,
+    ) -> UserProtocol:
+        """Create a new user.
+
+        Args:
+            email: The user's email address.
+            is_active: Whether the user account is active.
+            **extra_data: Additional data to store with the user.
+
+        Returns:
+            The newly created user.
+        """
+        ...
+
+    async def update(self, user_id: str, **data: Any) -> UserProtocol | None:
+        """Update a user's data.
+
+        Args:
+            user_id: The unique identifier of the user to update.
+            **data: The data to update.
+
+        Returns:
+            The updated user if found, None otherwise.
+        """
+        ...
+
+
+@runtime_checkable
+class LinkedIdentityRepositoryProtocol(Protocol):
+    """Protocol for linked identity data persistence.
+
+    Implementing applications must provide a repository that satisfies this protocol
+    for managing linked identity entities.
+    """
+
+    async def get_by_id(self, identity_id: str) -> LinkedIdentityProtocol | None:
+        """Retrieve a linked identity by its unique identifier.
+
+        Args:
+            identity_id: The unique identifier of the linked identity.
+
+        Returns:
+            The linked identity if found, None otherwise.
+        """
+        ...
+
+    async def get_by_provider(
+        self,
+        provider_name: str,
+        provider_subject: str,
+    ) -> LinkedIdentityProtocol | None:
+        """Retrieve a linked identity by provider name and subject.
+
+        Args:
+            provider_name: The name of the identity provider.
+            provider_subject: The subject identifier from the provider.
+
+        Returns:
+            The linked identity if found, None otherwise.
+        """
+        ...
+
+    async def get_by_user_id(self, user_id: str) -> list[LinkedIdentityProtocol]:
+        """Retrieve all linked identities for a user.
+
+        Args:
+            user_id: The unique identifier of the user.
+
+        Returns:
+            A list of linked identities for the user.
+        """
+        ...
+
+    async def create(
+        self,
+        user_id: str,
+        provider_name: str,
+        provider_subject: str,
+        email: str | None = None,
+        **extra_data: Any,
+    ) -> LinkedIdentityProtocol:
+        """Create a new linked identity.
+
+        Args:
+            user_id: The ID of the user to link.
+            provider_name: The name of the identity provider.
+            provider_subject: The subject identifier from the provider.
+            email: The email associated with this identity.
+            **extra_data: Additional data to store with the identity.
+
+        Returns:
+            The newly created linked identity.
+        """
+        ...
+
+    async def update_last_used(
+        self,
+        identity_id: str,
+        last_used_at: datetime,
+    ) -> LinkedIdentityProtocol | None:
+        """Update the last used timestamp for a linked identity.
+
+        Args:
+            identity_id: The unique identifier of the linked identity.
+            last_used_at: The timestamp to set.
+
+        Returns:
+            The updated linked identity if found, None otherwise.
+        """
+        ...
+
+    async def delete(self, identity_id: str) -> bool:
+        """Delete a linked identity.
+
+        Args:
+            identity_id: The unique identifier of the linked identity to delete.
+
+        Returns:
+            True if the identity was deleted, False if not found.
+        """
+        ...
+
+    async def delete_by_user_id(self, user_id: str) -> int:
+        """Delete all linked identities for a user.
+
+        Args:
+            user_id: The unique identifier of the user.
+
+        Returns:
+            The number of identities deleted.
+        """
+        ...
+
+
+@runtime_checkable
+class AuthStateRepositoryProtocol(Protocol):
+    """Protocol for auth state data persistence.
+
+    Implementing applications must provide a repository that satisfies this protocol
+    for managing OAuth flow state (PKCE, nonce, etc.).
+    """
+
+    async def create(
+        self,
+        state: str,
+        code_verifier: str,
+        nonce: str,
+        provider_name: str,
+        redirect_uri: str,
+        expires_at: datetime,
+        client_type: str = "web",
+        extra_data: dict[str, Any] | None = None,
+    ) -> AuthStateProtocol:
+        """Create a new auth state.
+
+        Args:
+            state: The OAuth state parameter.
+            code_verifier: The PKCE code verifier.
+            nonce: The nonce for ID token validation.
+            provider_name: The name of the provider.
+            redirect_uri: The redirect URI for this auth flow.
+            expires_at: When this auth state expires.
+            client_type: Type of client (web, desktop, mobile).
+            extra_data: Additional data to store.
+
+        Returns:
+            The newly created auth state.
+        """
+        ...
+
+    async def get_by_state(self, state: str) -> AuthStateProtocol | None:
+        """Retrieve an auth state by its state parameter.
+
+        Args:
+            state: The OAuth state parameter.
+
+        Returns:
+            The auth state if found, None otherwise.
+        """
+        ...
+
+    async def delete(self, state: str) -> bool:
+        """Delete an auth state.
+
+        Args:
+            state: The OAuth state parameter.
+
+        Returns:
+            True if the state was deleted, False if not found.
+        """
+        ...
+
+    async def delete_expired(self) -> int:
+        """Delete all expired auth states.
+
+        Returns:
+            The number of states deleted.
+        """
+        ...