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

mcp_security_framework/middleware/flask_auth_middleware.py

@@ -27,165 +27,167 @@ Last Modified: 2024-01-20
 
 import json
 import logging
-from typing import Any, Dict, Optional
+from typing import Any, Dict, List, Optional
 
-from flask import Request, Response, jsonify, current_app
+from flask import Request, Response, current_app, jsonify
 from werkzeug.exceptions import Unauthorized
 
 from mcp_security_framework.middleware.auth_middleware import AuthMiddleware
-from mcp_security_framework.schemas.models import AuthResult, AuthStatus, AuthMethod
 from mcp_security_framework.schemas.config import SecurityConfig
+from mcp_security_framework.schemas.models import AuthMethod, AuthResult, AuthStatus
 
 
 class FlaskAuthMiddleware(AuthMiddleware):
     """
     Flask Authentication Middleware Implementation
-    
+
     This class provides Flask-specific authentication middleware that
     handles authentication-only processing for Flask applications.
-    
+
     The middleware implements:
     - Flask request/response handling
     - Authentication caching and optimization
     - Framework-specific error responses
     - Security event logging
-    
+
     Attributes:
         config (SecurityConfig): Security configuration settings
         security_manager: Security manager instance
         logger (Logger): Logger instance for authentication operations
         _auth_cache (Dict): Authentication result cache
-        
+
     Example:
         >>> from mcp_security_framework.middleware import create_auth_middleware
         >>> middleware = create_auth_middleware(config, framework="flask")
         >>> app.wsgi_app = middleware(app.wsgi_app)
-        
+
     Raises:
         AuthMiddlewareError: When authentication processing fails
     """
-    
-    def __init__(self, config: SecurityConfig, security_manager: Any):
+
+    def __init__(self, security_manager: Any):
         """
         Initialize Flask Authentication Middleware.
-        
+
         Args:
-            config (SecurityConfig): Security configuration settings
-            security_manager: Security manager instance
+            security_manager: Security manager instance containing
+                all security components and configuration.
         """
-        super().__init__(config, security_manager)
+        super().__init__(security_manager)
         self.logger = logging.getLogger(__name__)
-    
+
     def __call__(self, environ: Dict[str, Any], start_response: Any) -> Any:
         """
         Process Flask request through authentication middleware.
-        
+
         This method implements the authentication-only processing
         pipeline for Flask requests, focusing solely on user
         authentication without authorization checks.
-        
+
         Args:
             environ (Dict[str, Any]): WSGI environment
             start_response: WSGI start_response function
-            
+
         Returns:
             Any: WSGI response
-            
+
         Raises:
             AuthMiddlewareError: If authentication processing fails
         """
         try:
             # Create Flask request object
             request = Request(environ)
-            
+
             # Check if path is public (bypasses authentication)
             if self._is_public_path(request):
                 return self._call_next(environ, start_response)
-            
+
             # Perform authentication
             auth_result = self._authenticate_only(request)
-            
+
             if not auth_result.is_valid:
                 return self._auth_error_response(auth_result, start_response)
-            
+
             # Add authentication info to request context
             request.auth_result = auth_result
             request.username = auth_result.username
             request.user_roles = auth_result.roles
             request.auth_method = auth_result.auth_method
-            
+
             # Process request
             response = self._call_next(environ, start_response)
-            
+
             # Log successful authentication
-            self._log_auth_event("authentication_successful", {
-                "ip_address": self._get_client_ip(request),
-                "username": auth_result.username,
-                "path": request.path,
-                "method": request.method,
-                "auth_method": auth_result.auth_method
-            })
-            
+            self._log_auth_event(
+                "authentication_successful",
+                {
+                    "ip_address": self._get_client_ip(request),
+                    "username": auth_result.username,
+                    "path": request.path,
+                    "method": request.method,
+                    "auth_method": auth_result.auth_method,
+                },
+            )
+
             return response
-            
+
         except Exception as e:
             self.logger.error(
                 "Flask authentication middleware processing failed",
                 extra={"error": str(e)},
-                exc_info=True
+                exc_info=True,
             )
             raise AuthMiddlewareError(
-                f"Authentication processing failed: {str(e)}",
-                error_code=-32035
+                f"Authentication processing failed: {str(e)}", error_code=-32035
             )
