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

mcp_security_framework/examples/flask_example.py

@@ -17,34 +17,39 @@ Version: 1.0.0
 License: MIT
 """
 
-import os
 import json
 import logging
-from typing import Dict, List, Any, Optional
+import os
 from datetime import datetime, timedelta, timezone
+from typing import Any, Dict, List, Optional
 
-from flask import Flask, request, jsonify, g, session, current_app
-from flask_cors import CORS
 import werkzeug
+from flask import Flask, current_app, g, jsonify, request, session
+from flask_cors import CORS
 
-from mcp_security_framework.core.security_manager import SecurityManager
+from mcp_security_framework.constants import (
+    AUTH_METHODS,
+    DEFAULT_CLIENT_IP,
+    DEFAULT_SECURITY_HEADERS,
+    HTTP_FORBIDDEN,
+    HTTP_TOO_MANY_REQUESTS,
+    HTTP_UNAUTHORIZED,
+    ErrorCodes,
+)
 from mcp_security_framework.core.auth_manager import AuthManager
-from mcp_security_framework.core.ssl_manager import SSLManager
 from mcp_security_framework.core.permission_manager import PermissionManager
 from mcp_security_framework.core.rate_limiter import RateLimiter
-from mcp_security_framework.schemas.config import SecurityConfig, AuthConfig, SSLConfig
-from mcp_security_framework.schemas.models import AuthResult, AuthStatus, AuthMethod
+from mcp_security_framework.core.security_manager import SecurityManager
+from mcp_security_framework.core.ssl_manager import SSLManager
 from mcp_security_framework.middleware.flask_middleware import FlaskSecurityMiddleware
-from mcp_security_framework.constants import (
-    DEFAULT_CLIENT_IP, DEFAULT_SECURITY_HEADERS, AUTH_METHODS,
-    ErrorCodes, HTTP_UNAUTHORIZED, HTTP_FORBIDDEN, HTTP_TOO_MANY_REQUESTS
-)
+from mcp_security_framework.schemas.config import AuthConfig, SecurityConfig, SSLConfig
+from mcp_security_framework.schemas.models import AuthMethod, AuthResult, AuthStatus
 
 
 class FlaskExample:
     """
     Complete Flask Example with Security Framework Implementation
-    
+
     This class demonstrates a production-ready Flask application
     with comprehensive security features including:
     - Multi-method authentication (API Key, JWT, Certificate)
@@ -54,11 +59,11 @@ class FlaskExample:
     - Security headers and CORS
     - Comprehensive logging and monitoring
     """
-    
+
     def __init__(self, config_path: Optional[str] = None):
         """
         Initialize Flask example with security configuration.
-        
+
         Args:
             config_path: Path to security configuration file
         """
@@ -68,37 +73,41 @@ class FlaskExample:
         self._setup_middleware()
         self._setup_routes()
         self._setup_error_handlers()
-    
+
     def _load_config(self, config_path: Optional[str]) -> SecurityConfig:
         """
         Load security configuration from file or create default.
-        
+
         Args:
             config_path: Path to configuration file
