byoi 0.1.0a1__py3-none-any.whl

byoi/service.py

@@ -0,0 +1,723 @@
+"""Authentication service for the BYOI library.
+
+This is the main service that orchestrates the OAuth/OIDC authentication flow,
+including PKCE, token exchange, and identity linking.
+"""
+
+from __future__ import annotations
+
+import logging
+from datetime import datetime, timedelta, timezone
+from typing import TYPE_CHECKING, Any, Self
+from urllib.parse import urlencode
+
+import httpx
+
+from byoi.errors import (
+    CannotUnlinkLastIdentityError,
+    IdentityAlreadyLinkedError,
+    IdentityNotFoundError,
+    InvalidCodeError,
+    InvalidStateError,
+    StateExpiredError,
+    TokenExchangeError,
+    TokenRefreshError,
+    UserNotFoundError,
+)
+from byoi.models import (
+    AuthenticatedUser,
+    AuthorizationRequest,
+    AuthorizationResponse,
+    LinkIdentityRequest,
+    LinkedIdentityInfo,
+    TokenExchangeRequest,
+    TokenExchangeResponse,
+    TokenRefreshRequest,
+    TokenRefreshResponse,
+    UnlinkIdentityRequest,
+    UserIdentities,
+)
+from byoi.pkce import generate_nonce, generate_pkce_pair, generate_state
+from byoi.providers import ProviderManager
+from byoi.repositories import (
+    AuthStateRepositoryProtocol,
+    LinkedIdentityRepositoryProtocol,
+    UserRepositoryProtocol,
+)
+from byoi.tokens import TokenValidator
+
+if TYPE_CHECKING:
+    from types import TracebackType
+
+logger = logging.getLogger("byoi.service")
+
+__all__ = ("AuthService",)
+
+
+class AuthService:
+    """Main authentication service for BYOI.
+
+    This service provides the core authentication functionality:
+    - Initiating OAuth authorization flows with PKCE
+    - Exchanging authorization codes for tokens
+    - Validating ID tokens
+    - Managing identity linking
+
+    The implementing application is responsible for:
+    - Providing repository implementations for data persistence
+    - Creating FastAPI routes that use this service
+    - Managing user sessions after authentication
+    """
+
+    # Default expiration time for auth state
+    DEFAULT_STATE_EXPIRATION_MINUTES = 10
+
+    def __init__(
+        self,
+        provider_manager: ProviderManager,
+        user_repository: UserRepositoryProtocol,
+        identity_repository: LinkedIdentityRepositoryProtocol,
+        auth_state_repository: AuthStateRepositoryProtocol,
+        state_expiration_minutes: int = DEFAULT_STATE_EXPIRATION_MINUTES,
+        http_client: httpx.AsyncClient | None = None,
+        http_timeout_seconds: float = 30.0,
+    ) -> None:
+        """Initialize the auth service.
+
+        Args:
+            provider_manager: Manager for OIDC providers.
+            user_repository: Repository for user data.
+            identity_repository: Repository for linked identities.
+            auth_state_repository: Repository for auth state (PKCE, nonce).
+            state_expiration_minutes: How long auth states are valid.
+            http_client: Optional HTTP client for making requests.
+            http_timeout_seconds: Timeout for HTTP requests in seconds.
+        """
+        self._provider_manager = provider_manager
+        self._user_repository = user_repository
+        self._identity_repository = identity_repository
+        self._auth_state_repository = auth_state_repository
+        self._state_expiration_minutes = state_expiration_minutes
+        self._http_client = http_client
+        self._owns_http_client = http_client is None
+        self._http_timeout_seconds = http_timeout_seconds
+        self._token_validator = TokenValidator(provider_manager)
+
+    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 service."""
+        await self.close()
+    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
+
+    # =========================================================================
+    # Authorization Flow
+    # =========================================================================
+
+    async def create_authorization_url(
+        self,
+        request: AuthorizationRequest,
+    ) -> AuthorizationResponse:
+        """Create an authorization URL for initiating OAuth flow.
+
+        This generates PKCE challenge, state, and nonce, stores them,
+        and returns the URL to redirect the user to.
+
+        Args:
+            request: The authorization request details.
+
+        Returns:
+            Response containing the authorization URL and state.
+
+        Raises:
+            ProviderNotFoundError: If the provider is not registered.
+        """
+        # Get provider config and authorization endpoint
+        config = self._provider_manager.get_provider(request.provider_name)
+        auth_endpoint = await self._provider_manager.get_authorization_endpoint(
+            request.provider_name
+        )
+
+        # Generate PKCE, state, and nonce
+        pkce = generate_pkce_pair()
+        state = generate_state()
+        nonce = generate_nonce()
+
+        # Calculate expiration
+        now = datetime.now(timezone.utc)
+        expires_at = now + timedelta(minutes=self._state_expiration_minutes)
+
+        # Store auth state
+        await self._auth_state_repository.create(
+            state=state,
+            code_verifier=pkce.code_verifier,
+            nonce=nonce,
+            provider_name=request.provider_name,
+            redirect_uri=request.redirect_uri,
+            expires_at=expires_at,
+            client_type=request.client_type.value,
+            extra_data=request.extra_data,
+        )
+
+        # Build authorization URL parameters
+        params = {
+            "response_type": "code",
+            "client_id": config.client_id,
+            "redirect_uri": request.redirect_uri,
+            "scope": " ".join(config.scopes),
+            "state": state,
+            "nonce": nonce,
+            "code_challenge": pkce.code_challenge,
+            "code_challenge_method": pkce.code_challenge_method,
+        }
+
+        # Add any extra auth params from config
+        params.update(config.extra_auth_params)
+
+        # Build the full authorization URL
+        authorization_url = f"{auth_endpoint}?{urlencode(params)}"
+
+        return AuthorizationResponse(
+            authorization_url=authorization_url,
+            state=state,
+            provider_name=request.provider_name,
+            expires_at=expires_at,
+        )
+
+    async def exchange_code(
+        self,
+        request: TokenExchangeRequest,
+    ) -> TokenExchangeResponse:
+        """Exchange an authorization code for tokens.
+
+        This validates the state, performs PKCE verification, exchanges
+        the code for tokens, and validates the ID token.
+
+        Args:
+            request: The token exchange request.
+
+        Returns:
+            Response containing tokens and identity information.
+
+        Raises:
+            InvalidStateError: If the state is invalid or not found.
+            StateExpiredError: If the state has expired.
+            InvalidCodeError: If the authorization code is invalid.
+            TokenExchangeError: If the token exchange fails.
+            TokenValidationError: If ID token validation fails.
+        """
+        # Retrieve and validate auth state
+        auth_state = await self._auth_state_repository.get_by_state(request.state)
+        if auth_state is None:
+            raise InvalidStateError(request.state)
+
+        # Check expiration
+        now = datetime.now(timezone.utc)
+        expires_at = auth_state.expires_at
+        # Ensure expires_at is timezone-aware for comparison
+        if expires_at.tzinfo is None:
+            expires_at = expires_at.replace(tzinfo=timezone.utc)
+
+        if now > expires_at:
+            # Clean up expired state
+            await self._auth_state_repository.delete(request.state)
+            raise StateExpiredError(request.state)
+
+        # Validate redirect URI matches
+        if auth_state.redirect_uri != request.redirect_uri:
+            raise InvalidStateError(request.state)
+
+        # Get provider config and token endpoint
+        provider_name = auth_state.provider_name
+        config = self._provider_manager.get_provider(provider_name)
+        token_endpoint = await self._provider_manager.get_token_endpoint(provider_name)
+
+        # Exchange code for tokens
+        token_response = await self._exchange_code_for_tokens(
+            code=request.code,
+            redirect_uri=request.redirect_uri,
+            code_verifier=auth_state.code_verifier,
+            client_id=config.client_id,
+            client_secret=config.client_secret,
+            token_endpoint=token_endpoint,
+        )
+
+        # Validate ID token
+        id_token = token_response.get("id_token")
+        if not id_token:
+            raise TokenExchangeError("No ID token in response")
+
+        identity = await self._token_validator.validate(
+            id_token=id_token,
+            provider_name=provider_name,
+            expected_nonce=auth_state.nonce,
+            verify_at_hash=False,  # Optional verification
+        )
+
+        # Clean up the auth state (one-time use)
+        await self._auth_state_repository.delete(request.state)
+
+        return TokenExchangeResponse(
+            identity=identity,
+            access_token=token_response["access_token"],
+            token_type=token_response.get("token_type", "Bearer"),
+            expires_in=token_response.get("expires_in"),
+            refresh_token=token_response.get("refresh_token"),
+            id_token=id_token,
+        )
+
+    async def _exchange_code_for_tokens(
+        self,
+        code: str,
+        redirect_uri: str,
+        code_verifier: str,
+        client_id: str,
+        client_secret: str | None,
+        token_endpoint: str,
+    ) -> dict[str, Any]:
+        """Exchange an authorization code for tokens.
+
+        Args:
+            code: The authorization code.
+            redirect_uri: The redirect URI.
+            code_verifier: The PKCE code verifier.
+            client_id: The OAuth client ID.
+            client_secret: The OAuth client secret (optional).
+            token_endpoint: The token endpoint URL.
+
+        Returns:
+            The token response as a dictionary.
+
+        Raises:
+            InvalidCodeError: If the code is invalid.
+            TokenExchangeError: If the exchange fails.
+        """
+        # Build request body
+        data = {
+            "grant_type": "authorization_code",
+            "code": code,
+            "redirect_uri": redirect_uri,
+            "client_id": client_id,
+            "code_verifier": code_verifier,
+        }
+
+        # Add client secret if provided (confidential clients)
+        if client_secret:
+            data["client_secret"] = client_secret
+
+        try:
+            client = await self._get_http_client()
+            response = await client.post(
+                token_endpoint,
+                data=data,
+                headers={"Content-Type": "application/x-www-form-urlencoded"},
+            )
+
+            if response.status_code == 400:
+                error_data = response.json()
+                error = error_data.get("error", "unknown_error")
+                error_description = error_data.get("error_description", "")
+
+                if error == "invalid_grant":
+                    raise InvalidCodeError(error_description or "Invalid or expired code")
+
+                raise TokenExchangeError(
+                    f"{error}: {error_description}" if error_description else error
+                )
+
+            response.raise_for_status()
+            return response.json()
+
+        except (InvalidCodeError, TokenExchangeError):
+            raise
+        except httpx.HTTPError as e:
+            raise TokenExchangeError(f"HTTP error: {e}") from e
+        except Exception as e:
+            raise TokenExchangeError(f"Unexpected error: {e}") from e
+
+    async def refresh_token(
+        self,
+        request: TokenRefreshRequest,
+    ) -> TokenRefreshResponse:
+        """Refresh an access token using a refresh token.
+
+        This method exchanges a refresh token for a new access token.
+        Some providers may also return a new refresh token (token rotation).
+
+        Args:
+            request: The token refresh request containing the refresh token.
+
+        Returns:
+            Response containing the new access token and optional refresh token.
+
+        Raises:
+            ProviderNotFoundError: If the provider is not registered.
+            TokenRefreshError: If the token refresh fails.
+        """
+        logger.debug("Refreshing token for provider=%s", request.provider_name)
+
+        # Get provider config and token endpoint
+        config = self._provider_manager.get_provider(request.provider_name)
+        token_endpoint = await self._provider_manager.get_token_endpoint(
+            request.provider_name
+        )
+
+        # Build request body
+        data: dict[str, str] = {
+            "grant_type": "refresh_token",
+            "refresh_token": request.refresh_token,
+            "client_id": config.client_id,
+        }
+
+        # Add client secret if provided (confidential clients)
+        if config.client_secret:
+            data["client_secret"] = config.client_secret
+
+        # Add scope if specified
+        if request.scope:
+            data["scope"] = request.scope
+
+        try:
+            client = await self._get_http_client()
+            response = await client.post(
+                token_endpoint,
+                data=data,
+                headers={"Content-Type": "application/x-www-form-urlencoded"},
+            )
+
+            if response.status_code == 400:
+                error_data = response.json()
+                error = error_data.get("error", "unknown_error")
+                error_description = error_data.get("error_description", "")
+
+                raise TokenRefreshError(
+                    f"{error}: {error_description}" if error_description else error
+                )
+
+            response.raise_for_status()
+            token_data = response.json()
+
+            logger.debug(
+                "Token refreshed successfully for provider=%s",
+                request.provider_name,
+            )
+
+            return TokenRefreshResponse(
+                access_token=token_data["access_token"],
+                token_type=token_data.get("token_type", "Bearer"),
+                expires_in=token_data.get("expires_in"),
+                refresh_token=token_data.get("refresh_token"),
+                scope=token_data.get("scope"),
+                id_token=token_data.get("id_token"),
+            )
+
+        except TokenRefreshError:
+            raise
+        except httpx.HTTPError as e:
+            logger.error("HTTP error during token refresh: %s", e)
+            raise TokenRefreshError(f"HTTP error: {e}") from e
+        except Exception as e:
+            logger.error("Unexpected error during token refresh: %s", e)
+            raise TokenRefreshError(f"Unexpected error: {e}") from e
+
+    # =========================================================================
+    # Authentication (Login/Register)
+    # =========================================================================
+
+    async def authenticate(
+        self,
+        request: TokenExchangeRequest,
+        auto_create_user: bool = True,
+    ) -> AuthenticatedUser:
+        """Authenticate a user using an authorization code.
+
+        This is the main authentication method. It:
+        1. Exchanges the code for tokens
+        2. Looks up or creates the user
+        3. Links/updates the identity
+
+        Args:
+            request: The token exchange request.
+            auto_create_user: Whether to create a new user if not found.
+
+        Returns:
+            The authenticated user information.
+
+        Raises:
+            InvalidStateError: If the state is invalid.
+            StateExpiredError: If the state has expired.
+            TokenExchangeError: If token exchange fails.
+            TokenValidationError: If ID token validation fails.
+            UserNotFoundError: If user not found and auto_create_user is False.
+        """
+        logger.debug("Authenticating user with state=%s", request.state)
+
+        # Exchange code for tokens and identity
+        token_response = await self.exchange_code(request)
+        identity = token_response.identity
+
+        # Check if this identity is already linked
+        existing_link = await self._identity_repository.get_by_provider(
+            provider_name=identity.provider_name,
+            provider_subject=identity.subject,
+        )
+
+        if existing_link:
+            # User exists, update last used time
+            await self._identity_repository.update_last_used(
+                identity_id=existing_link.id,
+                last_used_at=datetime.now(timezone.utc),
+            )
+
+            # Get the user
+            user = await self._user_repository.get_by_id(existing_link.user_id)
+            if user is None:
+                # This shouldn't happen, but handle it gracefully
+                logger.error(
+                    "Linked identity references non-existent user: identity_id=%s, user_id=%s",
+                    existing_link.id,
+                    existing_link.user_id,
+                )
+                raise UserNotFoundError(existing_link.user_id)
+
+            logger.info(
+                "User authenticated: user_id=%s, provider=%s, is_new=False",
+                user.id,
+                identity.provider_name,
+            )
+
+            return AuthenticatedUser(
+                user_id=user.id,
+                email=user.email,
+                is_new_user=False,
+                identity=identity,
+                linked_identity_id=existing_link.id,
+            )
+
+        # Identity not linked - create new user or fail
+        if not auto_create_user:
+            logger.warning(
+                "User not found and auto_create_user=False: provider=%s, subject=%s",
+                identity.provider_name,
+                identity.subject,
+            )
+            raise UserNotFoundError(
+                f"No user linked to {identity.provider_name}:{identity.subject}"
+            )
+
+        # Create new user
+        logger.debug("Creating new user for provider=%s", identity.provider_name)
+        user = await self._user_repository.create(
+            email=identity.email if identity.email_verified else None,
+        )
+
+        # Link the identity to the new user
+        linked_identity = await self._identity_repository.create(
+            user_id=user.id,
+            provider_name=identity.provider_name,
+            provider_subject=identity.subject,
+            email=identity.email,
+        )
+
+        logger.info(
+            "New user created and authenticated: user_id=%s, provider=%s, is_new=True",
+            user.id,
+            identity.provider_name,
+        )
+
+        return AuthenticatedUser(
+            user_id=user.id,
+            email=user.email,
+            is_new_user=True,
+            identity=identity,
+            linked_identity_id=linked_identity.id,
+        )
+
+    # =========================================================================
+    # Identity Linking
+    # =========================================================================
+
+    async def link_identity(
+        self,
+        request: LinkIdentityRequest,
+    ) -> LinkedIdentityInfo:
+        """Link a new identity to an existing user.
+
+        This allows users to add additional login methods to their account.
+
+        Args:
+            request: The link identity request.
+
+        Returns:
+            Information about the linked identity.
+
+        Raises:
+            UserNotFoundError: If the user doesn't exist.
+            IdentityAlreadyLinkedError: If the identity is already linked.
+            InvalidStateError: If the state is invalid.
+            TokenExchangeError: If token exchange fails.
+        """
+        # Verify user exists
+        user = await self._user_repository.get_by_id(request.user_id)
+        if user is None:
+            raise UserNotFoundError(request.user_id)
+
+        # Exchange code for tokens
+        token_response = await self.exchange_code(
+            TokenExchangeRequest(
+                code=request.code,
+                state=request.state,
+                redirect_uri=request.redirect_uri,
+            )
+        )
+        identity = token_response.identity
+
+        # Check if this identity is already linked to any user
+        existing_link = await self._identity_repository.get_by_provider(
+            provider_name=identity.provider_name,
+            provider_subject=identity.subject,
+        )
+
+        if existing_link:
+            raise IdentityAlreadyLinkedError(
+                provider_name=identity.provider_name,
+                provider_subject=identity.subject,
+                existing_user_id=existing_link.user_id,
+            )
+
+        # Create the link
+        linked_identity = await self._identity_repository.create(
+            user_id=request.user_id,
+            provider_name=identity.provider_name,
+            provider_subject=identity.subject,
+            email=identity.email,
+        )
+
+        # Get provider display name
+        provider_info = self._provider_manager.get_provider_info(identity.provider_name)
+
+        return LinkedIdentityInfo(
+            id=linked_identity.id,
+            provider_name=linked_identity.provider_name,
+            provider_display_name=provider_info.display_name,
+            email=linked_identity.email,
+            created_at=linked_identity.created_at,
+            last_used_at=linked_identity.last_used_at,
+        )
+
+    async def unlink_identity(
+        self,
+        request: UnlinkIdentityRequest,
+    ) -> bool:
+        """Unlink an identity from a user.
+
+        Args:
+            request: The unlink identity request.
+
+        Returns:
+            True if the identity was unlinked.
+
+        Raises:
+            UserNotFoundError: If the user doesn't exist.
+            IdentityNotFoundError: If the identity doesn't exist.
+            CannotUnlinkLastIdentityError: If this is the user's only identity.
+        """
+        # Verify user exists
+        user = await self._user_repository.get_by_id(request.user_id)
+        if user is None:
+            raise UserNotFoundError(request.user_id)
+
+        # Get the identity
+        identity = await self._identity_repository.get_by_id(request.identity_id)
+        if identity is None:
+            raise IdentityNotFoundError(request.identity_id)
+
+        # Verify the identity belongs to this user
+        if identity.user_id != request.user_id:
+            raise IdentityNotFoundError(request.identity_id)
+
+        # Check if this is the user's only identity
+        user_identities = await self._identity_repository.get_by_user_id(request.user_id)
+        if len(user_identities) <= 1:
+            raise CannotUnlinkLastIdentityError(request.user_id)
+
+        # Delete the identity
+        return await self._identity_repository.delete(request.identity_id)
+
+    async def get_user_identities(self, user_id: str) -> UserIdentities:
+        """Get all linked identities for a user.
+
+        Args:
+            user_id: The user's ID.
+
+        Returns:
+            List of linked identities.
+
+        Raises:
+            UserNotFoundError: If the user doesn't exist.
+        """
+        # Verify user exists
+        user = await self._user_repository.get_by_id(user_id)
+        if user is None:
+            raise UserNotFoundError(user_id)
+
+        # Get all identities
+        identities = await self._identity_repository.get_by_user_id(user_id)
+
+        # Convert to LinkedIdentityInfo
+        identity_infos = []
+        for identity in identities:
+            try:
+                provider_info = self._provider_manager.get_provider_info(
+                    identity.provider_name
+                )
+                display_name = provider_info.display_name
+            except Exception:
+                display_name = identity.provider_name
+
+            identity_infos.append(
+                LinkedIdentityInfo(
+                    id=identity.id,
+                    provider_name=identity.provider_name,
+                    provider_display_name=display_name,
+                    email=identity.email,
+                    created_at=identity.created_at,
+                    last_used_at=identity.last_used_at,
+                )
+            )
+
+        return UserIdentities(user_id=user_id, identities=identity_infos)
+
+    # =========================================================================
+    # Cleanup
+    # =========================================================================
+
+    async def cleanup_expired_states(self) -> int:
+        """Clean up expired auth states.
+
+        Should be called periodically to remove stale state data.
+
+        Returns:
+            Number of expired states removed.
+        """
+        return await self._auth_state_repository.delete_expired()
+
+    async def close(self) -> None:
+        """Close the auth service and release resources."""
+        if self._owns_http_client and self._http_client is not None:
+            await self._http_client.aclose()
+            self._http_client = None