awslabs.openapi-mcp-server 0.1.1__py3-none-any.whl

awslabs/openapi_mcp_server/auth/auth_cache.py

@@ -0,0 +1,190 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Authentication caching utilities.
+
+This module provides caching mechanisms for authentication tokens and other data.
+"""
+
+import time
+from typing import Any, Callable, Dict, Optional, TypeVar, cast
+
+
+# Type variable for cached function return types
+T = TypeVar('T')
+
+
+class TokenCache:
+    """Cache for authentication tokens and related data.
+
+    This class provides a simple time-based cache for authentication tokens
+    and other authentication-related data.
+    """
+
+    def __init__(self, max_size: int = 100, ttl: int = 300):
+        """Initialize the token cache.
+
+        Args:
+            max_size: Maximum number of items to store in the cache
+            ttl: Time-to-live in seconds for cached items
+
+        """
+        self._cache: Dict[str, Dict[str, Any]] = {}
+        self._max_size = max_size
+        self._ttl = ttl
+
+    def get(self, key: str) -> Optional[Any]:
+        """Get a value from the cache.
+
+        Args:
+            key: Cache key
+
+        Returns:
+            Any: Cached value or None if not found or expired
+
+        """
+        if key not in self._cache:
+            return None
+
+        item = self._cache[key]
+        if time.time() > item['expires_at']:
+            # Item has expired
+            del self._cache[key]
+            return None
+
+        return item['value']
+
+    def set(self, key: str, value: Any, ttl: Optional[int] = None) -> None:
+        """Set a value in the cache.
+
+        Args:
+            key: Cache key
+            value: Value to cache
+            ttl: Time-to-live in seconds (overrides default)
+
+        """
+        # Ensure we don't exceed max size
+        if len(self._cache) >= self._max_size and key not in self._cache:
+            # Remove oldest item (simple LRU implementation)
+            oldest_key = min(self._cache.items(), key=lambda x: x[1]['expires_at'])[0]
+            del self._cache[oldest_key]
+
+        # Calculate expiration time
+        expires_at = time.time() + (ttl if ttl is not None else self._ttl)
+
+        # Store the item
+        self._cache[key] = {
+            'value': value,
+            'expires_at': expires_at,
+        }
+
+    def delete(self, key: str) -> bool:
+        """Delete a value from the cache.
+
+        Args:
+            key: Cache key
+
+        Returns:
+            bool: True if the key was found and deleted, False otherwise
+
+        """
+        if key in self._cache:
+            del self._cache[key]
+            return True
+        return False
+
+    def clear(self) -> None:
+        """Clear the entire cache."""
+        self._cache.clear()
+
+    def cleanup(self) -> int:
+        """Remove expired items from the cache.
+
+        Returns:
+            int: Number of items removed
+
+        """
+        now = time.time()
+        expired_keys = [k for k, v in self._cache.items() if now > v['expires_at']]
+        for key in expired_keys:
+            del self._cache[key]
+        return len(expired_keys)
+
+
+# Global token cache instance
+_TOKEN_CACHE = TokenCache()
+
+
+def get_token_cache() -> TokenCache:
+    """Get the global token cache instance.
+
+    Returns:
+        TokenCache: Global token cache instance
+
+    """
+    return _TOKEN_CACHE
+
+
+def cached_auth_data(ttl: int = 300) -> Callable[[Callable[..., T]], Callable[..., T]]:
+    """Cache authentication data.
+
+    Args:
+        ttl: Time-to-live in seconds for cached items
+
+    Returns:
+        Callable: Decorator function
+
+    """
+
+    def decorator(func: Callable[..., T]) -> Callable[..., T]:
+        """Decorate function with caching.
+
+        Args:
+            func: Function to decorate
+
+        Returns:
+            Callable: Wrapped function
+
+        """
+        cache = get_token_cache()
+
+        def wrapper(*args: Any, **kwargs: Any) -> T:
+            """Wrap function with caching logic.
+
+            Args:
+                *args: Positional arguments
+                **kwargs: Keyword arguments
+
+            Returns:
+                T: Function result
+
+            """
+            # Create a cache key from the function name and arguments
+            key_parts = [func.__name__]
+            key_parts.extend(str(arg) for arg in args)
+            key_parts.extend(f'{k}={v}' for k, v in sorted(kwargs.items()))
+            cache_key = ':'.join(key_parts)
+
+            # Check if we have a cached result
+            cached_result = cache.get(cache_key)
+            if cached_result is not None:
+                return cast(T, cached_result)
+
+            # Call the function and cache the result
+            result = func(*args, **kwargs)
+            cache.set(cache_key, result, ttl)
+            return result
+
+        return wrapper
+
+    return decorator

awslabs/openapi_mcp_server/auth/auth_errors.py

@@ -0,0 +1,206 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Authentication error handling.
+
+This module provides centralized error handling for authentication providers.
+"""
+
+from enum import Enum
+from typing import Dict, Optional, Type
+
+
+class AuthErrorType(Enum):
+    """Authentication error types."""
+
+    MISSING_CREDENTIALS = 'missing_credentials'
+    INVALID_CREDENTIALS = 'invalid_credentials'
+    EXPIRED_TOKEN = 'expired_token'  # nosec B105
+    INSUFFICIENT_PERMISSIONS = 'insufficient_permissions'
+    CONFIGURATION_ERROR = 'configuration_error'
+    NETWORK_ERROR = 'network_error'
+    UNKNOWN_ERROR = 'unknown_error'
+
+
+class AuthError(Exception):
+    """Base class for authentication errors."""
+
+    def __init__(
+        self,
+        message: str,
+        error_type: AuthErrorType = AuthErrorType.UNKNOWN_ERROR,
+        details: Optional[Dict] = None,
+    ):
+        """Initialize the error.
+
+        Args:
+            message: Error message
+            error_type: Type of authentication error
+            details: Additional error details
+
+        """
+        self.message = message
+        self.error_type = error_type
+        self.details = details or {}
+        super().__init__(message)
+
+    def __str__(self) -> str:
+        """Get string representation of the error.
+
+        Returns:
+            str: Error message with type
+
+        """
+        return f'{self.error_type.value}: {self.message}'
+
+
+class MissingCredentialsError(AuthError):
+    """Error raised when required credentials are missing."""
+
+    def __init__(self, message: str, details: Optional[Dict] = None):
+        """Initialize the error.
+
+        Args:
+            message: Error message
+            details: Additional error details
+
+        """
+        super().__init__(
+            message=message, error_type=AuthErrorType.MISSING_CREDENTIALS, details=details
+        )
+
+
+class InvalidCredentialsError(AuthError):
+    """Error raised when credentials are invalid."""
+
+    def __init__(self, message: str, details: Optional[Dict] = None):
+        """Initialize the error.
+
+        Args:
+            message: Error message
+            details: Additional error details
+
+        """
+        super().__init__(
+            message=message, error_type=AuthErrorType.INVALID_CREDENTIALS, details=details
+        )
+
+
+class ExpiredTokenError(AuthError):
+    """Error raised when a token has expired."""
+
+    def __init__(self, message: str, details: Optional[Dict] = None):
+        """Initialize the error.
+
+        Args:
+            message: Error message
+            details: Additional error details
+
+        """
+        super().__init__(message=message, error_type=AuthErrorType.EXPIRED_TOKEN, details=details)
+
+
+class InsufficientPermissionsError(AuthError):
+    """Error raised when permissions are insufficient."""
+
+    def __init__(self, message: str, details: Optional[Dict] = None):
+        """Initialize the error.
+
+        Args:
+            message: Error message
+            details: Additional error details
+
+        """
+        super().__init__(
+            message=message, error_type=AuthErrorType.INSUFFICIENT_PERMISSIONS, details=details
+        )
+
+
+class ConfigurationError(AuthError):
+    """Error raised when there is a configuration issue."""
+
+    def __init__(self, message: str, details: Optional[Dict] = None):
+        """Initialize the error.
+
+        Args:
+            message: Error message
+            details: Additional error details
+
+        """
+        super().__init__(
+            message=message, error_type=AuthErrorType.CONFIGURATION_ERROR, details=details
+        )
+
+
+class NetworkError(AuthError):
+    """Error raised when there is a network issue."""
+
+    def __init__(self, message: str, details: Optional[Dict] = None):
+        """Initialize the error.
+
+        Args:
+            message: Error message
+            details: Additional error details
+
+        """
+        super().__init__(message=message, error_type=AuthErrorType.NETWORK_ERROR, details=details)
+
+
+# Map of error types to error classes
+ERROR_CLASSES: Dict[AuthErrorType, Type[AuthError]] = {
+    AuthErrorType.MISSING_CREDENTIALS: MissingCredentialsError,
+    AuthErrorType.INVALID_CREDENTIALS: InvalidCredentialsError,
+    AuthErrorType.EXPIRED_TOKEN: ExpiredTokenError,
+    AuthErrorType.INSUFFICIENT_PERMISSIONS: InsufficientPermissionsError,
+    AuthErrorType.CONFIGURATION_ERROR: ConfigurationError,
+    AuthErrorType.NETWORK_ERROR: NetworkError,
+    AuthErrorType.UNKNOWN_ERROR: AuthError,
+}
+
+
+def create_auth_error(
+    error_type: AuthErrorType, message: str, details: Optional[Dict] = None
+) -> AuthError:
+    """Create an authentication error of the specified type.
+
+    Args:
+        error_type: Type of authentication error
+        message: Error message
+        details: Additional error details
+
+    Returns:
+        AuthError: An instance of the appropriate error class
+
+    """
+    error_class = ERROR_CLASSES.get(error_type, AuthError)
+    if error_class == AuthError:
+        # For the base class, we need to pass the error_type explicitly
+        return AuthError(message=message, error_type=error_type, details=details)
+    else:
+        # For subclasses, the error_type is already set in the constructor
+        return error_class(message=message, details=details)
+
+
+def format_error_message(provider_name: str, error_type: AuthErrorType, message: str) -> str:
+    """Format an error message for consistent output.
+
+    Args:
+        provider_name: Name of the authentication provider
+        error_type: Type of authentication error
+        message: Error message
+
+    Returns:
+        str: Formatted error message
+
+    """
+    return f'[{provider_name.upper()}] {error_type.value}: {message}'

