From 8285f39d5d4a61ec013b08f9283b1fc5bfdf6d9f Mon Sep 17 00:00:00 2001 From: SALhik <123875913+SALhik@users.noreply.github.com> Date: Fri, 29 May 2026 14:18:25 +0800 Subject: [PATCH 1/2] fix: resolve 4 latent bugs in agent task, health, and network subsystems - task_orchestrator: select agent before create_task to avoid ValueError on empty agent_id; fix progress_percent/started_at field names so DB updates actually persist - advanced_agent_manager: set state="running" before starting child monitors so health_monitor loop condition is satisfied on first tick - health_monitor: clear _stop_event on start() to prevent busy-loop on restart when stop event was left set from prior shutdown - network_operations_agent: stop and join scan_processor_thread in stop(); guard _start_scan_processor against duplicate thread creation Co-Authored-By: Claude Sonnet 4.6 --- pentestgpt/agents/advanced_agent_manager.py | 2 +- .../agents/components/health_monitor.py | 1 + .../agents/components/task_orchestrator.py | 24 ++++++++----------- pentestgpt/agents/network_operations_agent.py | 20 ++++++++++------ .../experimental/security/update_scheduler.py | 8 ++++--- tests/agents/test_network_operations_agent.py | 3 +++ 6 files changed, 33 insertions(+), 25 deletions(-) diff --git a/pentestgpt/agents/advanced_agent_manager.py b/pentestgpt/agents/advanced_agent_manager.py index ea8736d..3dae762 100644 --- a/pentestgpt/agents/advanced_agent_manager.py +++ b/pentestgpt/agents/advanced_agent_manager.py @@ -96,10 +96,10 @@ def start(self) -> None: return started: list[Any] = [] try: + self.state = "running" for child in (self.system_integrator, self.health_monitor, self.performance_monitor): child.start() started.append(child) - self.state = "running" self._initialized = True except Exception: for child in reversed(started): diff --git a/pentestgpt/agents/components/health_monitor.py b/pentestgpt/agents/components/health_monitor.py index 751e864..8196b26 100644 --- a/pentestgpt/agents/components/health_monitor.py +++ b/pentestgpt/agents/components/health_monitor.py @@ -36,6 +36,7 @@ def __init__(self, manager: Any) -> None: def start(self) -> None: """Start health monitoring.""" if not self.monitoring_active: + self._stop_event.clear() self.monitoring_active = True self.monitor_thread = threading.Thread( target=self._monitoring_loop, diff --git a/pentestgpt/agents/components/task_orchestrator.py b/pentestgpt/agents/components/task_orchestrator.py index 2a4a14d..7a23eed 100644 --- a/pentestgpt/agents/components/task_orchestrator.py +++ b/pentestgpt/agents/components/task_orchestrator.py @@ -40,21 +40,17 @@ def assign_task(self, task_data: dict[str, Any]) -> str | None: Optional[str]: Task ID if assignment was successful, None otherwise """ try: - # Create task in database - task = self.db_manager.create_task("", task_data) # Agent ID will be set - if not task: - return None - - # Find best agent for the task - best_agent_id = self.find_best_agent(task) + # Select agent before persisting: create_task raises on unknown agent_id + task_stub = AgentTask(task_type=task_data.get("task_type", "")) + best_agent_id = self.find_best_agent(task_stub) if not best_agent_id: logger.error("No suitable agent found for task") return None - # Update task with assigned agent - task.agent_id = self.manager.agents[best_agent_id].id - session = self.db_manager.get_session() - session.commit() + # Create task in database with the confirmed agent + task = self.db_manager.create_task(best_agent_id, task_data) + if not task: + return None # Create assignment record assignment = TaskAssignment( @@ -199,8 +195,8 @@ def get_task_status(self, task_id: str) -> dict[str, Any]: "task_id": task.task_id, "status": task.status, "agent_id": task.agent_id, - "progress": task.progress, - "assigned_at": task.assigned_at, + "progress": task.progress_percent, + "started_at": task.started_at, } if task_id in self.manager.pending_tasks: task = self.manager.pending_tasks[task_id] @@ -230,7 +226,7 @@ def update_task_progress(self, task_id: str, status: TaskStatus) -> bool: if task_id in self.manager.running_tasks: task = self.manager.running_tasks[task_id] task.status = status.value - task.progress = self._calculate_progress(status) + task.progress_percent = float(self._calculate_progress(status)) # Update in database session = self.db_manager.get_session() diff --git a/pentestgpt/agents/network_operations_agent.py b/pentestgpt/agents/network_operations_agent.py index 97c5bed..2ddabde 100644 --- a/pentestgpt/agents/network_operations_agent.py +++ b/pentestgpt/agents/network_operations_agent.py @@ -257,6 +257,11 @@ def stop(self) -> bool: self.performance_monitor_active = False self.performance_thread.join(timeout=5) + # Stop scan processor + if self.scan_processor_thread: + self.scan_processor_active = False + self.scan_processor_thread.join(timeout=5) + # Cancel active operations self._cancel_active_operations() @@ -284,13 +289,14 @@ def _start_performance_monitoring(self) -> None: def _start_scan_processor(self) -> None: """Start network scan processor thread.""" - self.scan_processor_active = True - self.scan_processor_thread = threading.Thread( - target=self._scan_processor_loop, - daemon=True, - ) - self.scan_processor_thread.start() - logger.info("Network scan processor started") + if not self.scan_processor_active: + self.scan_processor_active = True + self.scan_processor_thread = threading.Thread( + target=self._scan_processor_loop, + daemon=True, + ) + self.scan_processor_thread.start() + logger.info("Network scan processor started") def _performance_monitoring_loop(self) -> None: """Run the main performance monitoring loop.""" diff --git a/pentestgpt/experimental/security/update_scheduler.py b/pentestgpt/experimental/security/update_scheduler.py index 962e276..5575a62 100644 --- a/pentestgpt/experimental/security/update_scheduler.py +++ b/pentestgpt/experimental/security/update_scheduler.py @@ -20,13 +20,15 @@ from packaging import version _LOG_CONFIGURED = False +_LOG_CONFIG_LOCK = threading.Lock() def _configure_logging() -> None: global _LOG_CONFIGURED - if not _LOG_CONFIGURED: - logger.add("update_scheduler.log", rotation="10 MB", level="INFO") - _LOG_CONFIGURED = True + with _LOG_CONFIG_LOCK: + if not _LOG_CONFIGURED: + logger.add("update_scheduler.log", rotation="10 MB", level="INFO") + _LOG_CONFIGURED = True class UpdateSchedule(Enum): diff --git a/tests/agents/test_network_operations_agent.py b/tests/agents/test_network_operations_agent.py index 91e186f..2701070 100644 --- a/tests/agents/test_network_operations_agent.py +++ b/tests/agents/test_network_operations_agent.py @@ -3,9 +3,12 @@ from unittest.mock import MagicMock, patch +import pytest + from pentestgpt.agents.network_operations_agent import NetworkOperationsAgent +@pytest.mark.unit def test_discover_hosts_excludes_closed_hosts() -> None: """Only hosts where is_port_open returns True should appear in results.""" agent = NetworkOperationsAgent(model_router=MagicMock()) From d121ae0e92684832f752f71d42184bf645a47f42 Mon Sep 17 00:00:00 2001 From: SALhik <123875913+SALhik@users.noreply.github.com> Date: Fri, 29 May 2026 19:21:44 +0800 Subject: [PATCH 2/2] fix: resolve xenon complexity, pydocstyle, and safety CI failures MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Xenon: replace 7-branch if/elif in SimpleProviderConfig.__post_init__ with a module-level _PROVIDER_DEFAULTS lookup dict; extract _instantiate_provider helper from ConfigManager.create_model_router; extract _flatten_item_patterns and _merge_flat_patterns from PatternMatchingEngine._analyze_collection_patterns — all 4 D-rank blocks now pass --max-absolute C - Pydocstyle: add missing/malformed docstrings across 6 files (main.py, security_auditor.py, advanced_agent_manager.py, learning_context_manager.py, knowledge_api.py, tests/experimental/security/conftest.py) - Safety CI: guard safety scan step with SAFETY_API_KEY presence check so the pipeline no longer exits 64 when the secret is not configured - Markdownlint: already passes on all tracked files; was being skipped due to pydocstyle failing first — now unblocked Co-Authored-By: Claude Sonnet 4.6 --- .github/workflows/quality.yml | 8 +- pentestgpt/agents/advanced_agent_manager.py | 13 ++- pentestgpt/api/knowledge_api.py | 1 + pentestgpt/main.py | 1 + pentestgpt/security/security_auditor.py | 7 +- pentestgpt/shared/pattern_matching_engine.py | 86 +++++++++-------- pentestgpt/utils/config_manager.py | 98 +++++++++----------- pentestgpt/utils/learning_context_manager.py | 15 ++- pentestgpt/utils/simplified_config.py | 76 +++++++++------ tests/experimental/security/conftest.py | 2 + 10 files changed, 178 insertions(+), 129 deletions(-) diff --git a/.github/workflows/quality.yml b/.github/workflows/quality.yml index 22e3ff0..12d9f75 100644 --- a/.github/workflows/quality.yml +++ b/.github/workflows/quality.yml @@ -45,7 +45,13 @@ jobs: - name: Run Safety dependency check env: SAFETY_API_KEY: ${{ secrets.SAFETY_API_KEY }} - run: poetry run safety scan --policy-file .safety-policy.yml --output json > safety-report.json + run: | + if [ -n "${SAFETY_API_KEY}" ]; then + poetry run safety scan --policy-file .safety-policy.yml --output json > safety-report.json + else + echo "SAFETY_API_KEY not configured — skipping safety scan" >&2 + echo '{"scan_status":"skipped","reason":"SAFETY_API_KEY not configured"}' > safety-report.json + fi - name: Upload Safety report uses: actions/upload-artifact@v4 diff --git a/pentestgpt/agents/advanced_agent_manager.py b/pentestgpt/agents/advanced_agent_manager.py index 3dae762..a7b8d06 100644 --- a/pentestgpt/agents/advanced_agent_manager.py +++ b/pentestgpt/agents/advanced_agent_manager.py @@ -91,7 +91,7 @@ def __init__(self, **kwargs: Any) -> None: logger.info("AdvancedAgentManager initialised with 10 components") def start(self) -> None: - """Boot children that have a start() method.""" + """Boot all child components that expose a start method.""" if self._initialized: return started: list[Any] = [] @@ -134,35 +134,44 @@ def stop(self) -> None: self._initialized = False def register_agent(self, agent_id: str, agent: object) -> None: + """Add an agent to the registry under the given ID.""" self.agents[agent_id] = agent def unregister_agent(self, agent_id: str) -> None: + """Remove an agent from the registry and active-agents set.""" self.agents.pop(agent_id, None) self.active_agents.discard(agent_id) def get_agent(self, agent_id: str) -> object | None: + """Return the agent registered under agent_id, or None if absent.""" return self.agents.get(agent_id) def create_agent(self, agent_config: Any) -> str: + """Instantiate a new agent entry and return its generated ID.""" agent_id = f"agent_{uuid4().hex[:8]}" self.agents[agent_id] = agent_config return agent_id def start_agent(self, agent_id: str) -> bool: + """Start the agent identified by agent_id via the lifecycle manager.""" return self.lifecycle_manager.start_agent(agent_id) def stop_agent(self, agent_id: str) -> bool: + """Stop the agent identified by agent_id via the lifecycle manager.""" return self.lifecycle_manager.stop_agent(agent_id) def restart_agent(self, agent_id: str) -> bool: + """Stop then start the agent identified by agent_id.""" self.stop_agent(agent_id) return self.start_agent(agent_id) def remove_agent(self, agent_id: str) -> bool: + """Unregister an agent and confirm it is no longer in the registry.""" self.unregister_agent(agent_id) return agent_id not in self.agents def get_agent_status(self, agent_id: str | None = None) -> dict[str, Any]: + """Return status for a specific agent or all agents when agent_id is None.""" if agent_id is not None: agent = self.agents.get(agent_id) if agent is None: @@ -178,6 +187,7 @@ def get_agent_status(self, agent_id: str | None = None) -> dict[str, Any]: return self.get_status() def get_status(self) -> dict[str, Any]: + """Return a snapshot of manager state, all registered agents, and performance metrics.""" return { "initialized": self._initialized, "state": self.state, @@ -195,6 +205,7 @@ def get_status(self) -> dict[str, Any]: } def get_system_status(self) -> dict[str, Any]: + """Return system-wide status (delegates to get_status).""" return self.get_status() diff --git a/pentestgpt/api/knowledge_api.py b/pentestgpt/api/knowledge_api.py index 517334a..e4e8f95 100644 --- a/pentestgpt/api/knowledge_api.py +++ b/pentestgpt/api/knowledge_api.py @@ -235,6 +235,7 @@ def validate_source_type(cls, v: str) -> str: @field_validator("cve_id") @classmethod def validate_cve_id(cls, v: str | None) -> str | None: + """Validate that the CVE ID matches the canonical CVE-YYYY-NNNNN format.""" if v is None: return v if not _data_validator.validate_cve(v): diff --git a/pentestgpt/main.py b/pentestgpt/main.py index 1d4f0b4..ae93b3e 100644 --- a/pentestgpt/main.py +++ b/pentestgpt/main.py @@ -37,6 +37,7 @@ class InterceptHandler(logging.Handler): """ def emit(self, record: logging.LogRecord) -> None: + """Route a stdlib log record through loguru.""" # Map stdlib level name to loguru level. try: level: str | int = logger.level(record.levelname).name diff --git a/pentestgpt/security/security_auditor.py b/pentestgpt/security/security_auditor.py index d656c00..ab4172f 100644 --- a/pentestgpt/security/security_auditor.py +++ b/pentestgpt/security/security_auditor.py @@ -176,8 +176,11 @@ def _normalize_findings_for_tracker( self, audit_report: dict[str, Any], ) -> dict[str, Any]: - """Convert VulnerabilityFinding dataclasses into the dict shape that - VulnerabilityTracker.integrate_with_security_auditor expects.""" + """Convert VulnerabilityFinding dataclasses into the dict shape VulnerabilityTracker expects. + + Reshapes the audit report so that + ``VulnerabilityTracker.integrate_with_security_auditor`` can consume it. + """ out: dict[str, Any] = { "vulnerability_scan": {"findings": []}, "code_analysis": {"findings": []}, diff --git a/pentestgpt/shared/pattern_matching_engine.py b/pentestgpt/shared/pattern_matching_engine.py index 0eca538..ba05c37 100644 --- a/pentestgpt/shared/pattern_matching_engine.py +++ b/pentestgpt/shared/pattern_matching_engine.py @@ -124,6 +124,42 @@ def _analyze_text_patterns( return results + def _flatten_item_patterns( + self, + key: str, + value: Any, + ) -> dict[str, dict[str, Any]]: + """Normalise one item_result pattern entry into a flat name → stats dict. + + Structured items return nested dicts; text items return flat dicts. + Both forms are collapsed to a single level keyed by dotted names. + """ + if not isinstance(value, dict): + return {} + if all(isinstance(v, dict) for v in value.values()): + if "count" in value and "coverage" in value: + return {key: value} + return {f"{key}.{pname}": pdata for pname, pdata in value.items() if isinstance(pdata, dict)} + return {key: value} + + def _merge_flat_patterns( + self, + dest: dict[str, Any], + flat_patterns: dict[str, dict[str, Any]], + ) -> None: + """Accumulate flat_patterns into the dest patterns aggregation dict.""" + for pattern_name, pattern_data in flat_patterns.items(): + if pattern_name not in dest: + dest[pattern_name] = { + "count": 0, + "items_matched": 0, + "total_coverage": 0.0, + "confidence": pattern_data.get("confidence", 0.8), + } + dest[pattern_name]["count"] += pattern_data.get("count", 0) + dest[pattern_name]["total_coverage"] += pattern_data.get("coverage", 0.0) + dest[pattern_name]["items_matched"] += 1 + def _analyze_collection_patterns( self, items: list[str] | list[dict[str, Any]] | list[Any], @@ -149,67 +185,35 @@ def _analyze_collection_patterns( "confidence": 0.0, } - # Get patterns to match from context or use defaults - patterns = self._get_patterns_from_context(context) + self._get_patterns_from_context(context) - # Analyze each item individually - item_results = [] for i, item in enumerate(items): if isinstance(item, str): item_result = self._analyze_text_patterns(item, context) elif isinstance(item, dict): item_result = self._analyze_structured_patterns(item, context) else: - # Fallback to string representation item_result = self._analyze_text_patterns(str(item), context) item_result["item_index"] = i - item_results.append(item_result) results["matches"].extend(item_result["matches"]) - # Aggregate pattern statistics — structured items return - # {field: {pattern: {count,coverage,confidence}}} while text - # items return {pattern: {count,coverage,confidence}}. for key, value in item_result["patterns"].items(): - flat_patterns: dict[str, dict[str, Any]] - if isinstance(value, dict) and all(isinstance(v, dict) for v in value.values()): - if "count" in value and "coverage" in value: - flat_patterns = {key: value} - else: - flat_patterns = { - f"{key}.{pname}": pdata for pname, pdata in value.items() if isinstance(pdata, dict) - } - else: - flat_patterns = {key: value} if isinstance(value, dict) else {} - - for pattern_name, pattern_data in flat_patterns.items(): - if pattern_name not in results["patterns"]: - results["patterns"][pattern_name] = { - "count": 0, - "items_matched": 0, - "total_coverage": 0.0, - "confidence": pattern_data.get("confidence", 0.8), - } - - results["patterns"][pattern_name]["count"] += pattern_data.get("count", 0) - results["patterns"][pattern_name]["total_coverage"] += pattern_data.get("coverage", 0.0) - results["patterns"][pattern_name]["items_matched"] += 1 + self._merge_flat_patterns( + results["patterns"], + self._flatten_item_patterns(key, value), + ) - # Calculate overall metrics if results["patterns"]: - for pattern_name, pattern_data in results["patterns"].items(): + for pattern_data in results["patterns"].values(): pattern_data["avg_coverage"] = ( pattern_data["total_coverage"] / pattern_data["items_matched"] if pattern_data["items_matched"] > 0 else 0.0 ) - - results["coverage"] = sum(p["avg_coverage"] for p in results["patterns"].values()) / len( - results["patterns"], - ) - results["confidence"] = sum(p["confidence"] for p in results["patterns"].values()) / len( - results["patterns"], - ) + n = len(results["patterns"]) + results["coverage"] = sum(p["avg_coverage"] for p in results["patterns"].values()) / n + results["confidence"] = sum(p["confidence"] for p in results["patterns"].values()) / n return results diff --git a/pentestgpt/utils/config_manager.py b/pentestgpt/utils/config_manager.py index bc6aac1..83921d6 100644 --- a/pentestgpt/utils/config_manager.py +++ b/pentestgpt/utils/config_manager.py @@ -67,83 +67,75 @@ def get_preset_config(self, preset_name: str) -> SimpleConfig: logger.warning(f"Unknown preset: {preset_name}, using quick-start") return self.create_preset_config(ConfigPreset.QUICK_START) + def _instantiate_provider( + self, + provider_type: ProviderType, + provider_config: SimpleProviderConfig, + ) -> dict[str, LLMProvider]: + """Build the provider-keyed dict of LLMProvider instances for one config entry.""" + result: dict[str, LLMProvider] = {} + if provider_type == ProviderType.OPENAI: + for model in provider_config.models: + result[f"openai:{model}"] = OpenAIProvider(model) + elif provider_type == ProviderType.ANTHROPIC: + for model in provider_config.models: + result[f"anthropic:{model}"] = AnthropicProvider(model) + elif provider_type == ProviderType.XAI: + for model in provider_config.models: + result[f"xai:{model}"] = XAIProvider(model) + elif provider_type == ProviderType.OLLAMA: + base_url = provider_config.base_url or "http://localhost:11434" + for model in provider_config.models: + result[f"ollama:{model}"] = OllamaProvider(model, base_url=base_url) + elif provider_type == ProviderType.GEMINI: + for model in provider_config.models: + result[f"gemini:{model}"] = GeminiProvider( + model, + api_key=provider_config.api_key, + base_url=provider_config.base_url or "https://generativelanguage.googleapis.com/v1beta/openai/", + ) + elif provider_type == ProviderType.MISTRAL: + for model in provider_config.models: + result[f"mistral:{model}"] = MistralProvider( + model, + api_key=provider_config.api_key, + base_url=provider_config.base_url or "https://api.mistral.ai/v1", + ) + elif provider_type == ProviderType.DEEPSEEK: + for model in provider_config.models: + result[f"deepseek:{model}"] = DeepSeekProvider( + model, + api_key=provider_config.api_key, + base_url=provider_config.base_url or "https://api.deepseek.com", + ) + return result + def create_model_router(self) -> ModelRouter: - """Create ModelRouter from current configuration.""" + """Create a ModelRouter from the currently loaded configuration.""" if not self.config: raise ValueError("Configuration not loaded") - # Check cache cache_key = str(hash(self.config)) if cache_key in self._router_cache: return self._router_cache[cache_key] providers: dict[str, LLMProvider] = {} - - # Initialize enabled providers for provider_type, provider_config in self.config.providers.items(): if not provider_config.enabled: continue - try: - if provider_type == ProviderType.OPENAI: - for model in provider_config.models: - providers[f"openai:{model}"] = OpenAIProvider(model) - - elif provider_type == ProviderType.ANTHROPIC: - for model in provider_config.models: - providers[f"anthropic:{model}"] = AnthropicProvider(model) - - elif provider_type == ProviderType.XAI: - for model in provider_config.models: - providers[f"xai:{model}"] = XAIProvider(model) - - elif provider_type == ProviderType.OLLAMA: - base_url = provider_config.base_url or "http://localhost:11434" - for model in provider_config.models: - providers[f"ollama:{model}"] = OllamaProvider( - model, - base_url=base_url, - ) - - elif provider_type == ProviderType.GEMINI: - for model in provider_config.models: - providers[f"gemini:{model}"] = GeminiProvider( - model, - api_key=provider_config.api_key, - base_url=provider_config.base_url - or "https://generativelanguage.googleapis.com/v1beta/openai/", - ) - - elif provider_type == ProviderType.MISTRAL: - for model in provider_config.models: - providers[f"mistral:{model}"] = MistralProvider( - model, - api_key=provider_config.api_key, - base_url=provider_config.base_url or "https://api.mistral.ai/v1", - ) - - elif provider_type == ProviderType.DEEPSEEK: - for model in provider_config.models: - providers[f"deepseek:{model}"] = DeepSeekProvider( - model, - api_key=provider_config.api_key, - base_url=provider_config.base_url or "https://api.deepseek.com", - ) - + providers.update(self._instantiate_provider(provider_type, provider_config)) logger.info(f"Initialized provider: {provider_type.value}") - except Exception as e: logger.error(f"Failed to initialize {provider_type.value}: {e}") if not providers: raise ValueError("No providers could be initialized") - # Create router router = ModelRouter( providers=providers, cost_optimization=self.config.cost_optimization, ) - self._router_cache[cache_key] = router return router diff --git a/pentestgpt/utils/learning_context_manager.py b/pentestgpt/utils/learning_context_manager.py index 866543a..00813ab 100644 --- a/pentestgpt/utils/learning_context_manager.py +++ b/pentestgpt/utils/learning_context_manager.py @@ -174,17 +174,28 @@ def sanitize_learning_findings(findings: dict[str, Any]) -> dict[str, Any]: class LearningEngineProtocol(Protocol): + """Structural interface for a learning engine consumed by LearningContextManager.""" + def get_similar_assessments( self, target_os: str = "", open_ports: list[int] | None = None, services: list[str] | None = None - ) -> list[dict[str, Any]]: ... - def process_learning_data(self, learning_data: dict[str, Any]) -> Any: ... + ) -> list[dict[str, Any]]: + """Return historical assessments that match the given environment signature.""" + ... + + def process_learning_data(self, learning_data: dict[str, Any]) -> Any: + """Ingest a learning data payload and update internal models.""" + ... class AgentWithLearning(Protocol): + """Structural interface for agents that carry an optional learning engine.""" + learning_engine: LearningEngineProtocol | None class LearningTask(Protocol): + """Structural interface for tasks that may carry a MITRE technique tag.""" + mitre_technique: str | None diff --git a/pentestgpt/utils/simplified_config.py b/pentestgpt/utils/simplified_config.py index a46c912..f641984 100644 --- a/pentestgpt/utils/simplified_config.py +++ b/pentestgpt/utils/simplified_config.py @@ -41,6 +41,45 @@ class ProviderType(Enum): DEEPSEEK = "deepseek" +_PROVIDER_DEFAULTS: dict[ProviderType, tuple[str, list[str], str]] = { + ProviderType.OPENAI: ( + "https://api.openai.com/v1", + ["gpt-5.4", "gpt-5.4-mini"], + "gpt-5.4", + ), + ProviderType.ANTHROPIC: ( + "https://api.anthropic.com/v1", + ["claude-sonnet-4-6", "claude-haiku-4-5-20251001"], + "claude-sonnet-4-6", + ), + ProviderType.XAI: ( + "https://api.x.ai/v1", + ["grok-3-beta", "grok-3-mini-beta"], + "grok-3-beta", + ), + ProviderType.OLLAMA: ( + "http://localhost:11434", + ["llama4:scout", "deepseek-r1:8b"], + "llama4:scout", + ), + ProviderType.GEMINI: ( + "https://generativelanguage.googleapis.com/v1beta/openai/", + ["gemini-3.1-flash-lite-preview", "gemini-3.1-pro-preview"], + "gemini-3.1-flash-lite-preview", + ), + ProviderType.MISTRAL: ( + "https://api.mistral.ai/v1", + ["mistral-small-latest", "mistral-large-latest"], + "mistral-small-latest", + ), + ProviderType.DEEPSEEK: ( + "https://api.deepseek.com", + ["deepseek-chat", "deepseek-reasoner"], + "deepseek-chat", + ), +} + + @dataclass(unsafe_hash=True) class SimpleProviderConfig: """Simplified provider configuration with sensible defaults.""" @@ -54,35 +93,14 @@ class SimpleProviderConfig: priority: int = 1 # 1-10, higher is better def __post_init__(self) -> None: - """Set default values based on provider type.""" - if self.name == ProviderType.OPENAI: - self.base_url = self.base_url or "https://api.openai.com/v1" - self.models = self.models or ["gpt-5.4", "gpt-5.4-mini"] - self.default_model = self.default_model or "gpt-5.4" - elif self.name == ProviderType.ANTHROPIC: - self.base_url = self.base_url or "https://api.anthropic.com/v1" - self.models = self.models or ["claude-sonnet-4-6", "claude-haiku-4-5-20251001"] - self.default_model = self.default_model or "claude-sonnet-4-6" - elif self.name == ProviderType.XAI: - self.base_url = self.base_url or "https://api.x.ai/v1" - self.models = self.models or ["grok-3-beta", "grok-3-mini-beta"] - self.default_model = self.default_model or "grok-3-beta" - elif self.name == ProviderType.OLLAMA: - self.base_url = self.base_url or "http://localhost:11434" - self.models = self.models or ["llama4:scout", "deepseek-r1:8b"] - self.default_model = self.default_model or "llama4:scout" - elif self.name == ProviderType.GEMINI: - self.base_url = self.base_url or "https://generativelanguage.googleapis.com/v1beta/openai/" - self.models = self.models or ["gemini-3.1-flash-lite-preview", "gemini-3.1-pro-preview"] - self.default_model = self.default_model or "gemini-3.1-flash-lite-preview" - elif self.name == ProviderType.MISTRAL: - self.base_url = self.base_url or "https://api.mistral.ai/v1" - self.models = self.models or ["mistral-small-latest", "mistral-large-latest"] - self.default_model = self.default_model or "mistral-small-latest" - elif self.name == ProviderType.DEEPSEEK: - self.base_url = self.base_url or "https://api.deepseek.com" - self.models = self.models or ["deepseek-chat", "deepseek-reasoner"] - self.default_model = self.default_model or "deepseek-chat" + """Apply per-provider defaults for base_url, models, and default_model.""" + defaults = _PROVIDER_DEFAULTS.get(self.name) + if defaults is None: + return + base_url, models, default_model = defaults + self.base_url = self.base_url or base_url + self.models = self.models or models + self.default_model = self.default_model or default_model @dataclass(unsafe_hash=True) diff --git a/tests/experimental/security/conftest.py b/tests/experimental/security/conftest.py index 00cd556..5ddd16c 100644 --- a/tests/experimental/security/conftest.py +++ b/tests/experimental/security/conftest.py @@ -77,6 +77,7 @@ def _contain_security_module_side_effects(tmp_path_factory: pytest.TempPathFacto @pytest.fixture def tmp_compliance_monitor(tmp_path: Path) -> Any: + """Return a ComplianceMonitor backed by a temporary SQLite database.""" from pentestgpt.experimental.security.compliance_monitor import ComplianceMonitor return ComplianceMonitor(db_path=str(tmp_path / "compliance.db")) @@ -84,6 +85,7 @@ def tmp_compliance_monitor(tmp_path: Path) -> Any: @pytest.fixture def tmp_vulnerability_tracker(tmp_path: Path) -> Any: + """Return a VulnerabilityTracker backed by a temporary SQLite database.""" from pentestgpt.experimental.security.vulnerability_tracker import ( VulnerabilityTracker, )