Improve security: redact sensitive evidence, tighten scope checks, secure persistence, and test coverage

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:

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)

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",

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 <redacted> example.com"
+
 
 # ---------------------------------------------------------------------------
 # Process kill on timeout

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

tests/test_nmap_tool.py

@@ -134,6 +134,15 @@ def test_malformed_xml_returns_empty(self, nmap):
         result = nmap.parse_output("<not valid xml <<<")
         assert result["open_ports"] == []
 
+    def test_dtd_is_rejected(self, nmap):
+        xml = '<!DOCTYPE nmaprun [<!ENTITY x "boom">]><nmaprun>&x;</nmaprun>'
+        result = nmap.parse_output(xml)
+        assert result["open_ports"] == []
+
+    def test_oversized_xml_is_rejected(self, nmap):
+        result = nmap.parse_output("<nmaprun>" + (" " * (10 * 1024 * 1024)) + "</nmaprun>")
+        assert result["open_ports"] == []
+
     def test_no_os_returns_none(self, nmap):
         xml = """<nmaprun><host><ports>
           <port portid="80" protocol="tcp">

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 "<redacted>" 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")

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
 

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,

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 "<!DOCTYPE" in upper_output or "<!ENTITY" in upper_output:
+            self.logger.warning("nmap: refusing XML containing a DTD or entity declaration")
+            return results
+
         try:
-            root = ET.fromstring(output)
+            root = ET.fromstring(output)  # noqa: S314 - guarded above
         except ET.ParseError as exc:
             self.logger.warning(f"nmap: failed to parse XML output: {exc}")
             return results

utils/redaction.py

@@ -0,0 +1,18 @@
+"""Central redaction helpers for evidence leaving a tool process."""
+
+import re
+
+
+_SENSITIVE_PATTERNS = (
+    re.compile(r"(?i)(authorization\s*:\s*(?:bearer|basic)\s+)[^\s,;]+"),
+    re.compile(r"(?i)((?:api[_-]?key|api[_-]?token|password|passwd|secret)\s*[=:]\s*)[^\s,;&]+"),
+    re.compile(r"(?i)((?:cookie|set-cookie)\s*:\s*)[^\r\n]+"),
+)
+
+
+def redact_sensitive_text(value: str) -> 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>", redacted)
+    return redacted

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