-    
+
     def _call_next(self, environ: Dict[str, Any], start_response: Any) -> Any:
         """
         Call the next WSGI application in the chain.
-        
+
         Args:
             environ (Dict[str, Any]): WSGI environment
             start_response: WSGI start_response function
-            
+
         Returns:
             Any: WSGI response
         """
         # This would typically call the Flask app
         # For now, return a simple response
-        status = '200 OK'
-        headers = [('Content-Type', 'application/json')]
+        status = "200 OK"
+        headers = [("Content-Type", "application/json")]
         start_response(status, headers)
         return [json.dumps({"message": "OK"}).encode()]
-    
+
     def _is_public_path(self, request: Request) -> bool:
         """
         Check if the request path is public (bypasses authentication).
-        
+
         Args:
             request (Request): Flask request object
-            
+
         Returns:
             bool: True if path is public, False otherwise
         """
         path = request.path
-        
+
         # Check configured public paths
-        if hasattr(self.config.auth, 'public_paths'):
+        if hasattr(self.config.auth, "public_paths"):
             for public_path in self.config.auth.public_paths:
                 if path.startswith(public_path):
                     return True
-        
+
         # Check common public paths
         public_paths = ["/health", "/status", "/metrics", "/docs", "/openapi.json"]
         return any(path.startswith(public_path) for public_path in public_paths)
-    
+
     def _get_client_ip(self, request: Request) -> str:
         """
         Get client IP address from Flask request.
-        
+
         Args:
             request (Request): Flask request object
-            
+
         Returns:
             str: Client IP address
         """
@@ -193,33 +195,35 @@ class FlaskAuthMiddleware(AuthMiddleware):
         forwarded_for = request.headers.get("X-Forwarded-For")
         if forwarded_for:
             return forwarded_for.split(",")[0].strip()
-        
+
         # Try X-Real-IP header
         real_ip = request.headers.get("X-Real-IP")
         if real_ip:
             return real_ip
-        
+
         # Use remote address
         if request.remote_addr:
             return request.remote_addr
-        
+
         # Fallback to default IP from config or environment
-        default_ip = getattr(self.config, 'default_client_ip', None)
+        default_ip = getattr(self.config, "default_client_ip", None)
         if default_ip:
             return default_ip
-        
+
         # Use environment variable or default
         import os
+
         from ..constants import DEFAULT_CLIENT_IP
-        return os.environ.get('DEFAULT_CLIENT_IP', DEFAULT_CLIENT_IP)
-    
+
+        return os.environ.get("DEFAULT_CLIENT_IP", DEFAULT_CLIENT_IP)
+
     def _get_cache_key(self, request: Request) -> str:
         """
         Generate cache key for authentication result.
-        
+
         Args:
             request (Request): Flask request object
-            
+
         Returns:
             str: Cache key
         """
@@ -227,15 +231,15 @@ class FlaskAuthMiddleware(AuthMiddleware):
         ip = self._get_client_ip(request)
         user_agent = request.headers.get("User-Agent", "")
         return f"auth:{ip}:{hash(user_agent)}"
-    
+
     def _try_auth_method(self, request: Request, method: str) -> AuthResult:
         """
         Try authentication using specific method with Flask request.
-        
+
         Args:
             request (Request): Flask request object
             method (str): Authentication method to try
-            
+
         Returns:
             AuthResult: Authentication result
         """
@@ -256,13 +260,13 @@ class FlaskAuthMiddleware(AuthMiddleware):
                     roles=[],
                     auth_method=None,
                     error_code=-32022,
-                    error_message=f"Unsupported authentication method: {method}"
+                    error_message=f"Unsupported authentication method: {method}",
                 )
         except Exception as e:
             self.logger.error(
                 f"Authentication method {method} failed",
                 extra={"error": str(e)},
-                exc_info=True
+                exc_info=True,
             )
             return AuthResult(
                 is_valid=False,
@@ -271,16 +275,16 @@ class FlaskAuthMiddleware(AuthMiddleware):
                 roles=[],
                 auth_method=None,
                 error_code=-32023,
-                error_message=f"Authentication method {method} failed: {str(e)}"
+                error_message=f"Authentication method {method} failed: {str(e)}",
             )
-    
+
     def _try_api_key_auth(self, request: Request) -> AuthResult:
         """
         Try API key authentication with Flask request.
-        
+
         Args:
             request (Request): Flask request object
-            
+
         Returns:
             AuthResult: Authentication result
         """
@@ -291,7 +295,7 @@ class FlaskAuthMiddleware(AuthMiddleware):
             auth_header = request.headers.get("Authorization")
             if auth_header and auth_header.startswith("Bearer "):
                 api_key = auth_header[7:]  # Remove "Bearer " prefix
