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

tests/test_middleware/test_flask_middleware.py

@@ -19,25 +19,35 @@ Version: 1.0.0
 License: MIT
 """
 
-import pytest
 import json
-from unittest.mock import Mock, patch, MagicMock
-from typing import Dict, List, Any
+from typing import Any, Dict, List
+from unittest.mock import MagicMock, Mock, patch
 
+import pytest
 from flask import Request, Response
 
+from mcp_security_framework.core.security_manager import SecurityManager
 from mcp_security_framework.middleware.flask_middleware import (
+    FlaskMiddlewareError,
     FlaskSecurityMiddleware,
-    FlaskMiddlewareError
 )
-from mcp_security_framework.core.security_manager import SecurityManager
-from mcp_security_framework.schemas.config import SecurityConfig, AuthConfig, RateLimitConfig
-from mcp_security_framework.schemas.models import AuthResult, ValidationResult, ValidationStatus, AuthStatus, AuthMethod
+from mcp_security_framework.schemas.config import (
+    AuthConfig,
+    RateLimitConfig,
+    SecurityConfig,
+)
+from mcp_security_framework.schemas.models import (
+    AuthMethod,
+    AuthResult,
+    AuthStatus,
+    ValidationResult,
+    ValidationStatus,
+)
 
 
 class TestFlaskSecurityMiddleware:
     """Test suite for FlaskSecurityMiddleware class."""
-    
+
     def setup_method(self):
         """Set up test fixtures before each test method."""
         # Create mock security manager
@@ -47,26 +57,26 @@ class TestFlaskSecurityMiddleware:
                 enabled=True,
                 methods=["api_key", "jwt"],
                 public_paths=["/health", "/docs"],
-                jwt_secret="test_jwt_secret_key"
+                jwt_secret="test_jwt_secret_key",
             ),
             rate_limit=RateLimitConfig(
-                enabled=True,
-                default_requests_per_minute=100,
-                window_size_seconds=60
-            )
+                enabled=True, default_requests_per_minute=100, window_size_seconds=60
+            ),
         )
-        
+
         # Create mock auth manager
         self.mock_auth_manager = Mock()
         self.mock_security_manager.auth_manager = self.mock_auth_manager
-        
+
         # Setup rate_limiter mock
         self.mock_security_manager.rate_limiter = Mock()
-        
+
         # Create middleware instance
         self.middleware = FlaskSecurityMiddleware(self.mock_security_manager)
-    
-    def create_mock_request(self, path: str = "/api/test", headers: Dict[str, str] = None) -> Mock:
+
+    def create_mock_request(
+        self, path: str = "/api/test", headers: Dict[str, str] = None
+    ) -> Mock:
         """Create a mock Flask request for testing."""
         mock_request = Mock(spec=Request)
         mock_request.path = path
@@ -74,112 +84,114 @@ class TestFlaskSecurityMiddleware:
         mock_request.headers = headers or {}
         mock_request.remote_addr = "127.0.0.1"
         return mock_request
-    
-    def create_mock_environ(self, path: str = "/api/test", headers: Dict[str, str] = None) -> Dict[str, Any]:
+
+    def create_mock_environ(
+        self, path: str = "/api/test", headers: Dict[str, str] = None
+    ) -> Dict[str, Any]:
         """Create a mock WSGI environ for testing."""
         environ = {
-            'REQUEST_METHOD': 'GET',
-            'PATH_INFO': path,
-            'QUERY_STRING': '',
-            'SERVER_NAME': 'localhost',
-            'SERVER_PORT': '5000',
-            'HTTP_HOST': 'localhost:5000',
-            'wsgi.url_scheme': 'http',
-            'wsgi.input': Mock(),
-            'wsgi.errors': Mock(),
-            'wsgi.version': (1, 0),
-            'wsgi.run_once': False,
-            'wsgi.multithread': False,
-            'wsgi.multiprocess': False,
+            "REQUEST_METHOD": "GET",
+            "PATH_INFO": path,
+            "QUERY_STRING": "",
+            "SERVER_NAME": "localhost",
+            "SERVER_PORT": "5000",
+            "HTTP_HOST": "localhost:5000",
+            "wsgi.url_scheme": "http",
+            "wsgi.input": Mock(),
+            "wsgi.errors": Mock(),
+            "wsgi.version": (1, 0),
+            "wsgi.run_once": False,
+            "wsgi.multithread": False,
+            "wsgi.multiprocess": False,
         }
-        
+
         # Add headers to environ
         if headers:
             for key, value in headers.items():
                 environ[f'HTTP_{key.upper().replace("-", "_")}'] = value
-        
+
         return environ
-    
+
     def test_initialization_success(self):
         """Test successful Flask middleware initialization."""
         assert isinstance(self.middleware, FlaskSecurityMiddleware)
         assert self.middleware.security_manager == self.mock_security_manager
         assert self.middleware.config == self.mock_security_manager.config
-    
+
     def test_call_success(self):
         """Test successful middleware call."""
         environ = self.create_mock_environ()
         mock_start_response = Mock()
-        
+
         # Mock successful authentication and authorization
         self.mock_security_manager.rate_limiter.check_rate_limit.return_value = True
-        
+
         # Mock successful authentication
         auth_result = AuthResult(
             is_valid=True,
             status=AuthStatus.SUCCESS,
             username="test_user",
             roles=["user"],
-            auth_method=AuthMethod.API_KEY
+            auth_method=AuthMethod.API_KEY,
         )
         self.middleware._authenticate_request = Mock(return_value=auth_result)
-        
+
         self.middleware._validate_permissions = Mock(return_value=True)
-        
+
         # Mock successful response
         mock_response = [b'{"message": "success"}']
         self.middleware._process_request = Mock(return_value=mock_response)
-        
+
         result = self.middleware(environ, mock_start_response)
-        
+
         # Verify middleware processed successfully
         assert result == mock_response
         self.middleware._process_request.assert_called_once()
-    
+
     def test_call_rate_limit_exceeded(self):
         """Test middleware call with rate limit exceeded."""
         environ = self.create_mock_environ()
         mock_start_response = Mock()
-        
+
         self.mock_security_manager.rate_limiter.check_rate_limit.return_value = False
-        
+
         result = self.middleware(environ, mock_start_response)
-        
+
         # Verify rate limit response
         assert isinstance(result, list)
         assert len(result) == 1
-        response_data = json.loads(result[0].decode('utf-8'))
+        response_data = json.loads(result[0].decode("utf-8"))
         assert response_data["error"] == "Rate limit exceeded"
-        
+
         # Verify start_response was called with correct status
         mock_start_response.assert_called_once()
         call_args = mock_start_response.call_args[0]
-        assert call_args[0] == '429 Too Many Requests'
-    
+        assert call_args[0] == "429 Too Many Requests"
+
     def test_call_public_path(self):
         """Test middleware call with public path."""
         environ = self.create_mock_environ(path="/health")
         mock_start_response = Mock()
-        
+
         self.mock_security_manager.rate_limiter.check_rate_limit.return_value = True
-        
+
         # Mock successful response for public path
         mock_response = [b'{"status": "healthy"}']
         self.middleware._process_request = Mock(return_value=mock_response)
-        
+
         result = self.middleware(environ, mock_start_response)
-        
+
         # Verify middleware processed successfully for public path
         assert result == mock_response
         self.middleware._process_request.assert_called_once()
-    
+
     def test_call_authentication_failed(self):
         """Test middleware call with authentication failure."""
         environ = self.create_mock_environ()
         mock_start_response = Mock()
-        
+
         self.mock_security_manager.rate_limiter.check_rate_limit.return_value = True
-        
+
         # Mock failed authentication
         auth_result = AuthResult(
             is_valid=False,
@@ -188,91 +200,91 @@ class TestFlaskSecurityMiddleware:
             roles=[],
             auth_method=AuthMethod.API_KEY,
             error_code=-32005,
-            error_message="Authentication failed"
+            error_message="Authentication failed",
         )
         self.middleware._authenticate_request = Mock(return_value=auth_result)
-        
+
         result = self.middleware(environ, mock_start_response)
-        
+
         # Verify authentication error response
         assert isinstance(result, list)
         assert len(result) == 1
-        response_data = json.loads(result[0].decode('utf-8'))
+        response_data = json.loads(result[0].decode("utf-8"))
         assert response_data["error"] == "Authentication failed"
-        
+
         # Verify start_response was called with correct status
         mock_start_response.assert_called_once()
         call_args = mock_start_response.call_args[0]
-        assert call_args[0] == '401 Unauthorized'
-    
+        assert call_args[0] == "401 Unauthorized"
+
     def test_call_permission_denied(self):
         """Test middleware call with permission denied."""
         environ = self.create_mock_environ()
         mock_start_response = Mock()
-        
+
         self.mock_security_manager.rate_limiter.check_rate_limit.return_value = True
-        
+
         # Mock successful authentication
         auth_result = AuthResult(
             is_valid=True,
             status=AuthStatus.SUCCESS,
             username="test_user",
             roles=["user"],
-            auth_method=AuthMethod.API_KEY
+            auth_method=AuthMethod.API_KEY,
         )
         self.middleware._authenticate_request = Mock(return_value=auth_result)
-        
+
         self.mock_security_manager.check_permissions.return_value = ValidationResult(
             is_valid=False,
             status=ValidationStatus.INVALID,
             error_code=-32007,
-            error_message="Permission denied"
+            error_message="Permission denied",
         )
-        
+
         # Mock permission error response
         mock_response = [b'{"error": "Permission denied"}']
         self.middleware._process_request = Mock(return_value=mock_response)
-        
+
         result = self.middleware(environ, mock_start_response)
-        
+
         # Verify permission error response
         assert result == mock_response
         self.middleware._process_request.assert_called_once()
-    
+
     def test_get_rate_limit_identifier(self):
         """Test getting rate limit identifier from request."""
         mock_request = self.create_mock_request()
         result = self.middleware._get_rate_limit_identifier(mock_request)
         assert result == "127.0.0.1"
-    
+
     def test_get_rate_limit_identifier_with_forwarded_for(self):
         """Test getting rate limit identifier with X-Forwarded-For header."""
         headers = {"X-Forwarded-For": "192.168.1.1, 10.0.0.1"}
         mock_request = self.create_mock_request(headers=headers)
         result = self.middleware._get_rate_limit_identifier(mock_request)
         assert result == "192.168.1.1"
-    
+
     def test_get_rate_limit_identifier_with_real_ip(self):
         """Test getting rate limit identifier with X-Real-IP header."""
         headers = {"X-Real-IP": "192.168.1.100"}
         mock_request = self.create_mock_request(headers=headers)
         result = self.middleware._get_rate_limit_identifier(mock_request)
         assert result == "192.168.1.100"
-    
+
     def test_get_request_path(self):
         """Test getting request path from Flask request."""
         mock_request = self.create_mock_request(path="/api/users")
         result = self.middleware._get_request_path(mock_request)
         assert result == "/api/users"
-    
+
     def test_get_required_permissions_from_request(self):
         """Test getting required permissions from request."""
         mock_request = self.create_mock_request()
         mock_request.required_permissions = ["read", "write"]
-        
+
         result = self.middleware._get_required_permissions(mock_request)
         assert result == ["read", "write"]
-    
+
     def test_get_required_permissions_default(self):
         """Test getting required permissions when not set."""
         mock_request = self.create_mock_request()
@@ -284,178 +296,183 @@ class TestFlaskSecurityMiddleware:
         mock_request.endpoint = mock_endpoint
         result = self.middleware._get_required_permissions(mock_request)
         assert result == []
-    
+
     def test_try_api_key_auth_success(self):
         """Test successful API key authentication."""
         headers = {"X-API-Key": "valid_key"}
         mock_request = self.create_mock_request(headers=headers)
-        
+
         expected_result = AuthResult(
             is_valid=True,
             status=AuthStatus.SUCCESS,
             username="test_user",
             roles=["user"],
-            auth_method=AuthMethod.API_KEY
+            auth_method=AuthMethod.API_KEY,
         )
         self.mock_auth_manager.authenticate_api_key.return_value = expected_result
-        
+
         result = self.middleware._try_api_key_auth(mock_request)
-        
+
         assert result == expected_result
         self.mock_auth_manager.authenticate_api_key.assert_called_once_with("valid_key")
-    
+
     def test_try_api_key_auth_from_authorization_header(self):
         """Test API key authentication from Authorization header."""
         headers = {"Authorization": "Bearer api_key_123"}
         mock_request = self.create_mock_request(headers=headers)
-        
+
         expected_result = AuthResult(
             is_valid=True,
             status=AuthStatus.SUCCESS,
             username="test_user",
             roles=["user"],
-            auth_method=AuthMethod.API_KEY
+            auth_method=AuthMethod.API_KEY,
         )
         self.mock_auth_manager.authenticate_api_key.return_value = expected_result
-        
+
         result = self.middleware._try_api_key_auth(mock_request)
-        
+
         assert result == expected_result
-        self.mock_auth_manager.authenticate_api_key.assert_called_once_with("api_key_123")
-    
+        self.mock_auth_manager.authenticate_api_key.assert_called_once_with(
+            "api_key_123"
+        )
+
     def test_try_api_key_auth_no_key(self):
         """Test API key authentication with no key provided."""
         mock_request = self.create_mock_request()
-        
+
         result = self.middleware._try_api_key_auth(mock_request)
-        
+
         assert result.is_valid is False
         assert result.error_code == -32024
         assert "API key not found" in result.error_message
-    
+
     def test_try_jwt_auth_success(self):
         """Test successful JWT authentication."""
         headers = {"Authorization": "Bearer jwt_token_123"}
         mock_request = self.create_mock_request(headers=headers)
-        
+
         expected_result = AuthResult(
             is_valid=True,
             status=AuthStatus.SUCCESS,
             username="test_user",
             roles=["user"],
-            auth_method=AuthMethod.JWT
+            auth_method=AuthMethod.JWT,
         )
         self.mock_auth_manager.authenticate_jwt_token.return_value = expected_result
-        
+
         result = self.middleware._try_jwt_auth(mock_request)
-        
+
         assert result == expected_result
-        self.mock_auth_manager.authenticate_jwt_token.assert_called_once_with("jwt_token_123")
-    
+        self.mock_auth_manager.authenticate_jwt_token.assert_called_once_with(
+            "jwt_token_123"
+        )
+
     def test_try_jwt_auth_no_token(self):
         """Test JWT authentication with no token provided."""
         mock_request = self.create_mock_request()
-        
+
         result = self.middleware._try_jwt_auth(mock_request)
-        
+
         assert result.is_valid is False
         assert result.error_code == -32025
         assert "JWT token not found" in result.error_message
-    
+
     def test_try_jwt_auth_invalid_header(self):
         """Test JWT authentication with invalid Authorization header."""
         headers = {"Authorization": "Invalid jwt_token_123"}
         mock_request = self.create_mock_request(headers=headers)
-        
+
         result = self.middleware._try_jwt_auth(mock_request)
-        
+
         assert result.is_valid is False
         assert result.error_code == -32025
         assert "JWT token not found" in result.error_message
-    
+
     def test_try_certificate_auth_not_implemented(self):
         """Test certificate authentication (not implemented)."""
         mock_request = self.create_mock_request()
-        
+
         result = self.middleware._try_certificate_auth(mock_request)
-        
+
         assert result.is_valid is False
         assert result.error_code == -32026
         assert "not implemented" in result.error_message
-    
+
     def test_try_basic_auth_not_implemented(self):
         """Test basic authentication (not implemented)."""
         mock_request = self.create_mock_request()
-        
+
         result = self.middleware._try_basic_auth(mock_request)
-        
+
         assert result.is_valid is False
         assert result.error_code == -32027
         assert "Basic authentication credentials not found" in result.error_message
-    
+
     def test_try_auth_method_unsupported(self):
         """Test authentication with unsupported method."""
         mock_request = self.create_mock_request()
-        
+
         result = self.middleware._try_auth_method(mock_request, "unsupported_method")
-        
+
         assert result.is_valid is False
         assert result.error_code == -32022
         assert "Unsupported authentication method" in result.error_message
-    
+
     def test_apply_security_headers(self):
         """Test applying security headers to Flask response."""
         mock_response = Mock(spec=Response)
         mock_response.headers = {}
-        
+
         headers = {
             "X-Content-Type-Options": "nosniff",
             "X-Frame-Options": "DENY",
-            "X-XSS-Protection": "1; mode=block"
+            "X-XSS-Protection": "1; mode=block",
         }
-        
+
         self.middleware._apply_security_headers(mock_response, headers)
-        
+
         assert mock_response.headers["X-Content-Type-Options"] == "nosniff"
         assert mock_response.headers["X-Frame-Options"] == "DENY"
         assert mock_response.headers["X-XSS-Protection"] == "1; mode=block"
-    
+
     def test_create_error_response(self):
         """Test creating error response."""
         from flask import Flask
+
         app = Flask(__name__)
-        
+
         with app.app_context():
             result = self.middleware._create_error_response(400, "Bad request")
-            
+
             assert isinstance(result, Response)
             assert result.status_code == 400
             response_data = json.loads(result.get_data(as_text=True))
             assert response_data["error"] == "Security violation"
             assert response_data["message"] == "Bad request"
-    
+
     def test_rate_limit_response(self):
         """Test creating rate limit response."""
         mock_start_response = Mock()
-        
+
         result = self.middleware._rate_limit_response(mock_start_response)
-        
+
         assert isinstance(result, list)
         assert len(result) == 1
-        response_data = json.loads(result[0].decode('utf-8'))
+        response_data = json.loads(result[0].decode("utf-8"))
         assert response_data["error"] == "Rate limit exceeded"
-        
+
         # Verify start_response was called
         mock_start_response.assert_called_once()
         call_args = mock_start_response.call_args[0]
-        assert call_args[0] == '429 Too Many Requests'
-        
+        assert call_args[0] == "429 Too Many Requests"
+
         # Check for Retry-After header
         headers = call_args[1]
-        retry_after_header = next((h for h in headers if h[0] == 'Retry-After'), None)
+        retry_after_header = next((h for h in headers if h[0] == "Retry-After"), None)
         assert retry_after_header is not None
-        assert retry_after_header[1] == '60'
-    
+        assert retry_after_header[1] == "60"
+
     def test_auth_error_response(self):
         """Test creating authentication error response."""
         auth_result = AuthResult(
@@ -465,118 +482,123 @@ class TestFlaskSecurityMiddleware:
             roles=[],
             auth_method=AuthMethod.API_KEY,
             error_code=-32005,
-            error_message="Invalid API key"
+            error_message="Invalid API key",
         )
-        
+
         mock_start_response = Mock()
         result = self.middleware._auth_error_response(auth_result, mock_start_response)
-        
+
         assert isinstance(result, list)
         assert len(result) == 1
-        response_data = json.loads(result[0].decode('utf-8'))
+        response_data = json.loads(result[0].decode("utf-8"))
         assert response_data["error"] == "Authentication failed"
         assert response_data["message"] == "Invalid API key"
-        
+
         # Verify start_response was called
         mock_start_response.assert_called_once()
         call_args = mock_start_response.call_args[0]
-        assert call_args[0] == '401 Unauthorized'
-        
+        assert call_args[0] == "401 Unauthorized"
+
         # Check for WWW-Authenticate header
         headers = call_args[1]
-        www_auth_header = next((h for h in headers if h[0] == 'WWW-Authenticate'), None)
+        www_auth_header = next((h for h in headers if h[0] == "WWW-Authenticate"), None)
         assert www_auth_header is not None
-        assert www_auth_header[1] == 'Bearer, ApiKey'
-    
+        assert www_auth_header[1] == "Bearer, ApiKey"
+
     def test_permission_error_response(self):
         """Test creating permission error response."""
         mock_start_response = Mock()
-        
+
         result = self.middleware._permission_error_response(mock_start_response)
-        
+
         assert isinstance(result, list)
         assert len(result) == 1
-        response_data = json.loads(result[0].decode('utf-8'))
+        response_data = json.loads(result[0].decode("utf-8"))
         assert response_data["error"] == "Permission denied"
-        assert response_data["message"] == "Insufficient permissions to access this resource"
-        
+        assert (
+            response_data["message"]
+            == "Insufficient permissions to access this resource"
+        )
+
         # Verify start_response was called
         mock_start_response.assert_called_once()
         call_args = mock_start_response.call_args[0]
-        assert call_args[0] == '403 Forbidden'
-    
+        assert call_args[0] == "403 Forbidden"
+
     def test_get_client_ip_from_forwarded_for(self):
         """Test getting client IP from X-Forwarded-For header."""
         headers = {"X-Forwarded-For": "192.168.1.1, 10.0.0.1"}
         mock_request = self.create_mock_request(headers=headers)
-        
+
         result = self.middleware._get_client_ip(mock_request)
         assert result == "192.168.1.1"
-    
+
     def test_get_client_ip_from_real_ip(self):
         """Test getting client IP from X-Real-IP header."""
         headers = {"X-Real-IP": "192.168.1.100"}
         mock_request = self.create_mock_request(headers=headers)
-        
+
         result = self.middleware._get_client_ip(mock_request)
         assert result == "192.168.1.100"
-    
+
     def test_get_client_ip_from_remote_addr(self):
         """Test getting client IP from remote_addr."""
         mock_request = self.create_mock_request()
         mock_request.remote_addr = "192.168.1.50"
-        
+
         result = self.middleware._get_client_ip(mock_request)
         assert result == "192.168.1.50"
-    
+
     def test_get_client_ip_fallback(self):
         """Test getting client IP with fallback."""
         mock_request = self.create_mock_request()
         mock_request.remote_addr = None
-        
+
         result = self.middleware._get_client_ip(mock_request)
         assert result == "127.0.0.1"
-    
+
     def test_get_security_headers(self):
         """Test getting security headers."""
         result = self.middleware._get_security_headers()
-        
+
         assert isinstance(result, list)
         assert all(isinstance(header, tuple) and len(header) == 2 for header in result)
-        
+
         # Check for standard security headers
         header_names = [header[0] for header in result]
-        assert 'X-Content-Type-Options' in header_names
-        assert 'X-Frame-Options' in header_names
-        assert 'X-XSS-Protection' in header_names
-        assert 'Strict-Transport-Security' in header_names
-        assert 'Content-Security-Policy' in header_names
-        assert 'Referrer-Policy' in header_names
-    
+        assert "X-Content-Type-Options" in header_names
+        assert "X-Frame-Options" in header_names
+        assert "X-XSS-Protection" in header_names
+        assert "Strict-Transport-Security" in header_names
+        assert "Content-Security-Policy" in header_names
+        assert "Referrer-Policy" in header_names
+
     def test_get_security_headers_with_custom_headers(self):
         """Test getting security headers with custom headers from config."""
         self.mock_security_manager.config.auth.security_headers = {
             "Custom-Security-Header": "custom_value"
         }
-        
+
         result = self.middleware._get_security_headers()
-        
+
         header_names = [header[0] for header in result]
-        assert 'Custom-Security-Header' in header_names
-        
-        custom_header = next(h for h in result if h[0] == 'Custom-Security-Header')
-        assert custom_header[1] == 'custom_value'
-    
+        assert "Custom-Security-Header" in header_names
+
+        custom_header = next(h for h in result if h[0] == "Custom-Security-Header")
+        assert custom_header[1] == "custom_value"
+
     def test_call_with_exception(self):
         """Test middleware call with exception."""
         environ = self.create_mock_environ()
         mock_start_response = Mock()
-        
+
         # Mock an exception during processing
-        self.mock_security_manager.rate_limiter.check_rate_limit.side_effect = Exception("Test error")
-        
+        self.mock_security_manager.rate_limiter.check_rate_limit.side_effect = (
+            Exception("Test error")
+        )
+
         with pytest.raises(FlaskMiddlewareError) as exc_info:
             self.middleware(environ, mock_start_response)
-        
+
         assert "Middleware processing failed" in str(exc_info.value)
         assert exc_info.value.error_code == -32003