appkit-assistant 0.14.1__py3-none-any.whl → 0.15.1__py3-none-any.whl

appkit_assistant/backend/mcp_auth_service.py

@@ -0,0 +1,796 @@
+"""MCP OAuth Authentication Service.
+
+Handles OAuth 2.0 flows for MCP servers including:
+- Discovery of OAuth metadata via RFC 8414
+- Authorization URL construction with PKCE support
+- Token exchange (authorization code for tokens)
+- Token refresh
+- Token storage and retrieval
+"""
+
+import base64
+import hashlib
+import logging
+import secrets
+from dataclasses import dataclass
+from datetime import UTC, datetime, timedelta
+from http import HTTPStatus
+from urllib.parse import urlencode, urlparse
+
+import httpx
+from sqlmodel import Session, select
+
+from appkit_assistant.backend.models import (
+    AssistantMCPUserToken,
+    MCPAuthType,
+    MCPServer,
+)
+from appkit_user.authentication.backend.entities import OAuthStateEntity
+
+logger = logging.getLogger(__name__)
+
+# Discovery paths per RFC 8414
+WELL_KNOWN_PATHS = [
+    "/.well-known/oauth-authorization-server",
+    "/.well-known/openid-configuration",
+]
+
+
+def _generate_pkce_pair() -> tuple[str, str]:
+    """Generate PKCE code verifier and challenge (S256).
+
+    Returns:
+        Tuple of (code_verifier, code_challenge)
+    """
+    # Generate code verifier (43-128 characters)
+    code_verifier = (
+        base64.urlsafe_b64encode(secrets.token_bytes(32)).decode("utf-8").rstrip("=")
+    )
+
+    # Generate code challenge (SHA256 hash of verifier)
+    code_challenge = (
+        base64.urlsafe_b64encode(hashlib.sha256(code_verifier.encode("utf-8")).digest())
+        .decode("utf-8")
+        .rstrip("=")
+    )
+
+    return code_verifier, code_challenge
+
+
+@dataclass
+class OAuthDiscoveryResult:
+    """Result of OAuth metadata discovery."""
+
+    issuer: str | None = None
+    authorization_endpoint: str | None = None
+    token_endpoint: str | None = None
+    registration_endpoint: str | None = None
+    scopes_supported: list[str] | None = None
+    error: str | None = None
+
+
+@dataclass
+class ClientRegistrationResult:
+    """Result of OAuth Dynamic Client Registration (RFC 7591)."""
+
+    client_id: str | None = None
+    client_secret: str | None = None
+    client_id_issued_at: int | None = None
+    client_secret_expires_at: int | None = None
+    error: str | None = None
+    error_description: str | None = None
+
+
+@dataclass
+class TokenResult:
+    """Result of token exchange or refresh."""
+
+    access_token: str | None = None
+    refresh_token: str | None = None
+    expires_in: int | None = None
+    token_type: str | None = None
+    scope: str | None = None
+    error: str | None = None
+    error_description: str | None = None
+
+
+class MCPAuthService:
+    """Service for handling MCP OAuth authentication flows."""
+
+    def __init__(self, redirect_uri: str) -> None:
+        """Initialize the service.
+
+        Args:
+            redirect_uri: The OAuth redirect URI for the callback endpoint.
+        """
+        self.redirect_uri = redirect_uri
+        self._http_client: httpx.AsyncClient | None = None
+
+    async def _get_client(self) -> httpx.AsyncClient:
+        """Get or create the HTTP client."""
+        if self._http_client is None:
+            self._http_client = httpx.AsyncClient(timeout=30.0)
+        return self._http_client
+
+    async def close(self) -> None:
+        """Close the HTTP client."""
+        if self._http_client is not None:
+            await self._http_client.aclose()
+            self._http_client = None
+
+    async def discover_oauth_config(self, server_url: str) -> OAuthDiscoveryResult:
+        """Discover OAuth metadata from the server.
+
+        Attempts to fetch OAuth metadata from well-known endpoints per RFC 8414.
+
+        Args:
+            server_url: The MCP server URL to discover OAuth config from.
+
+        Returns:
+            OAuthDiscoveryResult with discovered endpoints or error.
+        """
+        client = await self._get_client()
+        parsed = urlparse(server_url)
+        base_url = f"{parsed.scheme}://{parsed.netloc}"
+
+        for path in WELL_KNOWN_PATHS:
+            discovery_url = f"{base_url}{path}"
+            logger.debug("Attempting OAuth discovery at: %s", discovery_url)
+
+            try:
+                response = await client.get(discovery_url)
+                if response.status_code == HTTPStatus.OK:
+                    data = response.json()
+                    return OAuthDiscoveryResult(
+                        issuer=data.get("issuer"),
+                        authorization_endpoint=data.get("authorization_endpoint"),
+                        token_endpoint=data.get("token_endpoint"),
+                        registration_endpoint=data.get("registration_endpoint"),
+                        scopes_supported=data.get("scopes_supported"),
+                    )
+            except httpx.RequestError as e:
+                logger.debug("Discovery failed at %s: %s", discovery_url, str(e))
+                continue
+            except Exception as e:
+                logger.warning("Unexpected error during discovery: %s", str(e))
+                continue
+
+        return OAuthDiscoveryResult(
+            error="OAuth discovery failed: No valid metadata found at well-known paths"
+        )
+
+    async def register_client(
+        self,
+        registration_endpoint: str,
+        client_name: str = "AppKit Assistant",
+        additional_redirect_uris: list[str] | None = None,
+    ) -> ClientRegistrationResult:
+        """Register a new OAuth client via Dynamic Client Registration (RFC 7591).
+
+        Some OAuth providers like Atlassian MCP require clients to register
+        dynamically before they can authenticate.
+
+        Args:
+            registration_endpoint: The OAuth registration endpoint URL.
+            client_name: The name to register the client with.
+            additional_redirect_uris: Additional redirect URIs to register.
+
+        Returns:
+            ClientRegistrationResult with client_id and optionally client_secret.
+        """
+        client = await self._get_client()
+
+        redirect_uris = [self.redirect_uri]
+        if additional_redirect_uris:
+            redirect_uris.extend(additional_redirect_uris)
+
+        # Client metadata per RFC 7591
+        client_metadata = {
+            "client_name": client_name,
+            "redirect_uris": redirect_uris,
+            "grant_types": ["authorization_code", "refresh_token"],
+            "response_types": ["code"],
+            "token_endpoint_auth_method": "none",  # Public client (no secret)
+        }
+
+        logger.debug(
+            "Registering OAuth client at %s with metadata: %s",
+            registration_endpoint,
+            client_metadata,
+        )
+
+        try:
+            response = await client.post(
+                registration_endpoint,
+                json=client_metadata,
+                headers={"Content-Type": "application/json"},
+            )
+
+            if response.status_code in (HTTPStatus.OK, HTTPStatus.CREATED):
+                data = response.json()
+                logger.debug(
+                    "Successfully registered OAuth client: %s",
+                    data.get("client_id"),
+                )
+                return ClientRegistrationResult(
+                    client_id=data.get("client_id"),
+                    client_secret=data.get("client_secret"),
+                    client_id_issued_at=data.get("client_id_issued_at"),
+                    client_secret_expires_at=data.get("client_secret_expires_at"),
+                )
+
+            # Handle error response
+            try:
+                error_data = response.json()
+                return ClientRegistrationResult(
+                    error=error_data.get("error", "registration_failed"),
+                    error_description=error_data.get(
+                        "error_description",
+                        f"HTTP {response.status_code}: {response.text}",
+                    ),
+                )
+            except Exception:
+                return ClientRegistrationResult(
+                    error="registration_failed",
+                    error_description=f"HTTP {response.status_code}: {response.text}",
+                )
+
+        except httpx.RequestError as e:
+            logger.error("Client registration request failed: %s", str(e))
+            return ClientRegistrationResult(
+                error="request_failed",
+                error_description=str(e),
+            )
+
+    async def build_authorization_url_with_registration(
+        self,
+        server: MCPServer,
+        state: str | None = None,
+        session: Session | None = None,
+        user_id: int | None = None,
+    ) -> tuple[str, str]:
+        """Build the OAuth authorization URL, registering client via DCR if needed.
+
+        This method will automatically discover OAuth config and perform Dynamic
+        Client Registration (RFC 7591) if the server has no client_id configured
+        and a registration_endpoint is available.
+
+        Args:
+            server: The MCP server configuration.
+            state: Optional state parameter. If not provided, a random one is generated.
+            session: DB Session for storing PKCE state and updating server config.
+            user_id: User ID for binding state.
+
+        Returns:
+            Tuple of (authorization_url, state) where state should be stored
+            for verification.
+        """
+        # Ensure we work with a session-attached object to avoid StateProxy issues
+        # and to allow persisting changes if DCR happens.
+        if session and server.id:
+            db_server = session.get(MCPServer, server.id)
+            if db_server:
+                server = db_server
+
+        # If no client_id, attempt Dynamic Client Registration
+        if not server.oauth_client_id:
+            logger.debug(
+                "No client_id for server %s, attempting DCR",
+                server.name,
+            )
+
+            # Discover OAuth config to find registration endpoint
+            discovery_result = await self.discover_oauth_config(server.url)
+
+            if discovery_result.error:
+                msg = f"OAuth discovery failed: {discovery_result.error}"
+                raise ValueError(msg)
+
+            if not discovery_result.registration_endpoint:
+                msg = "Server has no client ID and no registration endpoint available"
+                raise ValueError(msg)
+
+            # Perform Dynamic Client Registration
+            reg_result = await self.register_client(
+                registration_endpoint=discovery_result.registration_endpoint,
+                client_name="AppKit Assistant",
+            )
+
+            if reg_result.error:
+                msg = (
+                    f"Client registration failed: "
+                    f"{reg_result.error}: {reg_result.error_description}"
+                )
+                raise ValueError(msg)
+
+            if not reg_result.client_id:
+                msg = "Client registration succeeded but no client_id returned"
+                raise ValueError(msg)
+
+            # Update server with registered client_id
+            server.oauth_client_id = reg_result.client_id
+            if reg_result.client_secret:
+                server.oauth_client_secret = reg_result.client_secret
+
+            # Also update discovered endpoints if not already set
+            auth_endpoint = discovery_result.authorization_endpoint
+            if not server.oauth_authorize_url and auth_endpoint:
+                server.oauth_authorize_url = auth_endpoint
+            if not server.oauth_token_url and discovery_result.token_endpoint:
+                server.oauth_token_url = discovery_result.token_endpoint
+
+            # Persist updated server configuration
+            if session:
+                session.add(server)
+                try:
+                    session.commit()
+                    logger.debug(
+                        "Persisted DCR client_id %s for server %s",
+                        reg_result.client_id,
+                        server.name,
+                    )
+                except Exception as e:
+                    logger.error("Failed to persist DCR client_id: %s", e)
+                    session.rollback()
+
+        # Now delegate to the synchronous URL builder
+        return self.build_authorization_url(
+            server=server,
+            state=state,
+            session=session,
+            user_id=user_id,
+        )
+
+    def build_authorization_url(
+        self,
+        server: MCPServer,
+        state: str | None = None,
+        session: Session | None = None,
+        user_id: int | None = None,
+    ) -> tuple[str, str]:
+        """Build the OAuth authorization URL for user login.
+
+        Supports PKCE by generating code_verifier and storing it if a session is
+        provided.
+
+        Args:
+            server: The MCP server configuration.
+            state: Optional state parameter. If not provided, a random one is generated.
+            session: DB Session for storing PKCE state.
+            user_id: User ID for binding state.
+
+        Returns:
+            Tuple of (authorization_url, state) where state should be stored
+            for verification.
+        """
+        if not server.oauth_authorize_url:
+            msg = "Server has no authorization URL configured"
+            raise ValueError(msg)
+
+        if state is None:
+            state = secrets.token_urlsafe(32)
+            logger.info("Generated new OAuth state: %s", state)
+
+        if not server.oauth_client_id:
+            msg = "Server has no client ID configured"
+            raise ValueError(msg)
+
+        # Generate PKCE parameters (required by OAuth 2.1 / MCP servers)
+        code_verifier, code_challenge = _generate_pkce_pair()
+
+        params = {
+            "response_type": "code",
+            "redirect_uri": self.redirect_uri,
+            "state": state,
+            "client_id": server.oauth_client_id,
+            "code_challenge": code_challenge,
+            "code_challenge_method": "S256",
+        }
+
+        # Store state in DB for CSRF protection, server mapping, and PKCE verifier
+        if session:
+            provider_key = f"mcp:{server.id}" if server.id else "mcp:unknown"
+            logger.info(
+                "Saving OAuth state to DB: state=%s, provider=%s, user_id=%s",
+                state,
+                provider_key,
+                user_id,
+            )
+            oauth_state = OAuthStateEntity(
+                session_id="mcp_auth_flow",
+                state=state,
+                provider=provider_key,
+                code_verifier=code_verifier,
+                expires_at=datetime.now(UTC) + timedelta(minutes=10),
+                user_id=user_id,
+            )
+            session.add(oauth_state)
+            try:
+                session.commit()
+                logger.info("OAuth state committed successfully: %s", state)
+            except Exception as e:
+                logger.error("Failed to commit OAuth state: %s", e)
+                session.rollback()
+        else:
+            logger.warning(
+                "No DB session provided to build_authorization_url. "
+                "PKCE and state check will fail."
+            )
+
+        if server.oauth_scopes:
+            params["scope"] = server.oauth_scopes
+
+        auth_url = f"{server.oauth_authorize_url}?{urlencode(params)}"
+        logger.debug(
+            "build_authorization_url: base=%s, result=%s",
+            server.oauth_authorize_url,
+            auth_url,
+        )
+        return auth_url, state
+
+    async def exchange_code_for_tokens(  # noqa: PLR0911
+        self,
+        server: MCPServer,
+        code: str,
+        state: str | None = None,
+        session: Session | None = None,
+    ) -> TokenResult:
+        """Exchange authorization code for access and refresh tokens.
+
+        Args:
+            server: The MCP server configuration.
+            code: The authorization code received from the callback.
+            state: The state parameter from the callback (required for PKCE).
+            session: DB Session for retrieving PKCE verifier.
+
+        Returns:
+            TokenResult with tokens or error information.
+        """
+        if not server.oauth_token_url:
+            return TokenResult(error="no_token_url", error_description="No token URL")
+
+        client = await self._get_client()
+
+        data = {
+            "grant_type": "authorization_code",
+            "code": code,
+            "redirect_uri": self.redirect_uri,
+        }
+
+        if not server.oauth_client_id:
+            logger.error("Missing client_id for server %s", server.name)
+            return TokenResult(
+                error="config_missing",
+                error_description="Client ID missing in server configuration",
+            )
+
+        data["client_id"] = server.oauth_client_id
+
+        # Helper to get session from manager if not provided
+        # (Though arguments type hint says session is Optional,
+        # we expect it for state check)
+
+        # Check OAuth state (CSRF) and retrieve PKCE code_verifier
+        code_verifier: str | None = None
+        if state and session:
+            provider_key = f"mcp:{server.id}" if server.id else "mcp:unknown"
+            oauth_state = session.exec(
+                select(OAuthStateEntity).where(
+                    OAuthStateEntity.state == state,
+                    OAuthStateEntity.provider == provider_key,
+                )
+            ).first()
+
+            if oauth_state:
+                if oauth_state.expires_at < datetime.now(UTC):
+                    logger.warning("OAuth state expired for state %s", state)
+                    return TokenResult(
+                        error="invalid_grant",
+                        error_description="OAuth state expired or invalid",
+                    )
+
+                # Retrieve PKCE code_verifier before cleanup
+                code_verifier = oauth_state.code_verifier
+
+                # Clean up used state
+                session.delete(oauth_state)
+                session.commit()
+            else:
+                logger.warning("No OAuth state found for state %s.", state)
+
+        # Add PKCE code_verifier if available (required by providers like Atlassian)
+        if code_verifier:
+            data["code_verifier"] = code_verifier
+
+        if server.oauth_client_secret:
+            data["client_secret"] = server.oauth_client_secret
+
+        try:
+            response = await client.post(
+                server.oauth_token_url,
+                data=data,
+                headers={"Content-Type": "application/x-www-form-urlencoded"},
+            )
+
+            if response.status_code == HTTPStatus.OK:
+                token_data = response.json()
+                return TokenResult(
+                    access_token=token_data.get("access_token"),
+                    refresh_token=token_data.get("refresh_token"),
+                    expires_in=token_data.get("expires_in"),
+                    token_type=token_data.get("token_type", "Bearer"),
+                    scope=token_data.get("scope"),
+                )
+
+            # Handle error response
+            try:
+                error_data = response.json()
+                return TokenResult(
+                    error=error_data.get("error", "unknown_error"),
+                    error_description=error_data.get(
+                        "error_description",
+                        f"HTTP {response.status_code}",
+                    ),
+                )
+            except Exception:
+                return TokenResult(
+                    error="http_error",
+                    error_description=f"HTTP {response.status_code}",
+                )
+
+        except httpx.RequestError as e:
+            logger.error("Token exchange request failed: %s", str(e))
+            return TokenResult(
+                error="request_failed",
+                error_description=str(e),
+            )
+
+    async def refresh_access_token(
+        self,
+        server: MCPServer,
+        refresh_token: str,
+    ) -> TokenResult:
+        """Refresh an access token using a refresh token.
+
+        Args:
+            server: The MCP server configuration.
+            refresh_token: The refresh token to use.
+
+        Returns:
+            TokenResult with new tokens or error information.
+        """
+        if not server.oauth_token_url:
+            return TokenResult(error="no_token_url", error_description="No token URL")
+
+        if not server.oauth_client_id:
+            return TokenResult(error="no_client_id", error_description="No client ID")
+
+        client = await self._get_client()
+
+        data = {
+            "grant_type": "refresh_token",
+            "refresh_token": refresh_token,
+            "client_id": server.oauth_client_id,
+        }
+
+        if server.oauth_client_secret:
+            data["client_secret"] = server.oauth_client_secret
+
+        try:
+            response = await client.post(
+                server.oauth_token_url,
+                data=data,
+                headers={"Content-Type": "application/x-www-form-urlencoded"},
+            )
+
+            if response.status_code == HTTPStatus.OK:
+                token_data = response.json()
+                return TokenResult(
+                    access_token=token_data.get("access_token"),
+                    refresh_token=token_data.get("refresh_token", refresh_token),
+                    expires_in=token_data.get("expires_in"),
+                    token_type=token_data.get("token_type", "Bearer"),
+                    scope=token_data.get("scope"),
+                )
+
+            try:
+                error_data = response.json()
+                return TokenResult(
+                    error=error_data.get("error", "unknown_error"),
+                    error_description=error_data.get(
+                        "error_description",
+                        f"HTTP {response.status_code}",
+                    ),
+                )
+            except Exception:
+                return TokenResult(
+                    error="http_error",
+                    error_description=f"HTTP {response.status_code}",
+                )
+
+        except httpx.RequestError as e:
+            logger.error("Token refresh request failed: %s", str(e))
+            return TokenResult(
+                error="request_failed",
+                error_description=str(e),
+            )
+
+    # Database operations
+
+    def get_user_token(
+        self,
+        session: Session,
+        user_id: int,
+        mcp_server_id: int,
+    ) -> AssistantMCPUserToken | None:
+        """Get a user's token for an MCP server.
+
+        Args:
+            session: Database session.
+            user_id: The user's ID.
+            mcp_server_id: The MCP server's ID.
+
+        Returns:
+            The token record or None if not found.
+        """
+        statement = select(AssistantMCPUserToken).where(
+            AssistantMCPUserToken.user_id == user_id,
+            AssistantMCPUserToken.mcp_server_id == mcp_server_id,
+        )
+        return session.exec(statement).first()
+
+    def save_user_token(
+        self,
+        session: Session,
+        user_id: int,
+        mcp_server_id: int,
+        token_result: TokenResult,
+    ) -> AssistantMCPUserToken:
+        """Save or update a user's token for an MCP server.
+
+        Args:
+            session: Database session.
+            user_id: The user's ID.
+            mcp_server_id: The MCP server's ID.
+            token_result: The token data from exchange or refresh.
+
+        Returns:
+            The saved token record.
+        """
+        # Calculate expiry
+        expires_in = token_result.expires_in or 3600  # Default 1 hour
+        expires_at = datetime.now(UTC) + timedelta(seconds=expires_in)
+
+        # Check for existing token
+        existing = self.get_user_token(session, user_id, mcp_server_id)
+
+        if existing:
+            existing.access_token = token_result.access_token or ""
+            if token_result.refresh_token:
+                existing.refresh_token = token_result.refresh_token
+            existing.expires_at = expires_at
+            existing.updated_at = datetime.now(UTC)
+            session.add(existing)
+            session.commit()
+            session.refresh(existing)
+            return existing
+
+        # Create new token
+        new_token = AssistantMCPUserToken(
+            user_id=user_id,
+            mcp_server_id=mcp_server_id,
+            access_token=token_result.access_token or "",
+            refresh_token=token_result.refresh_token,
+            expires_at=expires_at,
+        )
+        session.add(new_token)
+        session.commit()
+        session.refresh(new_token)
+        return new_token
+
+    def delete_user_token(
+        self,
+        session: Session,
+        user_id: int,
+        mcp_server_id: int,
+    ) -> bool:
+        """Delete a user's token for an MCP server.
+
+        Args:
+            session: Database session.
+            user_id: The user's ID.
+            mcp_server_id: The MCP server's ID.
+
+        Returns:
+            True if a token was deleted, False otherwise.
+        """
+        token = self.get_user_token(session, user_id, mcp_server_id)
+        if token:
+            session.delete(token)
+            session.commit()
+            return True
+        return False
+
+    def is_token_valid(self, token: AssistantMCPUserToken) -> bool:
+        """Check if a token is still valid (not expired).
+
+        Args:
+            token: The token to check.
+
+        Returns:
+            True if the token is valid, False if expired.
+        """
+        # Add 30 second buffer for clock skew
+        return token.expires_at > datetime.now(UTC) + timedelta(seconds=30)
+
+    async def ensure_valid_token(
+        self,
+        session: Session,
+        server: MCPServer,
+        token: AssistantMCPUserToken,
+    ) -> AssistantMCPUserToken | None:
+        """Ensure a token is valid, refreshing if necessary.
+
+        Args:
+            session: Database session.
+            server: The MCP server configuration.
+            token: The token to validate/refresh.
+
+        Returns:
+            A valid token or None if refresh failed.
+        """
+        if self.is_token_valid(token):
+            return token
+
+        # Token expired, try to refresh
+        if not token.refresh_token:
+            logger.warning("Token expired and no refresh token available")
+            return None
+
+        logger.debug("Refreshing expired token for server %s", server.name)
+        result = await self.refresh_access_token(server, token.refresh_token)
+
+        if result.error:
+            logger.error(
+                "Token refresh failed: %s - %s",
+                result.error,
+                result.error_description,
+            )
+            return None
+
+        # Save the refreshed token
+        return self.save_user_token(
+            session,
+            token.user_id,
+            token.mcp_server_id,
+            result,
+        )
+
+    def update_server_oauth_config(
+        self,
+        session: Session,
+        server: MCPServer,
+        discovery_result: OAuthDiscoveryResult,
+    ) -> MCPServer:
+        """Update server with discovered OAuth configuration.
+
+        Args:
+            session: Database session.
+            server: The MCP server to update.
+            discovery_result: The discovery result with OAuth metadata.
+
+        Returns:
+            The updated server.
+        """
+        server.oauth_issuer = discovery_result.issuer
+        server.oauth_authorize_url = discovery_result.authorization_endpoint
+        server.oauth_token_url = discovery_result.token_endpoint
+        if discovery_result.scopes_supported:
+            server.oauth_scopes = " ".join(discovery_result.scopes_supported)
+        server.oauth_discovered_at = datetime.now(UTC)
+        server.auth_type = MCPAuthType.OAUTH_DISCOVERY
+
+        session.add(server)
+        session.commit()
+        session.refresh(server)
+        return server