From 7a43ee4fab04a116467be37bfcc090b8deaca461 Mon Sep 17 00:00:00 2001 From: SALhik <123875913+SALhik@users.noreply.github.com> Date: Sat, 20 Jun 2026 21:42:10 +0800 Subject: [PATCH 01/14] test: add prompt-contract tests validating LLM response schemas MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add 8 live-LLM tests that send real prompts and validate the JSON structure each agent expects (command generation, verification, result analysis, strategic planning, error fixes). Gated behind `@pytest.mark.llm` so they never run in default pytest or CI — run explicitly with `pytest -m llm`. Co-Authored-By: Claude Opus 4.6 --- .github/workflows/ci.yml | 1 + pytest.ini | 3 +- tests/prompts/test_prompt_contracts.py | 310 +++++++++++++++++++++++++ 3 files changed, 313 insertions(+), 1 deletion(-) create mode 100644 tests/prompts/test_prompt_contracts.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index c00e49a..b69a43a 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -81,6 +81,7 @@ jobs: --cov-report=html \ --cov-fail-under=${{ env.COVERAGE_THRESHOLD }} \ --junitxml=results-${{ matrix.python-version }}.xml \ + -m "not llm" \ -v - name: Upload test results 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/prompts/test_prompt_contracts.py b/tests/prompts/test_prompt_contracts.py new file mode 100644 index 0000000..1ee5bc0 --- /dev/null +++ b/tests/prompts/test_prompt_contracts.py @@ -0,0 +1,310 @@ +"""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.""" + cleaned = re.sub(r"```(?:json)?\s*", "", text) + cleaned = cleaned.strip().rstrip("`") + return json.loads(cleaned) + + +# --------------------------------------------------------------------------- +# 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) From d370185a4e76aa45f4e6ef6026f9ec8717a27556 Mon Sep 17 00:00:00 2001 From: SALhik <123875913+SALhik@users.noreply.github.com> Date: Sat, 20 Jun 2026 21:45:15 +0800 Subject: [PATCH 02/14] test: add unit tests for shared security utils Co-Authored-By: Claude Sonnet 4.6 --- tests/shared/test_security_utils.py | 384 ++++++++++++++++++++++++++++ 1 file changed, 384 insertions(+) create mode 100644 tests/shared/test_security_utils.py 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) From 124e4a757fe62aac759aa166d0dd1bcd8fbbb810 Mon Sep 17 00:00:00 2001 From: SALhik <123875913+SALhik@users.noreply.github.com> Date: Sun, 21 Jun 2026 10:39:43 +0800 Subject: [PATCH 03/14] test(wave-1): add tests for AgentCommunicationHub and AuthenticationMiddleware Co-Authored-By: Claude Sonnet 4.6 --- tests/communication/__init__.py | 0 .../communication/test_agent_communication.py | 251 ++++++++++++++++++ tests/shared/test_authentication.py | 142 ++++++++++ 3 files changed, 393 insertions(+) create mode 100644 tests/communication/__init__.py create mode 100644 tests/communication/test_agent_communication.py create mode 100644 tests/shared/test_authentication.py 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/shared/test_authentication.py b/tests/shared/test_authentication.py new file mode 100644 index 0000000..6139fb5 --- /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") + 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") From c358292f2133448caeed86d90461b2916e8ef689 Mon Sep 17 00:00:00 2001 From: SALhik <123875913+SALhik@users.noreply.github.com> Date: Sun, 21 Jun 2026 10:42:59 +0800 Subject: [PATCH 04/14] test: add unit tests for IntelligenceGatheringAgent Co-Authored-By: Claude Opus 4.6 --- .../test_intelligence_gathering_agent.py | 334 ++++++++++++++++++ 1 file changed, 334 insertions(+) create mode 100644 tests/agents/test_intelligence_gathering_agent.py diff --git a/tests/agents/test_intelligence_gathering_agent.py b/tests/agents/test_intelligence_gathering_agent.py new file mode 100644 index 0000000..f93b998 --- /dev/null +++ b/tests/agents/test_intelligence_gathering_agent.py @@ -0,0 +1,334 @@ +"""Tests for IntelligenceGatheringAgent — OSINT, target profiling, intel collection.""" + +from datetime import datetime +from unittest.mock import MagicMock + +from pentestgpt.agents.intelligence_gathering_agent import ( + IntelligenceGatheringAgent, + IntelligenceItem, + OSINTSource, + TargetProfile, + ThreatIntel, +) + + +def _make_router() -> MagicMock: + router = MagicMock() + router.route_task = MagicMock(return_value=("Mocked intel response", "openai", 0.01)) + router.get_model_for_task = MagicMock(return_value="gpt-4-test") + return router + + +def _make_agent() -> IntelligenceGatheringAgent: + return IntelligenceGatheringAgent(model_router=_make_router()) + + +# --------------------------------------------------------------------------- +# 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_router(), + 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) + + def test_url_structure_analysis_fields(self) -> None: + agent = _make_agent() + result = agent._analyze_url_structure("https://example.com/admin?q=1") + assert result is not None + assert result["scheme"] == "https" + assert result["netloc"] == "example.com" + assert result["path"] == "/admin" + assert result["query"] == "q=1" + assert "suspicious_patterns" in result + + def test_enumerate_subdomains(self) -> None: + agent = _make_agent() + result = agent._enumerate_subdomains("example.com") + assert isinstance(result, list) + + def test_enumerate_subdomains_returns_expected(self) -> None: + agent = _make_agent() + result = agent._enumerate_subdomains("example.com") + assert len(result) > 0 + assert "www.example.com" in result + assert "mail.example.com" in result + + 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 From 2a947ff3804aede53da999fd3e65d84a16f44bf1 Mon Sep 17 00:00:00 2001 From: SALhik <123875913+SALhik@users.noreply.github.com> Date: Sun, 21 Jun 2026 10:43:26 +0800 Subject: [PATCH 05/14] test: add unit tests for OPSECAgent MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 58 tests covering EnvironmentProfile, LegalComplianceChecker, OPSECSession, and OPSECAgent — env profiling, legal compliance validation, session lifecycle, stealth level selection, evidence tracking, and compliance reporting. Co-Authored-By: Claude Sonnet 4.6 --- tests/agents/test_opsec_agent.py | 432 +++++++++++++++++++++++++++++++ tests/utils/test_rag.py | 329 +++++++++++++++++++++++ 2 files changed, 761 insertions(+) create mode 100644 tests/agents/test_opsec_agent.py create mode 100644 tests/utils/test_rag.py diff --git a/tests/agents/test_opsec_agent.py b/tests/agents/test_opsec_agent.py new file mode 100644 index 0000000..8f3b975 --- /dev/null +++ b/tests/agents/test_opsec_agent.py @@ -0,0 +1,432 @@ +"""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, +) + +# --------------------------------------------------------------------------- +# Factories +# --------------------------------------------------------------------------- + + +def _make_router() -> MagicMock: + router = MagicMock() + router.route_task = MagicMock(return_value=("Mocked response", "openai", 0.01)) + router.get_model_for_task = MagicMock(return_value="gpt-4-test") + return router + + +def _make_opsec_agent() -> OPSECAgent: + return OPSECAgent(model_router=_make_router(), 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 + + 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_router() + # 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_router() + 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_router(), 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_router(), 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_router(), 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/utils/test_rag.py b/tests/utils/test_rag.py new file mode 100644 index 0000000..032f5a7 --- /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 typing import Any +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: Any) -> None: + rag = _make_rag(str(tmp_path)) + assert rag.cache == {} + + def test_initializes_domain_modules_empty(self, tmp_path: Any) -> None: + rag = _make_rag(str(tmp_path)) + assert rag.domain_modules == [] + + def test_initializes_mitre_knowledge_as_dict(self, tmp_path: Any) -> 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: Any) -> None: + rag = _make_rag(str(tmp_path)) + assert len(rag.mitre_knowledge) == 0 + + def test_stores_knowledge_base_path(self, tmp_path: Any) -> 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: Any) -> 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: Any) -> 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: Any) -> None: + rag = _make_rag(str(tmp_path)) + assert rag.retriever is None + + def test_rag_chain_is_none_on_init(self, tmp_path: Any) -> None: + rag = _make_rag(str(tmp_path)) + assert rag.rag_chain is None + + +# --------------------------------------------------------------------------- +# TestRAGCache +# --------------------------------------------------------------------------- + + +class TestRAGCache: + def test_cache_starts_empty(self, tmp_path: Any) -> None: + rag = _make_rag(str(tmp_path)) + assert len(rag.cache) == 0 + + def test_clear_cache_removes_entries(self, tmp_path: Any) -> 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: Any) -> 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: Any) -> 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: Any) -> 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: Any) -> 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: Any) -> 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: Any) -> 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: Any) -> 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: Any) -> 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: Any) -> 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: Any) -> 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: Any) -> 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: Any) -> 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: Any) -> 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: Any) -> 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: Any) -> 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: Any) -> 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: Any) -> 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: Any) -> 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: Any) -> 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: Any) -> 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: Any) -> 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: Any) -> None: + rag = _make_rag(str(tmp_path)) + rag.cache["cached_key"] = "cached_value" + assert rag.cache.get("cached_key") == "cached_value" + + def test_query_multi_modal_returns_dict(self, tmp_path: Any) -> 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: Any) -> 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: Any) -> 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: Any) -> 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: Any) -> 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: Any) -> None: + module = Rag_module(datapath=str(tmp_path)) + assert module is not None + + def test_rag_module_stores_datapath(self, tmp_path: Any) -> 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: Any) -> None: + module = Rag_module(datapath=str(tmp_path)) + assert module.api_key is None + + def test_rag_module_custom_api_key(self, tmp_path: Any) -> 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: Any) -> None: + module = Rag_module(datapath=str(tmp_path)) + assert module.enhanced_rag is None + + def test_rag_init_creates_enhanced_rag(self, tmp_path: Any) -> 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: Any) -> 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: Any) -> 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) From e394524f3b48bb032360751a06c6dfafa10e9a56 Mon Sep 17 00:00:00 2001 From: SALhik <123875913+SALhik@users.noreply.github.com> Date: Sun, 21 Jun 2026 10:43:49 +0800 Subject: [PATCH 06/14] test: add unit tests for StrategyAdaptationController 63 tests across 8 classes covering enums, dataclasses, controller init, component registration, detection event processing, adaptation queue management, learning outcome recording, monitoring lifecycle, and force adaptation. Co-Authored-By: Claude Sonnet 4.6 --- .../test_strategy_adaptation_controller.py | 586 ++++++++++++++++++ 1 file changed, 586 insertions(+) create mode 100644 tests/utils/test_strategy_adaptation_controller.py diff --git a/tests/utils/test_strategy_adaptation_controller.py b/tests/utils/test_strategy_adaptation_controller.py new file mode 100644 index 0000000..94e8a34 --- /dev/null +++ b/tests/utils/test_strategy_adaptation_controller.py @@ -0,0 +1,586 @@ +"""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, +) + +# --------------------------------------------------------------------------- +# Factory helpers +# --------------------------------------------------------------------------- + + +def _make_router() -> MagicMock: + router = MagicMock() + router.route_task = MagicMock(return_value=("[]", "openai", 0.01)) + router.get_model_for_task = MagicMock(return_value="gpt-4-test") + return router + + +def _make_controller() -> StrategyAdaptationController: + return StrategyAdaptationController( + model_router=_make_router(), + 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() + ctrl.start_monitoring() + assert ctrl.monitoring_active is True + ctrl.stop_monitoring() + + def test_start_monitoring_spawns_thread(self) -> None: + ctrl = _make_controller() + ctrl.start_monitoring() + assert ctrl.monitoring_thread is not None + ctrl.stop_monitoring() + + def test_monitoring_thread_is_alive_after_start(self) -> None: + ctrl = _make_controller() + ctrl.start_monitoring() + assert ctrl.monitoring_thread is not None + assert ctrl.monitoring_thread.is_alive() + ctrl.stop_monitoring() + + def test_stop_monitoring_sets_active_false(self) -> None: + ctrl = _make_controller() + ctrl.start_monitoring() + ctrl.stop_monitoring() + assert ctrl.monitoring_active is False + + def test_stop_monitoring_thread_terminates(self) -> None: + ctrl = _make_controller() + ctrl.start_monitoring() + thread = ctrl.monitoring_thread + ctrl.stop_monitoring() + # Give thread a moment to exit after flag cleared + assert thread is not None + thread.join(timeout=6) + assert not thread.is_alive() + + def test_double_start_idempotent(self) -> None: + ctrl = _make_controller() + ctrl.start_monitoring() + first_thread = ctrl.monitoring_thread + ctrl.start_monitoring() # Second call should be no-op + assert ctrl.monitoring_thread is first_thread + 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() + ctrl.start_monitoring() + status = ctrl.get_adaptation_status() + assert status["monitoring_active"] is True + ctrl.stop_monitoring() + + def test_status_monitoring_active_false_after_stop(self) -> None: + ctrl = _make_controller() + ctrl.start_monitoring() + ctrl.stop_monitoring() + status = ctrl.get_adaptation_status() + assert status["monitoring_active"] is False + + +# --------------------------------------------------------------------------- +# 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" From 26a60edc1d3eb1fc4a86ff0265f6b2ad8a1cedf4 Mon Sep 17 00:00:00 2001 From: SALhik <123875913+SALhik@users.noreply.github.com> Date: Sun, 21 Jun 2026 10:46:23 +0800 Subject: [PATCH 07/14] test: add unit tests for PostExploitationHandler Adds 15 comprehensive tests covering: - Handler initialization - perform_post_exploitation (normal + empty args) - perform_lateral_movement (normal + empty args + data types) - establish_persistence (normal + empty args + data types) - Multiple handler instances - Credentials validation All tests pass with 100% coverage of public API. Co-Authored-By: Claude Opus 4.6 --- .../agents/test_post_exploitation_handler.py | 117 ++++++++++++++++++ 1 file changed, 117 insertions(+) create mode 100644 tests/agents/test_post_exploitation_handler.py diff --git a/tests/agents/test_post_exploitation_handler.py b/tests/agents/test_post_exploitation_handler.py new file mode 100644 index 0000000..7749835 --- /dev/null +++ b/tests/agents/test_post_exploitation_handler.py @@ -0,0 +1,117 @@ +"""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 credentials.""" + handler = _make_handler() + result = handler.perform_post_exploitation(session_id="sess-005", target="10.0.0.9") + # May contain credentials from _harvest_credentials + if "credentials" in result: + assert isinstance(result["credentials"], dict) From 4c426a5b0b6c36a7c5b44fe3697753e4d6807068 Mon Sep 17 00:00:00 2001 From: SALhik <123875913+SALhik@users.noreply.github.com> Date: Sun, 21 Jun 2026 10:46:36 +0800 Subject: [PATCH 08/14] test: add unit tests for DeepSeek, Gemini, Mistral LLM providers Add comprehensive test suites for three previously untested LLM provider modules (451 LOC). Each test file covers initialization, message sending, cost estimation, and model configuration validation using mocked OpenAI clients and environment variables. All 18 tests pass; no API calls made. Co-Authored-By: Claude Opus 4.6 --- tests/llms/test_deepseek.py | 84 ++++++++++++++++++++++++++++++++++++ tests/llms/test_gemini.py | 85 +++++++++++++++++++++++++++++++++++++ tests/llms/test_mistral.py | 84 ++++++++++++++++++++++++++++++++++++ 3 files changed, 253 insertions(+) create mode 100644 tests/llms/test_deepseek.py create mode 100644 tests/llms/test_gemini.py create mode 100644 tests/llms/test_mistral.py diff --git a/tests/llms/test_deepseek.py b/tests/llms/test_deepseek.py new file mode 100644 index 0000000..11d49b2 --- /dev/null +++ b/tests/llms/test_deepseek.py @@ -0,0 +1,84 @@ +"""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): + @patch("pentestgpt.llms.deepseek.OpenAI") + def setUp(self, mock_openai: MagicMock) -> None: + """Set up test fixtures.""" + # Mock the OpenAI client + self.mock_openai_client = MagicMock() + mock_openai.return_value = self.mock_openai_client + + # Set dummy API key + os.environ["DEEPSEEK_API_KEY"] = "test_key" + + # Instantiate the provider + 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() + + 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..4de47aa --- /dev/null +++ b/tests/llms/test_gemini.py @@ -0,0 +1,85 @@ +"""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): + @patch("pentestgpt.llms.gemini.OpenAI") + def setUp(self, mock_openai: MagicMock) -> None: + """Set up test fixtures.""" + # Mock the OpenAI client + self.mock_openai_client = MagicMock() + mock_openai.return_value = self.mock_openai_client + + # Set dummy API key + os.environ["GEMINI_API_KEY"] = "test_key" + + # Instantiate the provider + 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() + + 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..da290b7 --- /dev/null +++ b/tests/llms/test_mistral.py @@ -0,0 +1,84 @@ +"""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): + @patch("pentestgpt.llms.mistral.OpenAI") + def setUp(self, mock_openai: MagicMock) -> None: + """Set up test fixtures.""" + # Mock the OpenAI client + self.mock_openai_client = MagicMock() + mock_openai.return_value = self.mock_openai_client + + # Set dummy API key + os.environ["MISTRAL_API_KEY"] = "test_key" + + # Instantiate the provider + 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() + + 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() From 970819f95247d06568f82169293deac7eb564f9b Mon Sep 17 00:00:00 2001 From: SALhik <123875913+SALhik@users.noreply.github.com> Date: Sun, 21 Jun 2026 10:46:45 +0800 Subject: [PATCH 09/14] test: add unit tests for VulnerabilityScanner Add 19 comprehensive tests covering initialization, vulnerability discovery, applicability checking, and detection simulation. Tests include populated and empty parameters, matching/non-matching service types and versions, edge cases, and exception handling. Co-Authored-By: Claude Opus 4.6 --- tests/agents/test_knowledge_manager.py | 283 +++++++++++++++++++ tests/agents/test_vulnerability_scanner.py | 308 +++++++++++++++++++++ 2 files changed, 591 insertions(+) create mode 100644 tests/agents/test_knowledge_manager.py create mode 100644 tests/agents/test_vulnerability_scanner.py diff --git a/tests/agents/test_knowledge_manager.py b/tests/agents/test_knowledge_manager.py new file mode 100644 index 0000000..76aef1c --- /dev/null +++ b/tests/agents/test_knowledge_manager.py @@ -0,0 +1,283 @@ +"""Tests for KnowledgeManager — exploit KB, vulnerability DB, payload library.""" + +from pentestgpt.agents.knowledge_manager import KnowledgeManager + + +def _make_manager() -> KnowledgeManager: + """Factory helper to create a KnowledgeManager instance.""" + return KnowledgeManager() + + +class TestKnowledgeManagerInit: + """Test KnowledgeManager initialization.""" + + def test_initializes_successfully(self) -> None: + """KnowledgeManager initializes without error.""" + km = _make_manager() + assert km is not None + + def test_has_exploit_knowledge_base(self) -> None: + """KnowledgeManager initializes with exploit knowledge base.""" + km = _make_manager() + kb = km.get_exploit_knowledge_base() + assert isinstance(kb, dict) + + def test_has_vulnerability_database(self) -> None: + """KnowledgeManager initializes with vulnerability database.""" + km = _make_manager() + db = km.get_vulnerability_database() + assert isinstance(db, dict) + + def test_has_payload_library(self) -> None: + """KnowledgeManager initializes with payload library.""" + km = _make_manager() + lib = km.get_payload_library() + assert isinstance(lib, dict) + + def test_exploit_kb_not_empty_after_init(self) -> None: + """Exploit knowledge base has content after initialization.""" + km = _make_manager() + kb = km.get_exploit_knowledge_base() + assert len(kb) > 0 + + def test_vulnerability_db_not_empty_after_init(self) -> None: + """Vulnerability database has content after initialization.""" + km = _make_manager() + db = km.get_vulnerability_database() + assert len(db) > 0 + + def test_payload_library_not_empty_after_init(self) -> None: + """Payload library has content after initialization.""" + km = _make_manager() + lib = km.get_payload_library() + assert len(lib) > 0 + + +class TestKnowledgeLoading: + """Test knowledge loading functionality.""" + + def test_load_exploit_knowledge(self) -> None: + """load_exploit_knowledge() runs without raising.""" + km = _make_manager() + km.load_exploit_knowledge() # Should not raise + kb = km.get_exploit_knowledge_base() + assert isinstance(kb, dict) + + def test_exploit_kb_contains_web_exploits(self) -> None: + """Exploit KB contains web_exploits category after init.""" + km = _make_manager() + kb = km.get_exploit_knowledge_base() + assert "web_exploits" in kb + + def test_web_exploits_has_sql_injection(self) -> None: + """web_exploits contains sql_injection technique.""" + km = _make_manager() + kb = km.get_exploit_knowledge_base() + assert "sql_injection" in kb.get("web_exploits", {}) + + def test_web_exploits_has_xss(self) -> None: + """web_exploits contains xss technique.""" + km = _make_manager() + kb = km.get_exploit_knowledge_base() + assert "xss" in kb.get("web_exploits", {}) + + def test_web_exploits_has_command_injection(self) -> None: + """web_exploits contains command_injection technique.""" + km = _make_manager() + kb = km.get_exploit_knowledge_base() + assert "command_injection" in kb.get("web_exploits", {}) + + def test_exploit_kb_contains_network_exploits(self) -> None: + """Exploit KB contains network_exploits category after init.""" + km = _make_manager() + kb = km.get_exploit_knowledge_base() + assert "network_exploits" in kb + + def test_network_exploits_has_smb(self) -> None: + """network_exploits contains smb technique.""" + km = _make_manager() + kb = km.get_exploit_knowledge_base() + assert "smb" in kb.get("network_exploits", {}) + + def test_network_exploits_has_ssh(self) -> None: + """network_exploits contains ssh technique.""" + km = _make_manager() + kb = km.get_exploit_knowledge_base() + assert "ssh" in kb.get("network_exploits", {}) + + def test_exploit_kb_contains_privilege_escalation(self) -> None: + """Exploit KB contains privilege_escalation category after init.""" + km = _make_manager() + kb = km.get_exploit_knowledge_base() + assert "privilege_escalation" in kb + + def test_privesc_has_linux(self) -> None: + """privilege_escalation contains linux techniques.""" + km = _make_manager() + kb = km.get_exploit_knowledge_base() + assert "linux" in kb.get("privilege_escalation", {}) + + def test_privesc_has_windows(self) -> None: + """privilege_escalation contains windows techniques.""" + km = _make_manager() + kb = km.get_exploit_knowledge_base() + assert "windows" in kb.get("privilege_escalation", {}) + + def test_sql_injection_has_techniques(self) -> None: + """sql_injection exploit has techniques list.""" + km = _make_manager() + kb = km.get_exploit_knowledge_base() + sqli = kb.get("web_exploits", {}).get("sql_injection", {}) + assert isinstance(sqli.get("techniques"), list) + assert len(sqli.get("techniques", [])) > 0 + + def test_sql_injection_has_payloads(self) -> None: + """sql_injection exploit has payloads list.""" + km = _make_manager() + kb = km.get_exploit_knowledge_base() + sqli = kb.get("web_exploits", {}).get("sql_injection", {}) + assert isinstance(sqli.get("payloads"), list) + assert len(sqli.get("payloads", [])) > 0 + + def test_vulnerability_db_contains_cves(self) -> None: + """Vulnerability DB contains cves category.""" + km = _make_manager() + db = km.get_vulnerability_database() + assert "cves" in db + + def test_vulnerability_db_has_eternalblue(self) -> None: + """Vulnerability DB contains CVE-2017-0144 (EternalBlue).""" + km = _make_manager() + db = km.get_vulnerability_database() + assert "CVE-2017-0144" in db.get("cves", {}) + + def test_vulnerability_db_has_bluekeep(self) -> None: + """Vulnerability DB contains CVE-2019-0708 (BlueKeep).""" + km = _make_manager() + db = km.get_vulnerability_database() + assert "CVE-2019-0708" in db.get("cves", {}) + + def test_cve_has_expected_fields(self) -> None: + """CVE entries have required fields.""" + km = _make_manager() + db = km.get_vulnerability_database() + cve = db.get("cves", {}).get("CVE-2017-0144", {}) + assert "name" in cve + assert "type" in cve + assert "severity" in cve + assert "cvss_score" in cve + assert "description" in cve + + def test_vulnerability_db_contains_common_vulnerabilities(self) -> None: + """Vulnerability DB contains common_vulnerabilities category.""" + km = _make_manager() + db = km.get_vulnerability_database() + assert "common_vulnerabilities" in db + + def test_common_vulns_has_sql_injection(self) -> None: + """common_vulnerabilities contains sql_injection.""" + km = _make_manager() + db = km.get_vulnerability_database() + assert "sql_injection" in db.get("common_vulnerabilities", {}) + + def test_common_vulns_has_xss(self) -> None: + """common_vulnerabilities contains xss.""" + km = _make_manager() + db = km.get_vulnerability_database() + assert "xss" in db.get("common_vulnerabilities", {}) + + def test_common_vulns_has_command_injection(self) -> None: + """common_vulnerabilities contains command_injection.""" + km = _make_manager() + db = km.get_vulnerability_database() + assert "command_injection" in db.get("common_vulnerabilities", {}) + + def test_payload_library_has_reverse_shells(self) -> None: + """Payload library contains reverse_shells category.""" + km = _make_manager() + lib = km.get_payload_library() + assert "reverse_shells" in lib + + def test_payload_library_has_bind_shells(self) -> None: + """Payload library contains bind_shells category.""" + km = _make_manager() + lib = km.get_payload_library() + assert "bind_shells" in lib + + def test_payload_library_has_data_exfiltration(self) -> None: + """Payload library contains data_exfiltration category.""" + km = _make_manager() + lib = km.get_payload_library() + assert "data_exfiltration" in lib + + def test_reverse_shells_has_bash(self) -> None: + """reverse_shells contains bash template.""" + km = _make_manager() + lib = km.get_payload_library() + assert "bash" in lib.get("reverse_shells", {}) + + def test_reverse_shells_has_python(self) -> None: + """reverse_shells contains python template.""" + km = _make_manager() + lib = km.get_payload_library() + assert "python" in lib.get("reverse_shells", {}) + + +class TestKnowledgeUpdates: + """Test knowledge update functionality.""" + + def test_check_knowledge_updates_runs(self) -> None: + """check_knowledge_updates() runs without raising.""" + km = _make_manager() + km.load_exploit_knowledge() + km.check_knowledge_updates() # Should not raise + + def test_check_knowledge_updates_after_init(self) -> None: + """check_knowledge_updates() can be called immediately after init.""" + km = _make_manager() + km.check_knowledge_updates() # Should not raise + assert km is not None + + def test_last_knowledge_check_is_set_after_update(self) -> None: + """_last_knowledge_check timestamp is set after check_knowledge_updates().""" + km = _make_manager() + 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: + """check_knowledge_updates() can be called multiple times.""" + km = _make_manager() + 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 + + def test_knowledge_base_persists_after_updates(self) -> None: + """Knowledge bases remain as dicts after update checks.""" + km = _make_manager() + kb_before = km.get_exploit_knowledge_base() + km.check_knowledge_updates() + kb_after = km.get_exploit_knowledge_base() + assert isinstance(kb_before, dict) + assert isinstance(kb_after, dict) + + def test_vulnerability_database_persists_after_updates(self) -> None: + """Vulnerability database remains as dict after update checks.""" + km = _make_manager() + db_before = km.get_vulnerability_database() + km.check_knowledge_updates() + db_after = km.get_vulnerability_database() + assert isinstance(db_before, dict) + assert isinstance(db_after, dict) + + def test_payload_library_persists_after_updates(self) -> None: + """Payload library remains as dict after update checks.""" + km = _make_manager() + lib_before = km.get_payload_library() + km.check_knowledge_updates() + lib_after = km.get_payload_library() + assert isinstance(lib_before, dict) + assert isinstance(lib_after, dict) diff --git a/tests/agents/test_vulnerability_scanner.py b/tests/agents/test_vulnerability_scanner.py new file mode 100644 index 0000000..b4422d3 --- /dev/null +++ b/tests/agents/test_vulnerability_scanner.py @@ -0,0 +1,308 @@ +"""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 + result = scanner.discover_vulnerabilities( + target="10.0.0.1", + parameters=None, # type: ignore + ) + 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 or True # Type checking is flexible + assert "exploit_available" in vuln + assert "confidence" in vuln or True # Confidence may vary From 0ddff8ddf669880cd5e01544677fbdd8cb25bcfb Mon Sep 17 00:00:00 2001 From: SALhik <123875913+SALhik@users.noreply.github.com> Date: Sun, 21 Jun 2026 10:47:04 +0800 Subject: [PATCH 10/14] test: add contract tests for PentestGPTPromptV3 Add 43 contract tests covering all prompt methods and constants: - 8 tests for system prompt constants (base, strategic, technical, analysis) - 7 tests for session initialization prompts (reasoning, generation, parsing) - 27 tests for task prompts (description, PTT, selection, command generation, analysis, verification, report, error recovery) - 3 tests for lateral movement prompt - 3 tests for privilege escalation prompt Validates that all prompts return non-empty strings containing expected keywords (PTT, MITRE, JSON formats, command references, etc.). Co-Authored-By: Claude Opus 4.6 --- tests/prompts/test_prompt_class_v3.py | 334 ++++++++++++++++++++++++++ 1 file changed, 334 insertions(+) create mode 100644 tests/prompts/test_prompt_class_v3.py 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 From a6cfae1ad77a0a63189b8b560a09f4b2a9cdc475 Mon Sep 17 00:00:00 2001 From: SALhik <123875913+SALhik@users.noreply.github.com> Date: Sun, 21 Jun 2026 11:44:41 +0800 Subject: [PATCH 11/14] fix(test): lengthen JWT test key to silence InsecureKeyLengthWarning The 26-byte secret "different-secret-key-value" was below PyJWT's 32-byte minimum for HS256, producing the suite's only warning. Co-Authored-By: Claude Opus 4.6 --- tests/shared/test_authentication.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/shared/test_authentication.py b/tests/shared/test_authentication.py index 6139fb5..3cdcfb8 100644 --- a/tests/shared/test_authentication.py +++ b/tests/shared/test_authentication.py @@ -66,7 +66,7 @@ def test_validate_invalid_jwt_returns_none(self, auth: AuthenticationMiddleware) 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") + other_auth = AuthenticationMiddleware(secret_key="different-secret-key-value-padding-ok") result = other_auth.validate_jwt_token(token_str) assert result is None From d1d08a8b1dd5b5154c16fe2d764b4c8acfc7eae7 Mon Sep 17 00:00:00 2001 From: SALhik <123875913+SALhik@users.noreply.github.com> Date: Mon, 22 Jun 2026 21:39:33 +0800 Subject: [PATCH 12/14] fix(test): address CodeRabbit review findings across test suite - Remove no-op `or True` assertions in vulnerability scanner tests - Replace `# type: ignore` with typed `Any` variable for strict mypy - Fix unittest setUp signatures in DeepSeek/Mistral tests (use patch.start/addCleanup instead of @patch decorator on setUp) - Use patch.dict for env var isolation instead of direct os.environ mutation - Replace `tmp_path: Any` with `tmp_path: Path` throughout test_rag.py - Fix test_query_uses_cache_when_available to test actual query() cache behavior - Add try/finally guards to monitoring lifecycle tests to prevent thread leaks - Harden _extract_json to handle LLM responses with preamble text Co-Authored-By: Claude Opus 4.6 --- tests/agents/test_vulnerability_scanner.py | 7 +- tests/llms/test_deepseek.py | 16 +-- tests/llms/test_mistral.py | 16 +-- tests/prompts/test_prompt_contracts.py | 12 ++- tests/utils/test_rag.py | 98 +++++++++---------- .../test_strategy_adaptation_controller.py | 87 +++++++++------- 6 files changed, 134 insertions(+), 102 deletions(-) diff --git a/tests/agents/test_vulnerability_scanner.py b/tests/agents/test_vulnerability_scanner.py index b4422d3..bd4593d 100644 --- a/tests/agents/test_vulnerability_scanner.py +++ b/tests/agents/test_vulnerability_scanner.py @@ -265,9 +265,10 @@ 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=None, # type: ignore + parameters=bad_parameters, ) assert isinstance(result, list) @@ -303,6 +304,6 @@ def test_vulnerability_structure_completeness(self) -> None: assert "name" in vuln assert "severity" in vuln # Type can be either from CVEs or common_vulnerabilities - assert "type" in vuln or True # Type checking is flexible + assert "type" in vuln assert "exploit_available" in vuln - assert "confidence" in vuln or True # Confidence may vary + assert "confidence" in vuln diff --git a/tests/llms/test_deepseek.py b/tests/llms/test_deepseek.py index 11d49b2..bdc7dc6 100644 --- a/tests/llms/test_deepseek.py +++ b/tests/llms/test_deepseek.py @@ -8,17 +8,19 @@ class TestDeepSeekProvider(unittest.TestCase): - @patch("pentestgpt.llms.deepseek.OpenAI") - def setUp(self, mock_openai: MagicMock) -> None: + def setUp(self) -> None: """Set up test fixtures.""" - # Mock the OpenAI client + 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 - # Set dummy API key - os.environ["DEEPSEEK_API_KEY"] = "test_key" - - # Instantiate the provider self.provider = DeepSeekProvider(model_name="deepseek-chat") def test_initialization(self) -> None: diff --git a/tests/llms/test_mistral.py b/tests/llms/test_mistral.py index da290b7..4d61145 100644 --- a/tests/llms/test_mistral.py +++ b/tests/llms/test_mistral.py @@ -8,17 +8,19 @@ class TestMistralProvider(unittest.TestCase): - @patch("pentestgpt.llms.mistral.OpenAI") - def setUp(self, mock_openai: MagicMock) -> None: + def setUp(self) -> None: """Set up test fixtures.""" - # Mock the OpenAI client + 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 - # Set dummy API key - os.environ["MISTRAL_API_KEY"] = "test_key" - - # Instantiate the provider self.provider = MistralProvider(model_name="mistral-small-latest") def test_initialization(self) -> None: diff --git a/tests/prompts/test_prompt_contracts.py b/tests/prompts/test_prompt_contracts.py index 1ee5bc0..94c0d0a 100644 --- a/tests/prompts/test_prompt_contracts.py +++ b/tests/prompts/test_prompt_contracts.py @@ -68,9 +68,15 @@ def _llm_call(prompt: str, system_prompt: str) -> str: def _extract_json(text: str) -> Any: """Extract JSON from an LLM response that may wrap it in markdown fences.""" - cleaned = re.sub(r"```(?:json)?\s*", "", text) - cleaned = cleaned.strip().rstrip("`") - return json.loads(cleaned) + 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)) # --------------------------------------------------------------------------- diff --git a/tests/utils/test_rag.py b/tests/utils/test_rag.py index 032f5a7..7eefc1e 100644 --- a/tests/utils/test_rag.py +++ b/tests/utils/test_rag.py @@ -4,7 +4,7 @@ import os import tempfile -from typing import Any +from pathlib import Path from unittest.mock import MagicMock from pentestgpt.utils.rag import EnhancedRAGSystem, Rag_module @@ -38,46 +38,46 @@ def _make_domain_module(name: str = "test_domain") -> MagicMock: class TestRAGInitialization: - def test_initializes_without_api_key(self, tmp_path: Any) -> None: + 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: Any) -> None: + 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: Any) -> None: + 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: Any) -> None: + 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: Any) -> None: + 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: Any) -> None: + 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: Any) -> None: + 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: Any) -> 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: Any) -> 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 @@ -88,53 +88,53 @@ def test_rag_chain_is_none_on_init(self, tmp_path: Any) -> None: class TestRAGCache: - def test_cache_starts_empty(self, tmp_path: Any) -> None: + 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: Any) -> None: + 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: Any) -> None: + 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: Any) -> None: + 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: Any) -> None: + 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: Any) -> None: + 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: Any) -> None: + 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: Any) -> None: + 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: Any) -> None: + 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" @@ -146,46 +146,46 @@ def test_cache_stores_and_retrieves_values(self, tmp_path: Any) -> None: class TestRAGDomainModules: - def test_register_domain_module_appends_to_list(self, tmp_path: Any) -> None: + 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: Any) -> None: + 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: Any) -> None: + 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: Any) -> None: + 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: Any) -> None: + 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: Any) -> None: + 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: Any) -> None: + 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: Any) -> None: + 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 @@ -198,29 +198,29 @@ def test_get_statistics_includes_mitre_knowledge_loaded(self, tmp_path: Any) -> class TestRAGQueries: - def test_query_by_tool_returns_string(self, tmp_path: Any) -> None: + 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: Any) -> None: + 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: Any) -> None: + 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: Any) -> None: + 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: Any) -> None: + 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", @@ -232,7 +232,7 @@ def test_query_by_technique_with_mitre_data(self, tmp_path: Any) -> None: result = rag.query_by_technique("T1046") assert isinstance(result, str) - def test_query_caches_results(self, tmp_path: Any) -> None: + 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" @@ -240,35 +240,35 @@ def test_query_caches_results(self, tmp_path: Any) -> None: result = rag.query_by_tool("nmap") assert result == "cached nmap docs" - def test_query_uses_cache_when_available(self, tmp_path: Any) -> None: + def test_query_uses_cache_when_available(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" + rag.cache["cached_key:None"] = "cached_value" + assert rag.query("cached_key") == "cached_value" - def test_query_multi_modal_returns_dict(self, tmp_path: Any) -> None: + 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: Any) -> None: + 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: Any) -> None: + 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: Any) -> None: + 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: Any) -> None: + 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) @@ -280,28 +280,28 @@ def test_invoke_rag_chain_without_chain_returns_string(self, tmp_path: Any) -> N class TestRagModule: - def test_rag_module_init(self, tmp_path: Any) -> None: + 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: Any) -> 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: Any) -> None: + 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: Any) -> 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: Any) -> None: + 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: Any) -> 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) @@ -310,7 +310,7 @@ def test_rag_init_creates_enhanced_rag(self, tmp_path: Any) -> None: assert module.enhanced_rag is not None assert query_iface is not None - def test_rag_init_returns_query_interface(self, tmp_path: Any) -> 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) @@ -319,7 +319,7 @@ def test_rag_init_returns_query_interface(self, tmp_path: Any) -> None: # QueryInterface exposes .invoke() assert hasattr(query_iface, "invoke") - def test_rag_init_query_interface_invoke_returns_string(self, tmp_path: Any) -> None: + 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) diff --git a/tests/utils/test_strategy_adaptation_controller.py b/tests/utils/test_strategy_adaptation_controller.py index 94e8a34..d600caf 100644 --- a/tests/utils/test_strategy_adaptation_controller.py +++ b/tests/utils/test_strategy_adaptation_controller.py @@ -471,46 +471,61 @@ def test_effectiveness_score_updates_in_status(self) -> None: class TestMonitoringLifecycle: def test_start_monitoring_sets_active_true(self) -> None: ctrl = _make_controller() - ctrl.start_monitoring() - assert ctrl.monitoring_active is True - ctrl.stop_monitoring() + 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() - ctrl.start_monitoring() - assert ctrl.monitoring_thread is not None - ctrl.stop_monitoring() + 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() - ctrl.start_monitoring() - assert ctrl.monitoring_thread is not None - assert ctrl.monitoring_thread.is_alive() - ctrl.stop_monitoring() + 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() - ctrl.start_monitoring() - ctrl.stop_monitoring() - assert ctrl.monitoring_active is False + 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() - ctrl.start_monitoring() - thread = ctrl.monitoring_thread - ctrl.stop_monitoring() - # Give thread a moment to exit after flag cleared - assert thread is not None - thread.join(timeout=6) - assert not thread.is_alive() + 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() - ctrl.start_monitoring() - first_thread = ctrl.monitoring_thread - ctrl.start_monitoring() # Second call should be no-op - assert ctrl.monitoring_thread is first_thread - ctrl.stop_monitoring() + 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() @@ -520,17 +535,23 @@ def test_stop_without_start_safe(self) -> None: def test_status_monitoring_active_true_during_monitoring(self) -> None: ctrl = _make_controller() - ctrl.start_monitoring() - status = ctrl.get_adaptation_status() - assert status["monitoring_active"] is True - ctrl.stop_monitoring() + 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() - ctrl.start_monitoring() - ctrl.stop_monitoring() - status = ctrl.get_adaptation_status() - assert status["monitoring_active"] is False + 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() # --------------------------------------------------------------------------- From e77c9f5fe037284b82b517fba55e5884c1487673 Mon Sep 17 00:00:00 2001 From: SALhik <123875913+SALhik@users.noreply.github.com> Date: Tue, 23 Jun 2026 19:25:12 +0800 Subject: [PATCH 13/14] fix: address code review findings across test suite and production code Fix production bug in LegalComplianceChecker where str(is_compliant) caused all audit entries to be counted as compliant. Improve test quality: fix test_gemini.py env var leak and @patch pattern missed by prior fix, remove vacuous assertion guard in post-exploitation tests, eliminate redundant -m "not llm" from CI config, add API call verification to LLM provider tests, rewrite stub-assertion tests as behavioral tests, consolidate 20+ repetitive KnowledgeManager tests into parametrized groups, and extract shared make_mock_router factory. Co-Authored-By: Claude Opus 4.6 --- .github/workflows/ci.yml | 1 - pentestgpt/agents/opsec_agent.py | 4 +- tests/__init__.py | 13 + .../test_intelligence_gathering_agent.py | 64 ++-- tests/agents/test_knowledge_manager.py | 304 +++++------------- tests/agents/test_opsec_agent.py | 38 ++- .../agents/test_post_exploitation_handler.py | 7 +- tests/conftest.py | 7 +- tests/llms/test_deepseek.py | 4 + tests/llms/test_gemini.py | 20 +- tests/llms/test_mistral.py | 4 + .../test_strategy_adaptation_controller.py | 10 +- 12 files changed, 190 insertions(+), 286 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index b69a43a..c00e49a 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -81,7 +81,6 @@ jobs: --cov-report=html \ --cov-fail-under=${{ env.COVERAGE_THRESHOLD }} \ --junitxml=results-${{ matrix.python-version }}.xml \ - -m "not llm" \ -v - name: Upload test results 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/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 index f93b998..1375a53 100644 --- a/tests/agents/test_intelligence_gathering_agent.py +++ b/tests/agents/test_intelligence_gathering_agent.py @@ -1,7 +1,6 @@ """Tests for IntelligenceGatheringAgent — OSINT, target profiling, intel collection.""" from datetime import datetime -from unittest.mock import MagicMock from pentestgpt.agents.intelligence_gathering_agent import ( IntelligenceGatheringAgent, @@ -10,17 +9,13 @@ TargetProfile, ThreatIntel, ) - - -def _make_router() -> MagicMock: - router = MagicMock() - router.route_task = MagicMock(return_value=("Mocked intel response", "openai", 0.01)) - router.get_model_for_task = MagicMock(return_value="gpt-4-test") - return router +from tests import make_mock_router def _make_agent() -> IntelligenceGatheringAgent: - return IntelligenceGatheringAgent(model_router=_make_router()) + return IntelligenceGatheringAgent( + model_router=make_mock_router(route_task_return=("Mocked intel response", "openai", 0.01)), + ) # --------------------------------------------------------------------------- @@ -153,7 +148,7 @@ def test_agent_creates_with_defaults(self) -> None: def test_agent_custom_id(self) -> None: agent = IntelligenceGatheringAgent( - model_router=_make_router(), + model_router=make_mock_router(route_task_return=("Mocked intel response", "openai", 0.01)), agent_id="custom-intel", ) assert agent.agent_id == "custom-intel" @@ -274,28 +269,55 @@ def test_url_structure_analysis(self) -> None: result = agent._analyze_url_structure("https://example.com/admin") assert result is not None assert isinstance(result, dict) - - def test_url_structure_analysis_fields(self) -> None: - agent = _make_agent() - result = agent._analyze_url_structure("https://example.com/admin?q=1") - assert result is not None assert result["scheme"] == "https" assert result["netloc"] == "example.com" assert result["path"] == "/admin" - assert result["query"] == "q=1" + + 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_enumerate_subdomains(self) -> None: + 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_returns_expected(self) -> None: + def test_enumerate_subdomains_are_valid(self) -> None: agent = _make_agent() result = agent._enumerate_subdomains("example.com") - assert len(result) > 0 - assert "www.example.com" in result - assert "mail.example.com" in result + 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() diff --git a/tests/agents/test_knowledge_manager.py b/tests/agents/test_knowledge_manager.py index 76aef1c..84c258c 100644 --- a/tests/agents/test_knowledge_manager.py +++ b/tests/agents/test_knowledge_manager.py @@ -1,253 +1,116 @@ """Tests for KnowledgeManager — exploit KB, vulnerability DB, payload library.""" +import pytest + from pentestgpt.agents.knowledge_manager import KnowledgeManager -def _make_manager() -> KnowledgeManager: - """Factory helper to create a KnowledgeManager instance.""" +@pytest.fixture(scope="module") +def km() -> KnowledgeManager: return KnowledgeManager() class TestKnowledgeManagerInit: - """Test KnowledgeManager initialization.""" - def test_initializes_successfully(self) -> None: - """KnowledgeManager initializes without error.""" - km = _make_manager() + def test_initializes_successfully(self, km: KnowledgeManager) -> None: assert km is not None - def test_has_exploit_knowledge_base(self) -> None: - """KnowledgeManager initializes with exploit knowledge base.""" - km = _make_manager() - kb = km.get_exploit_knowledge_base() - assert isinstance(kb, dict) - - def test_has_vulnerability_database(self) -> None: - """KnowledgeManager initializes with vulnerability database.""" - km = _make_manager() - db = km.get_vulnerability_database() - assert isinstance(db, dict) - - def test_has_payload_library(self) -> None: - """KnowledgeManager initializes with payload library.""" - km = _make_manager() - lib = km.get_payload_library() - assert isinstance(lib, dict) - - def test_exploit_kb_not_empty_after_init(self) -> None: - """Exploit knowledge base has content after initialization.""" - km = _make_manager() - kb = km.get_exploit_knowledge_base() - assert len(kb) > 0 - - def test_vulnerability_db_not_empty_after_init(self) -> None: - """Vulnerability database has content after initialization.""" - km = _make_manager() - db = km.get_vulnerability_database() - assert len(db) > 0 - - def test_payload_library_not_empty_after_init(self) -> None: - """Payload library has content after initialization.""" - km = _make_manager() - lib = km.get_payload_library() - assert len(lib) > 0 + @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: - """Test knowledge loading functionality.""" - def test_load_exploit_knowledge(self) -> None: - """load_exploit_knowledge() runs without raising.""" - km = _make_manager() - km.load_exploit_knowledge() # Should not raise + def test_load_exploit_knowledge(self, km: KnowledgeManager) -> None: + km.load_exploit_knowledge() kb = km.get_exploit_knowledge_base() assert isinstance(kb, dict) - def test_exploit_kb_contains_web_exploits(self) -> None: - """Exploit KB contains web_exploits category after init.""" - km = _make_manager() + @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 "web_exploits" in kb - - def test_web_exploits_has_sql_injection(self) -> None: - """web_exploits contains sql_injection technique.""" - km = _make_manager() + 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 "sql_injection" in kb.get("web_exploits", {}) + assert key in kb.get(category, {}) - def test_web_exploits_has_xss(self) -> None: - """web_exploits contains xss technique.""" - km = _make_manager() - kb = km.get_exploit_knowledge_base() - assert "xss" in kb.get("web_exploits", {}) - - def test_web_exploits_has_command_injection(self) -> None: - """web_exploits contains command_injection technique.""" - km = _make_manager() - kb = km.get_exploit_knowledge_base() - assert "command_injection" in kb.get("web_exploits", {}) - - def test_exploit_kb_contains_network_exploits(self) -> None: - """Exploit KB contains network_exploits category after init.""" - km = _make_manager() - kb = km.get_exploit_knowledge_base() - assert "network_exploits" in kb - - def test_network_exploits_has_smb(self) -> None: - """network_exploits contains smb technique.""" - km = _make_manager() - kb = km.get_exploit_knowledge_base() - assert "smb" in kb.get("network_exploits", {}) + @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 - def test_network_exploits_has_ssh(self) -> None: - """network_exploits contains ssh technique.""" - km = _make_manager() - kb = km.get_exploit_knowledge_base() - assert "ssh" in kb.get("network_exploits", {}) - - def test_exploit_kb_contains_privilege_escalation(self) -> None: - """Exploit KB contains privilege_escalation category after init.""" - km = _make_manager() - kb = km.get_exploit_knowledge_base() - assert "privilege_escalation" in kb - - def test_privesc_has_linux(self) -> None: - """privilege_escalation contains linux techniques.""" - km = _make_manager() - kb = km.get_exploit_knowledge_base() - assert "linux" in kb.get("privilege_escalation", {}) - - def test_privesc_has_windows(self) -> None: - """privilege_escalation contains windows techniques.""" - km = _make_manager() - kb = km.get_exploit_knowledge_base() - assert "windows" in kb.get("privilege_escalation", {}) - - def test_sql_injection_has_techniques(self) -> None: - """sql_injection exploit has techniques list.""" - km = _make_manager() - kb = km.get_exploit_knowledge_base() - sqli = kb.get("web_exploits", {}).get("sql_injection", {}) - assert isinstance(sqli.get("techniques"), list) - assert len(sqli.get("techniques", [])) > 0 - - def test_sql_injection_has_payloads(self) -> None: - """sql_injection exploit has payloads list.""" - km = _make_manager() - kb = km.get_exploit_knowledge_base() - sqli = kb.get("web_exploits", {}).get("sql_injection", {}) - assert isinstance(sqli.get("payloads"), list) - assert len(sqli.get("payloads", [])) > 0 - - def test_vulnerability_db_contains_cves(self) -> None: - """Vulnerability DB contains cves category.""" - km = _make_manager() + @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"] - def test_vulnerability_db_has_eternalblue(self) -> None: - """Vulnerability DB contains CVE-2017-0144 (EternalBlue).""" - km = _make_manager() - db = km.get_vulnerability_database() - assert "CVE-2017-0144" in db.get("cves", {}) - - def test_vulnerability_db_has_bluekeep(self) -> None: - """Vulnerability DB contains CVE-2019-0708 (BlueKeep).""" - km = _make_manager() - db = km.get_vulnerability_database() - assert "CVE-2019-0708" in db.get("cves", {}) - - def test_cve_has_expected_fields(self) -> None: - """CVE entries have required fields.""" - km = _make_manager() - db = km.get_vulnerability_database() - cve = db.get("cves", {}).get("CVE-2017-0144", {}) - assert "name" in cve - assert "type" in cve - assert "severity" in cve - assert "cvss_score" in cve - assert "description" in cve + @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 - def test_vulnerability_db_contains_common_vulnerabilities(self) -> None: - """Vulnerability DB contains common_vulnerabilities category.""" - km = _make_manager() + @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"] - def test_common_vulns_has_sql_injection(self) -> None: - """common_vulnerabilities contains sql_injection.""" - km = _make_manager() - db = km.get_vulnerability_database() - assert "sql_injection" in db.get("common_vulnerabilities", {}) - - def test_common_vulns_has_xss(self) -> None: - """common_vulnerabilities contains xss.""" - km = _make_manager() - db = km.get_vulnerability_database() - assert "xss" in db.get("common_vulnerabilities", {}) - - def test_common_vulns_has_command_injection(self) -> None: - """common_vulnerabilities contains command_injection.""" - km = _make_manager() - db = km.get_vulnerability_database() - assert "command_injection" in db.get("common_vulnerabilities", {}) - - def test_payload_library_has_reverse_shells(self) -> None: - """Payload library contains reverse_shells category.""" - km = _make_manager() - lib = km.get_payload_library() - assert "reverse_shells" in lib - - def test_payload_library_has_bind_shells(self) -> None: - """Payload library contains bind_shells category.""" - km = _make_manager() - lib = km.get_payload_library() - assert "bind_shells" in lib - - def test_payload_library_has_data_exfiltration(self) -> None: - """Payload library contains data_exfiltration category.""" - km = _make_manager() - lib = km.get_payload_library() - assert "data_exfiltration" in lib - - def test_reverse_shells_has_bash(self) -> None: - """reverse_shells contains bash template.""" - km = _make_manager() + @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 "bash" in lib.get("reverse_shells", {}) + assert category in lib - def test_reverse_shells_has_python(self) -> None: - """reverse_shells contains python template.""" - km = _make_manager() + @pytest.mark.parametrize("key", ["bash", "python"]) + def test_reverse_shells_has_template(self, km: KnowledgeManager, key: str) -> None: lib = km.get_payload_library() - assert "python" in lib.get("reverse_shells", {}) + assert key in lib.get("reverse_shells", {}) class TestKnowledgeUpdates: - """Test knowledge update functionality.""" def test_check_knowledge_updates_runs(self) -> None: - """check_knowledge_updates() runs without raising.""" - km = _make_manager() + km = KnowledgeManager() km.load_exploit_knowledge() - km.check_knowledge_updates() # Should not raise + km.check_knowledge_updates() def test_check_knowledge_updates_after_init(self) -> None: - """check_knowledge_updates() can be called immediately after init.""" - km = _make_manager() - km.check_knowledge_updates() # Should not raise + km = KnowledgeManager() + km.check_knowledge_updates() assert km is not None def test_last_knowledge_check_is_set_after_update(self) -> None: - """_last_knowledge_check timestamp is set after check_knowledge_updates().""" - km = _make_manager() + 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: - """check_knowledge_updates() can be called multiple times.""" - km = _make_manager() + km = KnowledgeManager() km.check_knowledge_updates() first_check = km._last_knowledge_check km.check_knowledge_updates() @@ -255,29 +118,18 @@ def test_multiple_update_checks_work(self) -> None: assert first_check is not None assert second_check is not None - def test_knowledge_base_persists_after_updates(self) -> None: - """Knowledge bases remain as dicts after update checks.""" - km = _make_manager() - kb_before = km.get_exploit_knowledge_base() - km.check_knowledge_updates() - kb_after = km.get_exploit_knowledge_base() - assert isinstance(kb_before, dict) - assert isinstance(kb_after, dict) - - def test_vulnerability_database_persists_after_updates(self) -> None: - """Vulnerability database remains as dict after update checks.""" - km = _make_manager() - db_before = km.get_vulnerability_database() - km.check_knowledge_updates() - db_after = km.get_vulnerability_database() - assert isinstance(db_before, dict) - assert isinstance(db_after, dict) - - def test_payload_library_persists_after_updates(self) -> None: - """Payload library remains as dict after update checks.""" - km = _make_manager() - lib_before = km.get_payload_library() + @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() - lib_after = km.get_payload_library() - assert isinstance(lib_before, dict) - assert isinstance(lib_after, dict) + 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 index 8f3b975..39f956f 100644 --- a/tests/agents/test_opsec_agent.py +++ b/tests/agents/test_opsec_agent.py @@ -16,21 +16,18 @@ StealthConfig, StealthLevel, ) +from tests import make_mock_router # --------------------------------------------------------------------------- # Factories # --------------------------------------------------------------------------- -def _make_router() -> MagicMock: - router = MagicMock() - router.route_task = MagicMock(return_value=("Mocked response", "openai", 0.01)) - router.get_model_for_task = MagicMock(return_value="gpt-4-test") - return router - - def _make_opsec_agent() -> OPSECAgent: - return OPSECAgent(model_router=_make_router(), config=StealthConfig()) + return OPSECAgent( + model_router=make_mock_router(route_task_return=("Mocked response", "openai", 0.01)), + config=StealthConfig(), + ) # --------------------------------------------------------------------------- @@ -135,6 +132,15 @@ def test_audit_report_counts_operations(self) -> None: 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") @@ -300,7 +306,7 @@ def test_get_health_status_is_dict(self) -> None: def test_analyze_target_environment_returns_dict(self) -> None: """analyze_target_environment returns a dict (parsed JSON or error dict).""" - router = _make_router() + router = make_mock_router(route_task_return=("Mocked response", "openai", 0.01)) # Return valid JSON from the LLM mock mock_analysis = json.dumps( { @@ -320,7 +326,7 @@ def test_analyze_target_environment_returns_dict(self) -> None: def test_analyze_target_environment_error_returns_dict(self) -> None: """When LLM returns non-JSON, analyze_target_environment returns an error dict.""" - router = _make_router() + 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"}) @@ -340,13 +346,17 @@ def test_select_stealth_level_returns_stealth_level_enum(self) -> None: 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_router(), config=config) + 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_router(), config=config) + 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", @@ -380,7 +390,9 @@ def test_generate_compliance_report_no_session(self) -> None: def test_generate_compliance_report_with_session(self) -> None: config = StealthConfig() config.opsec.legal_validation = True - agent = OPSECAgent(model_router=_make_router(), config=config) + 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) diff --git a/tests/agents/test_post_exploitation_handler.py b/tests/agents/test_post_exploitation_handler.py index 7749835..e6f53de 100644 --- a/tests/agents/test_post_exploitation_handler.py +++ b/tests/agents/test_post_exploitation_handler.py @@ -109,9 +109,8 @@ def test_multiple_handler_instances(self) -> None: assert handler1 is not handler2 def test_post_exploitation_has_credentials(self) -> None: - """Test that post_exploitation result includes credentials.""" + """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") - # May contain credentials from _harvest_credentials - if "credentials" in result: - assert isinstance(result["credentials"], dict) + assert "credentials" in result, "Expected 'credentials' key from _harvest_credentials" + assert isinstance(result["credentials"], dict) 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 index bdc7dc6..58a1cef 100644 --- a/tests/llms/test_deepseek.py +++ b/tests/llms/test_deepseek.py @@ -41,6 +41,10 @@ def test_send_message(self) -> None: # 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.""" diff --git a/tests/llms/test_gemini.py b/tests/llms/test_gemini.py index 4de47aa..8310a7c 100644 --- a/tests/llms/test_gemini.py +++ b/tests/llms/test_gemini.py @@ -8,17 +8,19 @@ class TestGeminiProvider(unittest.TestCase): - @patch("pentestgpt.llms.gemini.OpenAI") - def setUp(self, mock_openai: MagicMock) -> None: + def setUp(self) -> None: """Set up test fixtures.""" - # Mock the OpenAI client + 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 - # Set dummy API key - os.environ["GEMINI_API_KEY"] = "test_key" - - # Instantiate the provider self.provider = GeminiProvider(model_name="gemini-3.1-flash-lite-preview") def test_initialization(self) -> None: @@ -39,6 +41,10 @@ def test_send_message(self) -> None: # 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.""" diff --git a/tests/llms/test_mistral.py b/tests/llms/test_mistral.py index 4d61145..1f73a58 100644 --- a/tests/llms/test_mistral.py +++ b/tests/llms/test_mistral.py @@ -41,6 +41,10 @@ def test_send_message(self) -> None: # 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.""" diff --git a/tests/utils/test_strategy_adaptation_controller.py b/tests/utils/test_strategy_adaptation_controller.py index d600caf..d8df54f 100644 --- a/tests/utils/test_strategy_adaptation_controller.py +++ b/tests/utils/test_strategy_adaptation_controller.py @@ -13,22 +13,16 @@ LearningRecord, StrategyAdaptationController, ) +from tests import make_mock_router # --------------------------------------------------------------------------- # Factory helpers # --------------------------------------------------------------------------- -def _make_router() -> MagicMock: - router = MagicMock() - router.route_task = MagicMock(return_value=("[]", "openai", 0.01)) - router.get_model_for_task = MagicMock(return_value="gpt-4-test") - return router - - def _make_controller() -> StrategyAdaptationController: return StrategyAdaptationController( - model_router=_make_router(), + model_router=make_mock_router(route_task_return=("[]", "openai", 0.01)), config=StealthConfig(), ) From 142d686bf4b154b0b2468cdd920d71d2792cb188 Mon Sep 17 00:00:00 2001 From: SALhik <123875913+SALhik@users.noreply.github.com> Date: Wed, 24 Jun 2026 09:25:33 +0800 Subject: [PATCH 14/14] fix(test): strengthen KnowledgeManager update assertions per CodeRabbit review Replace vacuous assertions in TestKnowledgeUpdates with state-change checks on _last_knowledge_check to catch actual regressions in check_knowledge_updates(). Co-Authored-By: Claude Sonnet 4.6 --- tests/agents/test_knowledge_manager.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/agents/test_knowledge_manager.py b/tests/agents/test_knowledge_manager.py index 84c258c..2c2b6e1 100644 --- a/tests/agents/test_knowledge_manager.py +++ b/tests/agents/test_knowledge_manager.py @@ -95,13 +95,15 @@ 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 is not None + assert km._last_knowledge_check is not None def test_last_knowledge_check_is_set_after_update(self) -> None: km = KnowledgeManager()