Merge pull request #105 from Agent-Hellboy/improve_typing

Change typing according to 3.10

Author: Agent-Hellboy

mcp_fuzzer/auth/loaders.py

@@ -1,6 +1,5 @@
 import json
 import os
-from typing import Dict
 
 from .manager import AuthManager
 from .providers import (
@@ -10,7 +9,6 @@
     create_custom_header_auth,
 )
 
-
 def setup_auth_from_env() -> AuthManager:
     auth_manager = AuthManager()
 
@@ -32,7 +30,7 @@ def setup_auth_from_env() -> AuthManager:
         try:
             headers_json = json.loads(custom_headers)
             if isinstance(headers_json, dict):
-                headers: Dict[str, str] = {
+                headers: dict[str, str] = {
                     str(k): str(v) for k, v in headers_json.items()
                 }
                 auth_manager.add_auth_provider(
@@ -55,7 +53,6 @@ def setup_auth_from_env() -> AuthManager:
 
     return auth_manager
 
-
 def load_auth_config(config_file: str) -> AuthManager:
     auth_manager = AuthManager()
 
@@ -95,7 +92,7 @@ def load_auth_config(config_file: str) -> AuthManager:
             headers = provider_config.get("headers")
             if not isinstance(headers, dict):
                 raise ValueError(f"Provider '{name}' custom headers must be a dict")
-            headers_str: Dict[str, str] = {str(k): str(v) for k, v in headers.items()}
+            headers_str: dict[str, str] = {str(k): str(v) for k, v in headers.items()}
             auth_manager.add_auth_provider(name, create_custom_header_auth(headers_str))
         else:
             raise ValueError(f"Unknown provider type: {provider_type}")

mcp_fuzzer/auth/manager.py

@@ -1,34 +1,33 @@
-from typing import Any, Dict, Optional
+from typing import Any
 
 from .providers import AuthProvider
 
-
 class AuthManager:
     """Manages authentication for different tools and services."""
 
     def __init__(self):
-        self.auth_providers: Dict[str, AuthProvider] = {}
-        self.tool_auth_mapping: Dict[str, str] = {}
+        self.auth_providers: dict[str, AuthProvider] = {}
+        self.tool_auth_mapping: dict[str, str] = {}
 
     def add_auth_provider(self, name: str, provider: AuthProvider):
         self.auth_providers[name] = provider
 
     def map_tool_to_auth(self, tool_name: str, auth_provider_name: str):
         self.tool_auth_mapping[tool_name] = auth_provider_name
 
-    def get_auth_for_tool(self, tool_name: str) -> Optional[AuthProvider]:
+    def get_auth_for_tool(self, tool_name: str) -> AuthProvider | None:
         auth_provider_name = self.tool_auth_mapping.get(tool_name)
         if auth_provider_name:
             return self.auth_providers.get(auth_provider_name)
         return None
 
-    def get_auth_headers_for_tool(self, tool_name: str) -> Dict[str, str]:
+    def get_auth_headers_for_tool(self, tool_name: str) -> dict[str, str]:
         provider = self.get_auth_for_tool(tool_name)
         if provider:
             return provider.get_auth_headers()
         return {}
 
-    def get_auth_params_for_tool(self, tool_name: str) -> Dict[str, Any]:
+    def get_auth_params_for_tool(self, tool_name: str) -> dict[str, Any]:
         provider = self.get_auth_for_tool(tool_name)
         if provider:
             return provider.get_auth_params()

mcp_fuzzer/auth/providers.py

@@ -1,78 +1,69 @@
 import base64
 from abc import ABC, abstractmethod
-from typing import Any, Dict
-
+from typing import Any
 
 class AuthProvider(ABC):
     @abstractmethod
-    def get_auth_headers(self) -> Dict[str, str]:
+    def get_auth_headers(self) -> dict[str, str]:
         pass
 
     @abstractmethod
-    def get_auth_params(self) -> Dict[str, Any]:
+    def get_auth_params(self) -> dict[str, Any]:
         pass
 
-
 class APIKeyAuth(AuthProvider):
     def __init__(self, api_key: str, header_name: str = "Authorization"):
         self.api_key = api_key
         self.header_name = header_name
 
-    def get_auth_headers(self) -> Dict[str, str]:
+    def get_auth_headers(self) -> dict[str, str]:
         return {self.header_name: f"Bearer {self.api_key}"}
 
-    def get_auth_params(self) -> Dict[str, Any]:
+    def get_auth_params(self) -> dict[str, Any]:
         return {}
 
-
 class BasicAuth(AuthProvider):
     def __init__(self, username: str, password: str):
         self.username = username
         self.password = password
 
-    def get_auth_headers(self) -> Dict[str, str]:
+    def get_auth_headers(self) -> dict[str, str]:
         credentials = f"{self.username}:{self.password}"
         encoded = base64.b64encode(credentials.encode()).decode()
         return {"Authorization": f"Basic {encoded}"}
 
-    def get_auth_params(self) -> Dict[str, Any]:
+    def get_auth_params(self) -> dict[str, Any]:
         return {}
 
-
 class OAuthTokenAuth(AuthProvider):
     def __init__(self, token: str, token_type: str = "Bearer"):
         self.token = token
         self.token_type = token_type
 
-    def get_auth_headers(self) -> Dict[str, str]:
+    def get_auth_headers(self) -> dict[str, str]:
         return {"Authorization": f"{self.token_type} {self.token}"}
 
-    def get_auth_params(self) -> Dict[str, Any]:
+    def get_auth_params(self) -> dict[str, Any]:
         return {}
 
-
 class CustomHeaderAuth(AuthProvider):
-    def __init__(self, headers: Dict[str, str]):
+    def __init__(self, headers: dict[str, str]):
         self.headers = dict(headers)
 
-    def get_auth_headers(self) -> Dict[str, str]:
+    def get_auth_headers(self) -> dict[str, str]:
         return dict(self.headers)
 
-    def get_auth_params(self) -> Dict[str, Any]:
+    def get_auth_params(self) -> dict[str, Any]:
         return {}
 
-
 def create_api_key_auth(api_key: str, header_name: str = "Authorization") -> APIKeyAuth:
     return APIKeyAuth(api_key, header_name)
 
-
 def create_basic_auth(username: str, password: str) -> BasicAuth:
     return BasicAuth(username, password)
 
-
 def create_oauth_auth(token: str, token_type: str = "Bearer") -> OAuthTokenAuth:
     return OAuthTokenAuth(token, token_type)
 
-
-def create_custom_header_auth(headers: Dict[str, str]) -> CustomHeaderAuth:
+def create_custom_header_auth(headers: dict[str, str]) -> CustomHeaderAuth:
     return CustomHeaderAuth(headers)

mcp_fuzzer/cli/args.py

@@ -2,7 +2,7 @@
 import argparse
 import logging
 import sys
-from typing import Any, Dict
+from typing import Any
 
 from rich.console import Console
 
@@ -13,7 +13,6 @@
 )
 from ..config import config
 
-
 def create_argument_parser() -> argparse.ArgumentParser:
     parser = argparse.ArgumentParser(
         description="MCP Fuzzer - Comprehensive fuzzing for MCP servers",
@@ -354,12 +353,10 @@ def create_argument_parser() -> argparse.ArgumentParser:
 
     return parser
 
-
 def parse_arguments() -> argparse.Namespace:
     parser = create_argument_parser()
     return parser.parse_args()
 
-
 def setup_logging(args: argparse.Namespace) -> None:
     if getattr(args, "log_level", None):
         level = getattr(logging, args.log_level)
@@ -381,8 +378,7 @@ def setup_logging(args: argparse.Namespace) -> None:
     logging.getLogger("urllib3").setLevel(logging.WARNING)
     logging.getLogger("psutil").setLevel(logging.WARNING)
 
-
-def build_unified_client_args(args: argparse.Namespace) -> Dict[str, Any]:
+def build_unified_client_args(args: argparse.Namespace) -> dict[str, Any]:
     client_args = {
         "mode": args.mode,
         "phase": args.phase,
@@ -442,7 +438,6 @@ def build_unified_client_args(args: argparse.Namespace) -> Dict[str, Any]:
 
     return client_args
 
-
 def print_startup_info(args: argparse.Namespace) -> None:
     # Resolve Console via package so tests can patch mcp_fuzzer.cli.Console
     cli_module = sys.modules.get("mcp_fuzzer.cli")
@@ -452,8 +447,7 @@ def print_startup_info(args: argparse.Namespace) -> None:
     console.print(f"Protocol: {args.protocol.upper()}")
     console.print(f"Endpoint: {args.endpoint}")
 
-
-def get_cli_config() -> Dict[str, Any]:
+def get_cli_config() -> dict[str, Any]:
     """Get CLI configuration as a dictionary for external callers/tests.
 
     Resolve helpers through the package namespace so unit tests patching
@@ -590,7 +584,6 @@ def get_cli_config() -> Dict[str, Any]:
         "export_markdown": getattr(args, "export_markdown", None),
     }
 
-
 def validate_arguments(args: argparse.Namespace) -> None:
     # Check if this is a utility command that doesn't need endpoint
     is_utility_command = (

mcp_fuzzer/cli/runner.py

@@ -3,16 +3,15 @@
 import os
 import signal
 import sys
-from typing import Any, Dict, List
+from typing import Any
 
 from rich.console import Console
 
 from ..transport import create_transport
 from ..safety_system.policy import configure_network_policy
 from ..safety_system import start_system_blocking, stop_system_blocking
 
-
-def create_transport_with_auth(args, client_args: Dict[str, Any]):
+def create_transport_with_auth(args, client_args: dict[str, Any]):
     try:
         auth_headers = None
         if client_args.get("auth_manager"):
@@ -33,9 +32,8 @@ def create_transport_with_auth(args, client_args: Dict[str, Any]):
         console.print(f"[bold red]Unexpected error:[/bold red] {transport_error}")
         sys.exit(1)
 
-
-def prepare_inner_argv(args) -> List[str]:
-    argv: List[str] = [sys.argv[0]]
+def prepare_inner_argv(args) -> list[str]:
+    argv: list[str] = [sys.argv[0]]
     mode = args.mode
     argv += ["--mode", mode]
     argv += ["--protocol", args.protocol]
@@ -59,22 +57,19 @@ def prepare_inner_argv(args) -> List[str]:
             argv += ["--allow-host", h]
     return argv
 
-
 def start_safety_if_enabled(args) -> bool:
     if getattr(args, "enable_safety_system", False):
         start_system_blocking()
         return True
     return False
 
-
 def stop_safety_if_started(started: bool) -> None:
     if started:
         try:
             stop_system_blocking()
         except Exception:
             pass
 
-
 def execute_inner_client(args, unified_client_main, argv):
     old_argv = sys.argv
     sys.argv = argv
@@ -177,7 +172,6 @@ def _cancel_all_tasks():  # pragma: no cover
         if should_exit:
             raise SystemExit(130)
 
-
 def run_with_retry_on_interrupt(args, unified_client_main, argv) -> None:
     try:
         execute_inner_client(args, unified_client_main, argv)

mcp_fuzzer/client/__init__.py

@@ -6,7 +6,6 @@
 """
 
 import logging
-from typing import List, Optional
 
 import emoji
 
@@ -17,8 +16,7 @@
 # For backward compatibility
 UnifiedMCPFuzzerClient = MCPFuzzerClient
 
-
-async def main(argv: Optional[List[str]] = None) -> int:
+async def main(argv: list[str] | None = None) -> int:
     """Main entry point for the CLI application.
 
     Args:
@@ -245,5 +243,4 @@ async def main(argv: Optional[List[str]] = None) -> int:
         # Ensure proper shutdown
         await client.cleanup()
 
-
 __all__ = ["MCPFuzzerClient", "UnifiedMCPFuzzerClient", "main"]

mcp_fuzzer/client/base.py

@@ -6,7 +6,6 @@
 """
 
 import logging
-from typing import Optional
 
 from ..auth import AuthManager
 from ..reports import FuzzerReporter
@@ -15,7 +14,6 @@
 from .tool_client import ToolClient
 from .protocol_client import ProtocolClient
 
-
 class MCPFuzzerClient:
     """
     Main client for fuzzing MCP servers.
@@ -27,10 +25,10 @@ class MCPFuzzerClient:
     def __init__(
         self,
         transport,
-        auth_manager: Optional[AuthManager] = None,
-        tool_timeout: Optional[float] = None,
-        reporter: Optional[FuzzerReporter] = None,
-        safety_system: Optional[SafetyProvider] = None,
+        auth_manager: AuthManager | None = None,
+        tool_timeout: float | None = None,
+        reporter: FuzzerReporter | None = None,
+        safety_system: SafetyProvider | None = None,
         max_concurrency: int = 5,
     ):
         """