-            
+
         Returns:
             SecurityConfig: Loaded configuration
         """
         if config_path and os.path.exists(config_path):
-            with open(config_path, 'r') as f:
+            with open(config_path, "r") as f:
                 config_data = json.load(f)
             return SecurityConfig(**config_data)
-        
+
         # Create production-ready default configuration
         return SecurityConfig(
             auth=AuthConfig(
                 enabled=True,
-                methods=[AUTH_METHODS["API_KEY"], AUTH_METHODS["JWT"], AUTH_METHODS["CERTIFICATE"]],
+                methods=[
+                    AUTH_METHODS["API_KEY"],
+                    AUTH_METHODS["JWT"],
+                    AUTH_METHODS["CERTIFICATE"],
+                ],
                 api_keys={
                     "admin_key_123": {"username": "admin", "roles": ["admin", "user"]},
                     "user_key_456": {"username": "user", "roles": ["user"]},
-                    "readonly_key_789": {"username": "readonly", "roles": ["readonly"]}
+                    "readonly_key_789": {"username": "readonly", "roles": ["readonly"]},
                 },
                 jwt_secret="your-super-secret-jwt-key-change-in-production",
                 jwt_algorithm="HS256",
                 jwt_expiry_hours=24,
                 public_paths=["/health", "/docs", "/metrics"],
-                security_headers=DEFAULT_SECURITY_HEADERS
+                security_headers=DEFAULT_SECURITY_HEADERS,
             ),
             ssl=SSLConfig(
                 enabled=False,  # Disable SSL for testing
@@ -106,7 +115,7 @@ class FlaskExample:
                 key_file=None,
                 ca_cert_file=None,
                 verify_mode="CERT_REQUIRED",
-                min_version="TLSv1.2"
+                min_version="TLSv1.2",
             ),
             rate_limit={
                 "enabled": True,
@@ -119,16 +128,16 @@ class FlaskExample:
                     "host": "localhost",
                     "port": 6379,
                     "db": 0,
-                    "password": None
+                    "password": None,
                 },
                 "exempt_paths": ["/health", "/metrics"],
-                "exempt_roles": ["admin"]
+                "exempt_roles": ["admin"],
             },
             permissions={
                 "enabled": True,
                 "roles_file": "config/roles.json",
                 "default_role": "user",
-                "hierarchy_enabled": True
+                "hierarchy_enabled": True,
             },
             logging={
                 "enabled": True,
@@ -138,100 +147,108 @@ class FlaskExample:
                 "max_file_size": 10,
                 "backup_count": 5,
                 "console_output": True,
-                "json_format": False
-            }
+                "json_format": False,
+            },
         )
-    
+
     def _create_flask_app(self) -> Flask:
         """
         Create Flask application with security features.
