workspace-mcp 1.1.7__py3-none-any.whl → 1.1.9__py3-none-any.whl

auth/oauth21/handler.py

@@ -0,0 +1,440 @@
+"""
+OAuth 2.1 Handler
+
+Main OAuth 2.1 authentication handler that integrates all components.
+Provides a unified interface for OAuth 2.1 functionality.
+"""
+
+import logging
+from typing import Dict, Any, Optional, List, Tuple
+
+from .config import OAuth2Config
+from .discovery import AuthorizationServerDiscovery
+from .oauth2 import OAuth2AuthorizationFlow
+from .tokens import TokenValidator, TokenValidationError
+from .sessions import SessionStore, Session
+from .middleware import AuthenticationMiddleware, AuthContext
+from .http import HTTPAuthHandler
+
+logger = logging.getLogger(__name__)
+
+
+class OAuth2Handler:
+    """Main OAuth 2.1 authentication handler."""
+
+    def __init__(self, config: OAuth2Config):
+        """
+        Initialize OAuth 2.1 handler.
+
+        Args:
+            config: OAuth 2.1 configuration
+        """
+        self.config = config
+        
+        # Initialize components
+        self.discovery = AuthorizationServerDiscovery(
+            resource_url=config.resource_url,
+            cache_ttl=config.discovery_cache_ttl,
+        )
+        
+        self.flow_handler = OAuth2AuthorizationFlow(
+            client_id=config.client_id,
+            client_secret=config.client_secret,
+            discovery_service=self.discovery,
+        )
+        
+        self.token_validator = TokenValidator(
+            discovery_service=self.discovery,
+            cache_ttl=config.jwks_cache_ttl,
+        )
+        
+        self.session_store = SessionStore(
+            default_session_timeout=config.session_timeout,
+            max_sessions_per_user=config.max_sessions_per_user,
+            cleanup_interval=config.session_cleanup_interval,
+            enable_persistence=config.enable_session_persistence,
+            persistence_file=str(config.get_session_persistence_path()) if config.get_session_persistence_path() else None,
+        )
+        
+        self.http_auth = HTTPAuthHandler()
+        
+        # Setup debug logging if enabled
+        if config.enable_debug_logging:
+            logging.getLogger("auth.oauth21").setLevel(logging.DEBUG)
+
+    async def start(self):
+        """Start the OAuth 2.1 handler and background tasks."""
+        await self.session_store.start_cleanup_task()
+        logger.info("OAuth 2.1 handler started")
+
+    async def stop(self):
+        """Stop the OAuth 2.1 handler and clean up resources."""
+        await self.session_store.stop_cleanup_task()
+        await self.flow_handler.close()
+        await self.token_validator.close()
+        logger.info("OAuth 2.1 handler stopped")
+
+    async def create_authorization_url(
+        self,
+        redirect_uri: str,
+        scopes: List[str],
+        state: Optional[str] = None,
+        session_id: Optional[str] = None,
+        additional_params: Optional[Dict[str, str]] = None,
+    ) -> Tuple[str, str, str]:
+        """
+        Create OAuth 2.1 authorization URL.
+
+        Args:
+            redirect_uri: OAuth redirect URI
+            scopes: Requested scopes
+            state: State parameter (generated if not provided)
+            session_id: Optional session ID to associate
+            additional_params: Additional authorization parameters
+
+        Returns:
+            Tuple of (authorization_url, state, code_verifier)
+
+        Raises:
+            ValueError: If configuration is invalid
+        """
+        if not self.config.authorization_server_url:
+            raise ValueError("Authorization server URL not configured")
+
+        # Build authorization URL
+        auth_url, final_state, code_verifier = await self.flow_handler.build_authorization_url(
+            authorization_server_url=self.config.authorization_server_url,
+            redirect_uri=redirect_uri,
+            scopes=scopes,
+            state=state,
+            resource=self.config.resource_url,
+            additional_params=additional_params,
+        )
+
+        # Store session association if provided
+        if session_id:
+            self._store_authorization_state(final_state, session_id, code_verifier)
+
+        logger.info(f"Created authorization URL for scopes: {scopes}")
+        return auth_url, final_state, code_verifier
+
+    async def exchange_code_for_session(
+        self,
+        authorization_code: str,
+        code_verifier: str,
+        redirect_uri: str,
+        state: Optional[str] = None,
+    ) -> Tuple[str, Session]:
+        """
+        Exchange authorization code for session.
+
+        Args:
+            authorization_code: Authorization code from callback
+            code_verifier: PKCE code verifier
+            redirect_uri: OAuth redirect URI
+            state: State parameter from authorization
+
+        Returns:
+            Tuple of (session_id, session)
+
+        Raises:
+            ValueError: If code exchange fails
+            TokenValidationError: If token validation fails
+        """
+        # Exchange code for tokens
+        token_response = await self.flow_handler.exchange_code_for_token(
+            authorization_server_url=self.config.authorization_server_url,
+            authorization_code=authorization_code,
+            code_verifier=code_verifier,
+            redirect_uri=redirect_uri,
+            resource=self.config.resource_url,
+        )
+
+        # Validate the received token
+        access_token = token_response["access_token"]
+        token_info = await self.token_validator.validate_token(
+            token=access_token,
+            expected_audience=self.config.expected_audience,
+            required_scopes=self.config.required_scopes,
+            authorization_server_url=self.config.authorization_server_url,
+        )
+
+        if not token_info["valid"]:
+            raise TokenValidationError("Received token is invalid")
+
+        # Extract user identity
+        user_id = token_info["user_identity"]
+
+        # Create session
+        session_id = self.session_store.create_session(
+            user_id=user_id,
+            token_info={
+                **token_response,
+                "validation_info": token_info,
+                "claims": token_info.get("claims", {}),
+            },
+            scopes=token_info.get("scopes", []),
+            authorization_server=self.config.authorization_server_url,
+            client_id=self.config.client_id,
+            metadata={
+                "auth_method": "oauth2_authorization_code",
+                "token_type": token_info.get("token_type"),
+                "created_via": "code_exchange",
+            }
+        )
+
+        session = self.session_store.get_session(session_id)
+        logger.info(f"Created session {session_id} for user {user_id}")
+
+        return session_id, session
+
+    async def authenticate_bearer_token(
+        self,
+        token: str,
+        required_scopes: Optional[List[str]] = None,
+        create_session: bool = True,
+    ) -> AuthContext:
+        """
+        Authenticate Bearer token and optionally create session.
+
+        Args:
+            token: Bearer token to authenticate
+            required_scopes: Required scopes (uses config default if not provided)
+            create_session: Whether to create a session for valid tokens
+
+        Returns:
+            Authentication context
+
+        Raises:
+            TokenValidationError: If token validation fails
+        """
+        auth_context = AuthContext()
+
+        try:
+            # Validate token
+            scopes_to_check = required_scopes or self.config.required_scopes
+            token_info = await self.token_validator.validate_token(
+                token=token,
+                expected_audience=self.config.expected_audience,
+                required_scopes=scopes_to_check,
+                authorization_server_url=self.config.authorization_server_url,
+            )
+
+            if token_info["valid"]:
+                auth_context.authenticated = True
+                auth_context.user_id = token_info["user_identity"]
+                auth_context.token_info = token_info
+                auth_context.scopes = token_info.get("scopes", [])
+
+                # Create session if requested
+                if create_session:
+                    session_id = self.session_store.create_session(
+                        user_id=auth_context.user_id,
+                        token_info=token_info,
+                        scopes=auth_context.scopes,
+                        authorization_server=self.config.authorization_server_url,
+                        client_id=self.config.client_id,
+                        metadata={
+                            "auth_method": "bearer_token",
+                            "created_via": "token_passthrough",
+                        }
+                    )
+                    
+                    auth_context.session_id = session_id
+                    auth_context.session = self.session_store.get_session(session_id)
+
+                logger.debug(f"Authenticated Bearer token for user {auth_context.user_id}")
+            else:
+                auth_context.error = "invalid_token"
+                auth_context.error_description = "Token validation failed"
+
+        except TokenValidationError as e:
+            auth_context.error = e.error_code
+            auth_context.error_description = str(e)
+            logger.warning(f"Bearer token validation failed: {e}")
+
+        return auth_context
+
+    async def refresh_session_token(self, session_id: str) -> bool:
+        """
+        Refresh tokens for a session.
+
+        Args:
+            session_id: Session identifier
+
+        Returns:
+            True if refresh was successful
+
+        Raises:
+            ValueError: If session not found or refresh fails
+        """
+        session = self.session_store.get_session(session_id)
+        if not session:
+            raise ValueError(f"Session {session_id} not found")
+
+        refresh_token = session.token_info.get("refresh_token")
+        if not refresh_token:
+            raise ValueError("Session has no refresh token")
+
+        try:
+            # Refresh the token
+            token_response = await self.flow_handler.refresh_access_token(
+                authorization_server_url=self.config.authorization_server_url,
+                refresh_token=refresh_token,
+                scopes=session.scopes,
+                resource=self.config.resource_url,
+            )
+
+            # Update session with new tokens
+            updated_token_info = {**session.token_info, **token_response}
+            success = self.session_store.update_session(
+                session_id=session_id,
+                token_info=updated_token_info,
+                extend_expiration=True,
+            )
+
+            if success:
+                logger.info(f"Refreshed tokens for session {session_id}")
+            
+            return success
+
+        except Exception as e:
+            logger.error(f"Failed to refresh token for session {session_id}: {e}")
+            raise ValueError(f"Token refresh failed: {str(e)}")
+
+    def create_middleware(self) -> AuthenticationMiddleware:
+        """
+        Create authentication middleware.
+
+        Returns:
+            Configured authentication middleware
+        """
+        return AuthenticationMiddleware(
+            app=None,  # Will be set when middleware is added to app
+            session_store=self.session_store,
+            token_validator=self.token_validator,
+            discovery_service=self.discovery,
+            http_auth_handler=self.http_auth,
+            required_scopes=self.config.required_scopes,
+            exempt_paths=self.config.exempt_paths,
+            authorization_server_url=self.config.authorization_server_url,
+            expected_audience=self.config.expected_audience,
+            enable_bearer_passthrough=self.config.enable_bearer_passthrough,
+        )
+
+    def get_session_info(self, session_id: str) -> Optional[Dict[str, Any]]:
+        """
+        Get session information.
+
+        Args:
+            session_id: Session identifier
+
+        Returns:
+            Session information dictionary or None
+        """
+        session = self.session_store.get_session(session_id)
+        if not session:
+            return None
+
+        return {
+            "session_id": session.session_id,
+            "user_id": session.user_id,
+            "scopes": session.scopes,
+            "created_at": session.created_at.isoformat(),
+            "last_accessed": session.last_accessed.isoformat(),
+            "expires_at": session.expires_at.isoformat() if session.expires_at else None,
+            "authorization_server": session.authorization_server,
+            "metadata": session.metadata,
+            "has_refresh_token": "refresh_token" in session.token_info,
+        }
+
+    def get_user_sessions(self, user_id: str) -> List[Dict[str, Any]]:
+        """
+        Get all sessions for a user.
+
+        Args:
+            user_id: User identifier
+
+        Returns:
+            List of session information
+        """
+        sessions = self.session_store.get_user_sessions(user_id)
+        return [
+            {
+                "session_id": session.session_id,
+                "created_at": session.created_at.isoformat(),
+                "last_accessed": session.last_accessed.isoformat(),
+                "expires_at": session.expires_at.isoformat() if session.expires_at else None,
+                "scopes": session.scopes,
+                "metadata": session.metadata,
+            }
+            for session in sessions
+        ]
+
+    def revoke_session(self, session_id: str) -> bool:
+        """
+        Revoke a session.
+
+        Args:
+            session_id: Session identifier
+
+        Returns:
+            True if session was revoked
+        """
+        success = self.session_store.remove_session(session_id)
+        if success:
+            logger.info(f"Revoked session {session_id}")
+        return success
+
+    def revoke_user_sessions(self, user_id: str) -> int:
+        """
+        Revoke all sessions for a user.
+
+        Args:
+            user_id: User identifier
+
+        Returns:
+            Number of sessions revoked
+        """
+        count = self.session_store.remove_user_sessions(user_id)
+        logger.info(f"Revoked {count} sessions for user {user_id}")
+        return count
+
+    def get_handler_stats(self) -> Dict[str, Any]:
+        """
+        Get OAuth 2.1 handler statistics.
+
+        Returns:
+            Handler statistics
+        """
+        session_stats = self.session_store.get_session_stats()
+        
+        return {
+            "config": {
+                "authorization_server": self.config.authorization_server_url,
+                "client_id": self.config.client_id,
+                "session_timeout": self.config.session_timeout,
+                "bearer_passthrough": self.config.enable_bearer_passthrough,
+            },
+            "sessions": session_stats,
+            "components": {
+                "discovery_cache_size": len(self.discovery.cache),
+                "token_validation_cache_size": len(self.token_validator.validation_cache),
+                "jwks_cache_size": len(self.token_validator.jwks_cache),
+            }
+        }
+
+    def _store_authorization_state(self, state: str, session_id: str, code_verifier: str):
+        """Store authorization state for later retrieval."""
+        # This could be enhanced with a proper state store
+        # For now, we'll use session metadata
+        pass
+
+    async def __aenter__(self):
+        """Async context manager entry."""
+        await self.start()
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        """Async context manager exit."""
+        await self.stop()

