mcp-proxy-adapter 6.1.1__py3-none-any.whl → 6.2.1__py3-none-any.whl

mcp_proxy_adapter/examples/security_test_client.py

@@ -1,18 +1,15 @@
 #!/usr/bin/env python3
 """
 Security Test Client for MCP Proxy Adapter
-
 This client tests various security configurations including:
 - Basic HTTP
 - HTTP + Token authentication
 - HTTPS
 - HTTPS + Token authentication
 - mTLS with certificate authentication
-
 Author: Vasiliy Zdanovskiy
 email: vasilyvz@gmail.com
 """
-
 import asyncio
 import json
 import os
@@ -22,20 +19,15 @@ import time
 from pathlib import Path
 from typing import Dict, List, Optional, Any
 from dataclasses import dataclass
-
 import aiohttp
 from aiohttp import ClientSession, ClientTimeout, TCPConnector
-
 # Add project root to path for imports
 project_root = Path(__file__).parent.parent.parent
 current_dir = Path(__file__).parent
 parent_dir = current_dir.parent
-
 sys.path.insert(0, str(project_root))
 sys.path.insert(0, str(current_dir))
 sys.path.insert(0, str(parent_dir))
-
-
 @dataclass
 class TestResult:
     """Test result data class."""
@@ -47,11 +39,8 @@ class TestResult:
     response_data: Optional[Dict] = None
     error_message: Optional[str] = None
     duration: float = 0.0
-
-
 class SecurityTestClient:
     """Security test client for comprehensive testing."""
-    
     def __init__(self, base_url: str = "http://localhost:8000"):
         """Initialize security test client."""
         self.base_url = base_url
@@ -62,7 +51,6 @@ class SecurityTestClient:
         self.ssl_manager = None
         self.auth_manager = None
         self.test_results: List[TestResult] = []
-        
         # Test tokens
         self.test_tokens = {
             "admin": "test-token-123",
@@ -72,7 +60,6 @@ class SecurityTestClient:
             "proxy": "proxy-token-123",
             "invalid": "invalid-token-999"
         }
-        
         # Test certificates
         self.test_certificates = {
             "admin": {
@@ -88,68 +75,51 @@ class SecurityTestClient:
                 "key": "mcp_proxy_adapter/examples/certs/readonly_key.pem"
             }
         }
-    
     async def __aenter__(self):
         """Async context manager entry."""
         timeout = ClientTimeout(total=30)
-        
         # Create SSL context for HTTPS connections
         ssl_context = self.create_ssl_context()
         connector = TCPConnector(ssl=ssl_context)
-        
         self.session = ClientSession(timeout=timeout, connector=connector)
         return self
-    
     def create_ssl_context_for_mtls(self) -> ssl.SSLContext:
         """Create SSL context for mTLS connections."""
         ssl_context = ssl.create_default_context()
-        
         # For mTLS testing
         ssl_context.check_hostname = False
         ssl_context.verify_mode = ssl.CERT_NONE
-        
         # Load client certificate and key
         cert_file = "./certs/user_cert.pem"
         key_file = "./certs/user_key.pem"
         ca_cert_file = "./certs/ca_cert.pem"
-        
         if os.path.exists(cert_file) and os.path.exists(key_file):
             ssl_context.load_cert_chain(certfile=cert_file, keyfile=key_file)
-        
         if os.path.exists(ca_cert_file):
             ssl_context.load_verify_locations(cafile=ca_cert_file)
-        
         return ssl_context
-    
     async def __aexit__(self, exc_type, exc_val, exc_tb):
         """Async context manager exit."""
         if self.session:
             await self.session.close()
-    
     def create_ssl_context(self, cert_file: Optional[str] = None,
                           key_file: Optional[str] = None,
                           ca_cert_file: Optional[str] = None) -> ssl.SSLContext:
         """Create SSL context for client."""
         ssl_context = ssl.create_default_context()
-        
         # For testing with self-signed certificates
         ssl_context.check_hostname = False
         ssl_context.verify_mode = ssl.CERT_NONE
-        
         if cert_file and key_file:
             ssl_context.load_cert_chain(certfile=cert_file, keyfile=key_file)
-        
         if ca_cert_file:
             ssl_context.load_verify_locations(cafile=ca_cert_file)
             # For testing, still don't verify
             ssl_context.verify_mode = ssl.CERT_NONE
-            
         return ssl_context
-    
     def create_auth_headers(self, auth_type: str, **kwargs) -> Dict[str, str]:
         """Create authentication headers."""
         headers = {"Content-Type": "application/json"}
-        
         if auth_type == "api_key":
             token = kwargs.get("token", "test-token-123")  # Use correct token
             headers["X-API-Key"] = token  # Use X-API-Key header
@@ -163,20 +133,15 @@ class SecurityTestClient:
             # For mTLS, we need to use client certificates
             # This is handled by SSL context, not headers
             pass