-        
+
         Returns:
             Flask: Configured Flask application
         """
         app = Flask(__name__)
-        app.config['SECRET_KEY'] = 'your-secret-key-change-in-production'
-        app.config['JSON_SORT_KEYS'] = False
-        
+        app.config["SECRET_KEY"] = "your-secret-key-change-in-production"
+        app.config["JSON_SORT_KEYS"] = False
+
         # Add CORS support
         CORS(app, origins=["https://trusted-domain.com"])
-        
+
         return app
-    
+
     def _setup_middleware(self):
         """Setup security middleware."""
         # For now, skip middleware setup to avoid WSGI issues
         # and use fallback authentication for testing
         print("Middleware setup skipped - using fallback authentication")
         self._setup_test_authentication()
-    
+
     def _setup_test_authentication(self):
         """Setup authentication for testing environment."""
         security_manager = self.security_manager
-        
+
         def get_current_user():
             """Get current user from request headers."""
             api_key = request.headers.get("X-API-Key")
             if api_key:
                 try:
-                    auth_result = security_manager.auth_manager.authenticate_api_key(api_key)
+                    auth_result = security_manager.auth_manager.authenticate_api_key(
+                        api_key
+                    )
                     if auth_result.is_valid:
                         return auth_result
                 except Exception as e:
                     print(f"Authentication error: {e}")
                     pass
-            
+
             # Check for JWT token
             auth_header = request.headers.get("Authorization")
             if auth_header and auth_header.startswith("Bearer "):
                 token = auth_header.split(" ")[1]
                 try:
-                    auth_result = security_manager.auth_manager.authenticate_jwt_token(token)
+                    auth_result = security_manager.auth_manager.authenticate_jwt_token(
+                        token
+                    )
                     if auth_result.is_valid:
                         return auth_result
                 except Exception as e:
                     print(f"JWT authentication error: {e}")
                     pass
-            
+
             return None
-        
+
         # Store function for use in routes
         self.get_current_user = get_current_user
-        
+
         # Make function available in app context
         self.app.get_current_user = get_current_user
-    
+
     def _setup_routes(self):
         """Setup application routes with security."""
-        
+
         # Get reference to authentication function
         get_current_user_func = self.get_current_user
-        
-        @self.app.route('/health', methods=['GET'])
+
+        @self.app.route("/health", methods=["GET"])
         def health_check():
             """Health check endpoint (public)."""
-            return jsonify({
-                "status": "healthy",
-                "timestamp": datetime.now(timezone.utc).isoformat(),
-                "version": "1.0.0"
-            })
-        
-        @self.app.route('/metrics', methods=['GET'])
+            return jsonify(
+                {
+                    "status": "healthy",
+                    "timestamp": datetime.now(timezone.utc).isoformat(),
+                    "version": "1.0.0",
+                }
+            )
+
+        @self.app.route("/metrics", methods=["GET"])
         def metrics():
             """Metrics endpoint (public)."""
-            return jsonify({
-                "requests_total": 1000,
-                "requests_per_minute": 60,
-                "active_connections": 25,
-                "uptime_seconds": 3600
-            })
-        
-        @self.app.route('/api/v1/users/me', methods=['GET'])
+            return jsonify(
+                {
+                    "requests_total": 1000,
+                    "requests_per_minute": 60,
+                    "active_connections": 25,
+                    "uptime_seconds": 3600,
+                }
+            )
+
+        @self.app.route("/api/v1/users/me", methods=["GET"])
         def get_current_user_route():
             """Get current user information (authenticated)."""
             # Try to get user info from middleware
-            user_info = getattr(g, 'user_info', None)
-            
+            user_info = getattr(g, "user_info", None)
+
             # If middleware didn't set user_info, try authentication
             if not user_info:
                 auth_result = get_current_user_func()
@@ -239,24 +256,26 @@ class FlaskExample:
                     user_info = {
                         "username": auth_result.username,
                         "roles": auth_result.roles,
-                        "permissions": auth_result.permissions
+                        "permissions": auth_result.permissions,
                     }
                 else:
                     return jsonify({"error": "Authentication required"}), 401
-            
-            return jsonify({
-                "username": user_info.get("username"),
-                "roles": user_info.get("roles", []),
-                "permissions": list(user_info.get("permissions", [])),
-                "last_login": datetime.now(timezone.utc).isoformat()
-            })
-        
-        @self.app.route('/api/v1/admin/users', methods=['GET'])
+
+            return jsonify(
+                {
+                    "username": user_info.get("username"),
+                    "roles": user_info.get("roles", []),
+                    "permissions": list(user_info.get("permissions", [])),
+                    "last_login": datetime.now(timezone.utc).isoformat(),
+                }
+            )
+
+        @self.app.route("/api/v1/admin/users", methods=["GET"])
         def get_all_users():
             """Get all users (admin only)."""
             # Try to get user info from middleware
-            user_info = getattr(g, 'user_info', None)
-            
+            user_info = getattr(g, "user_info", None)
+
             # If middleware didn't set user_info, try authentication
             if not user_info:
                 auth_result = get_current_user_func()
@@ -264,29 +283,35 @@ class FlaskExample:
                     user_info = {
                         "username": auth_result.username,
                         "roles": auth_result.roles,
-                        "permissions": auth_result.permissions
+                        "permissions": auth_result.permissions,
                     }
                 else:
                     return jsonify({"error": "Authentication required"}), 401
-            
+
             # Check admin permission
             if "admin" not in user_info.get("roles", []):
                 return jsonify({"error": "Admin access required"}), 403
-            
-            return jsonify({
-                "users": [
-                    {"username": "admin", "roles": ["admin"], "status": "active"},
-                    {"username": "user", "roles": ["user"], "status": "active"},
-                    {"username": "readonly", "roles": ["readonly"], "status": "active"}
-                ]
-            })
-        
-        @self.app.route('/api/v1/data', methods=['POST'])
+
+            return jsonify(
+                {
+                    "users": [
+                        {"username": "admin", "roles": ["admin"], "status": "active"},
+                        {"username": "user", "roles": ["user"], "status": "active"},
+                        {
+                            "username": "readonly",
+                            "roles": ["readonly"],
+                            "status": "active",
+                        },
+                    ]
+                }
+            )
+
+        @self.app.route("/api/v1/data", methods=["POST"])
         def create_data():
             """Create data (authenticated users)."""
             # Try to get user info from middleware
-            user_info = getattr(g, 'user_info', None)
-            
+            user_info = getattr(g, "user_info", None)
+
             # If middleware didn't set user_info, try authentication
             if not user_info:
                 auth_result = get_current_user_func()
@@ -294,30 +319,32 @@ class FlaskExample:
                     user_info = {
                         "username": auth_result.username,
                         "roles": auth_result.roles,
-                        "permissions": auth_result.permissions
+                        "permissions": auth_result.permissions,
                     }
                 else:
                     return jsonify({"error": "Authentication required"}), 401
-            
+
             # Check write permission
             if "readonly" in user_info.get("roles", []):
                 return jsonify({"error": "Write permission required"}), 403
-            
+
             # Process request data
             data = request.get_json()
-            return jsonify({
-                "id": "data_123",
-                "created_by": user_info.get("username"),
-                "data": data,
-                "created_at": datetime.now(timezone.utc).isoformat()
-            })
-        
-        @self.app.route('/api/v1/data/<data_id>', methods=['GET'])
+            return jsonify(
+                {
+                    "id": "data_123",
+                    "created_by": user_info.get("username"),
+                    "data": data,
+                    "created_at": datetime.now(timezone.utc).isoformat(),
+                }
+            )
+
+        @self.app.route("/api/v1/data/<data_id>", methods=["GET"])
         def get_data(data_id):
             """Get data by ID (authenticated users)."""
             # Try to get user info from middleware
-            user_info = getattr(g, 'user_info', None)
-            
+            user_info = getattr(g, "user_info", None)
+
             # If middleware didn't set user_info, try authentication
             if not user_info:
                 auth_result = get_current_user_func()
@@ -325,91 +352,129 @@ class FlaskExample:
                     user_info = {
                         "username": auth_result.username,
                         "roles": auth_result.roles,
-                        "permissions": auth_result.permissions
+                        "permissions": auth_result.permissions,
                     }
                 else:
                     return jsonify({"error": "Authentication required"}), 401
-            
-            return jsonify({
-                "id": data_id,
-                "data": {"example": "data"},
-                "created_by": user_info.get("username"),
-                "created_at": datetime.now(timezone.utc).isoformat()
-            })
-    
+
+            return jsonify(
+                {
+                    "id": data_id,
+                    "data": {"example": "data"},
+                    "created_by": user_info.get("username"),
+                    "created_at": datetime.now(timezone.utc).isoformat(),
+                }
+            )
+
     def _setup_error_handlers(self):
         """Setup custom error handlers."""
-        
+
         @self.app.errorhandler(401)
         def unauthorized(error):
             """Handle unauthorized errors."""
-            return jsonify({
-                "error": "Unauthorized",
-                "message": "Authentication required",
-                "status_code": 401,
-                "timestamp": datetime.now(timezone.utc).isoformat(),
-                "path": request.path
-            }), 401
-        
+            return (
+                jsonify(
+                    {
+                        "error": "Unauthorized",
+                        "message": "Authentication required",
+                        "status_code": 401,
+                        "timestamp": datetime.now(timezone.utc).isoformat(),
+                        "path": request.path,
+                    }
+                ),
+                401,
+            )
+
         @self.app.errorhandler(403)
         def forbidden(error):
             """Handle forbidden errors."""
-            return jsonify({
-                "error": "Forbidden",
-                "message": "Insufficient permissions",
-                "status_code": 403,
-                "timestamp": datetime.now(timezone.utc).isoformat(),
-                "path": request.path
-            }), 403
-        
+            return (
+                jsonify(
+                    {
+                        "error": "Forbidden",
+                        "message": "Insufficient permissions",
+                        "status_code": 403,
+                        "timestamp": datetime.now(timezone.utc).isoformat(),
+                        "path": request.path,
+                    }
+                ),
+                403,
+            )
+
         @self.app.errorhandler(404)
         def not_found(error):
             """Handle not found errors."""
-            return jsonify({
-                "error": "Not Found",
-                "message": "Resource not found",
-                "status_code": 404,
-                "timestamp": datetime.now(timezone.utc).isoformat(),
-                "path": request.path
-            }), 404
-        
+            return (
+                jsonify(
+                    {
+                        "error": "Not Found",
+                        "message": "Resource not found",
+                        "status_code": 404,
+                        "timestamp": datetime.now(timezone.utc).isoformat(),
+                        "path": request.path,
+                    }
+                ),
+                404,
+            )
+
         @self.app.errorhandler(429)
         def too_many_requests(error):
             """Handle rate limit errors."""
-            return jsonify({
-                "error": "Too Many Requests",
-                "message": "Rate limit exceeded",
-                "status_code": 429,
-                "timestamp": datetime.now(timezone.utc).isoformat(),
-                "path": request.path
-            }), 429
-        
+            return (
+                jsonify(
+                    {
+                        "error": "Too Many Requests",
+                        "message": "Rate limit exceeded",
+                        "status_code": 429,
+                        "timestamp": datetime.now(timezone.utc).isoformat(),
+                        "path": request.path,
+                    }
+                ),
+                429,
+            )
+
         @self.app.errorhandler(500)
         def internal_error(error):
             """Handle internal server errors."""
-            return jsonify({
-                "error": "Internal Server Error",
-                "message": "An unexpected error occurred",
-                "status_code": 500,
-                "timestamp": datetime.now(timezone.utc).isoformat(),
-                "path": request.path
-            }), 500
-        
+            return (
+                jsonify(
+                    {
+                        "error": "Internal Server Error",
+                        "message": "An unexpected error occurred",
+                        "status_code": 500,
+                        "timestamp": datetime.now(timezone.utc).isoformat(),
+                        "path": request.path,
+                    }
+                ),
+                500,
+            )
+
         @self.app.errorhandler(Exception)
         def handle_exception(error):
             """Handle general exceptions."""
-            return jsonify({
-                "error": "Internal Server Error",
-                "message": "An unexpected error occurred",
-                "status_code": 500,
-                "timestamp": datetime.now(timezone.utc).isoformat(),
-                "path": request.path
-            }), 500
-    
-    def run(self, host: str = "0.0.0.0", port: int = 5000, ssl_context: Optional[tuple] = None, debug: bool = False):
+            return (
+                jsonify(
+                    {
+                        "error": "Internal Server Error",
+                        "message": "An unexpected error occurred",
+                        "status_code": 500,
+                        "timestamp": datetime.now(timezone.utc).isoformat(),
+                        "path": request.path,
+                    }
+                ),
+                500,
+            )
+
+    def run(
+        self,
+        host: str = "0.0.0.0",
+        port: int = 5000,
+        ssl_context: Optional[tuple] = None,
+        debug: bool = False,
+    ):
         """
         Run the Flask application with security features.