auth/oauth21/http.py

@@ -0,0 +1,270 @@
+"""
+HTTP Authentication Handler
+
+Handles HTTP authentication headers and responses per RFC6750 (Bearer Token Usage)
+and OAuth 2.1 specifications.
+"""
+
+import logging
+import re
+from typing import Optional, Dict, Any
+
+logger = logging.getLogger(__name__)
+
+
+class HTTPAuthHandler:
+    """Handles HTTP authentication headers and responses."""
+
+    def __init__(self, resource_metadata_url: Optional[str] = None):
+        """
+        Initialize the HTTP authentication handler.
+
+        Args:
+            resource_metadata_url: URL for protected resource metadata discovery
+        """
+        self.resource_metadata_url = resource_metadata_url or "/.well-known/oauth-authorization-server"
+
+    def parse_authorization_header(self, header: str) -> Optional[str]:
+        """
+        Extract Bearer token from Authorization header per RFC6750.
+
+        Args:
+            header: Authorization header value
+
+        Returns:
+            Bearer token string or None if not found/invalid
+
+        Examples:
+            >>> handler = HTTPAuthHandler()
+            >>> handler.parse_authorization_header("Bearer abc123")
+            'abc123'
+            >>> handler.parse_authorization_header("Basic abc123")
+            None
+        """
+        if not header:
+            return None
+
+        # RFC6750 Section 2.1: Authorization Request Header Field
+        # Authorization: Bearer <token>
+        bearer_pattern = re.compile(r'^Bearer\s+([^\s]+)$', re.IGNORECASE)
+        match = bearer_pattern.match(header.strip())
+        
+        if match:
+            token = match.group(1)
+            # Basic validation - token should not be empty
+            if token:
+                logger.debug("Successfully extracted Bearer token from Authorization header")
+                return token
+            else:
+                logger.warning("Empty Bearer token in Authorization header")
+                return None
+        else:
+            logger.debug(f"Authorization header does not contain valid Bearer token: {header[:20]}...")
+            return None
+
+    def build_www_authenticate_header(
+        self,
+        realm: Optional[str] = None,
+        scope: Optional[str] = None,
+        error: Optional[str] = None,
+        error_description: Optional[str] = None,
+        error_uri: Optional[str] = None,
+    ) -> str:
+        """
+        Build WWW-Authenticate header for 401 responses per RFC6750.
+
+        Args:
+            realm: Authentication realm
+            scope: Required scope(s)
+            error: Error code (invalid_request, invalid_token, insufficient_scope)
+            error_description: Human-readable error description
+            error_uri: URI with error information
+
+        Returns:
+            WWW-Authenticate header value
+
+        Examples:
+            >>> handler = HTTPAuthHandler()
+            >>> handler.build_www_authenticate_header(realm="api")
+            'Bearer realm="api"'
+            >>> handler.build_www_authenticate_header(error="invalid_token")
+            'Bearer error="invalid_token"'
+        """
+        # Start with Bearer scheme
+        parts = ["Bearer"]
+        
+        # Add realm if provided
+        if realm:
+            parts.append(f'realm="{self._quote_attribute_value(realm)}"')
+        
+        # Add scope if provided
+        if scope:
+            parts.append(f'scope="{self._quote_attribute_value(scope)}"')
+        
+        # Add error information if provided
+        if error:
+            parts.append(f'error="{self._quote_attribute_value(error)}"')
+        
+        if error_description:
+            parts.append(f'error_description="{self._quote_attribute_value(error_description)}"')
+        
+        if error_uri:
+            parts.append(f'error_uri="{self._quote_attribute_value(error_uri)}"')
+
+        return " ".join(parts)
+
+    def build_resource_metadata_header(self) -> str:
+        """
+        Build WWW-Authenticate header with resource metadata URL for discovery.
+
+        Returns:
+            WWW-Authenticate header with AS_metadata_url parameter
+        """
+        return f'Bearer AS_metadata_url="{self.resource_metadata_url}"'
+
+    def _quote_attribute_value(self, value: str) -> str:
+        """
+        Quote attribute value for use in HTTP header per RFC7235.
+
+        Args:
+            value: Attribute value to quote
+
+        Returns:
+            Properly quoted value
+        """
+        # Escape quotes and backslashes
+        escaped = value.replace('\\', '\\\\').replace('"', '\\"')
+        return escaped
+
+    def extract_bearer_token_from_request(self, headers: Dict[str, str]) -> Optional[str]:
+        """
+        Extract Bearer token from HTTP request headers.
+
+        Args:
+            headers: HTTP request headers (case-insensitive dict)
+
+        Returns:
+            Bearer token or None
+        """
+        # Look for Authorization header (case-insensitive)
+        authorization = None
+        for key, value in headers.items():
+            if key.lower() == "authorization":
+                authorization = value
+                break
+        
+        if authorization:
+            return self.parse_authorization_header(authorization)
+        
+        return None
+
+    def is_bearer_token_request(self, headers: Dict[str, str]) -> bool:
+        """
+        Check if request contains Bearer token authentication.
+
+        Args:
+            headers: HTTP request headers
+
+        Returns:
+            True if request has Bearer token
+        """
+        token = self.extract_bearer_token_from_request(headers)
+        return token is not None
+
+    def build_error_response_headers(
+        self,
+        error: str,
+        error_description: Optional[str] = None,
+        realm: Optional[str] = None,
+        scope: Optional[str] = None,
+    ) -> Dict[str, str]:
+        """
+        Build complete error response headers for 401/403 responses.
+
+        Args:
+            error: OAuth error code
+            error_description: Human-readable error description
+            realm: Authentication realm
+            scope: Required scope
+
+        Returns:
+            Dictionary of response headers
+        """
+        headers = {
+            "WWW-Authenticate": self.build_www_authenticate_header(
+                realm=realm,
+                scope=scope,
+                error=error,
+                error_description=error_description,
+            ),
+            "Cache-Control": "no-store",
+            "Pragma": "no-cache",
+        }
+        
+        return headers
+
+    def validate_token_format(self, token: str) -> bool:
+        """
+        Validate Bearer token format per RFC6750.
+
+        Args:
+            token: Bearer token to validate
+
+        Returns:
+            True if token format is valid
+        """
+        if not token:
+            return False
+        
+        # RFC6750 - token should be ASCII and not contain certain characters
+        try:
+            # Check if token contains only valid characters
+            # Avoid control characters and certain special characters
+            invalid_chars = set('\x00\x01\x02\x03\x04\x05\x06\x07\x08\x09\x0a\x0b\x0c\x0d\x0e\x0f'
+                              '\x10\x11\x12\x13\x14\x15\x16\x17\x18\x19\x1a\x1b\x1c\x1d\x1e\x1f'
+                              '\x20\x7f"\\')
+            
+            if any(char in invalid_chars for char in token):
+                logger.warning("Bearer token contains invalid characters")
+                return False
+            
+            # Token should be ASCII
+            token.encode('ascii')
+            
+            return True
+            
+        except UnicodeEncodeError:
+            logger.warning("Bearer token contains non-ASCII characters")
+            return False
+
+    def get_token_info_from_headers(self, headers: Dict[str, str]) -> Dict[str, Any]:
+        """
+        Extract and validate token information from request headers.
+
+        Args:
+            headers: HTTP request headers
+
+        Returns:
+            Dictionary with token information
+        """
+        result = {
+            "has_bearer_token": False,
+            "token": None,
+            "valid_format": False,
+            "error": None,
+        }
+        
+        # Extract token
+        token = self.extract_bearer_token_from_request(headers)
+        
+        if token:
+            result["has_bearer_token"] = True
+            result["token"] = token
+            result["valid_format"] = self.validate_token_format(token)
+            
+            if not result["valid_format"]:
+                result["error"] = "invalid_token"
+        else:
+            result["error"] = "missing_token"
+        
+        return result