-        
         return headers
-    
     async def test_health_check(self, server_url: str, auth_type: str = "none", **kwargs) -> TestResult:
         """Test health check endpoint."""
         start_time = time.time()
         test_name = f"Health Check ({auth_type})"
-        
         try:
             headers = self.create_auth_headers(auth_type, **kwargs)
-            
             async with self.session.get(f"{server_url}/health", headers=headers) as response:
                 duration = time.time() - start_time
-                
                 if response.status == 200:
                     data = await response.json()
                     return TestResult(
@@ -199,7 +164,6 @@ class SecurityTestClient:
                         error_message=f"Health check failed: {error_text}",
                         duration=duration
                     )
-        
         except Exception as e:
             duration = time.time() - start_time
             return TestResult(
@@ -210,15 +174,12 @@ class SecurityTestClient:
                 error_message=f"Health check error: {str(e)}",
                 duration=duration
             )
-    
     async def test_echo_command(self, server_url: str, auth_type: str = "none", **kwargs) -> TestResult:
         """Test echo command."""
         start_time = time.time()
         test_name = f"Echo Command ({auth_type})"
-        
         try:
             headers = self.create_auth_headers(auth_type, **kwargs)
-            
             data = {
                 "jsonrpc": "2.0",
                 "method": "echo",
@@ -227,12 +188,10 @@ class SecurityTestClient:
                 },
                 "id": 1
             }
