mcp-security-framework 0.1.0__py3-none-any.whl → 1.1.1__py3-none-any.whl

mcp_security_framework/middleware/__init__.py

@@ -33,36 +33,35 @@ License: MIT
 
 from typing import Optional
 
-from .security_middleware import SecurityMiddleware, SecurityMiddlewareError
-from .fastapi_middleware import FastAPISecurityMiddleware, FastAPIMiddlewareError
-from .flask_middleware import FlaskSecurityMiddleware, FlaskMiddlewareError
-from .auth_middleware import AuthMiddleware, AuthMiddlewareError
-from .rate_limit_middleware import RateLimitMiddleware, RateLimitMiddlewareError
-from .mtls_middleware import MTLSMiddleware, MTLSMiddlewareError
-
 from ..core.security_manager import SecurityManager
 from ..schemas.config import SecurityConfig
+from .auth_middleware import AuthMiddleware, AuthMiddlewareError
+from .fastapi_middleware import FastAPIMiddlewareError, FastAPISecurityMiddleware
+from .flask_middleware import FlaskMiddlewareError, FlaskSecurityMiddleware
+from .mtls_middleware import MTLSMiddleware, MTLSMiddlewareError
+from .rate_limit_middleware import RateLimitMiddleware, RateLimitMiddlewareError
+from .security_middleware import SecurityMiddleware, SecurityMiddlewareError
 
 
 def create_fastapi_security_middleware(
-    security_manager: SecurityManager
+    security_manager: SecurityManager,
 ) -> FastAPISecurityMiddleware:
     """
     Create FastAPI security middleware.
-    
+
     Args:
         security_manager (SecurityManager): Security manager instance
-        
+
     Returns:
         FastAPISecurityMiddleware: Configured FastAPI security middleware
-        
+
     Raises:
         FastAPIMiddlewareError: If middleware creation fails
-        
+
     Example:
         >>> from mcp_security_framework.middleware import create_fastapi_security_middleware
         >>> from fastapi import FastAPI
-        >>> 
+        >>>
         >>> app = FastAPI()
         >>> security_manager = SecurityManager(config)
         >>> middleware = create_fastapi_security_middleware(security_manager)
@@ -72,24 +71,24 @@ def create_fastapi_security_middleware(
 
 
 def create_flask_security_middleware(
-    security_manager: SecurityManager
+    security_manager: SecurityManager,
 ) -> FlaskSecurityMiddleware:
     """
     Create Flask security middleware.
-    
+
     Args:
         security_manager (SecurityManager): Security manager instance
-        
+
     Returns:
         FlaskSecurityMiddleware: Configured Flask security middleware
-        
+
     Raises:
         FlaskMiddlewareError: If middleware creation fails
-        
+
     Example:
         >>> from mcp_security_framework.middleware import create_flask_security_middleware
         >>> from flask import Flask
-        >>> 
+        >>>
         >>> app = Flask(__name__)
         >>> security_manager = SecurityManager(config)
         >>> middleware = create_flask_security_middleware(security_manager)
