Fix types.

Author: jakozaur

src/cli/pentest.py

@@ -183,7 +183,7 @@ def main(
         return 1
 
     # Check model availability
-    backend_type = client.get_backend_type() if hasattr(client, 'get_backend_type') else "Ollama"
+    backend_type = client.get_backend_type() if hasattr(client, "get_backend_type") else "Ollama"
     click.echo(f"🔍 Checking {backend_type} model availability...")
     if not client.is_available():
         click.echo(f"❌ Model {client.get_model_name()} not available.")
@@ -193,7 +193,7 @@ def main(
     click.echo(f"✅ {backend_type} model {client.get_model_name()} ready")
 
     # Check if backend is busy before starting tests (Ollama only)
-    if not skip_busy_check and hasattr(client, 'check_status'):
+    if not skip_busy_check and hasattr(client, "check_status"):
         click.echo(f"🔍 Checking {backend_type} status...")
         try:
             status = client.check_status()
@@ -214,7 +214,9 @@ def main(
                         )
                         return 1
                 else:
-                    click.echo(f"🚫 {backend_type} is busy. Use --skip-busy-check to proceed anyway.")
+                    click.echo(
+                        f"🚫 {backend_type} is busy. Use --skip-busy-check to proceed anyway."
+                    )
                     return 1
             else:
                 click.echo(f"✅ {backend_type} status: Available")
@@ -223,7 +225,9 @@ def main(
         except Exception as e:
             click.echo(f"⚠️  Could not check {backend_type} status: {e}")
             if not quiet and not skip_busy_check:
-                if click.confirm(f"\nCould not determine if {backend_type} is busy. Continue anyway?"):
+                if click.confirm(
+                    f"\nCould not determine if {backend_type} is busy. Continue anyway?"
+                ):
                     click.echo("⚡ Proceeding with tests...")
                 else:
                     click.echo("🚫 Aborted. Use --skip-busy-check to proceed anyway.")

src/cli/setup.py

@@ -59,7 +59,7 @@ def interactive_backend_setup() -> bool:
     click.echo("Choose your preferred LLM backend:")
     click.echo("1. Ollama (local models)")
     click.echo("2. OpenRouter (cloud models)")
-    
+
     while True:
         try:
             choice = click.prompt("Select backend (1-2)", type=int)
@@ -71,7 +71,7 @@ def interactive_backend_setup() -> bool:
             return False
         except (ValueError, TypeError):
             click.echo("❌ Invalid input. Please enter 1 or 2.")
-    
+
     if choice == 1:
         return setup_ollama_backend()
     else:
@@ -81,21 +81,21 @@ def interactive_backend_setup() -> bool:
 def setup_ollama_backend() -> bool:
     """Setup Ollama backend configuration."""
     click.echo("\n🦙 Configuring Ollama Backend")
-    
+
     # Create default settings.yaml if it doesn't exist
     if not Path("settings.yaml").exists():
         settings_manager.create_default_settings()
-    
+
     host = click.prompt("Ollama host", default="localhost")
     port = click.prompt("Ollama port", default=11434, type=int)
     model = click.prompt("Model name", default="gpt-oss:20b")
-    
+
     # Save settings
     settings_manager.set("backend.provider", "ollama")
     settings_manager.set("ollama.host", host)
     settings_manager.set("ollama.port", port)
     settings_manager.set("ollama.model", model)
-    
+
     click.echo("✅ Ollama backend configured")
     return True
 
@@ -104,16 +104,16 @@ def setup_openrouter_backend() -> bool:
     """Setup OpenRouter backend configuration."""
     click.echo("\n🌐 Configuring OpenRouter Backend")
     click.echo("You'll need an OpenRouter API key from: https://openrouter.ai/keys")
-    
+
     # Create default settings.yaml if it doesn't exist
     if not Path("settings.yaml").exists():
         settings_manager.create_default_settings()
-    
+
     api_key = click.prompt("OpenRouter API key", hide_input=True)
     if not api_key.strip():
         click.echo("❌ API key is required for OpenRouter")
         return False
-    
+
     # Show popular models
     click.echo("\nPopular models:")
     models = [
@@ -124,30 +124,27 @@ def setup_openrouter_backend() -> bool:
         "google/gemini-pro-1.5",
         "mistral/mistral-large",
     ]
-    
+
     for i, model in enumerate(models, 1):
         click.echo(f"{i}. {model}")
-    
+
     try:
         choice = click.prompt("Select model (1-5, or enter custom)", default="1")
-        if choice.isdigit() and 1 <= int(choice) <= 5:
-            model = models[int(choice) - 1]
-        else:
-            model = choice
+        model = models[int(choice) - 1] if choice.isdigit() and 1 <= int(choice) <= 5 else choice
     except (ValueError, IndexError):
         model = "openai/gpt-oss-20b"
-    
+
     site_name = click.prompt("Site name (optional)", default="Red Team Testbed")
     site_url = click.prompt("Site URL (optional)", default="")
-    
+
     # Save settings
     settings_manager.set("backend.provider", "openrouter")
     settings_manager.set("openrouter.api_key", api_key)
     settings_manager.set("openrouter.model", model)
     settings_manager.set("openrouter.site_name", site_name)
     if site_url:
         settings_manager.set("openrouter.site_url", site_url)
-    
+
     click.echo("✅ OpenRouter backend configured")
     return True
 
@@ -159,32 +156,33 @@ def test_connection(config: dict[str, Any], verbose: bool = False) -> bool:
         if Path("settings.yaml").exists():
             settings = settings_manager.load_settings()
             backend = create_backend(settings)
-            provider = settings['backend']['provider']
+            provider = settings["backend"]["provider"]
             click.echo(f"🔗 Testing {provider} connection...")
         else:
             # Fallback to default Ollama configuration
             from src.utils.model_client import OllamaClient
+
             backend = OllamaClient()  # Uses default localhost:11434, gpt-oss:20b
             provider = "ollama"
             click.echo("🔗 Testing default Ollama connection...")
-        
-        if provider == 'ollama':
+
+        if provider == "ollama":
             return test_ollama_connection(backend, verbose)
         else:
             return test_openrouter_connection(backend, verbose)
-            
+
     except Exception as e:
         click.echo(f"❌ Backend setup failed: {e}")
         return False
 
 
-def test_ollama_connection(backend, verbose: bool = False) -> bool:
+def test_ollama_connection(backend: object, verbose: bool = False) -> bool:
     """Test connection to Ollama backend."""
     # Check if Ollama is busy before testing
     click.echo("🔍 Checking Ollama status...")
     try:
         # Handle both OllamaBackend and OllamaClient
-        if hasattr(backend, 'check_status'):
+        if hasattr(backend, "check_status"):
             status = backend.check_status()
         else:
             status = backend.check_ollama_status()
@@ -207,7 +205,7 @@ def test_ollama_connection(backend, verbose: bool = False) -> bool:
 
     try:
         # Handle both interfaces for availability check
-        if hasattr(backend, 'is_available'):
+        if hasattr(backend, "is_available"):
             model_available = backend.is_available()
             model_name = backend.get_model_name()
         else:
@@ -242,7 +240,7 @@ def test_ollama_connection(backend, verbose: bool = False) -> bool:
         return False
 
 
-def test_openrouter_connection(backend, verbose: bool = False) -> bool:
+def test_openrouter_connection(backend: object, verbose: bool = False) -> bool:
     """Test connection to OpenRouter backend."""
     try:
         if not backend.is_available():
@@ -282,7 +280,7 @@ def main(config: str, verbose: bool, configure: bool) -> int | None:
     This command will:
 
     \b
-    - Configure LLM backend (Ollama or OpenRouter) 
+    - Configure LLM backend (Ollama or OpenRouter)
     - Load and validate configuration
     - Set up logging and directories
     - Test backend connection and model availability
@@ -298,7 +296,7 @@ def main(config: str, verbose: bool, configure: bool) -> int | None:
             if not interactive_backend_setup():
                 return 1
             click.echo("")
-        
+
         # Check if settings.yaml exists
         if not Path("settings.yaml").exists():
             if configure:
@@ -307,7 +305,7 @@ def main(config: str, verbose: bool, configure: bool) -> int | None:
             else:
                 click.echo("💡 No settings.yaml found, using default Ollama configuration")
                 click.echo("   Run 'uv run setup --configure' to configure different backends")
-            
+
         # Load configuration (fallback compatibility)
         click.echo("📋 Loading configuration...")
         try:
@@ -317,7 +315,7 @@ def main(config: str, verbose: bool, configure: bool) -> int | None:
             # Use minimal config if config.yaml doesn't exist
             config_data = {
                 "output": {"results_dir": "results", "findings_dir": "findings"},
-                "logging": {"level": "INFO", "file": "testbed.log"}
+                "logging": {"level": "INFO", "file": "testbed.log"},
             }
 
         # Setup logging
@@ -358,4 +356,4 @@ def main(config: str, verbose: bool, configure: bool) -> int | None:
 
 
 if __name__ == "__main__":
-    main()
+    main()

src/utils/llm_backend.py

@@ -45,7 +45,7 @@ def get_model_name(self) -> str:
         """Get the current model name."""
         pass
 
-    @abstractmethod 
+    @abstractmethod
     def get_backend_type(self) -> str:
         """Get the backend type identifier."""
         pass
@@ -115,7 +115,7 @@ def get_backend_type(self) -> str:
         """Get the backend type identifier."""
         return "Ollama"
 
-    def check_status(self) -> Any:
+    def check_status(self) -> object:
         """Check Ollama status."""
         return self.client.check_ollama_status()
 
@@ -130,6 +130,7 @@ class OpenRouterBackend(LLMBackend):
     def __init__(self, config: dict[str, Any]) -> None:
         super().__init__(config)
         import logging
+
         import openai
 
         # Suppress noisy httpx logs from OpenAI client
@@ -181,13 +182,13 @@ def generate(
             )
 
             response_time = time.time() - start_time
-            
+
             choice = response.choices[0]
             content = choice.message.content or ""
-            
+
             # Extract thinking if available (for models that support it)
             thinking = None
-            if hasattr(choice.message, 'reasoning') and choice.message.reasoning:
+            if hasattr(choice.message, "reasoning") and choice.message.reasoning:
                 thinking = choice.message.reasoning
 
             return ModelResponse(
@@ -235,10 +236,10 @@ def chat(
 
             choice = response.choices[0]
             content = choice.message.content or ""
-            
+
             # Extract thinking if available (for models that support it)
             thinking = None
-            if hasattr(choice.message, 'reasoning') and choice.message.reasoning:
+            if hasattr(choice.message, "reasoning") and choice.message.reasoning:
                 thinking = choice.message.reasoning
 
             return ModelResponse(
@@ -301,4 +302,4 @@ def create_backend(settings: dict[str, Any]) -> LLMBackend:
         openrouter_config = settings.get("openrouter", {})
         return OpenRouterBackend(openrouter_config)
     else:
-        raise ValueError(f"Unsupported backend provider: {provider}")
+        raise ValueError(f"Unsupported backend provider: {provider}")

src/utils/model_client.py

@@ -11,7 +11,7 @@
 from src.utils.settings_manager import settings_manager
 
 
-def get_client():
+def get_client() -> object:
     """Factory function to get the configured LLM client."""
     try:
         settings = settings_manager.load_settings()