-        
+
         if not api_key:
             return AuthResult(
                 is_valid=False,
@@ -300,19 +304,19 @@ class FlaskAuthMiddleware(AuthMiddleware):
                 roles=[],
                 auth_method=AuthMethod.API_KEY,
                 error_code=-32012,
-                error_message="API key not found in request"
+                error_message="API key not found in request",
             )
-        
+
         # Authenticate using security manager
         return self.security_manager.auth_manager.authenticate_api_key(api_key)
-    
+
     def _try_jwt_auth(self, request: Request) -> AuthResult:
         """
         Try JWT authentication with Flask request.
-        
+
         Args:
             request (Request): Flask request object
-            
+
         Returns:
             AuthResult: Authentication result
         """
@@ -326,21 +330,21 @@ class FlaskAuthMiddleware(AuthMiddleware):
                 roles=[],
                 auth_method=AuthMethod.JWT,
                 error_code=-32013,
-                error_message="JWT token not found in Authorization header"
+                error_message="JWT token not found in Authorization header",
             )
-        
+
         token = auth_header[7:]  # Remove "Bearer " prefix
-        
+
         # Authenticate using security manager
         return self.security_manager.auth_manager.authenticate_jwt_token(token)
-    
+
     def _try_certificate_auth(self, request: Request) -> AuthResult:
         """
         Try certificate authentication with Flask request.
-        
+
         Args:
             request (Request): Flask request object
-            
+
         Returns:
             AuthResult: Authentication result
         """
@@ -354,16 +358,16 @@ class FlaskAuthMiddleware(AuthMiddleware):
             roles=[],
             auth_method=AuthMethod.CERTIFICATE,
             error_code=-32014,
-            error_message="Certificate authentication not implemented"
+            error_message="Certificate authentication not implemented",
         )
-    
+
     def _try_basic_auth(self, request: Request) -> AuthResult:
         """
         Try basic authentication with Flask request.
-        
+
         Args:
             request (Request): Flask request object
-            
+
         Returns:
             AuthResult: Authentication result
         """
@@ -377,9 +381,9 @@ class FlaskAuthMiddleware(AuthMiddleware):
                 roles=[],
                 auth_method=AuthMethod.BASIC,
                 error_code=-32015,
-                error_message="Basic authentication credentials not found"
+                error_message="Basic authentication credentials not found",
             )
-        
+
         # Basic auth is not implemented in this version
         return AuthResult(
             is_valid=False,
@@ -388,17 +392,17 @@ class FlaskAuthMiddleware(AuthMiddleware):
             roles=[],
             auth_method=AuthMethod.BASIC,
             error_code=-32016,
-            error_message="Basic authentication not implemented"
+            error_message="Basic authentication not implemented",
         )
-    
+
     def _auth_error_response(self, auth_result: AuthResult, start_response: Any) -> Any:
         """
         Create authentication error response for Flask.
-        
+
         Args:
             auth_result (AuthResult): Authentication result with error
             start_response: WSGI start_response function
-            
+
         Returns:
             Any: WSGI error response
         """
@@ -406,22 +410,21 @@ class FlaskAuthMiddleware(AuthMiddleware):
             "error": "Authentication failed",
             "error_code": auth_result.error_code,
             "error_message": auth_result.error_message,
-            "auth_method": auth_result.auth_method.value if auth_result.auth_method else None
+            "auth_method": (
+                auth_result.auth_method.value if auth_result.auth_method else None
+            ),
         }
-        
-        status = '401 Unauthorized'
-        headers = [
-            ('Content-Type', 'application/json'),
-            ('WWW-Authenticate', 'Bearer')
-        ]
-        
+
+        status = "401 Unauthorized"
+        headers = [("Content-Type", "application/json"), ("WWW-Authenticate", "Bearer")]
+
         start_response(status, headers)
         return [json.dumps(error_data).encode()]