@@ -99,83 +98,81 @@ def create_flask_security_middleware(
 
 
 def create_auth_middleware(
-    security_manager: SecurityManager,
-    framework: str = "fastapi",
-    cache_ttl: int = 300
+    security_manager: SecurityManager, framework: str = "fastapi", cache_ttl: int = 300
 ) -> AuthMiddleware:
     """
     Create authentication-only middleware for specified framework.
-    
+
     Args:
         security_manager (SecurityManager): Security manager instance
         framework (str): Target framework ("fastapi", "flask", "django")
         cache_ttl (int): Cache TTL in seconds (default: 300)
-        
+
     Returns:
         AuthMiddleware: Framework-specific authentication middleware
-        
+
     Raises:
         ValueError: If framework is not supported
         AuthMiddlewareError: If middleware creation fails
-        
+
     Example:
         >>> from mcp_security_framework.middleware import create_auth_middleware
-        >>> 
+        >>>
         >>> security_manager = SecurityManager(config)
         >>> auth_middleware = create_auth_middleware(security_manager, "fastapi")
     """
     if framework == "fastapi":
         from .fastapi_auth_middleware import FastAPIAuthMiddleware
+
         return FastAPIAuthMiddleware(security_manager.config, security_manager)
     elif framework == "flask":
         from .flask_auth_middleware import FlaskAuthMiddleware
+
         return FlaskAuthMiddleware(security_manager.config, security_manager)
     else:
         raise ValueError(f"Unsupported framework: {framework}")
 
 
 def create_rate_limit_middleware(
-    security_manager: SecurityManager
+    security_manager: SecurityManager,
 ) -> RateLimitMiddleware:
     """
     Create rate limiting middleware.
-    
+
     Args:
         security_manager (SecurityManager): Security manager instance
-        
+
     Returns:
         RateLimitMiddleware: Configured rate limiting middleware
-        
+
     Raises:
         RateLimitMiddlewareError: If middleware creation fails
-        
+
     Example:
         >>> from mcp_security_framework.middleware import create_rate_limit_middleware
-        >>> 
+        >>>
         >>> security_manager = SecurityManager(config)
         >>> rate_limit_middleware = create_rate_limit_middleware(security_manager)
     """
     return RateLimitMiddleware(security_manager)
 
 
-def create_mtls_middleware(
-    security_manager: SecurityManager
-) -> MTLSMiddleware:
+def create_mtls_middleware(security_manager: SecurityManager) -> MTLSMiddleware:
     """
     Create mTLS middleware.
-    
+
     Args:
         security_manager (SecurityManager): Security manager instance
-        
+
     Returns:
         MTLSMiddleware: Configured mTLS middleware
-        
+
     Raises:
         MTLSMiddlewareError: If middleware creation fails
-        
+
     Example:
         >>> from mcp_security_framework.middleware import create_mtls_middleware
-        >>> 
+        >>>
         >>> security_manager = SecurityManager(config)
         >>> mtls_middleware = create_mtls_middleware(security_manager)
     """
@@ -183,38 +180,36 @@ def create_mtls_middleware(
 
 
 def create_security_middleware(
-    security_manager: SecurityManager,
-    framework: str = "fastapi"
+    security_manager: SecurityManager, framework: str = "fastapi"
 ) -> SecurityMiddleware:
     """
     Create security middleware for specified framework.
-    
+
     Args:
         security_manager (SecurityManager): Security manager instance
         framework (str): Target framework ("fastapi", "flask", etc.)
-        
+
     Returns:
         SecurityMiddleware: Configured security middleware
-        
+
     Raises:
         SecurityMiddlewareError: If middleware creation fails
-        
+
     Example:
         >>> from mcp_security_framework.middleware import create_security_middleware
-        >>> 
+        >>>
         >>> security_manager = SecurityManager(config)
         >>> middleware = create_security_middleware(security_manager, "fastapi")
     """
     framework = framework.lower()
-    
+
     if framework == "fastapi":
         return create_fastapi_security_middleware(security_manager)
     elif framework == "flask":
         return create_flask_security_middleware(security_manager)
     else:
         raise SecurityMiddlewareError(
-            f"Unsupported framework: {framework}",
-            error_code=-32041
+            f"Unsupported framework: {framework}", error_code=-32041
         )
 
 
@@ -223,15 +218,12 @@ __all__ = [
     # Base classes
     "SecurityMiddleware",
     "SecurityMiddlewareError",
-    
     # FastAPI middleware
     "FastAPISecurityMiddleware",
     "FastAPIMiddlewareError",
-    
     # Flask middleware
     "FlaskSecurityMiddleware",
     "FlaskMiddlewareError",
-    
     # Specialized middleware
     "AuthMiddleware",
     "AuthMiddlewareError",
@@ -239,7 +231,6 @@ __all__ = [
     "RateLimitMiddlewareError",
     "MTLSMiddleware",
     "MTLSMiddlewareError",
-    
     # Factory functions
     "create_fastapi_security_middleware",
     "create_flask_security_middleware",

mcp_security_framework/middleware/auth_middleware.py

@@ -24,13 +24,13 @@ import logging
 from abc import ABC, abstractmethod
 from typing import Any, Dict, List, Optional
 
+from ..schemas.models import AuthResult, AuthStatus
 from .security_middleware import SecurityMiddleware, SecurityMiddlewareError
-from ..schemas.models import AuthResult
 
 
 class AuthMiddlewareError(SecurityMiddlewareError):
     """Raised when authentication middleware encounters an error."""
-    
+
     def __init__(self, message: str, error_code: int = -32032):
         self.message = message
         self.error_code = error_code
@@ -40,100 +40,99 @@ class AuthMiddlewareError(SecurityMiddlewareError):
 class AuthMiddleware(SecurityMiddleware):
     """
     Authentication-Only Middleware Class
-    
+
     This class provides authentication-only middleware that focuses
     solely on user authentication without performing authorization
     checks. It's useful for scenarios where authentication and
     authorization are handled separately.
-    
+
     The AuthMiddleware implements:
     - Authentication-only request processing
     - Multiple authentication method support
     - Authentication result caching
     - Authentication event logging
     - Framework-agnostic design
-    
+
     Key Responsibilities:
     - Process requests through authentication pipeline only
     - Handle multiple authentication methods
     - Cache authentication results for performance
     - Log authentication events and failures
     - Provide authentication status to downstream components
-    
+
     Attributes:
         Inherits all attributes from SecurityMiddleware
         _auth_cache (Dict): Cache for authentication results
         _cache_ttl (int): Time-to-live for cached results
-        
+
     Example:
         >>> from mcp_security_framework.middleware import AuthMiddleware
-        >>> 
+        >>>
         >>> security_manager = SecurityManager(config)
         >>> auth_middleware = AuthMiddleware(security_manager)
         >>> app.add_middleware(auth_middleware)
-        
+
     Note:
         This middleware only handles authentication. Authorization
         should be handled separately by other middleware or
         application logic.
     """
-    
+
     def __init__(self, security_manager, cache_ttl: int = 300):
         """
         Initialize Authentication-Only Middleware.
-        
+
         Args:
             security_manager: Security manager instance containing
                 all security components and configuration.
             cache_ttl (int): Time-to-live for cached authentication
                 results in seconds. Defaults to 300 seconds (5 minutes).
-                
+
         Raises:
             AuthMiddlewareError: If initialization fails
         """
         super().__init__(security_manager)
         self.logger = logging.getLogger(f"{__name__}.{self.__class__.__name__}")
         self._cache_ttl = cache_ttl
-        
+
         self.logger.info(
-            "Authentication middleware initialized",
-            extra={"cache_ttl": cache_ttl}
+            "Authentication middleware initialized", extra={"cache_ttl": cache_ttl}
         )
-    
+
     @abstractmethod
     def __call__(self, request: Any, call_next: Any) -> Any:
         """
         Process request through authentication middleware.
-        
+
         This method implements the authentication-only processing
         pipeline, focusing solely on user authentication.
-        
+
         Args:
             request: Framework-specific request object
             call_next: Framework-specific call_next function
-            
+
         Returns:
             Framework-specific response object
-            
+
         Raises:
             AuthMiddlewareError: If authentication processing fails
         """
         pass
-    
+
     def _authenticate_only(self, request: Any) -> AuthResult:
         """
         Perform authentication-only processing.
-        
+
         This method handles authentication without authorization
         checks, making it suitable for scenarios where auth and
         authz are separated.
-        
+
         Args:
             request: Framework-specific request object
-            
+
         Returns:
             AuthResult: Authentication result with user information
-            
+
         Raises:
             AuthMiddlewareError: If authentication process fails
         """
@@ -144,9 +143,9 @@ class AuthMiddleware(SecurityMiddleware):
                     status=AuthStatus.SUCCESS,
                     username="anonymous",
                     roles=[],
-                    auth_method=None
+                    auth_method=None,
                 )
-            
+
             # Check cache first
             cache_key = self._get_cache_key(request)
             if cache_key in self._auth_cache:
@@ -154,33 +153,33 @@ class AuthMiddleware(SecurityMiddleware):
                 if self._is_cache_valid(cached_result):
                     self.logger.debug(
                         "Using cached authentication result",
-                        extra={"username": cached_result.username}
+                        extra={"username": cached_result.username},
                     )
                     return cached_result
-            
+
             # Try each authentication method in order
             for method in self.config.auth.methods:
                 auth_result = self._try_auth_method(request, method)
                 if auth_result.is_valid:
                     # Cache successful authentication result
                     self._cache_auth_result(cache_key, auth_result)
-                    
+
                     self.logger.info(
                         "Authentication successful",
                         extra={
                             "username": auth_result.username,
                             "auth_method": auth_result.auth_method,
-                            "user_roles": auth_result.roles
-                        }
+                            "user_roles": auth_result.roles,
+                        },
                     )
                     return auth_result
-            
+
             # All authentication methods failed
             self.logger.warning(
                 "All authentication methods failed",
-                extra={"auth_methods": self.config.auth.methods}
+                extra={"auth_methods": self.config.auth.methods},
             )
-            
+
             failed_result = AuthResult(
                 is_valid=False,
                 status=AuthStatus.FAILED,
@@ -188,31 +187,28 @@ class AuthMiddleware(SecurityMiddleware):
                 roles=[],
                 auth_method=None,
                 error_code=-32033,
-                error_message="All authentication methods failed"
+                error_message="All authentication methods failed",
             )
-            
+
             # Cache failed result briefly to prevent repeated attempts
             self._cache_auth_result(cache_key, failed_result, ttl=60)
             return failed_result
-            
+
         except Exception as e:
             self.logger.error(
-                "Authentication process failed",
-                extra={"error": str(e)},
-                exc_info=True
+                "Authentication process failed", extra={"error": str(e)}, exc_info=True
             )
             raise AuthMiddlewareError(
-                f"Authentication process failed: {str(e)}",
-                error_code=-32034
+                f"Authentication process failed: {str(e)}", error_code=-32034
             )
-    
+
     def _get_cache_key(self, request: Any) -> str:
         """
         Generate cache key for authentication result.
-        
+
         Args:
             request: Framework-specific request object
-            
+
         Returns:
             str: Cache key for authentication result
         """
@@ -220,25 +216,27 @@ class AuthMiddleware(SecurityMiddleware):
         ip = self._get_rate_limit_identifier(request)
         user_agent = self._get_user_agent(request)
         return f"auth:{ip}:{user_agent}"
-    
+
     def _is_cache_valid(self, auth_result: AuthResult) -> bool:
         """
         Check if cached authentication result is still valid.
-        
+
         Args:
             auth_result (AuthResult): Cached authentication result
-            
+
         Returns:
             bool: True if cache is valid, False otherwise
         """
         # For now, assume cache is valid if it exists
         # In a real implementation, you might check timestamps
         return True
-    
-    def _cache_auth_result(self, cache_key: str, auth_result: AuthResult, ttl: int = None) -> None:
+
+    def _cache_auth_result(
+        self, cache_key: str, auth_result: AuthResult, ttl: int = None
+    ) -> None:
         """
         Cache authentication result.
-        
+
         Args:
             cache_key (str): Cache key for the result
             auth_result (AuthResult): Authentication result to cache
@@ -246,33 +244,34 @@ class AuthMiddleware(SecurityMiddleware):
         """
         if ttl is None:
             ttl = self._cache_ttl
-        
+
         # In a real implementation, you would store with timestamp
         self._auth_cache[cache_key] = auth_result
-    
+
     def _clear_auth_cache(self) -> None:
         """Clear all cached authentication results."""
         self._auth_cache.clear()
         self.logger.info("Authentication cache cleared")
-    
+
     def _get_user_agent(self, request: Any) -> str:
         """
         Get user agent from request.
-        
+
         Args:
             request: Framework-specific request object
-            
+
         Returns:
             str: User agent string
         """
         # This should be implemented by framework-specific subclasses
         return "unknown"
-    
-    def _log_auth_event(self, event_type: str, auth_result: AuthResult, 
-                       request_details: Dict[str, Any]) -> None:
+
+    def _log_auth_event(
+        self, event_type: str, auth_result: AuthResult, request_details: Dict[str, Any]
+    ) -> None:
         """
         Log authentication event.
-        
+
         Args:
             event_type (str): Type of authentication event
             auth_result (AuthResult): Authentication result
@@ -287,6 +286,6 @@ class AuthMiddleware(SecurityMiddleware):
                 "is_valid": auth_result.is_valid,
                 "error_code": auth_result.error_code,
                 "error_message": auth_result.error_message,
-                **request_details
-            }
+                **request_details,
+            },
         )