diff --git a/core/memory.py b/core/memory.py index 6b5d820..dbea64b 100644 --- a/core/memory.py +++ b/core/memory.py @@ -4,6 +4,7 @@ """ import json +import os from typing import Dict, List, Any, Optional from datetime import datetime from pathlib import Path @@ -158,8 +159,15 @@ def save_state(self, filepath: Path): "context": self.context } - filepath.parent.mkdir(parents=True, exist_ok=True) - with open(filepath, 'w') as f: + filepath.parent.mkdir(parents=True, exist_ok=True, mode=0o700) + try: + filepath.parent.chmod(0o700) + except OSError: + pass + flags = os.O_WRONLY | os.O_CREAT | os.O_TRUNC | getattr(os, "O_NOFOLLOW", 0) + fd = os.open(filepath, flags, 0o600) + os.fchmod(fd, 0o600) + with os.fdopen(fd, 'w') as f: json.dump(state, f, indent=2) def load_state(self, filepath: Path) -> bool: diff --git a/core/workflow.py b/core/workflow.py index 6932eae..d863156 100644 --- a/core/workflow.py +++ b/core/workflow.py @@ -3,17 +3,16 @@ Coordinates agents and manages pentest execution flow """ -import asyncio -from typing import Dict, Any, Optional, List +from typing import Dict, Any, List from pathlib import Path from datetime import datetime +import os -from core.agent import BaseAgent -from core.planner import PlannerAgent -from core.memory import PentestMemory, ToolExecution, Finding +from core.memory import PentestMemory from ai.gemini_client import GeminiClient from utils.logger import get_logger from utils.scope_validator import ScopeValidator +from utils.redaction import redact_sensitive_text class WorkflowEngine: @@ -166,10 +165,14 @@ async def _execute_step(self, step: Dict[str, Any]): step_type = step.get("type", "tool") if step_type == "tool": + # Re-resolve immediately before every external process. This narrows + # the DNS-rebinding/TOCTOU window and prevents a target that changed + # after workflow startup from reaching a prohibited address. + is_valid, reason = self.scope_validator.validate_target(self.target) + if not is_valid: + raise ValueError(f"Invalid target before tool execution: {reason}") # Use Tool Agent to select and execute tool tool_name = step["tool"] - objective = step.get("objective", f"Execute {tool_name}") - self.logger.info(f"Tool Agent selecting tool: {tool_name}") # Tool Agent executes the tool @@ -193,7 +196,7 @@ async def _execute_step(self, step: Dict[str, Any]): target=self.target, timestamp=datetime.now().isoformat(), exit_code=result.get("exit_code", 0), - output=result.get("raw_output", ""), # Store the FULL raw output + output=redact_sensitive_text(result.get("raw_output", "")), duration=result.get("duration", 0) ) self.memory.add_tool_execution(execution) @@ -204,7 +207,7 @@ async def _execute_step(self, step: Dict[str, Any]): tool=tool_name, target=self.target, command=result.get("command", ""), - output=result.get("raw_output", ""), + output=redact_sensitive_text(result.get("raw_output", "")), execution_id=execution_id # Pass execution ID to analyst ) @@ -228,14 +231,18 @@ async def _execute_step(self, step: Dict[str, Any]): # Save report output_dir = Path(self.config.get("output", {}).get("save_path", "./reports")) - output_dir.mkdir(parents=True, exist_ok=True) + output_dir.mkdir(parents=True, exist_ok=True, mode=0o700) + output_dir.chmod(0o700) # Use proper file extension extension_map = {"markdown": "md", "html": "html", "json": "json"} extension = extension_map.get(report_format, "md") report_file = output_dir / f"report_{self.memory.session_id}.{extension}" - with open(report_file, 'w', encoding='utf-8') as f: + flags = os.O_WRONLY | os.O_CREAT | os.O_TRUNC | getattr(os, "O_NOFOLLOW", 0) + fd = os.open(report_file, flags, 0o600) + os.fchmod(fd, 0o600) + with os.fdopen(fd, 'w', encoding='utf-8') as f: f.write(report["content"]) self.logger.info(f"Report saved to: {report_file}") @@ -250,6 +257,9 @@ async def _execute_ai_decision(self, decision: Dict[str, Any]): # Use Tool Agent to select appropriate tool try: + is_valid, reason = self.scope_validator.validate_target(self.target) + if not is_valid: + raise ValueError(f"Invalid target before tool execution: {reason}") tool_selection = await self.tool_agent.execute( objective=action, target=self.target @@ -277,7 +287,7 @@ async def _execute_ai_decision(self, decision: Dict[str, Any]): target=self.target, timestamp=datetime.now().isoformat(), exit_code=result.get("exit_code", 0), - output=result.get("raw_output", ""), + output=redact_sensitive_text(result.get("raw_output", "")), duration=result.get("duration", 0), ) self.memory.add_tool_execution(execution) @@ -287,7 +297,7 @@ async def _execute_ai_decision(self, decision: Dict[str, Any]): tool=tool_name, target=self.target, command=result.get("command", ""), - output=result.get("raw_output", ""), + output=redact_sensitive_text(result.get("raw_output", "")), execution_id=execution_id, ) @@ -391,7 +401,8 @@ def _maybe_advance_phase(self): def _save_session(self): """Save session state""" output_dir = Path(self.config.get("output", {}).get("save_path", "./reports")) - output_dir.mkdir(parents=True, exist_ok=True) + output_dir.mkdir(parents=True, exist_ok=True, mode=0o700) + output_dir.chmod(0o700) state_file = output_dir / f"session_{self.memory.session_id}.json" self.memory.save_state(state_file) diff --git a/pyproject.toml b/pyproject.toml index daa56c6..9ce4604 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,20 +25,23 @@ classifiers = [ ] dependencies = [ - "typer[all]>=0.24.1", - "rich>=15.0.0", - "langchain>=1.2.15", - "langchain-google-genai>=4.2.1", - "langchain-openai>=1.1.12", - "langchain-anthropic>=1.4.0", - "pyyaml>=6.0.3", - "python-dotenv>=1.2.2", - "pydantic>=2.12.5", - "asyncio>=4.0.0", - "aiofiles>=25.1.0", - "jinja2>=3.1.6", + "typer[all]>=0.24.1,<1", + "rich>=15.0.0,<16", + "langchain>=1.2.15,<2", + "langchain-google-genai>=4.2.1,<5", + "langchain-openai>=1.1.12,<2", + "langchain-anthropic>=1.4.0,<2", + "pyyaml>=6.0.3,<7", + "python-dotenv>=1.2.2,<2", + "pydantic>=2.12.5,<3", + "asyncio>=4.0.0,<5", + "aiofiles>=25.1.0,<26", + "jinja2>=3.1.6,<4", ] +# Release builds should be produced from a reviewed, hash-locked dependency +# set. Keep direct dependency bounds constrained to avoid unreviewed majors. + [project.optional-dependencies] dev = [ "pytest>=9.0.3", diff --git a/tests/test_base_tool.py b/tests/test_base_tool.py index debc6a1..1e60a36 100644 --- a/tests/test_base_tool.py +++ b/tests/test_base_tool.py @@ -73,6 +73,19 @@ def test_sensitive_flags_set_covers_key_flags(self): expected = {"--cookie", "--data", "--api-token", "--header", "-H"} assert expected.issubset(_SENSITIVE_FLAGS) + def test_execute_never_returns_secret_command(self, tool): + process = MagicMock(returncode=0) + process.communicate = AsyncMock(return_value=(b"ok", b"")) + tool.get_command = MagicMock( + return_value=["wpscan", "--api-token", "MY_SECRET_TOKEN", "example.com"] + ) + + with patch("asyncio.create_subprocess_exec", return_value=process): + result = asyncio.run(tool.execute("example.com")) + + assert "MY_SECRET_TOKEN" not in result["command"] + assert result["command"] == "wpscan --api-token example.com" + # --------------------------------------------------------------------------- # Process kill on timeout diff --git a/tests/test_memory.py b/tests/test_memory.py index 51cf53b..c4ed3f6 100644 --- a/tests/test_memory.py +++ b/tests/test_memory.py @@ -174,3 +174,9 @@ def test_load_wrong_schema_returns_false(self, memory, tmp_path): bad.write_text(json.dumps({"wrong": "schema"})) result = memory.load_state(bad) assert result is False + + def test_saved_state_is_private(self, memory, tmp_path): + state = tmp_path / "private" / "session.json" + memory.save_state(state) + assert state.stat().st_mode & 0o777 == 0o600 + assert state.parent.stat().st_mode & 0o777 == 0o700 diff --git a/tests/test_nmap_tool.py b/tests/test_nmap_tool.py index 687d6e1..cd1d104 100644 --- a/tests/test_nmap_tool.py +++ b/tests/test_nmap_tool.py @@ -134,6 +134,15 @@ def test_malformed_xml_returns_empty(self, nmap): result = nmap.parse_output("" + (" " * (10 * 1024 * 1024)) + "") + assert result["open_ports"] == [] + def test_no_os_returns_none(self, nmap): xml = """ diff --git a/tests/test_redaction.py b/tests/test_redaction.py new file mode 100644 index 0000000..9f0e9a2 --- /dev/null +++ b/tests/test_redaction.py @@ -0,0 +1,21 @@ +"""Tests for evidence redaction at the AI and persistence boundary.""" + +from utils.redaction import redact_sensitive_text + + +def test_redacts_authorization_header(): + value = redact_sensitive_text("Authorization: Bearer secret-token") + assert "secret-token" not in value + assert "" in value + + +def test_redacts_api_key_and_password(): + value = redact_sensitive_text("api_key=abc123&password=hunter2") + assert "abc123" not in value + assert "hunter2" not in value + + +def test_redacts_cookie_header(): + value = redact_sensitive_text("Cookie: session=secret; theme=dark\nbody") + assert "session=secret" not in value + assert value.endswith("\nbody") diff --git a/tests/test_scope_validator.py b/tests/test_scope_validator.py index 4dbefa0..fecd6e8 100644 --- a/tests/test_scope_validator.py +++ b/tests/test_scope_validator.py @@ -79,6 +79,12 @@ def test_unspecified_0000(self, validator): def test_ipv6_loopback(self, validator): assert validator._is_blacklisted("::1") is True + def test_ipv6_unique_local(self, validator): + assert validator._is_blacklisted("fd00::1") is True + + def test_ipv4_link_local(self, validator): + assert validator._is_blacklisted("169.254.169.254") is True + def test_private_class_a(self, validator): assert validator._is_blacklisted("10.0.0.1") is True diff --git a/tools/base_tool.py b/tools/base_tool.py index a7ba192..42058c5 100644 --- a/tools/base_tool.py +++ b/tools/base_tool.py @@ -6,7 +6,6 @@ import subprocess import shutil from typing import Dict, Any, Optional, List -from pathlib import Path from datetime import datetime from abc import ABC, abstractmethod @@ -70,7 +69,8 @@ async def execute(self, target: str, **kwargs) -> Dict[str, Any]: # Build command command = self.get_command(target, **kwargs) - self.logger.info(f"Executing: {self._sanitize_command_for_logging(command)}") + safe_command = self._sanitize_command_for_logging(command) + self.logger.info(f"Executing: {safe_command}") # Get timeout from config timeout = self.config.get("pentest", {}).get("tool_timeout", 300) @@ -104,7 +104,10 @@ async def execute(self, target: str, **kwargs) -> Dict[str, Any]: result = { "tool": self.tool_name, "target": target, - "command": " ".join(command), + # Commands may contain cookies, request bodies, passwords, or API + # tokens. Never let the executable form escape into memory, + # reports, logs, or AI prompts. + "command": safe_command, "exit_code": process.returncode, "duration": duration, "raw_output": output, diff --git a/tools/nmap.py b/tools/nmap.py index bb35b2a..f359b43 100644 --- a/tools/nmap.py +++ b/tools/nmap.py @@ -88,8 +88,20 @@ def parse_output(self, output: str) -> Dict[str, Any]: if not output.strip(): return results + # Nmap's expected document is small and never needs DTDs or entities. + # Reject those constructs and cap the input before invoking ElementTree. + # This prevents entity-expansion and memory-exhaustion payloads without + # adding a second XML implementation to the runtime dependency graph. + if len(output.encode("utf-8")) > 10 * 1024 * 1024: + self.logger.warning("nmap: refusing XML output larger than 10 MiB") + return results + upper_output = output.upper() + if " str: + """Remove common credential forms from tool evidence before storage or AI use.""" + redacted = value + for pattern in _SENSITIVE_PATTERNS: + redacted = pattern.sub(r"\1", redacted) + return redacted diff --git a/utils/scope_validator.py b/utils/scope_validator.py index 3ab1ca6..68b061f 100644 --- a/utils/scope_validator.py +++ b/utils/scope_validator.py @@ -4,7 +4,6 @@ """ import ipaddress -import re import socket from typing import List, Set, Optional from pathlib import Path @@ -97,7 +96,7 @@ def _is_blacklisted(self, host: str) -> bool: # Try to parse as a literal IP address first ip = ipaddress.ip_address(host) # Block loopback / link-local / unspecified regardless of CIDR list - if ip.is_loopback or ip.is_link_local or ip.is_unspecified: + if not ip.is_global: return True for network in self.blacklist_networks: if ip in network: @@ -109,7 +108,7 @@ def _is_blacklisted(self, host: str) -> bool: # Not a literal IP — check well-known loopback/special names _BLOCKED_NAMES = { 'localhost', '127.0.0.1', '::1', - 'ip6-localhost', 'ip6-loopback', '0.0.0.0', + 'ip6-localhost', 'ip6-loopback', '0.0.0.0', # noqa: S104 - denylist value } if host.lower() in _BLOCKED_NAMES: return True @@ -129,8 +128,7 @@ def _is_blacklisted(self, host: str) -> bool: if resolved_ip in network: return True # Also block loopback / link-local / unspecified explicitly - if (resolved_ip.is_loopback or resolved_ip.is_link_local - or resolved_ip.is_unspecified): + if not resolved_ip.is_global: return True except ValueError: continue