-    
+
     def _log_auth_event(self, event_type: str, details: Dict[str, Any]) -> None:
         """
         Log authentication event.
-        
+
         Args:
             event_type (str): Type of authentication event
             details (Dict[str, Any]): Event details
@@ -437,28 +440,185 @@ class FlaskAuthMiddleware(AuthMiddleware):
                     "path": details.get("path"),
                     "method": details.get("method"),
                     "auth_method": details.get("auth_method"),
-                    **details
-                }
+                    **details,
+                },
             )
         except Exception as e:
             self.logger.error(
                 "Failed to log authentication event",
                 extra={"error": str(e)},
-                exc_info=True
+                exc_info=True,
             )
 
+    def _apply_security_headers(self, response: Any) -> Any:
+        """
+        Apply security headers to response.
+
+        Args:
+            response: WSGI response
+
+        Returns:
+            Any: Response with security headers
+        """
+        # This would be implemented for actual Flask responses
+        # For now, return the response as-is
+        return response
+
+    def _create_error_response(
+        self, error_message: str, error_code: int, start_response: Any
+    ) -> Any:
+        """
+        Create error response for Flask.
+
+        Args:
+            error_message (str): Error message
+            error_code (int): Error code
+            start_response: WSGI start_response function
+
+        Returns:
+            Any: WSGI error response
+        """
+        error_data = {
+            "error": "Security error",
+            "error_code": error_code,
+            "error_message": error_message,
+        }
+
+        status = "400 Bad Request"
+        headers = [("Content-Type", "application/json")]
+
+        start_response(status, headers)
+        return [json.dumps(error_data).encode()]
+
+    def _get_rate_limit_identifier(self, request: Request) -> str:
+        """
+        Get rate limit identifier from request.
+
+        Args:
+            request (Request): Flask request object
+
+        Returns:
+            str: Rate limit identifier (IP address)
+        """
+        return self._get_client_ip(request)
+
+    def _get_request_path(self, request: Request) -> str:
+        """
+        Get request path from request.
+
+        Args:
+            request (Request): Flask request object
+
+        Returns:
+            str: Request path
+        """
+        return request.path
+
+    def _get_required_permissions(self, request: Request) -> List[str]:
+        """
+        Get required permissions for request.
+
+        Args:
+            request (Request): Flask request object
+
+        Returns:
+            List[str]: Required permissions
+        """
+        # This would be implemented based on your application's authorization logic
+        # For now, return an empty list
+        return []
+
+    def _validation_error_response(
+        self, error_message: str, error_code: int = -32000, start_response: Any = None
+    ) -> Any:
+        """
+        Create validation error response for Flask.
+
+        Args:
+            error_message (str): Error message
+            error_code (int): Error code
+            start_response: WSGI start_response function
+
+        Returns:
+            Any: WSGI error response
+        """
+        if start_response is None:
+            start_response = Mock()
+        return self._create_error_response(error_message, error_code, start_response)
+
+    def _rate_limit_error_response(
+        self, error_message: str, error_code: int = -32000, start_response: Any = None
+    ) -> Any:
+        """
+        Create rate limit error response for Flask.
+
+        Args:
+            error_message (str): Error message
+            error_code (int): Error code
+            start_response: WSGI start_response function
+
+        Returns:
+            Any: WSGI error response
+        """
+        if start_response is None:
+            start_response = Mock()
+        return self._create_error_response(error_message, error_code, start_response)
+
+    def _security_header_response(
+        self, error_message: str, error_code: int = -32000, start_response: Any = None
+    ) -> Any:
+        """
+        Create security header error response for Flask.
+
+        Args:
+            error_message (str): Error message
+            error_code (int): Error code
+            start_response: WSGI start_response function
+
+        Returns:
+            Any: WSGI error response
+        """
+        if start_response is None:
+            start_response = Mock()
+        return self._create_error_response(error_message, error_code, start_response)
+
+    def _authz_error_response(
+        self, auth_result: AuthResult, start_response: Any
+    ) -> Any:
+        """
+        Create authorization error response for Flask.
+
+        Args:
+            auth_result (AuthResult): Authentication result
+            start_response: WSGI start_response function
+
+        Returns:
+            Any: WSGI error response
+        """
+        error_data = {
+            "error": "Authorization failed",
+            "error_code": auth_result.error_code,
+            "error_message": auth_result.error_message,
+        }
+
+        status = "403 Forbidden"
+        headers = [("Content-Type", "application/json")]
+
+        start_response(status, headers)
+        return [json.dumps(error_data).encode()]
+
 
 class AuthMiddlewareError(Exception):
     """
     Authentication Middleware Error
-    
+
     This exception is raised when authentication middleware processing fails.
-    
+
     Attributes:
         message (str): Error message
         error_code (int): Error code for programmatic handling
     """
-    
+
     def __init__(self, message: str, error_code: int = -32030):
         self.message = message
         self.error_code = error_code