-        
+
         Args:
             host: Host to bind to
             port: Port to bind to
@@ -420,81 +485,84 @@ class FlaskExample:
         print(f"SSL Enabled: {self.config.ssl.enabled}")
         print(f"Authentication Methods: {self.config.auth.methods}")
         print(f"Rate Limiting: {self.config.rate_limit.enabled}")
-        
+
         if self.config.ssl.enabled and not ssl_context:
             ssl_context = (self.config.ssl.cert_file, self.config.ssl.key_file)
-        
-        self.app.run(
-            host=host,
-            port=port,
-            ssl_context=ssl_context,
-            debug=debug
-        )
+
+        self.app.run(host=host, port=port, ssl_context=ssl_context, debug=debug)
 
 
 # Example usage and testing
 class FlaskExampleTest:
     """Test class for Flask example functionality."""
-    
+
     @staticmethod
     def test_authentication():
         """Test authentication functionality."""
         example = FlaskExample()
-        
+
         # Test API key authentication
-        auth_result = example.security_manager.auth_manager.authenticate_api_key("admin_key_123")
+        auth_result = example.security_manager.auth_manager.authenticate_api_key(
+            "admin_key_123"
+        )
         assert auth_result.is_valid
         assert auth_result.username == "admin"
         assert "admin" in auth_result.roles
