diff --git a/pentestgpt/agents/opsec_agent.py b/pentestgpt/agents/opsec_agent.py
index 9318da3..0fc86fc 100644
--- a/pentestgpt/agents/opsec_agent.py
+++ b/pentestgpt/agents/opsec_agent.py
@@ -203,7 +203,7 @@ def validate_operation(
logger.info(f"Validating operation: {operation}")
# Log the validation attempt
- audit_entry = {
+ audit_entry: dict[str, Any] = {
"timestamp": datetime.now().isoformat(),
"operation": operation,
"context": context,
@@ -222,7 +222,7 @@ def validate_operation(
is_compliant = all(check[0] for check in compliance_checks)
reason = "; ".join(check[1] for check in compliance_checks if not check[0])
- audit_entry["result"] = str(is_compliant)
+ audit_entry["result"] = is_compliant
audit_entry["reason"] = reason
self.audit_log.append(audit_entry)
diff --git a/pytest.ini b/pytest.ini
index f07f9ca..fd62dd9 100644
--- a/pytest.ini
+++ b/pytest.ini
@@ -4,7 +4,7 @@ testpaths = tests
python_files = test_*.py
python_classes = Test*
python_functions = test_*
-addopts = -v --tb=short --strict-markers
+addopts = -v --tb=short --strict-markers -m "not llm"
markers =
unit: Unit tests (auto-applied to component tests)
integration: Integration tests (tests/integration/)
@@ -17,6 +17,7 @@ markers =
db: Database tests
config: Configuration tests
core: Core component tests
+ llm: Live LLM tests (require API key, run with: pytest -m llm)
asyncio_mode = auto
filterwarnings =
ignore::DeprecationWarning
diff --git a/tests/__init__.py b/tests/__init__.py
index 2a29fef..ec50e0c 100644
--- a/tests/__init__.py
+++ b/tests/__init__.py
@@ -1 +1,14 @@
"""Test suite for PentestGPT."""
+
+from typing import Any
+from unittest.mock import MagicMock
+
+
+def make_mock_router(
+ route_task_return: Any = ("Mocked LLM response", "openai", 0.01),
+) -> MagicMock:
+ """Create a mock model router. Usable from helper functions (not just fixtures)."""
+ router = MagicMock()
+ router.route_task = MagicMock(return_value=route_task_return)
+ router.get_model_for_task = MagicMock(return_value="gpt-4-test")
+ return router
diff --git a/tests/agents/test_intelligence_gathering_agent.py b/tests/agents/test_intelligence_gathering_agent.py
new file mode 100644
index 0000000..1375a53
--- /dev/null
+++ b/tests/agents/test_intelligence_gathering_agent.py
@@ -0,0 +1,356 @@
+"""Tests for IntelligenceGatheringAgent — OSINT, target profiling, intel collection."""
+
+from datetime import datetime
+
+from pentestgpt.agents.intelligence_gathering_agent import (
+ IntelligenceGatheringAgent,
+ IntelligenceItem,
+ OSINTSource,
+ TargetProfile,
+ ThreatIntel,
+)
+from tests import make_mock_router
+
+
+def _make_agent() -> IntelligenceGatheringAgent:
+ return IntelligenceGatheringAgent(
+ model_router=make_mock_router(route_task_return=("Mocked intel response", "openai", 0.01)),
+ )
+
+
+# ---------------------------------------------------------------------------
+# Dataclass tests
+# ---------------------------------------------------------------------------
+
+
+class TestDataclasses:
+ def test_osint_source_defaults(self) -> None:
+ source = OSINTSource(name="shodan", url="https://api.shodan.io")
+ assert source.rate_limit == 60
+ assert source.reliability_score == 1.0
+ assert source.error_count == 0
+
+ def test_osint_source_required_fields(self) -> None:
+ source = OSINTSource(name="custom-source", url="https://custom.example.com")
+ assert source.name == "custom-source"
+ assert source.url == "https://custom.example.com"
+ assert source.api_key_required is False
+ assert source.response_format == "json"
+ assert source.last_used is None
+
+ def test_intelligence_item_fields(self) -> None:
+ item = IntelligenceItem(
+ id="intel-001",
+ source="shodan",
+ type="host",
+ category="infrastructure",
+ content={"ip": "10.0.0.1"},
+ confidence=0.9,
+ relevance=0.8,
+ timestamp=datetime.now(),
+ )
+ assert item.validation_status == "pending"
+ assert item.tags == []
+
+ def test_intelligence_item_custom_tags(self) -> None:
+ item = IntelligenceItem(
+ id="intel-002",
+ source="virustotal",
+ type="malware",
+ category="threat",
+ content={"hash": "abc123"},
+ confidence=0.95,
+ relevance=0.9,
+ timestamp=datetime.now(),
+ tags=["critical", "ransomware"],
+ validation_status="validated",
+ )
+ assert item.tags == ["critical", "ransomware"]
+ assert item.validation_status == "validated"
+ assert item.expires_at is None
+
+ def test_target_profile_fields(self) -> None:
+ profile = TargetProfile(
+ target_id="t-001",
+ target_type="domain",
+ basic_info={"domain": "example.com"},
+ technical_info={},
+ social_info={},
+ threat_info={},
+ )
+ assert profile.intelligence_count == 0
+ assert profile.relationships == []
+
+ def test_target_profile_last_updated(self) -> None:
+ before = datetime.now()
+ profile = TargetProfile(
+ target_id="t-002",
+ target_type="ip",
+ basic_info={},
+ technical_info={},
+ social_info={},
+ threat_info={},
+ )
+ after = datetime.now()
+ assert before <= profile.last_updated <= after
+
+ def test_threat_intel_fields(self) -> None:
+ now = datetime.now()
+ intel = ThreatIntel(
+ threat_id="threat-001",
+ threat_type="malware",
+ indicators=["malware_hash_abc"],
+ ttps=["T1059"],
+ severity="high",
+ first_seen=now,
+ last_seen=now,
+ confidence=0.85,
+ source="internal",
+ context={"campaign": "test-campaign"},
+ )
+ assert intel.threat_type == "malware"
+ assert intel.severity == "high"
+ assert intel.confidence == 0.85
+ assert intel.source == "internal"
+ assert "malware_hash_abc" in intel.indicators
+ assert "T1059" in intel.ttps
+
+ def test_threat_intel_empty_lists(self) -> None:
+ now = datetime.now()
+ intel = ThreatIntel(
+ threat_id="threat-002",
+ threat_type="campaign",
+ indicators=[],
+ ttps=[],
+ severity="low",
+ first_seen=now,
+ last_seen=now,
+ confidence=0.5,
+ source="external",
+ context={},
+ )
+ assert intel.indicators == []
+ assert intel.ttps == []
+ assert intel.context == {}
+
+
+# ---------------------------------------------------------------------------
+# Agent initialization tests
+# ---------------------------------------------------------------------------
+
+
+class TestAgentInitialization:
+ def test_agent_creates_with_defaults(self) -> None:
+ agent = _make_agent()
+ assert agent.agent_id == "intelligence_gathering_agent"
+ assert isinstance(agent.osint_sources, dict)
+ assert len(agent.osint_sources) > 0
+
+ def test_agent_custom_id(self) -> None:
+ agent = IntelligenceGatheringAgent(
+ model_router=make_mock_router(route_task_return=("Mocked intel response", "openai", 0.01)),
+ agent_id="custom-intel",
+ )
+ assert agent.agent_id == "custom-intel"
+
+ def test_collection_stats_initialized(self) -> None:
+ agent = _make_agent()
+ assert agent.collection_stats["total_collections"] == 0
+ assert agent.collection_stats["successful_collections"] == 0
+ assert agent.collection_stats["failed_collections"] == 0
+ assert agent.collection_stats["intelligence_items_collected"] == 0
+ assert agent.collection_stats["threats_identified"] == 0
+
+ def test_osint_sources_populated(self) -> None:
+ agent = _make_agent()
+ # Verify known OSINT sources are loaded
+ source_names = list(agent.osint_sources.keys())
+ assert "Shodan" in source_names
+ assert "VirusTotal" in source_names
+ assert "Censys" in source_names
+
+ def test_osint_sources_are_osint_source_instances(self) -> None:
+ agent = _make_agent()
+ for source in agent.osint_sources.values():
+ assert isinstance(source, OSINTSource)
+
+ def test_agent_storage_initialized_empty(self) -> None:
+ agent = _make_agent()
+ assert agent.intelligence_items == {}
+ assert agent.target_profiles == {}
+ assert agent.threat_intel == {}
+ assert agent.active_collections == {}
+ assert agent.collection_history == []
+
+ def test_agent_integration_flags(self) -> None:
+ agent = _make_agent()
+ assert agent.knowledge_system_integration is True
+ assert agent.monitoring_integration is True
+ assert agent.strategy_agent_integration is True
+
+ def test_max_concurrent_collections(self) -> None:
+ agent = _make_agent()
+ assert agent.max_concurrent_collections == 5
+
+
+# ---------------------------------------------------------------------------
+# Intelligence operations tests
+# ---------------------------------------------------------------------------
+
+
+class TestIntelligenceOperations:
+ def test_get_intelligence_summary(self) -> None:
+ agent = _make_agent()
+ summary = agent.get_intelligence_summary()
+ assert isinstance(summary, dict)
+
+ def test_get_intelligence_summary_keys(self) -> None:
+ agent = _make_agent()
+ summary = agent.get_intelligence_summary()
+ assert "agent_id" in summary
+ assert "status" in summary
+ assert "intelligence_summary" in summary
+ assert "threat_intelligence" in summary
+ assert "collection_stats" in summary
+ assert "osint_sources" in summary
+
+ def test_get_intelligence_summary_agent_id(self) -> None:
+ agent = _make_agent()
+ summary = agent.get_intelligence_summary()
+ assert summary["agent_id"] == "intelligence_gathering_agent"
+
+ def test_get_intelligence_summary_osint_sources(self) -> None:
+ agent = _make_agent()
+ summary = agent.get_intelligence_summary()
+ osint = summary["osint_sources"]
+ assert osint["configured"] == len(agent.osint_sources)
+ assert isinstance(osint["available"], list)
+ assert len(osint["available"]) > 0
+
+ def test_dns_intelligence_helper(self) -> None:
+ agent = _make_agent()
+ result = agent._gather_dns_intelligence("example.com")
+ assert result is not None
+ assert isinstance(result, dict)
+
+ def test_dns_intelligence_structure(self) -> None:
+ agent = _make_agent()
+ result = agent._gather_dns_intelligence("example.com")
+ assert result is not None
+ assert result["domain"] == "example.com"
+ assert "a_records" in result
+ assert "mx_records" in result
+ assert "ns_records" in result
+ assert "dnssec_enabled" in result
+
+ def test_email_analysis_helper(self) -> None:
+ agent = _make_agent()
+ result = agent._analyze_email_format("test@example.com")
+ assert result is not None
+ assert isinstance(result, dict)
+
+ def test_email_analysis_structure(self) -> None:
+ agent = _make_agent()
+ result = agent._analyze_email_format("test@example.com")
+ assert result is not None
+ assert result["email"] == "test@example.com"
+ assert result["domain"] == "example.com"
+ assert result["local_part"] == "test"
+ assert result["is_valid_format"] is True
+
+ def test_email_analysis_invalid_email(self) -> None:
+ agent = _make_agent()
+ result = agent._analyze_email_format("not-an-email")
+ assert result is not None
+ assert result["is_valid_format"] is False
+
+ def test_url_structure_analysis(self) -> None:
+ agent = _make_agent()
+ result = agent._analyze_url_structure("https://example.com/admin")
+ assert result is not None
+ assert isinstance(result, dict)
+ assert result["scheme"] == "https"
+ assert result["netloc"] == "example.com"
+ assert result["path"] == "/admin"
+
+ def test_url_structure_analysis_with_query(self) -> None:
+ agent = _make_agent()
+ result = agent._analyze_url_structure("https://example.com/admin?q=1&page=2")
+ assert result is not None
+ assert result["query"] == "q=1&page=2"
+ assert "suspicious_patterns" in result
+
+ def test_url_structure_analysis_no_query(self) -> None:
+ agent = _make_agent()
+ result = agent._analyze_url_structure("https://example.com/page")
+ assert result is not None
+ assert result["query"] == ""
+
+ def test_url_structure_analysis_with_port(self) -> None:
+ agent = _make_agent()
+ result = agent._analyze_url_structure("https://example.com:8443/api")
+ assert result is not None
+ assert result["netloc"] == "example.com:8443"
+ assert result["path"] == "/api"
+
+ def test_url_structure_preserves_original(self) -> None:
+ agent = _make_agent()
+ url = "https://example.com/path?k=v"
+ result = agent._analyze_url_structure(url)
+ assert result is not None
+ assert result["original_url"] == url
+
+ def test_enumerate_subdomains_returns_list(self) -> None:
+ agent = _make_agent()
+ result = agent._enumerate_subdomains("example.com")
+ assert isinstance(result, list)
+ assert len(result) > 0
+
+ def test_enumerate_subdomains_are_valid(self) -> None:
+ agent = _make_agent()
+ result = agent._enumerate_subdomains("example.com")
+ for subdomain in result:
+ assert isinstance(subdomain, str)
+ assert subdomain.endswith("example.com")
+
+ def test_enumerate_subdomains_domain_specific(self) -> None:
+ agent = _make_agent()
+ result = agent._enumerate_subdomains("other.org")
+ for subdomain in result:
+ assert subdomain.endswith("other.org")
+
+ def test_start_collection_returns_string_id(self) -> None:
+ agent = _make_agent()
+ result = agent.start_collection(
+ target="example.com",
+ collection_type="domain",
+ )
+ assert isinstance(result, str)
+ assert result.startswith("col_")
+
+ def test_start_collection_increments_total_stat(self) -> None:
+ agent = _make_agent()
+ agent.start_collection(target="example.com", collection_type="domain")
+ # stat increments synchronously before the thread does its work
+ assert agent.collection_stats["total_collections"] == 1
+
+ def test_start_collection_multiple_types(self) -> None:
+ agent = _make_agent()
+ col_id_1 = agent.start_collection(target="10.0.0.1", collection_type="ip")
+ col_id_2 = agent.start_collection(target="user@example.com", collection_type="email")
+ assert col_id_1 != col_id_2
+ assert agent.collection_stats["total_collections"] == 2
+
+ def test_start_collection_returns_empty_when_at_capacity(self) -> None:
+ agent = _make_agent()
+ # Fill active_collections to the max capacity with fake entries
+ for i in range(agent.max_concurrent_collections):
+ agent.active_collections[f"col_fake_{i}"] = {"status": "active"}
+ result = agent.start_collection(target="example.com", collection_type="domain")
+ assert result == ""
+
+ def test_intelligence_summary_after_no_items(self) -> None:
+ agent = _make_agent()
+ summary = agent.get_intelligence_summary()
+ assert summary["intelligence_summary"]["total_items"] == 0
+ assert summary["threat_intelligence"]["total_threats"] == 0
diff --git a/tests/agents/test_knowledge_manager.py b/tests/agents/test_knowledge_manager.py
new file mode 100644
index 0000000..2c2b6e1
--- /dev/null
+++ b/tests/agents/test_knowledge_manager.py
@@ -0,0 +1,137 @@
+"""Tests for KnowledgeManager — exploit KB, vulnerability DB, payload library."""
+
+import pytest
+
+from pentestgpt.agents.knowledge_manager import KnowledgeManager
+
+
+@pytest.fixture(scope="module")
+def km() -> KnowledgeManager:
+ return KnowledgeManager()
+
+
+class TestKnowledgeManagerInit:
+
+ def test_initializes_successfully(self, km: KnowledgeManager) -> None:
+ assert km is not None
+
+ @pytest.mark.parametrize(
+ "getter",
+ [
+ "get_exploit_knowledge_base",
+ "get_vulnerability_database",
+ "get_payload_library",
+ ],
+ )
+ def test_knowledge_stores_are_non_empty_dicts(self, km: KnowledgeManager, getter: str) -> None:
+ result = getattr(km, getter)()
+ assert isinstance(result, dict)
+ assert len(result) > 0
+
+
+class TestKnowledgeLoading:
+
+ def test_load_exploit_knowledge(self, km: KnowledgeManager) -> None:
+ km.load_exploit_knowledge()
+ kb = km.get_exploit_knowledge_base()
+ assert isinstance(kb, dict)
+
+ @pytest.mark.parametrize("category", ["web_exploits", "network_exploits", "privilege_escalation"])
+ def test_exploit_kb_has_category(self, km: KnowledgeManager, category: str) -> None:
+ kb = km.get_exploit_knowledge_base()
+ assert category in kb
+
+ @pytest.mark.parametrize(
+ ("category", "key"),
+ [
+ ("web_exploits", "sql_injection"),
+ ("web_exploits", "xss"),
+ ("web_exploits", "command_injection"),
+ ("network_exploits", "smb"),
+ ("network_exploits", "ssh"),
+ ("privilege_escalation", "linux"),
+ ("privilege_escalation", "windows"),
+ ],
+ )
+ def test_exploit_kb_has_technique(self, km: KnowledgeManager, category: str, key: str) -> None:
+ kb = km.get_exploit_knowledge_base()
+ assert key in kb.get(category, {})
+
+ @pytest.mark.parametrize("field", ["techniques", "payloads"])
+ def test_sql_injection_has_list_field(self, km: KnowledgeManager, field: str) -> None:
+ sqli = km.get_exploit_knowledge_base().get("web_exploits", {}).get("sql_injection", {})
+ assert isinstance(sqli.get(field), list)
+ assert len(sqli.get(field, [])) > 0
+
+ @pytest.mark.parametrize("cve_id", ["CVE-2017-0144", "CVE-2019-0708"])
+ def test_vulnerability_db_has_cve(self, km: KnowledgeManager, cve_id: str) -> None:
+ db = km.get_vulnerability_database()
+ assert "cves" in db
+ assert cve_id in db["cves"]
+
+ @pytest.mark.parametrize("field", ["name", "type", "severity", "cvss_score", "description"])
+ def test_cve_has_expected_field(self, km: KnowledgeManager, field: str) -> None:
+ cve = km.get_vulnerability_database().get("cves", {}).get("CVE-2017-0144", {})
+ assert field in cve
+
+ @pytest.mark.parametrize("key", ["sql_injection", "xss", "command_injection"])
+ def test_common_vulns_has_key(self, km: KnowledgeManager, key: str) -> None:
+ db = km.get_vulnerability_database()
+ assert "common_vulnerabilities" in db
+ assert key in db["common_vulnerabilities"]
+
+ @pytest.mark.parametrize("category", ["reverse_shells", "bind_shells", "data_exfiltration"])
+ def test_payload_library_has_category(self, km: KnowledgeManager, category: str) -> None:
+ lib = km.get_payload_library()
+ assert category in lib
+
+ @pytest.mark.parametrize("key", ["bash", "python"])
+ def test_reverse_shells_has_template(self, km: KnowledgeManager, key: str) -> None:
+ lib = km.get_payload_library()
+ assert key in lib.get("reverse_shells", {})
+
+
+class TestKnowledgeUpdates:
+
+ def test_check_knowledge_updates_runs(self) -> None:
+ km = KnowledgeManager()
+ assert km._last_knowledge_check is None
+ km.load_exploit_knowledge()
+ km.check_knowledge_updates()
+ assert km._last_knowledge_check is not None
+
+ def test_check_knowledge_updates_after_init(self) -> None:
+ km = KnowledgeManager()
+ km.check_knowledge_updates()
+ assert km._last_knowledge_check is not None
+
+ def test_last_knowledge_check_is_set_after_update(self) -> None:
+ km = KnowledgeManager()
+ assert km._last_knowledge_check is None
+ km.check_knowledge_updates()
+ assert km._last_knowledge_check is not None
+
+ def test_multiple_update_checks_work(self) -> None:
+ km = KnowledgeManager()
+ km.check_knowledge_updates()
+ first_check = km._last_knowledge_check
+ km.check_knowledge_updates()
+ second_check = km._last_knowledge_check
+ assert first_check is not None
+ assert second_check is not None
+
+ @pytest.mark.parametrize(
+ "getter",
+ [
+ "get_exploit_knowledge_base",
+ "get_vulnerability_database",
+ "get_payload_library",
+ ],
+ )
+ def test_knowledge_persists_after_updates(self, getter: str) -> None:
+ km = KnowledgeManager()
+ before = getattr(km, getter)()
+ km.check_knowledge_updates()
+ after = getattr(km, getter)()
+ assert isinstance(before, dict)
+ assert isinstance(after, dict)
diff --git a/tests/agents/test_opsec_agent.py b/tests/agents/test_opsec_agent.py
new file mode 100644
index 0000000..39f956f
--- /dev/null
+++ b/tests/agents/test_opsec_agent.py
@@ -0,0 +1,444 @@
+"""Tests for OPSECAgent — environment profiling, compliance, session management."""
+
+import json
+from typing import Any
+from unittest.mock import MagicMock
+
+from pentestgpt.agents.opsec_agent import (
+ EnvironmentProfile,
+ LegalComplianceChecker,
+ OPSECAgent,
+ OPSECSession,
+)
+from pentestgpt.config.stealth_config import (
+ ComplianceLevel,
+ OPSECConfig,
+ StealthConfig,
+ StealthLevel,
+)
+from tests import make_mock_router
+
+# ---------------------------------------------------------------------------
+# Factories
+# ---------------------------------------------------------------------------
+
+
+def _make_opsec_agent() -> OPSECAgent:
+ return OPSECAgent(
+ model_router=make_mock_router(route_task_return=("Mocked response", "openai", 0.01)),
+ config=StealthConfig(),
+ )
+
+
+# ---------------------------------------------------------------------------
+# EnvironmentProfile tests
+# ---------------------------------------------------------------------------
+
+
+class TestEnvironmentProfile:
+ def test_profile_has_host_info(self) -> None:
+ profile = EnvironmentProfile()
+ assert isinstance(profile.host_info, dict)
+
+ def test_profile_has_system_info(self) -> None:
+ profile = EnvironmentProfile()
+ assert isinstance(profile.system_info, dict)
+
+ def test_profile_hash_is_deterministic(self) -> None:
+ profile = EnvironmentProfile()
+ assert isinstance(profile.profile_hash, str)
+ assert len(profile.profile_hash) > 0
+
+ def test_profile_hash_length(self) -> None:
+ """Profile hash is a 16-char hex substring of SHA-256."""
+ profile = EnvironmentProfile()
+ assert len(profile.profile_hash) == 16
+
+ def test_profile_has_security_tools_list(self) -> None:
+ profile = EnvironmentProfile()
+ assert isinstance(profile.security_tools, list)
+
+ def test_to_dict_contains_all_fields(self) -> None:
+ profile = EnvironmentProfile()
+ d = profile.to_dict()
+ assert "host_info" in d
+ assert "system_info" in d
+ assert "security_tools" in d
+ assert "network_info" in d
+ assert "profile_hash" in d
+ assert "timestamp" in d
+
+ def test_to_dict_returns_dict(self) -> None:
+ profile = EnvironmentProfile()
+ d = profile.to_dict()
+ assert isinstance(d, dict)
+
+ def test_system_info_contains_platform_keys(self) -> None:
+ profile = EnvironmentProfile()
+ assert "platform" in profile.system_info or "system" in profile.system_info
+
+
+# ---------------------------------------------------------------------------
+# LegalComplianceChecker tests
+# ---------------------------------------------------------------------------
+
+
+class TestLegalComplianceChecker:
+ def test_compliant_operation_with_authorization(self) -> None:
+ checker = LegalComplianceChecker(jurisdiction="US", compliance_level=ComplianceLevel.STRICT)
+ is_compliant, _reason = checker.validate_operation(
+ "port_scan",
+ {"authorized": True, "scope": {"targets": ["10.0.0.0/24"]}},
+ )
+ assert is_compliant is True
+
+ def test_noncompliant_without_authorization(self) -> None:
+ checker = LegalComplianceChecker()
+ is_compliant, reason = checker.validate_operation(
+ "port_scan",
+ {"authorized": False, "scope": {}},
+ )
+ assert is_compliant is False
+ assert "authorization" in reason.lower() or "scope" in reason.lower()
+
+ def test_noncompliant_without_scope(self) -> None:
+ checker = LegalComplianceChecker()
+ is_compliant, _reason = checker.validate_operation(
+ "port_scan",
+ {"authorized": True, "scope": {}},
+ )
+ assert is_compliant is False
+
+ def test_audit_log_records_validations(self) -> None:
+ checker = LegalComplianceChecker()
+ checker.validate_operation("scan", {"authorized": True, "scope": {"targets": ["10.0.0.1"]}})
+ checker.validate_operation("exploit", {"authorized": False, "scope": {}})
+ assert len(checker.audit_log) == 2
+
+ def test_generate_audit_report(self) -> None:
+ checker = LegalComplianceChecker()
+ checker.validate_operation("scan", {"authorized": True, "scope": {"targets": ["10.0.0.1"]}})
+ report = checker.generate_audit_report()
+ assert isinstance(report, dict)
+
+ def test_audit_report_contains_jurisdiction(self) -> None:
+ checker = LegalComplianceChecker(jurisdiction="EU")
+ report = checker.generate_audit_report()
+ assert report["jurisdiction"] == "EU"
+
+ def test_audit_report_counts_operations(self) -> None:
+ checker = LegalComplianceChecker()
+ checker.validate_operation("scan", {"authorized": True, "scope": {"targets": ["10.0.0.1"]}})
+ checker.validate_operation("exploit", {"authorized": False, "scope": {}})
+ report = checker.generate_audit_report()
+ assert report["total_operations"] == 2
+ assert report["compliant_operations"] + report["non_compliant_operations"] == report["total_operations"]
+
+ def test_audit_report_non_compliant_count(self) -> None:
+ checker = LegalComplianceChecker()
+ checker.validate_operation("scan", {"authorized": True, "scope": {"targets": ["10.0.0.1"]}})
+ checker.validate_operation("exploit", {"authorized": False, "scope": {}})
+ report = checker.generate_audit_report()
+ assert report["compliant_operations"] >= 1
+ assert report["non_compliant_operations"] >= 1
+
+ def test_illegal_operation_is_noncompliant(self) -> None:
+ checker = LegalComplianceChecker(jurisdiction="US")
+ is_compliant, _reason = checker.validate_operation(
+ "unauthorized_access",
+ {"authorized": True, "scope": {"targets": ["10.0.0.1"]}},
+ )
+ assert is_compliant is False
+
+ def test_default_jurisdiction_is_us(self) -> None:
+ checker = LegalComplianceChecker()
+ assert checker.jurisdiction == "US"
+
+ def test_default_compliance_level_is_strict(self) -> None:
+ checker = LegalComplianceChecker()
+ assert checker.compliance_level == ComplianceLevel.STRICT
+
+ def test_audit_log_starts_empty(self) -> None:
+ checker = LegalComplianceChecker()
+ assert checker.audit_log == []
+
+ def test_validate_operation_returns_tuple(self) -> None:
+ checker = LegalComplianceChecker()
+ result = checker.validate_operation("scan", {"authorized": True, "scope": {"targets": ["10.0.0.1"]}})
+ assert isinstance(result, tuple)
+ assert len(result) == 2
+
+
+# ---------------------------------------------------------------------------
+# OPSECSession tests
+# ---------------------------------------------------------------------------
+
+
+class TestOPSECSession:
+ def test_session_stores_session_id(self) -> None:
+ session = OPSECSession("sess-001", OPSECConfig())
+ assert session.session_id == "sess-001"
+
+ def test_start_session_creates_environment_profile(self) -> None:
+ session = OPSECSession("sess-002", OPSECConfig())
+ session.start_session()
+ assert session.environment_profile is not None
+ assert isinstance(session.environment_profile, EnvironmentProfile)
+
+ def test_start_session_creates_compliance_checker(self) -> None:
+ config = OPSECConfig(legal_validation=True)
+ session = OPSECSession("sess-003", config)
+ session.start_session()
+ assert session.compliance_checker is not None
+
+ def test_start_session_no_compliance_checker_when_disabled(self) -> None:
+ config = OPSECConfig(legal_validation=False)
+ session = OPSECSession("sess-004", config)
+ session.start_session()
+ assert session.compliance_checker is None
+
+ def test_end_session_logs_activity(self) -> None:
+ config = OPSECConfig(cleanup_automation=False)
+ session = OPSECSession("sess-005", config)
+ session.start_session()
+ session.end_session()
+ activity_types = [a["activity_type"] for a in session.activity_log]
+ assert "session_ended" in activity_types
+
+ def test_get_session_stats_returns_dict(self) -> None:
+ session = OPSECSession("sess-006", OPSECConfig())
+ session.start_session()
+ stats = session.get_session_stats()
+ assert isinstance(stats, dict)
+
+ def test_get_session_stats_contains_session_id(self) -> None:
+ session = OPSECSession("sess-007", OPSECConfig())
+ session.start_session()
+ stats = session.get_session_stats()
+ assert "session_id" in stats
+ assert stats["session_id"] == "sess-007"
+
+ def test_get_session_stats_contains_duration(self) -> None:
+ session = OPSECSession("sess-008", OPSECConfig())
+ session.start_session()
+ stats = session.get_session_stats()
+ assert "duration_seconds" in stats
+ assert stats["duration_seconds"] >= 0
+
+ def test_get_session_stats_without_start(self) -> None:
+ """Session stats work even before start_session is called."""
+ session = OPSECSession("sess-009", OPSECConfig())
+ stats = session.get_session_stats()
+ assert isinstance(stats, dict)
+
+ def test_validate_operation_no_compliance_checker(self) -> None:
+ config = OPSECConfig(legal_validation=False)
+ session = OPSECSession("sess-010", config)
+ session.start_session()
+ is_ok, _reason = session.validate_operation("scan", {"authorized": True, "scope": {}})
+ assert is_ok is True
+
+ def test_validate_operation_with_compliance_checker(self) -> None:
+ config = OPSECConfig(legal_validation=True)
+ session = OPSECSession("sess-011", config)
+ session.start_session()
+ is_ok, _reason = session.validate_operation(
+ "scan",
+ {"authorized": True, "scope": {"targets": ["10.0.0.1"]}},
+ )
+ assert is_ok is True
+
+ def test_activity_log_is_initially_empty(self) -> None:
+ session = OPSECSession("sess-012", OPSECConfig())
+ assert session.activity_log == []
+
+
+# ---------------------------------------------------------------------------
+# OPSECAgent tests
+# ---------------------------------------------------------------------------
+
+
+class TestOPSECAgent:
+ def test_initialization(self) -> None:
+ agent = _make_opsec_agent()
+ assert agent.metadata.agent_name == "OPSECAgent"
+
+ def test_initialization_metadata_type(self) -> None:
+ agent = _make_opsec_agent()
+ assert agent.metadata.agent_type == "OPSEC"
+
+ def test_create_session_returns_id(self) -> None:
+ agent = _make_opsec_agent()
+ session_id = agent.create_session()
+ assert isinstance(session_id, str)
+ assert len(session_id) > 0
+
+ def test_create_session_sets_current_session(self) -> None:
+ agent = _make_opsec_agent()
+ agent.create_session()
+ assert agent.current_session is not None
+
+ def test_create_session_adds_to_history(self) -> None:
+ agent = _make_opsec_agent()
+ agent.create_session()
+ assert len(agent.session_history) == 1
+
+ def test_end_current_session_clears_current(self) -> None:
+ agent = _make_opsec_agent()
+ agent.create_session()
+ agent.end_current_session()
+ assert agent.current_session is None
+
+ def test_end_current_session_without_session(self) -> None:
+ """Calling end_current_session with no session should not raise."""
+ agent = _make_opsec_agent()
+ agent.end_current_session() # should not raise
+
+ def test_get_health_status(self) -> None:
+ agent = _make_opsec_agent()
+ status = agent.get_health_status()
+ assert status["is_healthy"] is True
+
+ def test_get_health_status_is_dict(self) -> None:
+ agent = _make_opsec_agent()
+ status = agent.get_health_status()
+ assert isinstance(status, dict)
+
+ def test_analyze_target_environment_returns_dict(self) -> None:
+ """analyze_target_environment returns a dict (parsed JSON or error dict)."""
+ router = make_mock_router(route_task_return=("Mocked response", "openai", 0.01))
+ # Return valid JSON from the LLM mock
+ mock_analysis = json.dumps(
+ {
+ "security_measures": ["firewall"],
+ "detection_risks": ["IDS"],
+ "evasion_strategies": ["timing"],
+ "legal_considerations": ["authorization required"],
+ "recommended_opsec_level": "balanced",
+ "precautionary_measures": ["vpn"],
+ "risk_assessment": "medium",
+ }
+ )
+ router.route_task = MagicMock(return_value=(mock_analysis, "openai", 0.01))
+ agent = OPSECAgent(model_router=router, config=StealthConfig())
+ result = agent.analyze_target_environment({"target": "10.0.0.1", "type": "network"})
+ assert isinstance(result, dict)
+
+ def test_analyze_target_environment_error_returns_dict(self) -> None:
+ """When LLM returns non-JSON, analyze_target_environment returns an error dict."""
+ router = make_mock_router(route_task_return=("Mocked response", "openai", 0.01))
+ router.route_task = MagicMock(return_value=("not json at all", "openai", 0.01))
+ agent = OPSECAgent(model_router=router, config=StealthConfig())
+ result = agent.analyze_target_environment({"target": "10.0.0.1"})
+ assert isinstance(result, dict)
+ # Error path returns error key and default opsec level
+ assert "error" in result or "recommended_opsec_level" in result
+
+ def test_select_stealth_level_returns_stealth_level_enum(self) -> None:
+ agent = _make_opsec_agent()
+ target_analysis: dict[str, Any] = {
+ "risk_assessment": "medium",
+ "recommended_opsec_level": "balanced",
+ }
+ result = agent.select_stealth_level(target_analysis=target_analysis, context={})
+ assert isinstance(result, StealthLevel)
+
+ def test_select_stealth_level_uses_config_when_not_balanced(self) -> None:
+ """If config.stealth_level != BALANCED, select_stealth_level returns that level."""
+ config = StealthConfig(stealth_level=StealthLevel.MAXIMUM)
+ agent = OPSECAgent(
+ model_router=make_mock_router(route_task_return=("Mocked response", "openai", 0.01)), config=config
+ )
+ result = agent.select_stealth_level(target_analysis={}, context={})
+ assert result == StealthLevel.MAXIMUM
+
+ def test_select_stealth_level_maps_aggressive(self) -> None:
+ config = StealthConfig(stealth_level=StealthLevel.BALANCED)
+ agent = OPSECAgent(
+ model_router=make_mock_router(route_task_return=("Mocked response", "openai", 0.01)), config=config
+ )
+ target_analysis: dict[str, Any] = {
+ "risk_assessment": "medium",
+ "recommended_opsec_level": "aggressive",
+ }
+ result = agent.select_stealth_level(target_analysis=target_analysis, context={})
+ assert result == StealthLevel.AGGRESSIVE
+
+ def test_get_opsec_status_returns_dict(self) -> None:
+ agent = _make_opsec_agent()
+ status = agent.get_opsec_status()
+ assert isinstance(status, dict)
+
+ def test_get_opsec_status_no_active_session(self) -> None:
+ agent = _make_opsec_agent()
+ status = agent.get_opsec_status()
+ assert status["active_session"] is None
+
+ def test_get_opsec_status_with_active_session(self) -> None:
+ agent = _make_opsec_agent()
+ agent.create_session()
+ status = agent.get_opsec_status()
+ assert status["active_session"] is not None
+
+ def test_generate_compliance_report_no_session(self) -> None:
+ agent = _make_opsec_agent()
+ report = agent.generate_compliance_report()
+ # No active session → error dict
+ assert isinstance(report, dict)
+ assert "error" in report
+
+ def test_generate_compliance_report_with_session(self) -> None:
+ config = StealthConfig()
+ config.opsec.legal_validation = True
+ agent = OPSECAgent(
+ model_router=make_mock_router(route_task_return=("Mocked response", "openai", 0.01)), config=config
+ )
+ agent.create_session()
+ report = agent.generate_compliance_report()
+ assert isinstance(report, dict)
+ # Should have real report fields from LegalComplianceChecker
+ assert "jurisdiction" in report or "error" not in report
+
+ def test_validate_operation_no_session(self) -> None:
+ agent = _make_opsec_agent()
+ is_ok, reason = agent.validate_operation("scan", {"authorized": True, "scope": {}})
+ assert is_ok is True
+ assert "No session" in reason
+
+ def test_log_activity_with_no_session_does_not_raise(self) -> None:
+ agent = _make_opsec_agent()
+ agent.log_activity("test_activity", {"detail": "value"}) # should not raise
+
+ def test_log_activity_with_session_records_activity(self) -> None:
+ agent = _make_opsec_agent()
+ agent.create_session()
+ agent.log_activity("test_event", {"key": "value"})
+ assert agent.current_session is not None
+ activity_types = [a["activity_type"] for a in agent.current_session.activity_log]
+ assert "test_event" in activity_types
+
+ def test_add_evidence_file_with_session(self) -> None:
+ agent = _make_opsec_agent()
+ agent.create_session()
+ agent.add_evidence_file("/tmp/test.log")
+ assert agent.current_session is not None
+ assert "/tmp/test.log" in agent.current_session.evidence_files
+
+ def test_add_cleanup_task_with_session(self) -> None:
+ agent = _make_opsec_agent()
+ agent.create_session()
+ agent.add_cleanup_task("clear_history")
+ assert agent.current_session is not None
+ assert "clear_history" in agent.current_session.cleanup_tasks
+
+ def test_cleanup_evidence_no_session(self) -> None:
+ agent = _make_opsec_agent()
+ result = agent.cleanup_evidence()
+ assert "error" in result
+
+ def test_multiple_sessions_tracked_in_history(self) -> None:
+ agent = _make_opsec_agent()
+ agent.create_session()
+ agent.end_current_session()
+ agent.create_session()
+ assert len(agent.session_history) == 2
diff --git a/tests/agents/test_post_exploitation_handler.py b/tests/agents/test_post_exploitation_handler.py
new file mode 100644
index 0000000..e6f53de
--- /dev/null
+++ b/tests/agents/test_post_exploitation_handler.py
@@ -0,0 +1,116 @@
+"""Tests for PostExploitationHandler — post-exploit workflows."""
+
+from pentestgpt.agents.post_exploitation_handler import PostExploitationHandler
+
+
+def _make_handler() -> PostExploitationHandler:
+ """Factory helper for PostExploitationHandler."""
+ return PostExploitationHandler()
+
+
+class TestPostExploitationHandler:
+ """Tests for PostExploitationHandler class."""
+
+ def test_initialization(self) -> None:
+ """Test handler initialization."""
+ handler = _make_handler()
+ assert handler is not None
+ assert isinstance(handler, PostExploitationHandler)
+
+ def test_perform_post_exploitation(self) -> None:
+ """Test perform_post_exploitation returns dict."""
+ handler = _make_handler()
+ result = handler.perform_post_exploitation(session_id="sess-001", target="10.0.0.1")
+ assert isinstance(result, dict)
+ assert "system_information" in result
+ assert "network_information" in result
+ assert "file_system_access" in result
+
+ def test_perform_post_exploitation_with_different_target(self) -> None:
+ """Test perform_post_exploitation with different target."""
+ handler = _make_handler()
+ result = handler.perform_post_exploitation(session_id="sess-002", target="192.168.1.1")
+ assert isinstance(result, dict)
+
+ def test_perform_lateral_movement(self) -> None:
+ """Test perform_lateral_movement returns dict."""
+ handler = _make_handler()
+ result = handler.perform_lateral_movement(session_id="sess-001", target="10.0.0.2")
+ assert isinstance(result, dict)
+ assert "paths" in result
+ assert "successful_pivots" in result
+ assert "credentials_used" in result
+
+ def test_perform_lateral_movement_paths_list(self) -> None:
+ """Test that lateral movement returns valid paths list."""
+ handler = _make_handler()
+ result = handler.perform_lateral_movement(session_id="sess-003", target="10.0.0.5")
+ assert isinstance(result.get("paths"), list)
+ assert isinstance(result.get("successful_pivots"), int)
+ assert isinstance(result.get("credentials_used"), list)
+
+ def test_establish_persistence(self) -> None:
+ """Test establish_persistence returns dict."""
+ handler = _make_handler()
+ result = handler.establish_persistence(session_id="sess-001", target="10.0.0.1")
+ assert isinstance(result, dict)
+ assert "mechanisms" in result
+ assert "persistence_established" in result
+ assert "access_method" in result
+
+ def test_establish_persistence_mechanisms(self) -> None:
+ """Test that persistence returns valid mechanisms list."""
+ handler = _make_handler()
+ result = handler.establish_persistence(session_id="sess-004", target="10.0.0.8")
+ assert isinstance(result.get("mechanisms"), list)
+ assert isinstance(result.get("persistence_established"), bool)
+ assert isinstance(result.get("access_method"), str)
+
+ def test_perform_post_exploitation_empty_session_id(self) -> None:
+ """Test perform_post_exploitation with empty session_id."""
+ handler = _make_handler()
+ result = handler.perform_post_exploitation(session_id="", target="10.0.0.1")
+ assert isinstance(result, dict)
+
+ def test_perform_post_exploitation_empty_target(self) -> None:
+ """Test perform_post_exploitation with empty target."""
+ handler = _make_handler()
+ result = handler.perform_post_exploitation(session_id="sess-001", target="")
+ assert isinstance(result, dict)
+
+ def test_perform_lateral_movement_empty_session_id(self) -> None:
+ """Test perform_lateral_movement with empty session_id."""
+ handler = _make_handler()
+ result = handler.perform_lateral_movement(session_id="", target="10.0.0.2")
+ assert isinstance(result, dict)
+
+ def test_perform_lateral_movement_empty_target(self) -> None:
+ """Test perform_lateral_movement with empty target."""
+ handler = _make_handler()
+ result = handler.perform_lateral_movement(session_id="sess-001", target="")
+ assert isinstance(result, dict)
+
+ def test_establish_persistence_empty_session_id(self) -> None:
+ """Test establish_persistence with empty session_id."""
+ handler = _make_handler()
+ result = handler.establish_persistence(session_id="", target="10.0.0.1")
+ assert isinstance(result, dict)
+
+ def test_establish_persistence_empty_target(self) -> None:
+ """Test establish_persistence with empty target."""
+ handler = _make_handler()
+ result = handler.establish_persistence(session_id="sess-001", target="")
+ assert isinstance(result, dict)
+
+ def test_multiple_handler_instances(self) -> None:
+ """Test that multiple handler instances are independent."""
+ handler1 = _make_handler()
+ handler2 = _make_handler()
+ assert handler1 is not handler2
+
+ def test_post_exploitation_has_credentials(self) -> None:
+ """Test that post_exploitation result includes harvested credentials."""
+ handler = _make_handler()
+ result = handler.perform_post_exploitation(session_id="sess-005", target="10.0.0.9")
+ assert "credentials" in result, "Expected 'credentials' key from _harvest_credentials"
+ assert isinstance(result["credentials"], dict)
diff --git a/tests/agents/test_vulnerability_scanner.py b/tests/agents/test_vulnerability_scanner.py
new file mode 100644
index 0000000..bd4593d
--- /dev/null
+++ b/tests/agents/test_vulnerability_scanner.py
@@ -0,0 +1,309 @@
+"""Tests for VulnerabilityScanner — discovery, applicability, simulation."""
+
+from typing import Any
+
+from pentestgpt.agents.vulnerability_scanner import VulnerabilityScanner
+
+
+def _make_vuln_db() -> dict[str, Any]:
+ """Create a test vulnerability database."""
+ return {
+ "cves": {
+ "CVE-2021-44228": {
+ "name": "Log4Shell",
+ "severity": "critical",
+ "type": "rce",
+ "cvss_score": 9.8,
+ "description": "Critical RCE in Log4j",
+ "affected_versions": ["2.0.0", "2.0.1", "2.14.1"],
+ "exploitation_technique": "JNDI injection",
+ },
+ "CVE-2021-3156": {
+ "name": "Baron Samedit",
+ "severity": "high",
+ "type": "privilege_escalation",
+ "cvss_score": 8.2,
+ "description": "Heap overflow in sudo",
+ "affected_versions": ["1.8.2", "1.8.30"],
+ "exploitation_technique": "Heap spray",
+ },
+ },
+ "common_vulnerabilities": {
+ "sql_injection": {
+ "detection_methods": ["parameter_fuzzing", "time_based_blind"],
+ "exploitation_techniques": ["union_based", "time_based"],
+ },
+ "cross_site_scripting": {
+ "detection_methods": ["payload_injection", "dom_analysis"],
+ "exploitation_techniques": ["cookie_stealing", "phishing"],
+ },
+ },
+ }
+
+
+def _make_scanner() -> VulnerabilityScanner:
+ """Create a test VulnerabilityScanner instance."""
+ return VulnerabilityScanner(vulnerability_database=_make_vuln_db())
+
+
+class TestVulnerabilityScanner:
+ """Tests for VulnerabilityScanner initialization and discovery."""
+
+ def test_initialization(self) -> None:
+ """Test scanner initializes correctly."""
+ scanner = _make_scanner()
+ assert scanner is not None
+ assert scanner.vulnerability_database is not None
+ assert "cves" in scanner.vulnerability_database
+ assert "common_vulnerabilities" in scanner.vulnerability_database
+
+ def test_initialization_empty_db(self) -> None:
+ """Test scanner initializes with empty database."""
+ scanner = VulnerabilityScanner(vulnerability_database={})
+ assert scanner is not None
+ assert scanner.vulnerability_database == {}
+
+ def test_discover_vulnerabilities_populated_params(self) -> None:
+ """Test discovering vulnerabilities with populated parameters."""
+ scanner = _make_scanner()
+ result = scanner.discover_vulnerabilities(
+ target="10.0.0.1",
+ parameters={
+ "target_type": "web_service",
+ "target_version": "2.0.0",
+ },
+ )
+ assert isinstance(result, list)
+ # Result may be empty or populated depending on applicability checks
+ for vuln in result:
+ assert isinstance(vuln, dict)
+ assert "vulnerability_id" in vuln
+ assert "name" in vuln
+ assert "severity" in vuln
+
+ def test_discover_vulnerabilities_empty_parameters(self) -> None:
+ """Test discovering vulnerabilities with empty parameters dict."""
+ scanner = _make_scanner()
+ result = scanner.discover_vulnerabilities(
+ target="10.0.0.1",
+ parameters={},
+ )
+ assert isinstance(result, list)
+
+ def test_discover_vulnerabilities_with_applicable_version(self) -> None:
+ """Test discovery when target version matches affected versions."""
+ scanner = _make_scanner()
+ result = scanner.discover_vulnerabilities(
+ target="10.0.0.1",
+ parameters={
+ "target_type": "unknown",
+ "target_version": "2.0.0", # Matches CVE-2021-44228
+ },
+ )
+ assert isinstance(result, list)
+ # May contain matching vulnerabilities
+
+ def test_discover_vulnerabilities_no_matching_cves(self) -> None:
+ """Test discovery with non-matching target type and version."""
+ scanner = _make_scanner()
+ result = scanner.discover_vulnerabilities(
+ target="10.0.0.1",
+ parameters={
+ "target_type": "unknown",
+ "target_version": "9.9.9", # No matching affected versions
+ },
+ )
+ assert isinstance(result, list)
+
+ def test_discover_vulnerabilities_with_web_service(self) -> None:
+ """Test discovery with web service target type."""
+ scanner = _make_scanner()
+ result = scanner.discover_vulnerabilities(
+ target="example.com",
+ parameters={
+ "target_type": "web_service",
+ "target_version": "1.0.0",
+ },
+ )
+ assert isinstance(result, list)
+
+ def test_discover_vulnerabilities_with_network_service(self) -> None:
+ """Test discovery with network service target type."""
+ scanner = _make_scanner()
+ result = scanner.discover_vulnerabilities(
+ target="10.0.0.5",
+ parameters={
+ "target_type": "network_service",
+ "target_version": "1.0.0",
+ },
+ )
+ assert isinstance(result, list)
+
+ def test_check_vulnerability_applicability_matching_type(self) -> None:
+ """Test applicability check for matching vulnerability type."""
+ scanner = _make_scanner()
+ vuln_data = {
+ "name": "Test RCE",
+ "type": "web_exploitation",
+ "severity": "high",
+ "affected_versions": [],
+ }
+ result = scanner._check_vulnerability_applicability(
+ vuln_data=vuln_data,
+ target="10.0.0.1",
+ parameters={"target_type": "web_service", "target_version": "1.0.0"},
+ )
+ assert isinstance(result, bool)
+ assert result is True # Should match web_service → web_exploitation
+
+ def test_check_vulnerability_applicability_non_matching_type(self) -> None:
+ """Test applicability check for non-matching vulnerability type."""
+ scanner = _make_scanner()
+ vuln_data = {
+ "name": "Test RCE",
+ "type": "web_exploitation",
+ "severity": "high",
+ "affected_versions": [],
+ }
+ result = scanner._check_vulnerability_applicability(
+ vuln_data=vuln_data,
+ target="10.0.0.1",
+ parameters={"target_type": "network_service", "target_version": "1.0.0"},
+ )
+ assert isinstance(result, bool)
+ assert result is False # Non-matching types
+
+ def test_check_vulnerability_applicability_matching_version(self) -> None:
+ """Test applicability check for matching affected version."""
+ scanner = _make_scanner()
+ vuln_data = {
+ "name": "Test Priv Esc",
+ "type": "privilege_escalation",
+ "severity": "high",
+ "affected_versions": ["1.8.2", "1.8.30"],
+ }
+ result = scanner._check_vulnerability_applicability(
+ vuln_data=vuln_data,
+ target="10.0.0.1",
+ parameters={"target_type": "unknown", "target_version": "1.8.2"},
+ )
+ assert isinstance(result, bool)
+ assert result is True # Version matches affected_versions
+
+ def test_check_vulnerability_applicability_non_matching_version(self) -> None:
+ """Test applicability check for non-matching affected version."""
+ scanner = _make_scanner()
+ vuln_data = {
+ "name": "Test Priv Esc",
+ "type": "privilege_escalation",
+ "severity": "high",
+ "affected_versions": ["1.8.2", "1.8.30"],
+ }
+ result = scanner._check_vulnerability_applicability(
+ vuln_data=vuln_data,
+ target="10.0.0.1",
+ parameters={"target_type": "unknown", "target_version": "1.9.0"},
+ )
+ assert isinstance(result, bool)
+ assert result is False # Version does not match
+
+ def test_check_vulnerability_applicability_empty_affected_versions(self) -> None:
+ """Test applicability check with empty affected_versions list."""
+ scanner = _make_scanner()
+ vuln_data = {
+ "name": "Test Vuln",
+ "type": "other_type",
+ "severity": "medium",
+ "affected_versions": [],
+ }
+ result = scanner._check_vulnerability_applicability(
+ vuln_data=vuln_data,
+ target="10.0.0.1",
+ parameters={"target_type": "unknown", "target_version": "1.0.0"},
+ )
+ assert isinstance(result, bool)
+ assert result is False # No match
+
+ def test_simulate_vulnerability_detection(self) -> None:
+ """Test vulnerability detection simulation returns boolean."""
+ scanner = _make_scanner()
+ result = scanner._simulate_vulnerability_detection(
+ vuln_type="sql_injection",
+ target="10.0.0.1",
+ )
+ assert isinstance(result, bool)
+
+ def test_simulate_vulnerability_detection_multiple_calls(self) -> None:
+ """Test vulnerability detection can vary across multiple calls."""
+ scanner = _make_scanner()
+ results = [
+ scanner._simulate_vulnerability_detection(
+ vuln_type="xss",
+ target="10.0.0.1",
+ )
+ for _ in range(20)
+ ]
+ assert all(isinstance(r, bool) for r in results)
+ # With 40% detection rate, we expect some mix of True/False
+ # But with low sample size, could be all True or all False
+
+ def test_discover_vulnerabilities_returns_list(self) -> None:
+ """Test that discover_vulnerabilities always returns a list."""
+ scanner = _make_scanner()
+ result1 = scanner.discover_vulnerabilities(
+ target="10.0.0.1",
+ parameters={"target_type": "web_service", "target_version": "1.0.0"},
+ )
+ result2 = scanner.discover_vulnerabilities(
+ target="10.0.0.2",
+ parameters={},
+ )
+ assert isinstance(result1, list)
+ assert isinstance(result2, list)
+
+ def test_discover_vulnerabilities_handles_exception(self) -> None:
+ """Test that discover_vulnerabilities returns empty list on exception."""
+ scanner = _make_scanner()
+ # None as parameters should not crash, returns empty list
+ bad_parameters: Any = None
+ result = scanner.discover_vulnerabilities(
+ target="10.0.0.1",
+ parameters=bad_parameters,
+ )
+ assert isinstance(result, list)
+
+ def test_check_applicability_handles_missing_keys(self) -> None:
+ """Test applicability check handles missing optional keys."""
+ scanner = _make_scanner()
+ vuln_data = {
+ "name": "Test",
+ "type": "web_exploitation",
+ # Missing 'affected_versions' key
+ }
+ result = scanner._check_vulnerability_applicability(
+ vuln_data=vuln_data,
+ target="10.0.0.1",
+ parameters={"target_type": "web_service"},
+ )
+ assert isinstance(result, bool)
+
+ def test_vulnerability_structure_completeness(self) -> None:
+ """Test returned vulnerabilities have expected structure."""
+ scanner = _make_scanner()
+ result = scanner.discover_vulnerabilities(
+ target="10.0.0.1",
+ parameters={
+ "target_type": "web_service",
+ "target_version": "2.0.0",
+ },
+ )
+ assert isinstance(result, list)
+ for vuln in result:
+ # Check that basic fields exist
+ assert "vulnerability_id" in vuln
+ assert "name" in vuln
+ assert "severity" in vuln
+ # Type can be either from CVEs or common_vulnerabilities
+ assert "type" in vuln
+ assert "exploit_available" in vuln
+ assert "confidence" in vuln
diff --git a/tests/communication/__init__.py b/tests/communication/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/tests/communication/test_agent_communication.py b/tests/communication/test_agent_communication.py
new file mode 100644
index 0000000..f87aadd
--- /dev/null
+++ b/tests/communication/test_agent_communication.py
@@ -0,0 +1,251 @@
+"""Tests for inter-agent communication framework."""
+
+from typing import Any
+
+from pentestgpt.communication.agent_communication import (
+ AgentCommunicationHub,
+ AgentEndpoint,
+ AgentRole,
+ CommunicationStatus,
+ Message,
+ MessagePriority,
+ MessageQueue,
+ MessageType,
+)
+
+
+def _make_message(
+ sender: str = "agent-a",
+ receiver: str = "agent-b",
+ msg_type: MessageType = MessageType.REQUEST,
+ priority: MessagePriority = MessagePriority.NORMAL,
+ content: dict[str, Any] | None = None,
+) -> Message:
+ return Message(
+ message_id="msg-001",
+ message_type=msg_type,
+ priority=priority,
+ sender_id=sender,
+ receiver_id=receiver,
+ content=content or {"action": "test"},
+ )
+
+
+def _make_endpoint(
+ agent_id: str = "agent-a",
+ agent_type: str = "strategy",
+ status: CommunicationStatus = CommunicationStatus.CONNECTED,
+) -> AgentEndpoint:
+ return AgentEndpoint(
+ agent_id=agent_id,
+ agent_type=agent_type,
+ endpoint_url=f"http://localhost/{agent_id}",
+ status=status,
+ roles=[AgentRole.PUBLISHER, AgentRole.SUBSCRIBER],
+ capabilities=["send", "receive"],
+ )
+
+
+def _make_hub() -> AgentCommunicationHub:
+ return AgentCommunicationHub(hub_id="test-hub")
+
+
+class TestMessageQueue:
+ def test_put_and_get_message(self) -> None:
+ q = MessageQueue(max_size=10)
+ msg = _make_message()
+ assert q.put(msg) is True
+ retrieved = q.get(timeout=1.0)
+ assert retrieved is not None
+ assert retrieved.message_id == "msg-001"
+
+ def test_empty_queue_returns_none(self) -> None:
+ q = MessageQueue(max_size=10)
+ result = q.get(timeout=0.1)
+ assert result is None
+
+ def test_qsize_reflects_items(self) -> None:
+ q = MessageQueue(max_size=10)
+ assert q.qsize() == 0
+ q.put(_make_message())
+ assert q.qsize() == 1
+
+ def test_empty_check(self) -> None:
+ q = MessageQueue(max_size=10)
+ assert q.empty() is True
+ q.put(_make_message())
+ assert q.empty() is False
+
+ def test_priority_ordering(self) -> None:
+ q = MessageQueue(max_size=10)
+ low = _make_message(sender="low", priority=MessagePriority.LOW)
+ low.message_id = "low"
+ high = _make_message(sender="high", priority=MessagePriority.CRITICAL)
+ high.message_id = "high"
+ q.put(low)
+ q.put(high)
+ first = q.get(timeout=1.0)
+ assert first is not None
+ assert first.message_id == "high"
+
+ def test_shutdown_sets_flag(self) -> None:
+ q = MessageQueue(max_size=10)
+ q.shutdown()
+ assert q._shutdown is True
+
+
+class TestHubRegistration:
+ def test_register_agent_succeeds(self) -> None:
+ hub = _make_hub()
+ ep = _make_endpoint(agent_id="agent-1")
+ result = hub.register_agent(ep)
+ assert result is True
+ assert "agent-1" in hub.registered_agents
+
+ def test_register_duplicate_agent_updates(self) -> None:
+ hub = _make_hub()
+ ep = _make_endpoint(agent_id="agent-1")
+ hub.register_agent(ep)
+ result = hub.register_agent(ep)
+ assert result is True
+
+ def test_register_invalid_endpoint_url_fails(self) -> None:
+ hub = _make_hub()
+ ep = _make_endpoint(agent_id="agent-1")
+ ep.endpoint_url = "internal://agent-1"
+ result = hub.register_agent(ep)
+ assert result is False
+
+ def test_unregister_agent_succeeds(self) -> None:
+ hub = _make_hub()
+ ep = _make_endpoint(agent_id="agent-1")
+ hub.register_agent(ep)
+ result = hub.unregister_agent("agent-1")
+ assert result is True
+ assert "agent-1" not in hub.registered_agents
+
+ def test_unregister_nonexistent_agent_fails(self) -> None:
+ hub = _make_hub()
+ result = hub.unregister_agent("ghost-agent")
+ assert result is False
+
+
+class TestHubMessaging:
+ def test_send_message_returns_message_id(self) -> None:
+ hub = _make_hub()
+ hub.register_agent(_make_endpoint(agent_id="sender"))
+ hub.register_agent(_make_endpoint(agent_id="receiver"))
+ result = hub.send_message(
+ sender_id="sender",
+ receiver_id="receiver",
+ message_type=MessageType.REQUEST,
+ content={"action": "test"},
+ )
+ assert isinstance(result, str)
+ assert hub.communication_stats["messages_sent"] >= 1
+
+ def test_send_message_from_unregistered_sender_fails(self) -> None:
+ hub = _make_hub()
+ hub.register_agent(_make_endpoint(agent_id="receiver"))
+ result = hub.send_message(
+ sender_id="ghost",
+ receiver_id="receiver",
+ message_type=MessageType.REQUEST,
+ content={"action": "test"},
+ )
+ assert result is None
+
+ def test_broadcast_message_to_subscribers(self) -> None:
+ hub = _make_hub()
+ hub.register_agent(_make_endpoint(agent_id="broadcaster"))
+ hub.register_agent(_make_endpoint(agent_id="sub-1"))
+ hub.register_agent(_make_endpoint(agent_id="sub-2"))
+ hub.subscribe_to_topic("sub-1", "alerts")
+ hub.subscribe_to_topic("sub-2", "alerts")
+ result = hub.broadcast_message(
+ sender_id="broadcaster",
+ topic="alerts",
+ content={"alert": "test"},
+ )
+ assert isinstance(result, str)
+
+ def test_broadcast_to_empty_topic_returns_none(self) -> None:
+ hub = _make_hub()
+ hub.register_agent(_make_endpoint(agent_id="broadcaster"))
+ result = hub.broadcast_message(
+ sender_id="broadcaster",
+ topic="empty-topic",
+ content={"alert": "test"},
+ )
+ assert result is None
+
+
+class TestHubTopics:
+ def test_subscribe_to_topic(self) -> None:
+ hub = _make_hub()
+ hub.register_agent(_make_endpoint(agent_id="sub-1"))
+ result = hub.subscribe_to_topic("sub-1", "findings")
+ assert result is True
+ assert "sub-1" in hub.broadcast_topics["findings"]
+
+ def test_subscribe_unregistered_agent_fails(self) -> None:
+ hub = _make_hub()
+ result = hub.subscribe_to_topic("ghost", "findings")
+ assert result is False
+
+ def test_unsubscribe_from_topic(self) -> None:
+ hub = _make_hub()
+ hub.register_agent(_make_endpoint(agent_id="sub-1"))
+ hub.subscribe_to_topic("sub-1", "findings")
+ result = hub.unsubscribe_from_topic("sub-1", "findings")
+ assert result is True
+ assert "sub-1" not in hub.broadcast_topics.get("findings", set())
+
+
+class TestHubSessions:
+ def test_create_session_returns_id(self) -> None:
+ hub = _make_hub()
+ hub.register_agent(_make_endpoint(agent_id="init"))
+ hub.register_agent(_make_endpoint(agent_id="resp"))
+ session_id = hub.create_session("init", "resp", [MessageType.REQUEST])
+ assert isinstance(session_id, str)
+ assert session_id in hub.agent_sessions
+
+ def test_create_session_unregistered_agent_fails(self) -> None:
+ hub = _make_hub()
+ hub.register_agent(_make_endpoint(agent_id="init"))
+ result = hub.create_session("init", "ghost", [MessageType.REQUEST])
+ assert result is None
+
+ def test_close_session(self) -> None:
+ hub = _make_hub()
+ hub.register_agent(_make_endpoint(agent_id="init"))
+ hub.register_agent(_make_endpoint(agent_id="resp"))
+ session_id = hub.create_session("init", "resp", [MessageType.REQUEST])
+ assert session_id is not None
+ result = hub.close_session(session_id)
+ assert result is True
+ assert session_id not in hub.agent_sessions
+
+ def test_close_nonexistent_session_fails(self) -> None:
+ hub = _make_hub()
+ result = hub.close_session("nonexistent-session")
+ assert result is False
+
+
+class TestHubStatus:
+ def test_get_hub_status_structure(self) -> None:
+ hub = _make_hub()
+ status = hub.get_hub_status()
+ assert isinstance(status, dict)
+ assert status["hub_id"] == "test-hub"
+ assert "agents" in status
+ assert "communication" in status
+ assert "topics" in status
+ assert "performance" in status
+
+ def test_get_hub_status_reflects_agents(self) -> None:
+ hub = _make_hub()
+ hub.register_agent(_make_endpoint(agent_id="a1"))
+ status = hub.get_hub_status()
+ assert status["agents"]["total_registered"] == 1
diff --git a/tests/conftest.py b/tests/conftest.py
index 3e1467b..0ec2357 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -187,10 +187,9 @@ def mock_model_router():
``ModelRouter`` instance backed by mock providers should use
``mock_model_router_with_providers`` instead.
"""
- router = MagicMock()
- router.route_task.return_value = "Mocked LLM response"
- router.get_model_for_task.return_value = "gpt-4-test"
- return router
+ from tests import make_mock_router
+
+ return make_mock_router(route_task_return="Mocked LLM response")
@pytest.fixture
diff --git a/tests/llms/test_deepseek.py b/tests/llms/test_deepseek.py
new file mode 100644
index 0000000..58a1cef
--- /dev/null
+++ b/tests/llms/test_deepseek.py
@@ -0,0 +1,90 @@
+"""Tests for DeepSeek LLM provider."""
+
+import os
+import unittest
+from unittest.mock import MagicMock, patch
+
+from pentestgpt.llms.deepseek import DeepSeekProvider
+
+
+class TestDeepSeekProvider(unittest.TestCase):
+ def setUp(self) -> None:
+ """Set up test fixtures."""
+ self._openai_patcher = patch("pentestgpt.llms.deepseek.OpenAI")
+ mock_openai = self._openai_patcher.start()
+ self.addCleanup(self._openai_patcher.stop)
+
+ self._env_patcher = patch.dict(os.environ, {"DEEPSEEK_API_KEY": "test_key"})
+ self._env_patcher.start()
+ self.addCleanup(self._env_patcher.stop)
+
+ self.mock_openai_client = MagicMock()
+ mock_openai.return_value = self.mock_openai_client
+
+ self.provider = DeepSeekProvider(model_name="deepseek-chat")
+
+ def test_initialization(self) -> None:
+ """Test DeepSeekProvider initialization."""
+ assert self.provider is not None
+ assert self.provider.config.model_name == "deepseek-chat"
+
+ def test_send_message(self) -> None:
+ """Test sending a message to DeepSeek."""
+ # Mock the API response
+ mock_response = MagicMock()
+ mock_response.choices[0].message.content = "Hello, world!"
+ self.mock_openai_client.chat.completions.create.return_value = mock_response
+
+ # Call the method
+ response = self.provider.send_message("Hello", "test_conversation")
+
+ # Assert the response
+ assert response == "Hello, world!"
+ self.mock_openai_client.chat.completions.create.assert_called_once()
+ call_kwargs = self.mock_openai_client.chat.completions.create.call_args
+ assert call_kwargs is not None
+ assert call_kwargs.kwargs["model"] == "deepseek-chat"
+ assert any(m["role"] == "user" and m["content"] == "Hello" for m in call_kwargs.kwargs["messages"])
+
+ def test_send_new_conversation(self) -> None:
+ """Test starting a new conversation with DeepSeek."""
+ # Mock the API response
+ mock_response = MagicMock()
+ mock_response.choices[0].message.content = "New conversation started."
+ self.mock_openai_client.chat.completions.create.return_value = mock_response
+
+ # Call the method
+ response, conversation_id = self.provider.send_new_conversation("Start")
+
+ # Assert the response
+ assert response == "New conversation started."
+ assert conversation_id is not None
+
+ def test_count_tokens(self) -> None:
+ """Test token counting for DeepSeek."""
+ # Call the method
+ token_count = self.provider.count_tokens("some text")
+
+ # Assert the count (approximately len(text) // 4)
+ assert token_count == 2
+
+ def test_get_cost_estimate(self) -> None:
+ """Test cost estimation for DeepSeek."""
+ # Call the method
+ cost = self.provider.get_cost_estimate(1000, 2000)
+
+ # Assert the cost
+ # deepseek-chat: 0.00028 per 1k input, 0.00042 per 1k output
+ expected_cost = (1000 / 1000) * 0.00028 + (2000 / 1000) * 0.00042
+ assert cost == expected_cost
+
+ def test_deepseek_reasoner_model_config(self) -> None:
+ """Test DeepSeek reasoner model configuration."""
+ assert "deepseek-reasoner" in DeepSeekProvider.MODELS
+ config = DeepSeekProvider.MODELS["deepseek-reasoner"]
+ assert config.model_name == "deepseek-reasoner"
+ assert config.context_window == 128000
+
+
+if __name__ == "__main__":
+ unittest.main()
diff --git a/tests/llms/test_gemini.py b/tests/llms/test_gemini.py
new file mode 100644
index 0000000..8310a7c
--- /dev/null
+++ b/tests/llms/test_gemini.py
@@ -0,0 +1,91 @@
+"""Tests for Gemini LLM provider."""
+
+import os
+import unittest
+from unittest.mock import MagicMock, patch
+
+from pentestgpt.llms.gemini import GeminiProvider
+
+
+class TestGeminiProvider(unittest.TestCase):
+ def setUp(self) -> None:
+ """Set up test fixtures."""
+ self._openai_patcher = patch("pentestgpt.llms.gemini.OpenAI")
+ mock_openai = self._openai_patcher.start()
+ self.addCleanup(self._openai_patcher.stop)
+
+ self._env_patcher = patch.dict(os.environ, {"GEMINI_API_KEY": "test_key"})
+ self._env_patcher.start()
+ self.addCleanup(self._env_patcher.stop)
+
+ self.mock_openai_client = MagicMock()
+ mock_openai.return_value = self.mock_openai_client
+
+ self.provider = GeminiProvider(model_name="gemini-3.1-flash-lite-preview")
+
+ def test_initialization(self) -> None:
+ """Test GeminiProvider initialization."""
+ assert self.provider is not None
+ assert self.provider.config.model_name == "gemini-3.1-flash-lite-preview"
+
+ def test_send_message(self) -> None:
+ """Test sending a message to Gemini."""
+ # Mock the API response
+ mock_response = MagicMock()
+ mock_response.choices[0].message.content = "Hello, world!"
+ self.mock_openai_client.chat.completions.create.return_value = mock_response
+
+ # Call the method
+ response = self.provider.send_message("Hello", "test_conversation")
+
+ # Assert the response
+ assert response == "Hello, world!"
+ self.mock_openai_client.chat.completions.create.assert_called_once()
+ call_kwargs = self.mock_openai_client.chat.completions.create.call_args
+ assert call_kwargs is not None
+ assert call_kwargs.kwargs["model"] == "gemini-3.1-flash-lite-preview"
+ assert any(m["role"] == "user" and m["content"] == "Hello" for m in call_kwargs.kwargs["messages"])
+
+ def test_send_new_conversation(self) -> None:
+ """Test starting a new conversation with Gemini."""
+ # Mock the API response
+ mock_response = MagicMock()
+ mock_response.choices[0].message.content = "New conversation started."
+ self.mock_openai_client.chat.completions.create.return_value = mock_response
+
+ # Call the method
+ response, conversation_id = self.provider.send_new_conversation("Start")
+
+ # Assert the response
+ assert response == "New conversation started."
+ assert conversation_id is not None
+
+ def test_count_tokens(self) -> None:
+ """Test token counting for Gemini."""
+ # Call the method
+ token_count = self.provider.count_tokens("some text")
+
+ # Assert the count (approximately len(text) // 4)
+ assert token_count == 2
+
+ def test_get_cost_estimate(self) -> None:
+ """Test cost estimation for Gemini."""
+ # Call the method
+ cost = self.provider.get_cost_estimate(1000, 2000)
+
+ # Assert the cost
+ # gemini-3.1-flash-lite-preview: 0.00025 per 1k input, 0.0015 per 1k output
+ expected_cost = (1000 / 1000) * 0.00025 + (2000 / 1000) * 0.0015
+ assert cost == expected_cost
+
+ def test_gemini_pro_model_config(self) -> None:
+ """Test Gemini pro model configuration."""
+ assert "gemini-3.1-pro-preview" in GeminiProvider.MODELS
+ config = GeminiProvider.MODELS["gemini-3.1-pro-preview"]
+ assert config.model_name == "gemini-3.1-pro-preview"
+ assert config.context_window == 2000000
+ assert config.supports_vision is True
+
+
+if __name__ == "__main__":
+ unittest.main()
diff --git a/tests/llms/test_mistral.py b/tests/llms/test_mistral.py
new file mode 100644
index 0000000..1f73a58
--- /dev/null
+++ b/tests/llms/test_mistral.py
@@ -0,0 +1,90 @@
+"""Tests for Mistral LLM provider."""
+
+import os
+import unittest
+from unittest.mock import MagicMock, patch
+
+from pentestgpt.llms.mistral import MistralProvider
+
+
+class TestMistralProvider(unittest.TestCase):
+ def setUp(self) -> None:
+ """Set up test fixtures."""
+ self._openai_patcher = patch("pentestgpt.llms.mistral.OpenAI")
+ mock_openai = self._openai_patcher.start()
+ self.addCleanup(self._openai_patcher.stop)
+
+ self._env_patcher = patch.dict(os.environ, {"MISTRAL_API_KEY": "test_key"})
+ self._env_patcher.start()
+ self.addCleanup(self._env_patcher.stop)
+
+ self.mock_openai_client = MagicMock()
+ mock_openai.return_value = self.mock_openai_client
+
+ self.provider = MistralProvider(model_name="mistral-small-latest")
+
+ def test_initialization(self) -> None:
+ """Test MistralProvider initialization."""
+ assert self.provider is not None
+ assert self.provider.config.model_name == "mistral-small-latest"
+
+ def test_send_message(self) -> None:
+ """Test sending a message to Mistral."""
+ # Mock the API response
+ mock_response = MagicMock()
+ mock_response.choices[0].message.content = "Hello, world!"
+ self.mock_openai_client.chat.completions.create.return_value = mock_response
+
+ # Call the method
+ response = self.provider.send_message("Hello", "test_conversation")
+
+ # Assert the response
+ assert response == "Hello, world!"
+ self.mock_openai_client.chat.completions.create.assert_called_once()
+ call_kwargs = self.mock_openai_client.chat.completions.create.call_args
+ assert call_kwargs is not None
+ assert call_kwargs.kwargs["model"] == "mistral-small-latest"
+ assert any(m["role"] == "user" and m["content"] == "Hello" for m in call_kwargs.kwargs["messages"])
+
+ def test_send_new_conversation(self) -> None:
+ """Test starting a new conversation with Mistral."""
+ # Mock the API response
+ mock_response = MagicMock()
+ mock_response.choices[0].message.content = "New conversation started."
+ self.mock_openai_client.chat.completions.create.return_value = mock_response
+
+ # Call the method
+ response, conversation_id = self.provider.send_new_conversation("Start")
+
+ # Assert the response
+ assert response == "New conversation started."
+ assert conversation_id is not None
+
+ def test_count_tokens(self) -> None:
+ """Test token counting for Mistral."""
+ # Call the method
+ token_count = self.provider.count_tokens("some text")
+
+ # Assert the count (approximately len(text) // 4)
+ assert token_count == 2
+
+ def test_get_cost_estimate(self) -> None:
+ """Test cost estimation for Mistral."""
+ # Call the method
+ cost = self.provider.get_cost_estimate(1000, 2000)
+
+ # Assert the cost
+ # mistral-small-latest: 0.00015 per 1k input, 0.00046 per 1k output
+ expected_cost = (1000 / 1000) * 0.00015 + (2000 / 1000) * 0.00046
+ assert cost == expected_cost
+
+ def test_mistral_large_model_config(self) -> None:
+ """Test Mistral large model configuration."""
+ assert "mistral-large-latest" in MistralProvider.MODELS
+ config = MistralProvider.MODELS["mistral-large-latest"]
+ assert config.model_name == "mistral-large-latest"
+ assert config.context_window == 131072
+
+
+if __name__ == "__main__":
+ unittest.main()
diff --git a/tests/prompts/test_prompt_class_v3.py b/tests/prompts/test_prompt_class_v3.py
new file mode 100644
index 0000000..52e6839
--- /dev/null
+++ b/tests/prompts/test_prompt_class_v3.py
@@ -0,0 +1,334 @@
+"""Tests for PentestGPTPromptV3 — prompt contract tests ensuring structure and keywords."""
+
+from pentestgpt.prompts.prompt_class_v3 import PentestGPTPromptV3
+
+
+class TestSystemPrompts:
+ """Test system prompt constants are non-empty strings."""
+
+ def test_base_system_prompt_nonempty(self) -> None:
+ """Base system prompt must be non-empty."""
+ assert isinstance(PentestGPTPromptV3.SYSTEM_PROMPT_BASE, str)
+ assert len(PentestGPTPromptV3.SYSTEM_PROMPT_BASE) > 50
+
+ def test_base_system_prompt_contains_key_content(self) -> None:
+ """Base system prompt contains expert guidance."""
+ base = PentestGPTPromptV3.SYSTEM_PROMPT_BASE
+ assert "penetration tester" in base.lower() or "hacking" in base.lower()
+
+ def test_strategic_prompt_extends_base(self) -> None:
+ """Strategic prompt includes and extends base prompt."""
+ assert PentestGPTPromptV3.SYSTEM_PROMPT_BASE in PentestGPTPromptV3.SYSTEM_PROMPT_STRATEGIC
+ assert len(PentestGPTPromptV3.SYSTEM_PROMPT_STRATEGIC) > len(PentestGPTPromptV3.SYSTEM_PROMPT_BASE)
+
+ def test_strategic_prompt_contains_strategic_content(self) -> None:
+ """Strategic prompt mentions planning and strategy."""
+ strategic = PentestGPTPromptV3.SYSTEM_PROMPT_STRATEGIC
+ assert "strategic" in strategic.lower() or "planning" in strategic.lower()
+
+ def test_technical_prompt_extends_base(self) -> None:
+ """Technical prompt includes and extends base prompt."""
+ assert PentestGPTPromptV3.SYSTEM_PROMPT_BASE in PentestGPTPromptV3.SYSTEM_PROMPT_TECHNICAL
+ assert len(PentestGPTPromptV3.SYSTEM_PROMPT_TECHNICAL) > len(PentestGPTPromptV3.SYSTEM_PROMPT_BASE)
+
+ def test_technical_prompt_contains_technical_content(self) -> None:
+ """Technical prompt mentions command generation."""
+ technical = PentestGPTPromptV3.SYSTEM_PROMPT_TECHNICAL
+ assert "command" in technical.lower() or "tool" in technical.lower()
+
+ def test_analysis_prompt_extends_base(self) -> None:
+ """Analysis prompt includes and extends base prompt."""
+ assert PentestGPTPromptV3.SYSTEM_PROMPT_BASE in PentestGPTPromptV3.SYSTEM_PROMPT_ANALYSIS
+ assert len(PentestGPTPromptV3.SYSTEM_PROMPT_ANALYSIS) > len(PentestGPTPromptV3.SYSTEM_PROMPT_BASE)
+
+ def test_analysis_prompt_contains_analysis_content(self) -> None:
+ """Analysis prompt mentions results analysis."""
+ analysis = PentestGPTPromptV3.SYSTEM_PROMPT_ANALYSIS
+ assert "analyz" in analysis.lower() or "vulnerab" in analysis.lower() or "severity" in analysis.lower()
+
+
+class TestSessionInitPrompts:
+ """Test session initialization prompt methods."""
+
+ def test_reasoning_session_init_returns_string(self) -> None:
+ """Reasoning session init returns a string."""
+ result = PentestGPTPromptV3.get_reasoning_session_init()
+ assert isinstance(result, str)
+ assert len(result) > 0
+
+ def test_reasoning_session_init_contains_ptt(self) -> None:
+ """Reasoning session mentions PTT."""
+ result = PentestGPTPromptV3.get_reasoning_session_init()
+ assert "PTT" in result or "Penetration Testing Tree" in result
+
+ def test_reasoning_session_init_mentions_mitre(self) -> None:
+ """Reasoning session mentions MITRE."""
+ result = PentestGPTPromptV3.get_reasoning_session_init()
+ assert "MITRE" in result or "ATT&CK" in result
+
+ def test_generation_session_init_returns_string(self) -> None:
+ """Command generation session init returns a string."""
+ result = PentestGPTPromptV3.get_generation_session_init()
+ assert isinstance(result, str)
+ assert len(result) > 0
+
+ def test_generation_session_init_contains_command_keyword(self) -> None:
+ """Generation session mentions command generation."""
+ result = PentestGPTPromptV3.get_generation_session_init()
+ assert "command" in result.lower()
+
+ def test_parsing_session_init_returns_string(self) -> None:
+ """Parsing session init returns a string."""
+ result = PentestGPTPromptV3.get_parsing_session_init()
+ assert isinstance(result, str)
+ assert len(result) > 0
+
+ def test_parsing_session_init_contains_parsing_keyword(self) -> None:
+ """Parsing session mentions parsing/output handling."""
+ result = PentestGPTPromptV3.get_parsing_session_init()
+ assert "pars" in result.lower() or "output" in result.lower() or "extract" in result.lower()
+
+
+class TestTaskPrompts:
+ """Test task-related prompt methods."""
+
+ def test_task_description_prompt_returns_string(self) -> None:
+ """Task description prompt returns a string."""
+ result = PentestGPTPromptV3.get_task_description_prompt("10.0.0.1 — Linux web server")
+ assert isinstance(result, str)
+ assert len(result) > 0
+
+ def test_task_description_prompt_includes_target(self) -> None:
+ """Task description prompt includes the target info."""
+ target = "10.0.0.1 — Linux web server"
+ result = PentestGPTPromptV3.get_task_description_prompt(target)
+ assert "10.0.0.1" in result
+
+ def test_task_description_prompt_mentions_ptt(self) -> None:
+ """Task description prompt mentions PTT initialization."""
+ result = PentestGPTPromptV3.get_task_description_prompt("10.0.0.1")
+ assert "PTT" in result or "Penetration Testing Tree" in result
+
+ def test_ptt_update_prompt_returns_string(self) -> None:
+ """PTT update prompt returns a string."""
+ result = PentestGPTPromptV3.get_ptt_update_prompt(
+ current_ptt="Root\n Recon\n Port scan",
+ results="Found ports 22, 80, 443",
+ )
+ assert isinstance(result, str)
+ assert len(result) > 20
+
+ def test_ptt_update_prompt_includes_inputs(self) -> None:
+ """PTT update prompt includes current PTT and results."""
+ ptt = "Root\n Recon\n Port scan"
+ results = "Found ports 22, 80, 443"
+ result = PentestGPTPromptV3.get_ptt_update_prompt(current_ptt=ptt, results=results)
+ assert "Recon" in result or "Port scan" in result or "Found ports" in result
+
+ def test_task_selection_prompt_returns_string(self) -> None:
+ """Task selection prompt returns a string."""
+ result = PentestGPTPromptV3.get_task_selection_prompt("Root\n Recon\n Port scan [done]")
+ assert isinstance(result, str)
+ assert len(result) > 0
+
+ def test_task_selection_prompt_includes_ptt(self) -> None:
+ """Task selection prompt mentions the provided PTT."""
+ ptt = "Root\n Recon"
+ result = PentestGPTPromptV3.get_task_selection_prompt(ptt)
+ assert "Recon" in result or "Root" in result
+
+ def test_command_generation_prompt_returns_string(self) -> None:
+ """Command generation prompt returns a string."""
+ result = PentestGPTPromptV3.get_command_generation_prompt_with_examples(
+ task="Scan for open ports",
+ target="10.0.0.1",
+ technique_id="T1046",
+ )
+ assert isinstance(result, str)
+ assert len(result) > 0
+
+ def test_command_generation_prompt_includes_task_and_target(self) -> None:
+ """Command generation prompt includes task and target."""
+ task = "Scan for open ports"
+ target = "10.0.0.1"
+ result = PentestGPTPromptV3.get_command_generation_prompt_with_examples(
+ task=task,
+ target=target,
+ )
+ assert "Scan for open ports" in result
+ assert "10.0.0.1" in result
+
+ def test_command_generation_prompt_includes_technique_when_provided(self) -> None:
+ """Command generation prompt includes technique ID when provided."""
+ result = PentestGPTPromptV3.get_command_generation_prompt_with_examples(
+ task="Scan for open ports",
+ target="10.0.0.1",
+ technique_id="T1046",
+ )
+ assert "T1046" in result
+
+ def test_command_generation_prompt_mentions_json_format(self) -> None:
+ """Command generation prompt mentions JSON format."""
+ result = PentestGPTPromptV3.get_command_generation_prompt_with_examples(
+ task="Scan for open ports",
+ target="10.0.0.1",
+ )
+ assert "JSON" in result or "json" in result or "{" in result
+
+ def test_results_analysis_prompt_returns_string(self) -> None:
+ """Results analysis prompt returns a string."""
+ result = PentestGPTPromptV3.get_results_analysis_prompt(
+ command="nmap -sV 10.0.0.1",
+ output="22/tcp open ssh OpenSSH 8.2",
+ )
+ assert isinstance(result, str)
+ assert len(result) > 0
+
+ def test_results_analysis_prompt_includes_command_and_output(self) -> None:
+ """Results analysis prompt includes command and output."""
+ cmd = "nmap -sV 10.0.0.1"
+ output = "22/tcp open ssh OpenSSH 8.2"
+ result = PentestGPTPromptV3.get_results_analysis_prompt(command=cmd, output=output)
+ assert "nmap" in result or "22" in result
+
+ def test_results_analysis_prompt_includes_technique_when_provided(self) -> None:
+ """Results analysis prompt includes technique ID when provided."""
+ result = PentestGPTPromptV3.get_results_analysis_prompt(
+ command="nmap -sV 10.0.0.1",
+ output="22/tcp open ssh",
+ technique_id="T1595.001",
+ )
+ assert "T1595.001" in result
+
+ def test_command_verification_prompt_returns_string(self) -> None:
+ """Command verification prompt returns a string."""
+ result = PentestGPTPromptV3.get_command_verification_prompt(
+ commands=["nmap -sV 10.0.0.1", "gobuster dir -u http://10.0.0.1"],
+ )
+ assert isinstance(result, str)
+ assert len(result) > 0
+
+ def test_command_verification_prompt_includes_commands(self) -> None:
+ """Command verification prompt includes the commands to verify."""
+ result = PentestGPTPromptV3.get_command_verification_prompt(
+ commands=["nmap -sV 10.0.0.1"],
+ )
+ assert "nmap" in result
+
+ def test_command_verification_prompt_mentions_json(self) -> None:
+ """Command verification prompt mentions JSON output format."""
+ result = PentestGPTPromptV3.get_command_verification_prompt(
+ commands=["nmap -sV 10.0.0.1"],
+ )
+ assert "JSON" in result or "json" in result or "{" in result
+
+ def test_report_generation_prompt_returns_string(self) -> None:
+ """Report generation prompt returns a string."""
+ result = PentestGPTPromptV3.get_report_generation_prompt(
+ findings=[{"vuln": "SSH weak cipher"}],
+ mitre_coverage={"tactics": ["reconnaissance"]},
+ )
+ assert isinstance(result, str)
+ assert len(result) > 0
+
+ def test_report_generation_prompt_includes_findings(self) -> None:
+ """Report generation prompt includes findings."""
+ findings = [{"vuln": "SSH weak cipher"}]
+ result = PentestGPTPromptV3.get_report_generation_prompt(
+ findings=findings,
+ mitre_coverage={},
+ )
+ assert "SSH" in result or "cipher" in result or "Finding" in result
+
+ def test_report_generation_prompt_mentions_mitre_coverage(self) -> None:
+ """Report generation prompt mentions MITRE coverage."""
+ result = PentestGPTPromptV3.get_report_generation_prompt(
+ findings=[],
+ mitre_coverage={"tactics": ["reconnaissance"]},
+ )
+ assert "MITRE" in result or "ATT&CK" in result
+
+ def test_error_recovery_prompt_returns_string(self) -> None:
+ """Error recovery prompt returns a string."""
+ result = PentestGPTPromptV3.get_error_recovery_prompt(
+ failed_command="nmap -sV 10.0.0.1",
+ error_message="Connection timed out",
+ context={"phase": "reconnaissance"},
+ )
+ assert isinstance(result, str)
+ assert len(result) > 0
+
+ def test_error_recovery_prompt_includes_error_context(self) -> None:
+ """Error recovery prompt includes command and error message."""
+ result = PentestGPTPromptV3.get_error_recovery_prompt(
+ failed_command="nmap -sV 10.0.0.1",
+ error_message="Connection timed out",
+ context={},
+ )
+ assert "error" in result.lower() or "timed out" in result.lower() or "Connection" in result or "nmap" in result
+
+
+class TestLateralMovementPrompt:
+ """Test lateral movement prompt method."""
+
+ def test_lateral_movement_prompt_returns_string(self) -> None:
+ """Lateral movement prompt returns a string."""
+ result = PentestGPTPromptV3.get_lateral_movement_prompt(
+ current_access={"host": "10.0.0.1", "user": "admin"},
+ discovered_hosts=["10.0.0.2", "10.0.0.3"],
+ credentials=[{"user": "admin", "pass": "secret"}],
+ )
+ assert isinstance(result, str)
+ assert len(result) > 0
+
+ def test_lateral_movement_prompt_includes_inputs(self) -> None:
+ """Lateral movement prompt includes access and host info."""
+ result = PentestGPTPromptV3.get_lateral_movement_prompt(
+ current_access={"host": "10.0.0.1"},
+ discovered_hosts=["10.0.0.2"],
+ credentials=[],
+ )
+ assert "10.0.0.1" in result or "10.0.0.2" in result or "movement" in result.lower()
+
+ def test_lateral_movement_prompt_mentions_mitre(self) -> None:
+ """Lateral movement prompt mentions MITRE techniques."""
+ result = PentestGPTPromptV3.get_lateral_movement_prompt(
+ current_access={},
+ discovered_hosts=[],
+ credentials=[],
+ )
+ assert "MITRE" in result or "T1021" in result or "ATT&CK" in result
+
+
+class TestPrivilegeEscalationPrompt:
+ """Test privilege escalation prompt method."""
+
+ def test_privilege_escalation_prompt_returns_string(self) -> None:
+ """Privilege escalation prompt returns a string."""
+ result = PentestGPTPromptV3.get_privilege_escalation_prompt(
+ current_user="www-data",
+ system_info={"os": "Linux", "kernel": "5.10"},
+ enum_results={"sudo": True},
+ )
+ assert isinstance(result, str)
+ assert len(result) > 0
+
+ def test_privilege_escalation_prompt_includes_user_info(self) -> None:
+ """Privilege escalation prompt includes user and system info."""
+ result = PentestGPTPromptV3.get_privilege_escalation_prompt(
+ current_user="www-data",
+ system_info={"os": "Linux"},
+ enum_results={},
+ )
+ assert (
+ "www-data" in result or "Linux" in result or "privilege" in result.lower() or "escalation" in result.lower()
+ )
+
+ def test_privilege_escalation_prompt_mentions_mitre(self) -> None:
+ """Privilege escalation prompt mentions MITRE techniques."""
+ result = PentestGPTPromptV3.get_privilege_escalation_prompt(
+ current_user="www-data",
+ system_info={},
+ enum_results={},
+ )
+ assert "MITRE" in result or "T1068" in result or "ATT&CK" in result
diff --git a/tests/prompts/test_prompt_contracts.py b/tests/prompts/test_prompt_contracts.py
new file mode 100644
index 0000000..94c0d0a
--- /dev/null
+++ b/tests/prompts/test_prompt_contracts.py
@@ -0,0 +1,316 @@
+"""Prompt-contract tests: validate that real LLM responses match expected schemas.
+
+These tests send actual prompts to a real LLM provider and verify the response
+can be parsed into the JSON structure each agent expects. They catch prompt
+drift, model migration regressions, and format changes that mocked tests miss.
+
+Run explicitly with:
+ ./venv/bin/pytest -m llm tests/prompts/test_prompt_contracts.py
+
+Requires at least one of OPENAI_API_KEY or ANTHROPIC_API_KEY in the environment.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import re
+from typing import Any
+
+import pytest
+
+from pentestgpt.mitre import MITRETactic, MITRETechnique
+from pentestgpt.prompts.mitre_prompts import MITREPrompts
+from pentestgpt.utils.APIs.base_provider import LLMProvider
+
+pytestmark = pytest.mark.llm
+
+# ---------------------------------------------------------------------------
+# Provider helpers
+# ---------------------------------------------------------------------------
+
+_OPENAI_KEY = os.environ.get("OPENAI_API_KEY")
+_ANTHROPIC_KEY = os.environ.get("ANTHROPIC_API_KEY")
+
+skip_no_llm = pytest.mark.skipif(
+ not (_OPENAI_KEY or _ANTHROPIC_KEY),
+ reason="No LLM API key found (set OPENAI_API_KEY or ANTHROPIC_API_KEY)",
+)
+
+
+def _get_provider() -> tuple[LLMProvider, str]:
+ """Return (provider_instance, provider_key) for the first available API key."""
+ if _OPENAI_KEY:
+ from pentestgpt.utils.APIs.openai_provider import OpenAIProvider
+
+ provider: LLMProvider = OpenAIProvider(model_name="gpt-4o-mini", api_key=_OPENAI_KEY)
+ return provider, "openai:gpt-4o-mini"
+ if _ANTHROPIC_KEY:
+ from pentestgpt.utils.APIs.anthropic_provider import AnthropicProvider
+
+ provider = AnthropicProvider(model_name="claude-haiku-4-5-20251001", api_key=_ANTHROPIC_KEY)
+ return provider, "anthropic:claude-haiku-4-5-20251001"
+ pytest.skip("No LLM API key available")
+ raise AssertionError("unreachable")
+
+
+def _llm_call(prompt: str, system_prompt: str) -> str:
+ """Send a single prompt to the cheapest available LLM and return the text."""
+ provider, _ = _get_provider()
+ response, _conv_id = provider.send_new_conversation(
+ message=prompt,
+ system_prompt=system_prompt,
+ temperature=0.0,
+ max_tokens=2048,
+ )
+ return str(response)
+
+
+def _extract_json(text: str) -> Any:
+ """Extract JSON from an LLM response that may wrap it in markdown fences."""
+ fenced = re.search(r"```(?:json)?\s*([\s\S]*?)```", text, flags=re.IGNORECASE)
+ candidate = fenced.group(1).strip() if fenced else text.strip()
+ try:
+ return json.loads(candidate)
+ except json.JSONDecodeError:
+ match = re.search(r"(\{[\s\S]*\}|\[[\s\S]*\])", candidate)
+ if not match:
+ raise
+ return json.loads(match.group(1))
+
+
+# ---------------------------------------------------------------------------
+# Shared fixtures
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture(scope="module")
+def technique() -> MITRETechnique:
+ return MITRETechnique(
+ id="T1595",
+ name="Active Scanning",
+ tactic=MITRETactic.RECONNAISSANCE,
+ description="Adversaries may execute active reconnaissance scans.",
+ platforms=["PRE"],
+ tools=["nmap", "masscan"],
+ detection_methods=["Network monitoring"],
+ )
+
+
+@pytest.fixture(scope="module")
+def target_info() -> dict[str, Any]:
+ return {
+ "target": "192.168.1.100",
+ "platform": "Linux",
+ "open_ports": [22, 80, 443],
+ "services": ["ssh", "http", "https"],
+ "os": "Ubuntu 22.04",
+ }
+
+
+# ---------------------------------------------------------------------------
+# 1. Command generation prompt → parseable {"commands": [...]}
+# ---------------------------------------------------------------------------
+
+
+@skip_no_llm
+class TestCommandGenerationContract:
+
+ def test_response_is_valid_json_with_commands_key(
+ self, technique: MITRETechnique, target_info: dict[str, Any]
+ ) -> None:
+ prompt = MITREPrompts.generate_command_prompt(technique, target_info)
+ response = _llm_call(prompt, MITREPrompts.get_system_prompt())
+
+ data = _extract_json(response)
+ assert isinstance(data, dict), f"Expected dict, got {type(data).__name__}"
+ assert "commands" in data, f"Missing 'commands' key. Keys: {list(data.keys())}"
+ assert isinstance(data["commands"], list), "'commands' must be a list"
+ assert len(data["commands"]) >= 1, "Expected at least one command"
+
+ def test_each_command_has_required_fields(self, technique: MITRETechnique, target_info: dict[str, Any]) -> None:
+ prompt = MITREPrompts.generate_command_prompt(technique, target_info)
+ response = _llm_call(prompt, MITREPrompts.get_system_prompt())
+
+ data = _extract_json(response)
+ for i, cmd in enumerate(data["commands"]):
+ assert isinstance(cmd, dict), f"commands[{i}] is not a dict"
+ assert "command" in cmd, f"commands[{i}] missing 'command' field"
+ assert isinstance(cmd["command"], str), f"commands[{i}].command is not a string"
+ assert len(cmd["command"].strip()) > 0, f"commands[{i}].command is empty"
+
+
+# ---------------------------------------------------------------------------
+# 2. Command verification prompt → {"verified": bool, "issues": [...]}
+# ---------------------------------------------------------------------------
+
+
+@skip_no_llm
+class TestCommandVerificationContract:
+
+ def test_verification_response_schema(self, technique: MITRETechnique) -> None:
+ prompt = MITREPrompts.generate_verification_prompt(
+ technique=technique,
+ command="nmap -sV -p 1-1000 192.168.1.100",
+ )
+ response = _llm_call(prompt, MITREPrompts.get_system_prompt())
+
+ data = _extract_json(response)
+ assert isinstance(data, dict)
+ assert "verified" in data, f"Missing 'verified'. Keys: {list(data.keys())}"
+ assert isinstance(data["verified"], bool), "'verified' must be a boolean"
+ assert "issues" in data, f"Missing 'issues'. Keys: {list(data.keys())}"
+ assert isinstance(data["issues"], list), "'issues' must be a list"
+
+
+# ---------------------------------------------------------------------------
+# 3. Result analysis prompt → {"status": str, "findings": [...]}
+# ---------------------------------------------------------------------------
+
+
+@skip_no_llm
+class TestResultAnalysisContract:
+
+ def test_analysis_response_schema(self, technique: MITRETechnique) -> None:
+ nmap_output = (
+ "Starting Nmap 7.94 ( https://nmap.org )\n"
+ "Nmap scan report for 192.168.1.100\n"
+ "PORT STATE SERVICE VERSION\n"
+ "22/tcp open ssh OpenSSH 8.9p1\n"
+ "80/tcp open http Apache httpd 2.4.52\n"
+ "443/tcp open ssl/http Apache httpd 2.4.52\n"
+ "Nmap done: 1 IP address (1 host up) scanned in 12.34 seconds"
+ )
+ prompt = MITREPrompts.generate_analysis_prompt(
+ technique=technique,
+ command_output=nmap_output,
+ command_executed="nmap -sV 192.168.1.100",
+ )
+ response = _llm_call(prompt, MITREPrompts.get_system_prompt())
+
+ data = _extract_json(response)
+ assert isinstance(data, dict)
+ assert "status" in data, f"Missing 'status'. Keys: {list(data.keys())}"
+ assert data["status"] in ("success", "failed"), f"Unexpected status: {data['status']}"
+ assert "findings" in data, f"Missing 'findings'. Keys: {list(data.keys())}"
+ assert isinstance(data["findings"], list)
+
+ def test_findings_have_required_fields(self, technique: MITRETechnique) -> None:
+ nmap_output = (
+ "PORT STATE SERVICE VERSION\n"
+ "22/tcp open ssh OpenSSH 8.9p1\n"
+ "80/tcp open http Apache httpd 2.4.52\n"
+ )
+ prompt = MITREPrompts.generate_analysis_prompt(
+ technique=technique,
+ command_output=nmap_output,
+ command_executed="nmap -sV 192.168.1.100",
+ )
+ response = _llm_call(prompt, MITREPrompts.get_system_prompt())
+
+ data = _extract_json(response)
+ for i, finding in enumerate(data.get("findings", [])):
+ assert isinstance(finding, dict), f"findings[{i}] is not a dict"
+ assert "type" in finding or "description" in finding, f"findings[{i}] lacks both 'type' and 'description'"
+
+
+# ---------------------------------------------------------------------------
+# 4. Strategic planning prompt → {"continue_current_tactic": bool, ...}
+# ---------------------------------------------------------------------------
+
+
+@skip_no_llm
+class TestStrategicPlanningContract:
+
+ def test_planning_response_schema(self, target_info: dict[str, Any]) -> None:
+ prompt = MITREPrompts.generate_planning_prompt(
+ current_tactic=MITRETactic.RECONNAISSANCE,
+ completed_techniques=["T1595"],
+ findings={"open_ports": [22, 80, 443], "services": ["ssh", "http", "https"]},
+ target_info=target_info,
+ )
+ response = _llm_call(prompt, MITREPrompts.get_system_prompt())
+
+ data = _extract_json(response)
+ assert isinstance(data, dict)
+ assert "continue_current_tactic" in data, f"Missing 'continue_current_tactic'. Keys: {list(data.keys())}"
+ assert isinstance(data["continue_current_tactic"], bool)
+ assert "recommended_techniques" in data, f"Missing 'recommended_techniques'. Keys: {list(data.keys())}"
+ assert isinstance(data["recommended_techniques"], list)
+
+
+# ---------------------------------------------------------------------------
+# 5. Error fix suggestion prompt → [{error_index, fix_type, suggestion}]
+# ---------------------------------------------------------------------------
+
+
+@skip_no_llm
+class TestErrorFixContract:
+
+ def test_fix_suggestions_schema(self, target_info: dict[str, Any]) -> None:
+ errors = [
+ {
+ "command": "nmap -sV 192.168.1.100",
+ "error": "Connection timed out",
+ "exit_code": 1,
+ },
+ ]
+ prompt = f"""The following errors were detected during penetration testing:
+
+Errors:
+{json.dumps(errors, indent=2)}
+
+Target Information:
+{json.dumps(target_info, indent=2)}
+
+For each error, suggest a fix or alternative approach.
+
+Output as JSON array:
+[
+ {{
+ "error_index": 0,
+ "fix_type": "retry/modify/alternative",
+ "suggestion": "specific fix or alternative command",
+ "rationale": "why this should work"
+ }}
+]"""
+ response = _llm_call(prompt, MITREPrompts.get_system_prompt())
+
+ data = _extract_json(response)
+ assert isinstance(data, list), f"Expected list, got {type(data).__name__}"
+ assert len(data) >= 1, "Expected at least one fix suggestion"
+ for i, fix in enumerate(data):
+ assert isinstance(fix, dict), f"fixes[{i}] is not a dict"
+ assert "fix_type" in fix, f"fixes[{i}] missing 'fix_type'"
+ assert "suggestion" in fix, f"fixes[{i}] missing 'suggestion'"
+
+
+# ---------------------------------------------------------------------------
+# 6. Generic verification prompt (no MITRE technique) → {"verified": bool}
+# ---------------------------------------------------------------------------
+
+
+@skip_no_llm
+class TestGenericVerificationContract:
+
+ def test_generic_verification_schema(self) -> None:
+ prompt = """Verify the following penetration testing commands for correctness:
+
+Commands:
+nmap -sV -p 22,80,443 192.168.1.100
+
+Check for:
+1. Syntax errors
+2. Missing required arguments
+3. Potential issues
+4. Security concerns
+
+If issues found, provide corrected commands.
+Output as JSON: {"verified": true/false, "commands": [...], "issues": [...]}"""
+
+ response = _llm_call(prompt, MITREPrompts.get_system_prompt())
+
+ data = _extract_json(response)
+ assert isinstance(data, dict)
+ assert "verified" in data, f"Missing 'verified'. Keys: {list(data.keys())}"
+ assert isinstance(data["verified"], bool)
diff --git a/tests/shared/test_authentication.py b/tests/shared/test_authentication.py
new file mode 100644
index 0000000..3cdcfb8
--- /dev/null
+++ b/tests/shared/test_authentication.py
@@ -0,0 +1,142 @@
+"""Tests for AuthenticationMiddleware — JWT, API keys, permissions."""
+
+from datetime import datetime, timedelta
+
+import pytest
+
+from pentestgpt.shared.security.authentication import (
+ AuthenticationMiddleware,
+ AuthToken,
+ SecurityError,
+ check_sql_injection_patterns,
+ safe_file_path,
+ validate_input_length,
+)
+
+
+def _make_auth() -> AuthenticationMiddleware:
+ return AuthenticationMiddleware(secret_key="test-secret-key-for-unit-tests-only")
+
+
+@pytest.fixture
+def auth() -> AuthenticationMiddleware:
+ return _make_auth()
+
+
+class TestAuthToken:
+ def test_auth_token_fields(self) -> None:
+ token = AuthToken(
+ token="abc123",
+ user_id="user-1",
+ expires_at=datetime.now() + timedelta(hours=1),
+ permissions=["read", "write"],
+ )
+ assert token.user_id == "user-1"
+ assert "read" in token.permissions
+
+ def test_auth_token_default_permissions(self) -> None:
+ token = AuthToken(
+ token="abc123",
+ user_id="user-1",
+ expires_at=datetime.now() + timedelta(hours=1),
+ )
+ assert token.permissions == []
+
+
+class TestAuthenticationMiddleware:
+ def test_init_with_explicit_secret(self, auth: AuthenticationMiddleware) -> None:
+ assert auth.secret_key == "test-secret-key-for-unit-tests-only"
+
+ def test_init_generates_key_when_missing(self) -> None:
+ a = AuthenticationMiddleware(secret_key=None)
+ assert a.secret_key is not None
+ assert len(a.secret_key) > 0
+
+ def test_create_and_validate_jwt_token(self, auth: AuthenticationMiddleware) -> None:
+ token_str = auth.create_jwt_token(user_id="test-user", permissions=["read"])
+ assert isinstance(token_str, str)
+ validated = auth.validate_jwt_token(token_str)
+ assert validated is not None
+ assert validated.user_id == "test-user"
+ assert "read" in validated.permissions
+
+ def test_validate_invalid_jwt_returns_none(self, auth: AuthenticationMiddleware) -> None:
+ result = auth.validate_jwt_token("invalid.token.here")
+ assert result is None
+
+ def test_validate_jwt_wrong_secret_returns_none(self, auth: AuthenticationMiddleware) -> None:
+ token_str = auth.create_jwt_token(user_id="user-1", permissions=["read"])
+ other_auth = AuthenticationMiddleware(secret_key="different-secret-key-value-padding-ok")
+ result = other_auth.validate_jwt_token(token_str)
+ assert result is None
+
+ def test_revoke_token_in_active_sessions(self, auth: AuthenticationMiddleware) -> None:
+ token = AuthToken(
+ token="session-token",
+ user_id="user-1",
+ expires_at=datetime.now() + timedelta(hours=1),
+ permissions=["read"],
+ )
+ auth.active_sessions["session-token"] = token
+ result = auth.revoke_token("session-token")
+ assert result is True
+ assert "session-token" not in auth.active_sessions
+
+ def test_revoke_missing_token_returns_false(self, auth: AuthenticationMiddleware) -> None:
+ result = auth.revoke_token("nonexistent-token")
+ assert result is False
+
+ def test_has_permission_present(self, auth: AuthenticationMiddleware) -> None:
+ token = AuthToken(
+ token="abc",
+ user_id="user-1",
+ expires_at=datetime.now() + timedelta(hours=1),
+ permissions=["admin", "read"],
+ )
+ assert auth.has_permission(token, "read") is True
+
+ def test_has_permission_missing(self, auth: AuthenticationMiddleware) -> None:
+ token = AuthToken(
+ token="abc",
+ user_id="user-1",
+ expires_at=datetime.now() + timedelta(hours=1),
+ permissions=["read"],
+ )
+ assert auth.has_permission(token, "nonexistent") is False
+
+ def test_validate_api_key_known(self, auth: AuthenticationMiddleware) -> None:
+ auth.valid_api_keys["test-key-123"] = {"name": "test", "permissions": ["service"]}
+ result = auth.validate_api_key("test-key-123")
+ assert result is not None
+ assert result["name"] == "test"
+
+ def test_validate_api_key_unknown(self, auth: AuthenticationMiddleware) -> None:
+ result = auth.validate_api_key("unknown-key")
+ assert result is None
+
+
+class TestInputValidationHelpers:
+ def test_validate_input_length_passes(self) -> None:
+ result = validate_input_length("short input", max_length=100)
+ assert result == "short input"
+
+ def test_validate_input_length_truncates(self) -> None:
+ result = validate_input_length("x" * 200, max_length=100)
+ assert len(result) == 100
+
+ def test_check_sql_injection_detects_drop(self) -> None:
+ assert check_sql_injection_patterns("'; DROP TABLE users; --") is True
+
+ def test_check_sql_injection_detects_union(self) -> None:
+ assert check_sql_injection_patterns("1 UNION 1=1") is True
+
+ def test_check_sql_injection_clean_input(self) -> None:
+ assert check_sql_injection_patterns("normal search query") is False
+
+ def test_safe_file_path_blocks_traversal(self) -> None:
+ with pytest.raises(SecurityError):
+ safe_file_path("../../etc/passwd")
+
+ def test_safe_file_path_blocks_absolute(self) -> None:
+ with pytest.raises(SecurityError):
+ safe_file_path("/etc/passwd")
diff --git a/tests/shared/test_security_utils.py b/tests/shared/test_security_utils.py
new file mode 100644
index 0000000..92cd1fe
--- /dev/null
+++ b/tests/shared/test_security_utils.py
@@ -0,0 +1,384 @@
+"""Tests for shared security utilities -- path traversal, sanitization, input safety."""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+from typing import Any
+from unittest.mock import patch
+
+import pytest
+
+from pentestgpt.shared.utils.security import (
+ SecureConfiguration,
+ SecurePath,
+ SecurityViolationError,
+ check_input_safety,
+ sanitize_filename,
+ sanitize_sql_input,
+ validate_file_extension,
+)
+
+# ---------------------------------------------------------------------------
+# Factory helpers
+# ---------------------------------------------------------------------------
+
+
+def _make_secure_path(base_dir: str | Path) -> SecurePath:
+ """Create a SecurePath rooted at *base_dir*."""
+ return SecurePath(str(base_dir))
+
+
+def _make_secure_config() -> SecureConfiguration:
+ """Create a SecureConfiguration instance."""
+ return SecureConfiguration()
+
+
+# ---------------------------------------------------------------------------
+# SecurePath
+# ---------------------------------------------------------------------------
+
+
+class TestSecurePath:
+ """Tests for SecurePath -- directory traversal prevention."""
+
+ def test_validate_relative_path_succeeds(self, tmp_path: Path) -> None:
+ sp = _make_secure_path(tmp_path)
+ result = sp.validate_path("subdir/file.txt")
+ assert str(result).startswith(str(tmp_path))
+
+ def test_reject_absolute_path(self, tmp_path: Path) -> None:
+ sp = _make_secure_path(tmp_path)
+ with pytest.raises(SecurityViolationError, match="Absolute paths not allowed"):
+ sp.validate_path("/etc/passwd")
+
+ def test_reject_directory_traversal(self, tmp_path: Path) -> None:
+ sp = _make_secure_path(tmp_path)
+ with pytest.raises(SecurityViolationError, match="traversal"):
+ sp.validate_path("../../etc/passwd")
+
+ def test_reject_null_bytes(self, tmp_path: Path) -> None:
+ sp = _make_secure_path(tmp_path)
+ with pytest.raises(SecurityViolationError, match="Null bytes"):
+ sp.validate_path("file\x00.txt")
+
+ def test_validate_path_returns_resolved_path(self, tmp_path: Path) -> None:
+ sp = _make_secure_path(tmp_path)
+ result = sp.validate_path("a/../b/file.txt")
+ # resolved path must be inside base_dir and have no ".."
+ assert ".." not in str(result)
+ assert str(result).startswith(str(tmp_path))
+
+ # -- read / write -------------------------------------------------------
+
+ def test_read_file_safely(self, tmp_path: Path) -> None:
+ sp = _make_secure_path(tmp_path)
+ test_file = tmp_path / "test.txt"
+ test_file.write_text("hello world")
+ content = sp.read_file_safely("test.txt")
+ assert content == "hello world"
+
+ def test_read_file_safely_missing_file(self, tmp_path: Path) -> None:
+ sp = _make_secure_path(tmp_path)
+ with pytest.raises(FileNotFoundError):
+ sp.read_file_safely("nonexistent.txt")
+
+ def test_read_file_safely_rejects_traversal(self, tmp_path: Path) -> None:
+ sp = _make_secure_path(tmp_path)
+ with pytest.raises(SecurityViolationError):
+ sp.read_file_safely("../../etc/passwd")
+
+ def test_write_file_safely(self, tmp_path: Path) -> None:
+ sp = _make_secure_path(tmp_path)
+ sp.write_file_safely("output.txt", "secure content")
+ assert (tmp_path / "output.txt").read_text() == "secure content"
+
+ def test_write_file_safely_creates_dirs(self, tmp_path: Path) -> None:
+ sp = _make_secure_path(tmp_path)
+ sp.write_file_safely("nested/dir/file.txt", "data", create_dirs=True)
+ assert (tmp_path / "nested" / "dir" / "file.txt").read_text() == "data"
+
+ def test_write_file_safely_rejects_traversal(self, tmp_path: Path) -> None:
+ sp = _make_secure_path(tmp_path)
+ with pytest.raises(SecurityViolationError):
+ sp.write_file_safely("../../evil.txt", "payload")
+
+ def test_ensure_directory_exists(self, tmp_path: Path) -> None:
+ sp = _make_secure_path(tmp_path)
+ target = tmp_path / "a" / "b" / "file.txt"
+ sp.ensure_directory_exists(target)
+ assert (tmp_path / "a" / "b").is_dir()
+
+ def test_base_directory_created_if_missing(self, tmp_path: Path) -> None:
+ new_base = tmp_path / "brand_new_dir"
+ sp = _make_secure_path(new_base)
+ assert new_base.is_dir()
+ assert sp.base_directory == new_base.resolve()
+
+
+# ---------------------------------------------------------------------------
+# sanitize_filename
+# ---------------------------------------------------------------------------
+
+
+class TestSanitizeFilename:
+ """Tests for sanitize_filename."""
+
+ def test_removes_path_separators(self) -> None:
+ result = sanitize_filename("../../etc/passwd")
+ assert "/" not in result
+ assert "\\" not in result
+
+ def test_removes_null_bytes(self) -> None:
+ result = sanitize_filename("file\x00name.txt")
+ assert "\x00" not in result
+
+ def test_preserves_valid_filename(self) -> None:
+ result = sanitize_filename("report_2024.pdf")
+ assert result == "report_2024.pdf"
+
+ def test_replaces_dangerous_chars(self) -> None:
+ result = sanitize_filename('file<>:"|?*.txt')
+ # all dangerous chars should be replaced with underscore
+ for ch in '<>:"|?*':
+ assert ch not in result
+ assert result.endswith(".txt")
+
+ def test_strips_leading_trailing_spaces_and_dots(self) -> None:
+ result = sanitize_filename(" ..hidden.. ")
+ assert not result.startswith(" ")
+ assert not result.startswith(".")
+ assert not result.endswith(".")
+ assert not result.endswith(" ")
+
+ def test_empty_becomes_unnamed(self) -> None:
+ result = sanitize_filename("...")
+ assert result == "unnamed_file"
+
+ def test_truncates_long_filenames(self) -> None:
+ long_name = "a" * 300 + ".txt"
+ result = sanitize_filename(long_name)
+ assert len(result) <= 255
+ assert result.endswith(".txt")
+
+ def test_basename_extraction(self) -> None:
+ # os.path.basename should strip directory components
+ result = sanitize_filename("/some/dir/file.txt")
+ assert result == "file.txt"
+
+
+# ---------------------------------------------------------------------------
+# validate_file_extension
+# ---------------------------------------------------------------------------
+
+
+class TestValidateFileExtension:
+ """Tests for validate_file_extension."""
+
+ def test_allowed_extension_passes(self) -> None:
+ assert validate_file_extension("report.pdf", [".pdf", ".txt"]) is True
+
+ def test_disallowed_extension_fails(self) -> None:
+ assert validate_file_extension("malware.exe", [".pdf", ".txt"]) is False
+
+ def test_case_insensitive_match(self) -> None:
+ assert validate_file_extension("IMAGE.PNG", [".png"]) is True
+
+ def test_extensions_without_leading_dot(self) -> None:
+ assert validate_file_extension("data.csv", ["csv", "tsv"]) is True
+
+ def test_empty_filename_returns_false(self) -> None:
+ assert validate_file_extension("", [".pdf"]) is False
+
+ def test_no_extension_rejected(self) -> None:
+ assert validate_file_extension("Makefile", [".pdf", ".txt"]) is False
+
+
+# ---------------------------------------------------------------------------
+# check_input_safety
+# ---------------------------------------------------------------------------
+
+
+class TestCheckInputSafety:
+ """Tests for check_input_safety -- XSS, eval, exec detection (NOT SQL)."""
+
+ def test_safe_input_passes(self) -> None:
+ assert check_input_safety("hello world") is True
+
+ def test_empty_input_passes(self) -> None:
+ assert check_input_safety("") is True
+
+ def test_rejects_oversized_input(self) -> None:
+ assert check_input_safety("x" * 2000, max_length=1000) is False
+
+ def test_rejects_oversized_input_default_limit(self) -> None:
+ # default max_length is 1000
+ assert check_input_safety("a" * 1001) is False
+
+ def test_raises_on_null_bytes(self) -> None:
+ with pytest.raises(SecurityViolationError, match="Null bytes"):
+ check_input_safety("hello\x00world")
+
+ def test_rejects_xss_script_tag(self) -> None:
+ assert check_input_safety("") is False
+
+ def test_rejects_javascript_protocol(self) -> None:
+ assert check_input_safety("javascript:alert(1)") is False
+
+ def test_rejects_data_uri(self) -> None:
+ assert check_input_safety("data:text/html,hi
") is False
+
+ def test_rejects_eval_call(self) -> None:
+ assert check_input_safety("eval(something)") is False
+
+ def test_rejects_exec_call(self) -> None:
+ assert check_input_safety("exec(cmd)") is False
+
+ def test_rejects_system_call(self) -> None:
+ assert check_input_safety("system('ls')") is False
+
+ def test_rejects_shell_exec_call(self) -> None:
+ assert check_input_safety("shell_exec('whoami')") is False
+
+ def test_rejects_directory_traversal_pattern(self) -> None:
+ assert check_input_safety("../../../etc/passwd") is False
+
+ def test_case_insensitive_detection(self) -> None:
+ assert check_input_safety("") is False
+
+ def test_normal_sql_not_blocked(self) -> None:
+ # check_input_safety does NOT detect SQL patterns
+ assert check_input_safety("SELECT * FROM users") is True
+
+
+# ---------------------------------------------------------------------------
+# sanitize_sql_input
+# ---------------------------------------------------------------------------
+
+
+class TestSanitizeSqlInput:
+ """Tests for sanitize_sql_input -- quote escaping, semicolon/comment stripping."""
+
+ def test_escapes_single_quotes(self) -> None:
+ result = sanitize_sql_input("O'Brien")
+ assert "''" in result
+ # The single quote should be doubled
+ assert result == "O''Brien"
+
+ def test_strips_semicolons(self) -> None:
+ result = sanitize_sql_input("1; DROP TABLE users")
+ assert ";" not in result
+
+ def test_strips_line_comments(self) -> None:
+ result = sanitize_sql_input("value -- this is a comment")
+ assert "--" not in result
+
+ def test_strips_block_comments(self) -> None:
+ result = sanitize_sql_input("value /* block comment */ rest")
+ assert "/*" not in result
+ assert "*/" not in result
+
+ def test_removes_null_bytes(self) -> None:
+ result = sanitize_sql_input("hello\x00world")
+ assert "\x00" not in result
+
+ def test_empty_input_returns_empty(self) -> None:
+ assert sanitize_sql_input("") == ""
+
+ def test_does_not_strip_sql_keywords(self) -> None:
+ # sanitize_sql_input does NOT strip SQL keywords like DROP/SELECT
+ result = sanitize_sql_input("DROP TABLE users")
+ assert "DROP" in result
+
+ def test_combined_attack_sanitized(self) -> None:
+ result = sanitize_sql_input("admin'; DROP TABLE users; --")
+ assert ";" not in result
+ assert "--" not in result
+ assert "''" in result # single quote doubled
+
+
+# ---------------------------------------------------------------------------
+# SecureConfiguration
+# ---------------------------------------------------------------------------
+
+
+class TestSecureConfiguration:
+ """Tests for SecureConfiguration -- load_from_environment and validate."""
+
+ @patch("pentestgpt.database.engine.get_default_url", return_value="sqlite:///test.db")
+ def test_load_from_environment(self, mock_get_url: Any) -> None:
+ config = _make_secure_config()
+ result = config.load_from_environment()
+ assert isinstance(result, dict)
+ assert result["database_url"] == "sqlite:///test.db"
+ assert "openai_api_key" in result
+ assert "rate_limit" in result
+ mock_get_url.assert_called_once()
+
+ @patch("pentestgpt.database.engine.get_default_url", return_value="sqlite:///test.db")
+ def test_load_sets_config_attribute(self, mock_get_url: Any) -> None:
+ config = _make_secure_config()
+ config.load_from_environment()
+ assert config.config["database_url"] == "sqlite:///test.db"
+
+ @patch("pentestgpt.database.engine.get_default_url", return_value="sqlite:///test.db")
+ def test_load_reads_env_vars(self, mock_get_url: Any) -> None:
+ config = _make_secure_config()
+ with patch.dict(os.environ, {"OPENAI_API_KEY": "sk-test-key", "DEBUG": "true"}):
+ result = config.load_from_environment()
+ assert result["openai_api_key"] == "sk-test-key"
+ assert result["debug_mode"] is True
+
+ def test_validate_configuration_empty_config(self) -> None:
+ config = _make_secure_config()
+ # config.config is empty dict -- validate should report issues
+ warnings = config.validate_configuration()
+ assert isinstance(warnings, list)
+ # Should warn about missing secret_key and api key
+ assert any("SECRET_KEY" in w for w in warnings)
+
+ @patch("pentestgpt.database.engine.get_default_url", return_value="sqlite:///test.db")
+ def test_validate_after_load_warns_missing_secret_key(self, mock_get_url: Any) -> None:
+ config = _make_secure_config()
+ with patch.dict(os.environ, {}, clear=False):
+ # Ensure no secret key is set
+ os.environ.pop("PENTESTGPT_SECRET_KEY", None)
+ config.load_from_environment()
+ warnings = config.validate_configuration()
+ assert any("SECRET_KEY" in w for w in warnings)
+
+ @patch("pentestgpt.database.engine.get_default_url", return_value="sqlite:///test.db")
+ def test_validate_with_valid_config(self, mock_get_url: Any) -> None:
+ config = _make_secure_config()
+ with patch.dict(
+ os.environ,
+ {
+ "PENTESTGPT_SECRET_KEY": "mysecret",
+ "OPENAI_API_KEY": "sk-key",
+ "RATE_LIMIT": "100",
+ "SESSION_TIMEOUT": "3600",
+ },
+ ):
+ config.load_from_environment()
+ warnings = config.validate_configuration()
+ # With valid config, fewer warnings expected
+ assert not any("RATE_LIMIT" in w for w in warnings)
+ assert not any("SESSION_TIMEOUT" in w for w in warnings)
+
+ def test_validate_low_rate_limit(self) -> None:
+ config = _make_secure_config()
+ config.config = {"rate_limit": 0, "session_timeout": 3600, "database_url": "sqlite:///x.db"}
+ warnings = config.validate_configuration()
+ assert any("RATE_LIMIT" in w for w in warnings)
+
+ def test_validate_low_session_timeout(self) -> None:
+ config = _make_secure_config()
+ config.config = {"rate_limit": 100, "session_timeout": 10, "database_url": "sqlite:///x.db"}
+ warnings = config.validate_configuration()
+ assert any("SESSION_TIMEOUT" in w for w in warnings)
+
+ def test_validate_invalid_db_url(self) -> None:
+ config = _make_secure_config()
+ config.config = {"database_url": "badurl://something"}
+ warnings = config.validate_configuration()
+ assert any("database url" in w.lower() for w in warnings)
diff --git a/tests/utils/test_rag.py b/tests/utils/test_rag.py
new file mode 100644
index 0000000..7eefc1e
--- /dev/null
+++ b/tests/utils/test_rag.py
@@ -0,0 +1,329 @@
+"""Tests for EnhancedRAGSystem — caching, domain modules, query interface."""
+
+from __future__ import annotations
+
+import os
+import tempfile
+from pathlib import Path
+from unittest.mock import MagicMock
+
+from pentestgpt.utils.rag import EnhancedRAGSystem, Rag_module
+
+# ---------------------------------------------------------------------------
+# Factory helpers
+# ---------------------------------------------------------------------------
+
+
+def _make_rag(tmp_path: str | None = None) -> EnhancedRAGSystem:
+ """Create an EnhancedRAGSystem rooted under *tmp_path* (no LangChain required)."""
+ kb_path = tmp_path or tempfile.mkdtemp()
+ vs_path = os.path.join(kb_path, "vectorstore")
+ os.makedirs(vs_path, exist_ok=True)
+ return EnhancedRAGSystem(
+ knowledge_base_path=kb_path,
+ vector_store_path=vs_path,
+ )
+
+
+def _make_domain_module(name: str = "test_domain") -> MagicMock:
+ """Return a MagicMock that satisfies the DomainKnowledgeModule interface."""
+ module = MagicMock()
+ module.get_domain_name.return_value = name
+ return module
+
+
+# ---------------------------------------------------------------------------
+# TestRAGInitialization
+# ---------------------------------------------------------------------------
+
+
+class TestRAGInitialization:
+ def test_initializes_without_api_key(self, tmp_path: Path) -> None:
+ rag = _make_rag(str(tmp_path))
+ assert rag.cache == {}
+
+ def test_initializes_domain_modules_empty(self, tmp_path: Path) -> None:
+ rag = _make_rag(str(tmp_path))
+ assert rag.domain_modules == []
+
+ def test_initializes_mitre_knowledge_as_dict(self, tmp_path: Path) -> None:
+ rag = _make_rag(str(tmp_path))
+ assert isinstance(rag.mitre_knowledge, dict)
+
+ def test_initializes_mitre_knowledge_empty_without_mitre_path(self, tmp_path: Path) -> None:
+ rag = _make_rag(str(tmp_path))
+ assert len(rag.mitre_knowledge) == 0
+
+ def test_stores_knowledge_base_path(self, tmp_path: Path) -> None:
+ kb = str(tmp_path)
+ vs = os.path.join(kb, "vectorstore")
+ os.makedirs(vs, exist_ok=True)
+ rag = EnhancedRAGSystem(knowledge_base_path=kb, vector_store_path=vs)
+ assert rag.knowledge_base_path == kb
+
+ def test_stores_vector_store_path(self, tmp_path: Path) -> None:
+ kb = str(tmp_path)
+ vs = os.path.join(kb, "vectorstore")
+ os.makedirs(vs, exist_ok=True)
+ rag = EnhancedRAGSystem(knowledge_base_path=kb, vector_store_path=vs)
+ assert rag.vector_store_path == vs
+
+ def test_vectorstore_is_none_without_documents(self, tmp_path: Path) -> None:
+ # Empty knowledge base → no vectorstore created
+ rag = _make_rag(str(tmp_path))
+ assert rag.vectorstore is None
+
+ def test_retriever_is_none_without_vectorstore(self, tmp_path: Path) -> None:
+ rag = _make_rag(str(tmp_path))
+ assert rag.retriever is None
+
+ def test_rag_chain_is_none_on_init(self, tmp_path: Path) -> None:
+ rag = _make_rag(str(tmp_path))
+ assert rag.rag_chain is None
+
+
+# ---------------------------------------------------------------------------
+# TestRAGCache
+# ---------------------------------------------------------------------------
+
+
+class TestRAGCache:
+ def test_cache_starts_empty(self, tmp_path: Path) -> None:
+ rag = _make_rag(str(tmp_path))
+ assert len(rag.cache) == 0
+
+ def test_clear_cache_removes_entries(self, tmp_path: Path) -> None:
+ rag = _make_rag(str(tmp_path))
+ rag.cache["query_key"] = "cached result"
+ rag.clear_cache()
+ assert len(rag.cache) == 0
+
+ def test_clear_cache_on_empty_is_safe(self, tmp_path: Path) -> None:
+ rag = _make_rag(str(tmp_path))
+ rag.clear_cache() # must not raise
+ assert rag.cache == {}
+
+ def test_clear_cache_removes_multiple_entries(self, tmp_path: Path) -> None:
+ rag = _make_rag(str(tmp_path))
+ for i in range(5):
+ rag.cache[f"key_{i}"] = f"value_{i}"
+ rag.clear_cache()
+ assert rag.cache == {}
+
+ def test_get_cache_statistics_returns_dict(self, tmp_path: Path) -> None:
+ rag = _make_rag(str(tmp_path))
+ stats = rag.get_cache_statistics()
+ assert isinstance(stats, dict)
+
+ def test_get_cache_statistics_reports_size(self, tmp_path: Path) -> None:
+ rag = _make_rag(str(tmp_path))
+ rag.cache["k1"] = "v1"
+ rag.cache["k2"] = "v2"
+ stats = rag.get_cache_statistics()
+ assert stats["cache_size"] == 2
+
+ def test_get_cache_statistics_empty_cache_size_zero(self, tmp_path: Path) -> None:
+ rag = _make_rag(str(tmp_path))
+ stats = rag.get_cache_statistics()
+ assert stats["cache_size"] == 0
+
+ def test_get_cache_statistics_includes_cache_keys(self, tmp_path: Path) -> None:
+ rag = _make_rag(str(tmp_path))
+ rag.cache["sample_key"] = "val"
+ stats = rag.get_cache_statistics()
+ assert "cache_keys" in stats
+ assert "sample_key" in stats["cache_keys"]
+
+ def test_cache_stores_and_retrieves_values(self, tmp_path: Path) -> None:
+ rag = _make_rag(str(tmp_path))
+ rag.cache["cached_key"] = "cached_value"
+ assert rag.cache.get("cached_key") == "cached_value"
+
+
+# ---------------------------------------------------------------------------
+# TestRAGDomainModules
+# ---------------------------------------------------------------------------
+
+
+class TestRAGDomainModules:
+ def test_register_domain_module_appends_to_list(self, tmp_path: Path) -> None:
+ rag = _make_rag(str(tmp_path))
+ module = _make_domain_module("web")
+ rag.register_domain_module(module)
+ assert len(rag.domain_modules) == 1
+
+ def test_register_domain_module_stores_exact_object(self, tmp_path: Path) -> None:
+ rag = _make_rag(str(tmp_path))
+ module = _make_domain_module("network")
+ rag.register_domain_module(module)
+ assert rag.domain_modules[0] is module
+
+ def test_register_multiple_domain_modules(self, tmp_path: Path) -> None:
+ rag = _make_rag(str(tmp_path))
+ for name in ("web", "network", "cloud"):
+ rag.register_domain_module(_make_domain_module(name))
+ assert len(rag.domain_modules) == 3
+
+ def test_get_statistics_returns_dict(self, tmp_path: Path) -> None:
+ rag = _make_rag(str(tmp_path))
+ stats = rag.get_statistics()
+ assert isinstance(stats, dict)
+
+ def test_get_statistics_includes_vectorstore_key(self, tmp_path: Path) -> None:
+ rag = _make_rag(str(tmp_path))
+ stats = rag.get_statistics()
+ assert "vectorstore_initialized" in stats
+
+ def test_get_statistics_vectorstore_false_without_docs(self, tmp_path: Path) -> None:
+ rag = _make_rag(str(tmp_path))
+ stats = rag.get_statistics()
+ assert stats["vectorstore_initialized"] is False
+
+ def test_get_statistics_includes_cache_size(self, tmp_path: Path) -> None:
+ rag = _make_rag(str(tmp_path))
+ rag.cache["x"] = "y"
+ stats = rag.get_statistics()
+ assert stats["cache_size"] == 1
+
+ def test_get_statistics_includes_mitre_knowledge_loaded(self, tmp_path: Path) -> None:
+ rag = _make_rag(str(tmp_path))
+ stats = rag.get_statistics()
+ assert "mitre_knowledge_loaded" in stats
+ assert stats["mitre_knowledge_loaded"] is False
+
+
+# ---------------------------------------------------------------------------
+# TestRAGQueries
+# ---------------------------------------------------------------------------
+
+
+class TestRAGQueries:
+ def test_query_by_tool_returns_string(self, tmp_path: Path) -> None:
+ rag = _make_rag(str(tmp_path))
+ result = rag.query_by_tool("nmap")
+ assert isinstance(result, str)
+
+ def test_query_by_tool_without_vectorstore_returns_str(self, tmp_path: Path) -> None:
+ rag = _make_rag(str(tmp_path))
+ # No vectorstore → empty string fallback, still a str
+ result = rag.query_by_tool("masscan")
+ assert isinstance(result, str)
+
+ def test_query_by_technique_returns_string(self, tmp_path: Path) -> None:
+ rag = _make_rag(str(tmp_path))
+ # T1046 not in mitre_knowledge → returns ""
+ result = rag.query_by_technique("T1046")
+ assert isinstance(result, str)
+
+ def test_query_by_technique_unknown_id_returns_empty(self, tmp_path: Path) -> None:
+ rag = _make_rag(str(tmp_path))
+ result = rag.query_by_technique("T9999")
+ assert result == ""
+
+ def test_query_by_technique_with_mitre_data(self, tmp_path: Path) -> None:
+ rag = _make_rag(str(tmp_path))
+ rag.mitre_knowledge["T1046"] = {
+ "name": "Network Service Discovery",
+ "description": "Scanning network services",
+ "tools": ["nmap"],
+ "platforms": ["Linux"],
+ }
+ # With mitre data but no vectorstore, falls back to query() → ""
+ result = rag.query_by_technique("T1046")
+ assert isinstance(result, str)
+
+ def test_query_caches_results(self, tmp_path: Path) -> None:
+ rag = _make_rag(str(tmp_path))
+ # Seed the cache directly
+ cache_key = "How to use nmap for penetration testing:None"
+ rag.cache[cache_key] = "cached nmap docs"
+ result = rag.query_by_tool("nmap")
+ assert result == "cached nmap docs"
+
+ def test_query_uses_cache_when_available(self, tmp_path: Path) -> None:
+ rag = _make_rag(str(tmp_path))
+ rag.cache["cached_key:None"] = "cached_value"
+ assert rag.query("cached_key") == "cached_value"
+
+ def test_query_multi_modal_returns_dict(self, tmp_path: Path) -> None:
+ rag = _make_rag(str(tmp_path))
+ result = rag.query_multi_modal("web attack")
+ assert isinstance(result, dict)
+
+ def test_query_multi_modal_default_modes(self, tmp_path: Path) -> None:
+ rag = _make_rag(str(tmp_path))
+ result = rag.query_multi_modal("sql injection")
+ assert "tools" in result
+ assert "techniques" in result
+ assert "exploits" in result
+
+ def test_query_multi_modal_custom_modes(self, tmp_path: Path) -> None:
+ rag = _make_rag(str(tmp_path))
+ result = rag.query_multi_modal("ssh brute force", modes=["tools"])
+ assert "tools" in result
+ assert "techniques" not in result
+
+ def test_search_similar_without_vectorstore_returns_empty_list(self, tmp_path: Path) -> None:
+ rag = _make_rag(str(tmp_path))
+ result = rag.search_similar("nmap scan")
+ assert result == []
+
+ def test_invoke_rag_chain_without_chain_returns_string(self, tmp_path: Path) -> None:
+ rag = _make_rag(str(tmp_path))
+ result = rag.invoke_rag_chain("How to enumerate SMB?")
+ assert isinstance(result, str)
+
+
+# ---------------------------------------------------------------------------
+# TestRagModule
+# ---------------------------------------------------------------------------
+
+
+class TestRagModule:
+ def test_rag_module_init(self, tmp_path: Path) -> None:
+ module = Rag_module(datapath=str(tmp_path))
+ assert module is not None
+
+ def test_rag_module_stores_datapath(self, tmp_path: Path) -> None:
+ path = str(tmp_path)
+ module = Rag_module(datapath=path)
+ assert module.datapath == path
+
+ def test_rag_module_default_api_key_is_none(self, tmp_path: Path) -> None:
+ module = Rag_module(datapath=str(tmp_path))
+ assert module.api_key is None
+
+ def test_rag_module_custom_api_key(self, tmp_path: Path) -> None:
+ module = Rag_module(datapath=str(tmp_path), api_key="sk-test")
+ assert module.api_key == "sk-test"
+
+ def test_rag_module_enhanced_rag_initially_none(self, tmp_path: Path) -> None:
+ module = Rag_module(datapath=str(tmp_path))
+ assert module.enhanced_rag is None
+
+ def test_rag_init_creates_enhanced_rag(self, tmp_path: Path) -> None:
+ path = str(tmp_path)
+ vs_path = os.path.join(path, "vs")
+ os.makedirs(vs_path, exist_ok=True)
+ module = Rag_module(datapath=path)
+ query_iface = module.rag_init(vs_path)
+ assert module.enhanced_rag is not None
+ assert query_iface is not None
+
+ def test_rag_init_returns_query_interface(self, tmp_path: Path) -> None:
+ path = str(tmp_path)
+ vs_path = os.path.join(path, "vs")
+ os.makedirs(vs_path, exist_ok=True)
+ module = Rag_module(datapath=path)
+ query_iface = module.rag_init(vs_path)
+ # QueryInterface exposes .invoke()
+ assert hasattr(query_iface, "invoke")
+
+ def test_rag_init_query_interface_invoke_returns_string(self, tmp_path: Path) -> None:
+ path = str(tmp_path)
+ vs_path = os.path.join(path, "vs")
+ os.makedirs(vs_path, exist_ok=True)
+ module = Rag_module(datapath=path)
+ query_iface = module.rag_init(vs_path)
+ result = query_iface.invoke("What is nmap?")
+ assert isinstance(result, str)
diff --git a/tests/utils/test_strategy_adaptation_controller.py b/tests/utils/test_strategy_adaptation_controller.py
new file mode 100644
index 0000000..d8df54f
--- /dev/null
+++ b/tests/utils/test_strategy_adaptation_controller.py
@@ -0,0 +1,601 @@
+"""Tests for StrategyAdaptationController — detection events, adaptation rules, monitoring."""
+
+from datetime import datetime
+from unittest.mock import MagicMock
+
+from pentestgpt.config.stealth_config import StealthConfig
+from pentestgpt.utils.strategy_adaptation_controller import (
+ AdaptationAction,
+ AdaptationPriority,
+ AdaptationType,
+ DetectionEvent,
+ DetectionEventType,
+ LearningRecord,
+ StrategyAdaptationController,
+)
+from tests import make_mock_router
+
+# ---------------------------------------------------------------------------
+# Factory helpers
+# ---------------------------------------------------------------------------
+
+
+def _make_controller() -> StrategyAdaptationController:
+ return StrategyAdaptationController(
+ model_router=make_mock_router(route_task_return=("[]", "openai", 0.01)),
+ config=StealthConfig(),
+ )
+
+
+def _make_detection_event(
+ event_type: DetectionEventType = DetectionEventType.SIGNATURE_DETECTION,
+ severity: str = "medium",
+ response_required: bool = True,
+) -> DetectionEvent:
+ return DetectionEvent(
+ id="evt-001",
+ event_type=event_type,
+ timestamp=datetime.now(),
+ severity=severity,
+ description="Test detection event",
+ source="test-monitor",
+ indicators=["T1055", "nmap"],
+ context={"technique": "T1055", "tool": "nmap"},
+ response_required=response_required,
+ )
+
+
+def _make_adaptation_action(
+ adaptation_type: AdaptationType = AdaptationType.DETECTION_RESPONSE,
+ priority: AdaptationPriority = AdaptationPriority.MEDIUM,
+ target_component: str = "orchestrator",
+) -> AdaptationAction:
+ return AdaptationAction(
+ id="act-001",
+ adaptation_type=adaptation_type,
+ priority=priority,
+ description="Test adaptation action",
+ target_component=target_component,
+ parameters={"operation_delay": 2.0},
+ )
+
+
+# ---------------------------------------------------------------------------
+# Enum and dataclass tests
+# ---------------------------------------------------------------------------
+
+
+class TestEnumsAndDataclasses:
+ def test_detection_event_type_values(self) -> None:
+ assert DetectionEventType.SIGNATURE_DETECTION.value == "signature_detection"
+ assert DetectionEventType.BEHAVIORAL_ANOMALY.value == "behavioral_anomaly"
+ assert DetectionEventType.THRESHOLD_EXCEEDED.value == "threshold_exceeded"
+ assert DetectionEventType.NMAP_SCAN_DETECTED.value == "nmap_scan_detected"
+ assert DetectionEventType.NETWORK_ANOMALY.value == "network_anomaly"
+ assert DetectionEventType.USER_REPORT.value == "user_report"
+ assert DetectionEventType.MANUAL_REPORT.value == "manual_report"
+
+ def test_adaptation_type_values(self) -> None:
+ assert AdaptationType.DETECTION_RESPONSE.value == "detection_response"
+ assert AdaptationType.PERFORMANCE_OPTIMIZATION.value == "performance_optimization"
+ assert AdaptationType.ENVIRONMENT_CHANGE.value == "environment_change"
+ assert AdaptationType.FEEDBACK_INTEGRATION.value == "feedback_integration"
+ assert AdaptationType.RISK_MITIGATION.value == "risk_mitigation"
+ assert AdaptationType.RESOURCE_MANAGEMENT.value == "resource_management"
+
+ def test_adaptation_priority_ordering(self) -> None:
+ # Lower numeric value == higher urgency
+ assert AdaptationPriority.CRITICAL.value < AdaptationPriority.HIGH.value
+ assert AdaptationPriority.HIGH.value < AdaptationPriority.MEDIUM.value
+ assert AdaptationPriority.MEDIUM.value < AdaptationPriority.LOW.value
+
+ def test_detection_event_fields(self) -> None:
+ event = _make_detection_event()
+ assert event.id == "evt-001"
+ assert event.severity == "medium"
+ assert event.source == "test-monitor"
+ assert event.event_type == DetectionEventType.SIGNATURE_DETECTION
+ assert event.response_required is True
+ assert isinstance(event.indicators, list)
+ assert isinstance(event.context, dict)
+
+ def test_detection_event_defaults(self) -> None:
+ event = DetectionEvent(
+ id="evt-min",
+ event_type=DetectionEventType.USER_REPORT,
+ timestamp=datetime.now(),
+ severity="low",
+ description="Minimal event",
+ source="user",
+ )
+ assert event.indicators == []
+ assert event.context == {}
+ assert event.false_positive_probability == 0.0
+ assert event.response_required is True
+
+ def test_adaptation_action_fields(self) -> None:
+ action = _make_adaptation_action()
+ assert action.id == "act-001"
+ assert action.implementation_status == "pending"
+ assert action.estimated_impact == "medium"
+ assert action.implemented_at is None
+
+ def test_learning_record_fields(self) -> None:
+ record = LearningRecord(
+ id="rec-001",
+ event_id="evt-001",
+ action_id="act-001",
+ outcome="success",
+ effectiveness_score=0.9,
+ performance_impact="positive",
+ )
+ assert record.id == "rec-001"
+ assert record.outcome == "success"
+ assert record.effectiveness_score == 0.9
+ assert record.lessons_learned == []
+
+
+# ---------------------------------------------------------------------------
+# Controller initialization tests
+# ---------------------------------------------------------------------------
+
+
+class TestControllerInitialization:
+ def test_creates_with_valid_arguments(self) -> None:
+ ctrl = _make_controller()
+ assert ctrl is not None
+
+ def test_monitoring_inactive_on_init(self) -> None:
+ ctrl = _make_controller()
+ assert ctrl.monitoring_active is False
+
+ def test_monitoring_thread_none_on_init(self) -> None:
+ ctrl = _make_controller()
+ assert ctrl.monitoring_thread is None
+
+ def test_adaptation_rules_populated(self) -> None:
+ ctrl = _make_controller()
+ assert len(ctrl.adaptation_rules) > 0
+
+ def test_signature_detection_rule_exists(self) -> None:
+ ctrl = _make_controller()
+ assert DetectionEventType.SIGNATURE_DETECTION in ctrl.adaptation_rules
+
+ def test_behavioral_anomaly_rule_exists(self) -> None:
+ ctrl = _make_controller()
+ assert DetectionEventType.BEHAVIORAL_ANOMALY in ctrl.adaptation_rules
+
+ def test_threshold_exceeded_rule_exists(self) -> None:
+ ctrl = _make_controller()
+ assert DetectionEventType.THRESHOLD_EXCEEDED in ctrl.adaptation_rules
+
+ def test_detection_events_initially_empty(self) -> None:
+ ctrl = _make_controller()
+ assert len(ctrl.detection_events) == 0
+
+ def test_adaptation_queue_initially_empty(self) -> None:
+ ctrl = _make_controller()
+ assert len(ctrl.adaptation_queue) == 0
+
+ def test_learning_records_initially_empty(self) -> None:
+ ctrl = _make_controller()
+ assert len(ctrl.learning_records) == 0
+
+ def test_detection_patterns_initialized(self) -> None:
+ ctrl = _make_controller()
+ assert len(ctrl.detection_patterns) > 0
+ assert "nmap_detection" in ctrl.detection_patterns
+
+ def test_component_references_none_on_init(self) -> None:
+ ctrl = _make_controller()
+ assert ctrl.opsec_agent is None
+ assert ctrl.evasion_engine is None
+ assert ctrl.detection_monitor is None
+
+
+# ---------------------------------------------------------------------------
+# Component registration tests
+# ---------------------------------------------------------------------------
+
+
+class TestRegisterComponent:
+ def test_register_opsec_agent(self) -> None:
+ ctrl = _make_controller()
+ mock_component = MagicMock()
+ ctrl.register_component("opsec_agent", mock_component)
+ assert ctrl.opsec_agent is mock_component
+
+ def test_register_evasion_engine(self) -> None:
+ ctrl = _make_controller()
+ mock_component = MagicMock()
+ ctrl.register_component("evasion_engine", mock_component)
+ assert ctrl.evasion_engine is mock_component
+
+ def test_register_detection_monitor(self) -> None:
+ ctrl = _make_controller()
+ mock_component = MagicMock()
+ ctrl.register_component("detection_monitor", mock_component)
+ assert ctrl.detection_monitor is mock_component
+
+ def test_register_custom_component(self) -> None:
+ ctrl = _make_controller()
+ mock_component = MagicMock()
+ ctrl.register_component("custom_handler", mock_component)
+ assert ctrl.custom_handler is mock_component # type: ignore[attr-defined]
+
+ def test_register_overwrites_existing(self) -> None:
+ ctrl = _make_controller()
+ first = MagicMock()
+ second = MagicMock()
+ ctrl.register_component("opsec_agent", first)
+ ctrl.register_component("opsec_agent", second)
+ assert ctrl.opsec_agent is second
+
+
+# ---------------------------------------------------------------------------
+# Detection event processing tests
+# ---------------------------------------------------------------------------
+
+
+class TestDetectionEventProcessing:
+ def test_process_event_returns_bool(self) -> None:
+ ctrl = _make_controller()
+ event = _make_detection_event()
+ result = ctrl.process_detection_event(event)
+ assert isinstance(result, bool)
+
+ def test_process_event_stores_in_history(self) -> None:
+ ctrl = _make_controller()
+ event = _make_detection_event()
+ ctrl.process_detection_event(event)
+ assert len(ctrl.detection_events) == 1
+
+ def test_process_multiple_events_all_stored(self) -> None:
+ ctrl = _make_controller()
+ for i in range(3):
+ event = DetectionEvent(
+ id=f"evt-{i}",
+ event_type=DetectionEventType.SIGNATURE_DETECTION,
+ timestamp=datetime.now(),
+ severity="medium",
+ description=f"Event {i}",
+ source="test",
+ )
+ ctrl.process_detection_event(event)
+ assert len(ctrl.detection_events) == 3
+
+ def test_signature_detection_triggers_adaptations(self) -> None:
+ ctrl = _make_controller()
+ event = _make_detection_event(event_type=DetectionEventType.SIGNATURE_DETECTION)
+ result = ctrl.process_detection_event(event)
+ # SIGNATURE_DETECTION has a rule with immediate_actions, so should return True
+ assert result is True
+
+ def test_behavioral_anomaly_triggers_adaptations(self) -> None:
+ ctrl = _make_controller()
+ event = _make_detection_event(event_type=DetectionEventType.BEHAVIORAL_ANOMALY)
+ result = ctrl.process_detection_event(event)
+ assert result is True
+
+ def test_threshold_exceeded_triggers_adaptations(self) -> None:
+ ctrl = _make_controller()
+ event = _make_detection_event(event_type=DetectionEventType.THRESHOLD_EXCEEDED)
+ result = ctrl.process_detection_event(event)
+ assert result is True
+
+ def test_no_response_required_returns_false(self) -> None:
+ ctrl = _make_controller()
+ event = _make_detection_event(response_required=False)
+ result = ctrl.process_detection_event(event)
+ assert result is False
+
+ def test_no_response_event_still_stored(self) -> None:
+ ctrl = _make_controller()
+ event = _make_detection_event(response_required=False)
+ ctrl.process_detection_event(event)
+ assert len(ctrl.detection_events) == 1
+
+ def test_event_type_without_rule_may_return_false(self) -> None:
+ ctrl = _make_controller()
+ # USER_REPORT has no rule configured
+ event = _make_detection_event(event_type=DetectionEventType.USER_REPORT)
+ result = ctrl.process_detection_event(event)
+ # Return value depends on AI mock; with mock returning "[]" it should be False
+ assert isinstance(result, bool)
+
+ def test_signature_detection_queues_adaptation_actions(self) -> None:
+ ctrl = _make_controller()
+ event = _make_detection_event(event_type=DetectionEventType.SIGNATURE_DETECTION)
+ ctrl.process_detection_event(event)
+ # Rule has 3 immediate_actions; AI mock returns "[]"
+ assert len(ctrl.adaptation_queue) == 3
+
+ def test_adaptation_queue_priority_ordering(self) -> None:
+ ctrl = _make_controller()
+ low_event = _make_detection_event(event_type=DetectionEventType.THRESHOLD_EXCEEDED)
+ high_event = _make_detection_event(event_type=DetectionEventType.SIGNATURE_DETECTION)
+ ctrl.process_detection_event(low_event)
+ ctrl.process_detection_event(high_event)
+ # Queue should be non-empty and ordered (CRITICAL before MEDIUM)
+ assert len(ctrl.adaptation_queue) > 0
+ assert ctrl.adaptation_queue[0].priority.value <= ctrl.adaptation_queue[-1].priority.value
+
+
+# ---------------------------------------------------------------------------
+# Adaptation status tests
+# ---------------------------------------------------------------------------
+
+
+class TestAdaptationStatus:
+ def test_get_adaptation_status_returns_dict(self) -> None:
+ ctrl = _make_controller()
+ status = ctrl.get_adaptation_status()
+ assert isinstance(status, dict)
+
+ def test_status_contains_expected_keys(self) -> None:
+ ctrl = _make_controller()
+ status = ctrl.get_adaptation_status()
+ expected_keys = {
+ "monitoring_active",
+ "pending_adaptations",
+ "active_adaptations",
+ "completed_adaptations",
+ "recent_events",
+ "performance_metrics",
+ "learning_records",
+ "success_patterns",
+ "effectiveness_score",
+ }
+ assert expected_keys.issubset(status.keys())
+
+ def test_monitoring_active_reflects_state(self) -> None:
+ ctrl = _make_controller()
+ status = ctrl.get_adaptation_status()
+ assert status["monitoring_active"] is False
+
+ def test_pending_adaptations_zero_initially(self) -> None:
+ ctrl = _make_controller()
+ status = ctrl.get_adaptation_status()
+ assert status["pending_adaptations"] == 0
+
+ def test_pending_adaptations_increments_after_event(self) -> None:
+ ctrl = _make_controller()
+ event = _make_detection_event(event_type=DetectionEventType.SIGNATURE_DETECTION)
+ ctrl.process_detection_event(event)
+ status = ctrl.get_adaptation_status()
+ assert status["pending_adaptations"] >= 1
+
+ def test_effectiveness_score_zero_initially(self) -> None:
+ ctrl = _make_controller()
+ status = ctrl.get_adaptation_status()
+ assert status["effectiveness_score"] == 0.0
+
+ def test_learning_records_zero_initially(self) -> None:
+ ctrl = _make_controller()
+ status = ctrl.get_adaptation_status()
+ assert status["learning_records"] == 0
+
+
+# ---------------------------------------------------------------------------
+# record_adaptation_outcome tests
+# ---------------------------------------------------------------------------
+
+
+class TestRecordAdaptationOutcome:
+ def test_records_success_outcome(self) -> None:
+ ctrl = _make_controller()
+ ctrl.record_adaptation_outcome(
+ action_id="act-001",
+ outcome="success",
+ effectiveness_score=0.85,
+ performance_impact="positive",
+ )
+ assert len(ctrl.learning_records) == 1
+
+ def test_records_failure_outcome(self) -> None:
+ ctrl = _make_controller()
+ ctrl.record_adaptation_outcome(
+ action_id="act-002",
+ outcome="failure",
+ effectiveness_score=0.1,
+ performance_impact="negative",
+ )
+ assert len(ctrl.learning_records) == 1
+ assert ctrl.learning_records[0].outcome == "failure"
+
+ def test_records_partial_outcome(self) -> None:
+ ctrl = _make_controller()
+ ctrl.record_adaptation_outcome(
+ action_id="act-003",
+ outcome="partial",
+ effectiveness_score=0.5,
+ performance_impact="neutral",
+ )
+ assert ctrl.learning_records[0].outcome == "partial"
+
+ def test_multiple_records_accumulate(self) -> None:
+ ctrl = _make_controller()
+ for i in range(5):
+ ctrl.record_adaptation_outcome(
+ action_id=f"act-{i:03d}",
+ outcome="success",
+ effectiveness_score=0.7 + i * 0.05,
+ performance_impact="positive",
+ )
+ assert len(ctrl.learning_records) == 5
+
+ def test_effectiveness_score_stored_in_history(self) -> None:
+ ctrl = _make_controller()
+ ctrl.record_adaptation_outcome(
+ action_id="act-001",
+ outcome="success",
+ effectiveness_score=0.9,
+ performance_impact="positive",
+ )
+ assert 0.9 in list(ctrl.effectiveness_history)
+
+ def test_status_reflects_learning_record(self) -> None:
+ ctrl = _make_controller()
+ ctrl.record_adaptation_outcome(
+ action_id="act-001",
+ outcome="success",
+ effectiveness_score=0.75,
+ performance_impact="positive",
+ )
+ status = ctrl.get_adaptation_status()
+ assert status["learning_records"] == 1
+
+ def test_effectiveness_score_updates_in_status(self) -> None:
+ ctrl = _make_controller()
+ ctrl.record_adaptation_outcome(
+ action_id="act-001",
+ outcome="success",
+ effectiveness_score=0.8,
+ performance_impact="positive",
+ )
+ status = ctrl.get_adaptation_status()
+ assert abs(status["effectiveness_score"] - 0.8) < 1e-6
+
+
+# ---------------------------------------------------------------------------
+# Monitoring lifecycle tests
+# ---------------------------------------------------------------------------
+
+
+class TestMonitoringLifecycle:
+ def test_start_monitoring_sets_active_true(self) -> None:
+ ctrl = _make_controller()
+ try:
+ ctrl.start_monitoring()
+ assert ctrl.monitoring_active is True
+ finally:
+ ctrl.stop_monitoring()
+
+ def test_start_monitoring_spawns_thread(self) -> None:
+ ctrl = _make_controller()
+ try:
+ ctrl.start_monitoring()
+ assert ctrl.monitoring_thread is not None
+ finally:
+ ctrl.stop_monitoring()
+
+ def test_monitoring_thread_is_alive_after_start(self) -> None:
+ ctrl = _make_controller()
+ try:
+ ctrl.start_monitoring()
+ assert ctrl.monitoring_thread is not None
+ assert ctrl.monitoring_thread.is_alive()
+ finally:
+ ctrl.stop_monitoring()
+
+ def test_stop_monitoring_sets_active_false(self) -> None:
+ ctrl = _make_controller()
+ try:
+ ctrl.start_monitoring()
+ ctrl.stop_monitoring()
+ assert ctrl.monitoring_active is False
+ finally:
+ if ctrl.monitoring_active:
+ ctrl.stop_monitoring()
+
+ def test_stop_monitoring_thread_terminates(self) -> None:
+ ctrl = _make_controller()
+ try:
+ ctrl.start_monitoring()
+ thread = ctrl.monitoring_thread
+ ctrl.stop_monitoring()
+ assert thread is not None
+ thread.join(timeout=6)
+ assert not thread.is_alive()
+ finally:
+ if ctrl.monitoring_active:
+ ctrl.stop_monitoring()
+
+ def test_double_start_idempotent(self) -> None:
+ ctrl = _make_controller()
+ try:
+ ctrl.start_monitoring()
+ first_thread = ctrl.monitoring_thread
+ ctrl.start_monitoring() # Second call should be no-op
+ assert ctrl.monitoring_thread is first_thread
+ finally:
+ ctrl.stop_monitoring()
+
+ def test_stop_without_start_safe(self) -> None:
+ ctrl = _make_controller()
+ # Should not raise even when monitoring was never started
+ ctrl.stop_monitoring()
+ assert ctrl.monitoring_active is False
+
+ def test_status_monitoring_active_true_during_monitoring(self) -> None:
+ ctrl = _make_controller()
+ try:
+ ctrl.start_monitoring()
+ status = ctrl.get_adaptation_status()
+ assert status["monitoring_active"] is True
+ finally:
+ ctrl.stop_monitoring()
+
+ def test_status_monitoring_active_false_after_stop(self) -> None:
+ ctrl = _make_controller()
+ try:
+ ctrl.start_monitoring()
+ ctrl.stop_monitoring()
+ status = ctrl.get_adaptation_status()
+ assert status["monitoring_active"] is False
+ finally:
+ if ctrl.monitoring_active:
+ ctrl.stop_monitoring()
+
+
+# ---------------------------------------------------------------------------
+# Force adaptation tests
+# ---------------------------------------------------------------------------
+
+
+class TestForceAdaptation:
+ def test_force_adaptation_returns_true(self) -> None:
+ ctrl = _make_controller()
+ result = ctrl.force_adaptation(
+ {
+ "description": "Manual override",
+ "target_component": "orchestrator",
+ "parameters": {"operation_delay": 5.0},
+ }
+ )
+ assert result is True
+
+ def test_force_adaptation_queues_action(self) -> None:
+ ctrl = _make_controller()
+ ctrl.force_adaptation(
+ {
+ "description": "Manual override",
+ "target_component": "orchestrator",
+ "parameters": {},
+ }
+ )
+ assert len(ctrl.adaptation_queue) == 1
+
+ def test_force_adaptation_uses_feedback_integration_type(self) -> None:
+ ctrl = _make_controller()
+ ctrl.force_adaptation({"description": "test"})
+ action = ctrl.adaptation_queue[0]
+ assert action.adaptation_type == AdaptationType.FEEDBACK_INTEGRATION
+
+ def test_force_adaptation_defaults_target_to_orchestrator(self) -> None:
+ ctrl = _make_controller()
+ ctrl.force_adaptation({"description": "test"})
+ action = ctrl.adaptation_queue[0]
+ assert action.target_component == "orchestrator"
+
+ def test_force_adaptation_uses_provided_target(self) -> None:
+ ctrl = _make_controller()
+ ctrl.force_adaptation(
+ {
+ "description": "evasion override",
+ "target_component": "evasion_engine",
+ "parameters": {},
+ }
+ )
+ action = ctrl.adaptation_queue[0]
+ assert action.target_component == "evasion_engine"