-            
-            async with self.session.post(f"{server_url}/cmd", 
-                                        headers=headers, 
+            async with self.session.post(f"{server_url}/cmd",
+                                        headers=headers,
                                         json=data) as response:
                 duration = time.time() - start_time
-                
                 if response.status == 200:
                     data = await response.json()
                     return TestResult(
@@ -255,7 +214,6 @@ class SecurityTestClient:
                         error_message=f"Echo command failed: {error_text}",
                         duration=duration
                     )
-        
         except Exception as e:
             duration = time.time() - start_time
             return TestResult(
@@ -266,27 +224,22 @@ class SecurityTestClient:
                 error_message=f"Echo command error: {str(e)}",
                 duration=duration
             )
-    
     async def test_security_command(self, server_url: str, auth_type: str = "none", **kwargs) -> TestResult:
         """Test security command."""
         start_time = time.time()
         test_name = f"Security Command ({auth_type})"
-        
         try:
             headers = self.create_auth_headers(auth_type, **kwargs)
-            
             data = {
                 "jsonrpc": "2.0",
                 "method": "health",
                 "params": {},
                 "id": 2
             }
-            
-            async with self.session.post(f"{server_url}/cmd", 
-                                        headers=headers, 
+            async with self.session.post(f"{server_url}/cmd",
+                                        headers=headers,
                                         json=data) as response:
                 duration = time.time() - start_time
-                
                 if response.status == 200:
                     data = await response.json()
                     return TestResult(
@@ -309,7 +262,6 @@ class SecurityTestClient:
                         error_message=f"Security command failed: {error_text}",
                         duration=duration
                     )
-        
         except Exception as e:
             duration = time.time() - start_time
             return TestResult(
@@ -320,15 +272,12 @@ class SecurityTestClient:
                 error_message=f"Security command error: {str(e)}",
                 duration=duration
             )
-    
     async def test_health(self) -> TestResult:
         """Test health endpoint."""
         return await self.test_health_check(self.base_url, "none")
-    
     async def test_command_execution(self) -> TestResult:
         """Test command execution."""
         return await self.test_echo_command(self.base_url, "none")
-    
     async def test_authentication(self) -> TestResult:
         """Test authentication."""
         if "api_key" in self.auth_methods:
@@ -346,38 +295,29 @@ class SecurityTestClient:
                 success=False,
                 error_message="No authentication method available"
             )
-    
     async def test_negative_authentication(self) -> TestResult:
         """Test negative authentication (should fail)."""
         return await self.test_echo_command(self.base_url, "api_key", token="invalid-token")
-    
     async def test_no_auth_required(self) -> TestResult:
         """Test that no authentication is required."""
         return await self.test_echo_command(self.base_url, "none")
-    
-
-    
     async def test_negative_auth(self, server_url: str, auth_type: str = "none", **kwargs) -> TestResult:
         """Test negative authentication scenarios."""
         start_time = time.time()
         test_name = f"Negative Auth ({auth_type})"
-        
         try:
             # Use invalid token
             headers = self.create_auth_headers("api_key", token="invalid-token-999")
-            
             data = {
                 "jsonrpc": "2.0",
                 "method": "echo",
                 "params": {"message": "Should fail"},
                 "id": 3
             }
-            
-            async with self.session.post(f"{server_url}/cmd", 
-                                        headers=headers, 
+            async with self.session.post(f"{server_url}/cmd",
+                                        headers=headers,
                                         json=data) as response:
                 duration = time.time() - start_time
-                
                 # Expected to fail with 401
                 if response.status == 401:
                     return TestResult(
@@ -399,7 +339,6 @@ class SecurityTestClient:
                         error_message=f"Expected 401, got {response.status}",
                         duration=duration
                     )
-        
         except Exception as e:
             duration = time.time() - start_time
             return TestResult(
@@ -410,30 +349,25 @@ class SecurityTestClient:
                 error_message=f"Negative auth error: {str(e)}",
                 duration=duration
             )
-    
     async def test_role_based_access(self, server_url: str, auth_type: str = "none", **kwargs) -> TestResult:
         """Test role-based access control."""
         start_time = time.time()
         test_name = f"Role-Based Access ({auth_type})"
-        
         try:
             # Test with different roles
             role = kwargs.get("role", "user")
             token = self.test_tokens.get(role, self.test_tokens["user"])
             headers = self.create_auth_headers("api_key", token=token)
-            
             data = {
                 "jsonrpc": "2.0",
                 "method": "echo",
                 "params": {"message": f"Testing {role} role"},
                 "id": 4
             }
-            
-            async with self.session.post(f"{server_url}/cmd", 
-                                        headers=headers, 
+            async with self.session.post(f"{server_url}/cmd",
+                                        headers=headers,
                                         json=data) as response:
                 duration = time.time() - start_time
-                
                 if response.status == 200:
                     data = await response.json()
                     return TestResult(
@@ -456,7 +390,6 @@ class SecurityTestClient:
                         error_message=f"Role-based access failed: {error_text}",
                         duration=duration
                     )
-        
         except Exception as e:
             duration = time.time() - start_time
             return TestResult(
@@ -467,31 +400,26 @@ class SecurityTestClient:
                 error_message=f"Role-based access error: {str(e)}",
                 duration=duration
             )
-    
     async def test_role_permissions(self, server_url: str, auth_type: str = "none", **kwargs) -> TestResult:
         """Test role permissions with role_test command."""
         start_time = time.time()
         test_name = f"Role Permissions Test ({auth_type})"
-        
         try:
             # Test with different roles and actions
             role = kwargs.get("role", "user")
             action = kwargs.get("action", "read")
             token = self.test_tokens.get(role, self.test_tokens["user"])
             headers = self.create_auth_headers("api_key", token=token)
-            
             data = {
                 "jsonrpc": "2.0",
                 "method": "role_test",
                 "params": {"action": action},
                 "id": 5
             }
-            
-            async with self.session.post(f"{server_url}/cmd", 
-                                        headers=headers, 
+            async with self.session.post(f"{server_url}/cmd",
+                                        headers=headers,
                                         json=data) as response:
                 duration = time.time() - start_time
-                
                 if response.status == 200:
                     data = await response.json()
                     return TestResult(
@@ -514,7 +442,6 @@ class SecurityTestClient:
                         error_message=f"Role permissions test failed: {error_text}",
                         duration=duration
                     )
-        
         except Exception as e:
             duration = time.time() - start_time
             return TestResult(
@@ -525,26 +452,22 @@ class SecurityTestClient:
                 error_message=f"Role permissions test error: {str(e)}",
                 duration=duration
             )
-    
     async def test_multiple_roles(self, server_url: str, auth_type: str = "none", **kwargs) -> TestResult:
         """Test multiple roles with different permissions."""
         start_time = time.time()
         test_name = f"Multiple Roles Test ({auth_type})"
-        
         try:
             # Test admin role (should have all permissions)
             admin_token = self.test_tokens.get("admin", "admin-token-123")
             admin_headers = self.create_auth_headers("api_key", token=admin_token)
-            
             admin_data = {
                 "jsonrpc": "2.0",
                 "method": "role_test",
                 "params": {"action": "manage"},
                 "id": 6
             }
-            
-            async with self.session.post(f"{server_url}/cmd", 
-                                        headers=admin_headers, 
+            async with self.session.post(f"{server_url}/cmd",
+                                        headers=admin_headers,
                                         json=admin_data) as response:
                 if response.status != 200:
                     return TestResult(
@@ -556,22 +479,18 @@ class SecurityTestClient:
                         error_message="Admin role test failed",
                         duration=time.time() - start_time
                     )
-                
                 # Test readonly role (should only have read permission)
                 readonly_token = self.test_tokens.get("readonly", "readonly-token-123")
                 readonly_headers = self.create_auth_headers("api_key", token=readonly_token)
-                
                 readonly_data = {
                     "jsonrpc": "2.0",
                     "method": "role_test",
                     "params": {"action": "write"},
                 }
-                
-                async with self.session.post(f"{server_url}/cmd", 
-                                            headers=readonly_headers, 
+                async with self.session.post(f"{server_url}/cmd",
+                                            headers=readonly_headers,
                                             json=readonly_data) as response:
                     duration = time.time() - start_time
-                    
                     # Readonly should be denied write access
                     if response.status == 403:
                         return TestResult(
@@ -593,7 +512,6 @@ class SecurityTestClient:
                             error_message="Readonly role incorrectly allowed write access",
                             duration=duration
                         )
-        
         except Exception as e:
             duration = time.time() - start_time
             return TestResult(
@@ -604,12 +522,10 @@ class SecurityTestClient:
                 error_message=f"Multiple roles test error: {str(e)}",
                 duration=duration
             )
-    
     async def run_security_tests(self, server_url: str, auth_type: str = "none", **kwargs) -> List[TestResult]:
         """Run comprehensive security tests."""
         print(f"\n🔒 Running security tests for {server_url} ({auth_type})")
         print("=" * 60)
-        
         tests = [
             self.test_health_check(server_url, auth_type, **kwargs),
             self.test_echo_command(server_url, auth_type, **kwargs),
@@ -617,13 +533,11 @@ class SecurityTestClient:
             self.test_negative_auth(server_url, auth_type, **kwargs),
             self.test_role_based_access(server_url, auth_type, **kwargs)
         ]
-        
         results = []
         for test in tests:
             result = await test
             results.append(result)
             self.test_results.append(result)
-            
             # Print result
             status = "✅ PASS" if result.success else "❌ FAIL"
             print(f"{status} {result.test_name}")
@@ -633,9 +547,7 @@ class SecurityTestClient:
             if result.error_message:
                 print(f"   Error: {result.error_message}")
             print()
-        
         return results
-    
     async def test_all_scenarios(self) -> Dict[str, List[TestResult]]:
         """Test all security scenarios."""
         scenarios = {
@@ -660,41 +572,32 @@ class SecurityTestClient:
                 "auth": "certificate"
             }
         }
-        
         all_results = {}
-        
         for scenario_name, config in scenarios.items():
             print(f"\n🚀 Testing scenario: {scenario_name.upper()}")
             print("=" * 60)
-            
             try:
                 results = await self.run_security_tests(
-                    config["url"], 
+                    config["url"],
                     config["auth"]
                 )
                 all_results[scenario_name] = results
-                
             except Exception as e:
                 print(f"❌ Failed to test {scenario_name}: {e}")
                 all_results[scenario_name] = []
-        
         return all_results
-    
     def print_summary(self):
         """Print test summary."""
         print("\n" + "=" * 80)
         print("📊 SECURITY TEST SUMMARY")
         print("=" * 80)
-        
         total_tests = len(self.test_results)
         passed_tests = sum(1 for r in self.test_results if r.success)
         failed_tests = total_tests - passed_tests
-        
         print(f"Total Tests: {total_tests}")
         print(f"Passed: {passed_tests} ✅")
         print(f"Failed: {failed_tests} ❌")
         print(f"Success Rate: {(passed_tests/total_tests*100):.1f}%")
-        
         if failed_tests > 0:
             print("\n❌ Failed Tests:")
             for result in self.test_results:
@@ -702,31 +605,25 @@ class SecurityTestClient:
                     print(f"  - {result.test_name} ({result.server_url})")
                     if result.error_message:
                         print(f"    Error: {result.error_message}")
-        
         print("\n✅ Passed Tests:")
         for result in self.test_results:
             if result.success:
                 print(f"  - {result.test_name} ({result.server_url})")
-
-
 async def main():
     """Main function."""
     import argparse
-    
     parser = argparse.ArgumentParser(description="Security Test Client for MCP Proxy Adapter")
-    parser.add_argument("--server", default="http://localhost:8000", 
+    parser.add_argument("--server", default="http://localhost:8000",
                        help="Server URL to test")
-    parser.add_argument("--auth", choices=["none", "api_key", "basic", "certificate"], 
+    parser.add_argument("--auth", choices=["none", "api_key", "basic", "certificate"],
                        default="none", help="Authentication type")
-    parser.add_argument("--all-scenarios", action="store_true", 
+    parser.add_argument("--all-scenarios", action="store_true",
                        help="Test all security scenarios")
     parser.add_argument("--token", help="API token for authentication")
     parser.add_argument("--cert", help="Client certificate file")
     parser.add_argument("--key", help="Client private key file")
     parser.add_argument("--ca-cert", help="CA certificate file")
-    
     args = parser.parse_args()
-    
     if args.all_scenarios:
         # Test all scenarios
         async with SecurityTestClient() as client:
@@ -737,7 +634,5 @@ async def main():
         async with SecurityTestClient(args.server) as client:
             await client.run_security_tests(args.server, args.auth, token=args.token)
             client.print_summary()
-
-
 if __name__ == "__main__":
     asyncio.run(main())