-        
+
         print("✅ API Key authentication test passed")
-    
+
     @staticmethod
     def test_rate_limiting():
         """Test rate limiting functionality."""
         example = FlaskExample()
-        
+
         # Test rate limiting
         identifier = "test_user"
         for i in range(5):
-            is_allowed = example.security_manager.rate_limiter.check_rate_limit(identifier)
+            is_allowed = example.security_manager.rate_limiter.check_rate_limit(
+                identifier
+            )
             print(f"Request {i+1}: {'Allowed' if is_allowed else 'Blocked'}")
-        
+
         print("✅ Rate limiting test completed")
-    
+
     @staticmethod
     def test_permissions():
         """Test permission checking."""
         example = FlaskExample()
-        
+
         # Test admin permissions
         admin_roles = ["admin"]
         user_roles = ["user"]
         readonly_roles = ["readonly"]
-        
+
         # Admin should have all permissions
         admin_result = example.security_manager.permission_manager.validate_access(
             admin_roles, ["read", "write", "delete"]
         )
         assert admin_result.is_valid
-        
+
         # User should have read and write permissions
         user_result = example.security_manager.permission_manager.validate_access(
             user_roles, ["read", "write"]
         )
         assert user_result.is_valid
-        
+
         # Readonly should only have read permission
-        readonly_read_result = example.security_manager.permission_manager.validate_access(
-            readonly_roles, ["read"]
+        readonly_read_result = (
+            example.security_manager.permission_manager.validate_access(
+                readonly_roles, ["read"]
+            )
         )
         assert readonly_read_result.is_valid
-        
-        readonly_write_result = example.security_manager.permission_manager.validate_access(
-            readonly_roles, ["write"]
+
+        readonly_write_result = (
+            example.security_manager.permission_manager.validate_access(
+                readonly_roles, ["write"]
+            )
         )
         assert not readonly_write_result.is_valid
-        
+
         print("✅ Permission checking test passed")
 
 
@@ -504,43 +572,46 @@ if __name__ == "__main__":
     FlaskExampleTest.test_authentication()
     FlaskExampleTest.test_rate_limiting()
     FlaskExampleTest.test_permissions()
-    
+
     # Start server in background thread for testing
     print("\nStarting Flask Example Server in background...")
     example = FlaskExample()
-    
+
     import threading
     import time
+
     import requests
-    
+
     # Start server in background thread
     server_thread = threading.Thread(target=example.run, daemon=True)
     server_thread.start()
-    
+
     # Wait for server to start
     time.sleep(3)
-    
+
     try:
         # Test server endpoints
         print("Testing server endpoints...")
-        
+
         # Test health endpoint
         response = requests.get("http://localhost:5000/health", timeout=5)
         print(f"Health endpoint: {response.status_code}")
-        
+
         # Test metrics endpoint
         response = requests.get("http://localhost:5000/metrics", timeout=5)
         print(f"Metrics endpoint: {response.status_code}")
-        
+
         # Test protected endpoint with API key
         headers = {"X-API-Key": "admin_key_123"}
-        response = requests.get("http://localhost:5000/api/v1/users/me", headers=headers, timeout=5)
+        response = requests.get(
+            "http://localhost:5000/api/v1/users/me", headers=headers, timeout=5
+        )
         print(f"Protected endpoint: {response.status_code}")
-        
+
         print("✅ Server testing completed successfully")
-        
+
     except requests.exceptions.RequestException as e:
         print(f"⚠️  Server testing failed: {e}")
-    
+
     # Server will automatically stop when main thread exits
     print("Flask example completed")