awslabs/openapi_mcp_server/auth/auth_factory.py

@@ -0,0 +1,146 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Authentication provider factory."""
+
+from awslabs.openapi_mcp_server import logger
+from awslabs.openapi_mcp_server.api.config import Config
+from awslabs.openapi_mcp_server.auth.auth_protocol import AuthProviderProtocol
+from awslabs.openapi_mcp_server.auth.auth_provider import NullAuthProvider
+from typing import Any, Dict, Type
+
+
+# Registry of authentication providers
+_AUTH_PROVIDERS: Dict[str, Type[Any]] = {'none': NullAuthProvider}
+
+# Cache for provider instances
+_PROVIDER_CACHE: Dict[int, AuthProviderProtocol] = {}
+
+
+def register_auth_provider(auth_type: str, provider_class: Type[Any]) -> None:
+    """Register an authentication provider.
+
+    Args:
+        auth_type: The authentication type identifier
+        provider_class: The provider class to register
+
+    Raises:
+        ValueError: If auth_type is already registered
+
+    """
+    auth_type = auth_type.lower()
+    if auth_type in _AUTH_PROVIDERS:
+        raise ValueError(f"Authentication provider for type '{auth_type}' already registered")
+
+    _AUTH_PROVIDERS[auth_type] = provider_class
+    logger.debug(
+        f"Registered authentication provider for type '{auth_type}': {provider_class.__name__}"
+    )
+
+
+def _get_provider_instance(
+    auth_type: str, config_hash: int, config: Config
+) -> AuthProviderProtocol:
+    """Get a cached provider instance or create a new one.
+
+    Args:
+        auth_type: The authentication type
+        config_hash: Hash of the configuration to differentiate instances
+        config: The configuration object
+
+    Returns:
+        AuthProviderProtocol: The authentication provider instance
+
+    """
+    # Check if we have a cached instance
+    if config_hash in _PROVIDER_CACHE:
+        logger.debug(f'Using cached authentication provider for {auth_type}')
+        return _PROVIDER_CACHE[config_hash]
+
+    # Create a new instance
+    provider_class = _AUTH_PROVIDERS[auth_type]
+    provider = provider_class(config)
+
+    # Cache the instance
+    _PROVIDER_CACHE[config_hash] = provider
+
+    logger.debug(f'Created new authentication provider: {provider.provider_name}')
+    return provider
+
+
+def get_auth_provider(config: Config) -> AuthProviderProtocol:
+    """Get an authentication provider based on configuration.
+
+    Args:
+        config: The application configuration
+
+    Returns:
+        AuthProviderProtocol: An authentication provider instance
+
+    Notes:
+        If the specified auth_type is not registered, falls back to NullAuthProvider
+        Uses caching to avoid creating duplicate provider instances
+
+    """
+    auth_type = config.auth_type.lower()
+
+    if auth_type not in _AUTH_PROVIDERS:
+        logger.warning(f"Unknown authentication type '{auth_type}'. Falling back to 'none'.")
+        auth_type = 'none'
+
+    # Create a hash of the relevant config parts for caching
+    config_hash = hash(
+        (
+            auth_type,
+            getattr(config, 'auth_token', None),
+            getattr(config, 'auth_username', None),
+            getattr(config, 'auth_password', None),
+            getattr(config, 'auth_api_key', None),
+            getattr(config, 'auth_api_key_name', None),
+            getattr(config, 'auth_api_key_in', None),
+        )
+    )
+
+    # Get or create provider instance
+    provider = _get_provider_instance(auth_type, config_hash, config)
+
+    logger.info(f'Created authentication provider: {provider.provider_name}')
+
+    if not provider.is_configured() and auth_type != 'none':
+        logger.warning(
+            f"Authentication provider '{provider.provider_name}' is not properly configured"
+        )
+
+    return provider
+
+
+def is_auth_type_available(auth_type: str) -> bool:
+    """Check if an authentication type is available.
+
+    Args:
+        auth_type: The authentication type to check
+
+    Returns:
+        bool: True if available, False otherwise
+
+    """
+    return auth_type.lower() in _AUTH_PROVIDERS
+
+
+def clear_provider_cache() -> None:
+    """Clear the provider instance cache.
+
+    This is useful for testing or when configuration changes.
+    """
+    _PROVIDER_CACHE.clear()
+    logger.debug('Authentication provider cache cleared')

awslabs/openapi_mcp_server/auth/auth_protocol.py

@@ -0,0 +1,63 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Authentication provider protocols and type definitions."""
+
+import httpx
+from awslabs.openapi_mcp_server.api.config import Config
+from typing import Dict, Optional, Protocol, TypeVar, runtime_checkable
+
+
+@runtime_checkable
+class AuthProviderProtocol(Protocol):
+    """Protocol defining the interface for authentication providers.
+
+    This protocol allows for better type checking and removes the need for casting.
+    """
+
+    @property
+    def provider_name(self) -> str:
+        """Get the name of the authentication provider."""
+        ...
+
+    def is_configured(self) -> bool:
+        """Check if the authentication provider is properly configured."""
+        ...
+
+    def get_auth_headers(self) -> Dict[str, str]:
+        """Get authentication headers for HTTP requests."""
+        ...
+
+    def get_auth_params(self) -> Dict[str, str]:
+        """Get authentication query parameters for HTTP requests."""
+        ...
+
+    def get_auth_cookies(self) -> Dict[str, str]:
+        """Get authentication cookies for HTTP requests."""
+        ...
+
+    def get_httpx_auth(self) -> Optional[httpx.Auth]:
+        """Get authentication object for HTTPX."""
+        ...
+
+
+# Type variable for auth provider classes that can be instantiated with a Config
+T = TypeVar('T', bound=AuthProviderProtocol)
+
+
+class AuthProviderFactory(Protocol):
+    """Protocol for auth provider factory functions."""
+
+    def __call__(self, config: Config) -> AuthProviderProtocol:
+        """Create an authentication provider instance."""
+        ...