diff --git a/pentestagent/mcp/hexstrike_adapter.py b/pentestagent/mcp/hexstrike_adapter.py new file mode 100644 index 0000000..dc74ef5 --- /dev/null +++ b/pentestagent/mcp/hexstrike_adapter.py @@ -0,0 +1,177 @@ +"""Adapter to manage a vendored HexStrike MCP server. + +This adapter provides a simple programmatic API to start/stop the vendored +HexStrike server (expected under ``third_party/hexstrike``) and to perform a +health check before returning control to the caller. + +The adapter is intentionally lightweight (no Docker) and uses an async +subprocess so the server can run in the background while the TUI/runtime +operates. +""" + +import asyncio +import os +import shutil +import sys +from pathlib import Path +from typing import Optional + +try: + import aiohttp +except Exception: + aiohttp = None + + +LOOT_DIR = Path("loot/artifacts") +LOOT_DIR.mkdir(parents=True, exist_ok=True) +LOG_FILE = LOOT_DIR / "hexstrike.log" + + +class HexstrikeAdapter: + """Manage a vendored HexStrike server under `third_party/hexstrike`. + + Usage: + adapter = HexstrikeAdapter() + await adapter.start() + # ... use MCPManager to connect to the server + await adapter.stop() + """ + + def __init__( + self, + host: str = "127.0.0.1", + port: int = 8888, + python_cmd: str = "python3", + server_script: Optional[Path] = None, + cwd: Optional[Path] = None, + env: Optional[dict] = None, + ) -> None: + self.host = host + self.port = int(port) + self.python_cmd = python_cmd + self.server_script = ( + server_script + or Path("third_party/hexstrike/hexstrike_server.py") + ) + self.cwd = cwd or Path.cwd() + self.env = {**os.environ, **(env or {})} + + self._process: Optional[asyncio.subprocess.Process] = None + self._reader_task: Optional[asyncio.Task] = None + + def _build_command(self): + return [self.python_cmd, str(self.server_script), "--port", str(self.port)] + + async def start(self, background: bool = True, timeout: int = 30) -> bool: + """Start the vendored HexStrike server. + + Returns True if the server started and passed health check within + `timeout` seconds. + """ + if not self.server_script.exists(): + raise FileNotFoundError( + f"HexStrike server script not found at {self.server_script}." + ) + + if self._process and self._process.returncode is None: + return await self.health_check(timeout=1) + + cmd = self._build_command() + + # Resolve python command if possible + resolved = shutil.which(self.python_cmd) or self.python_cmd + + self._process = await asyncio.create_subprocess_exec( + resolved, + *cmd[1:], + cwd=str(self.cwd), + env=self.env, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.STDOUT, + ) + + # Start a background reader task to capture logs + loop = asyncio.get_running_loop() + self._reader_task = loop.create_task(self._capture_output()) + + # Wait for health check + try: + return await self.health_check(timeout=timeout) + except Exception: + return False + + async def _capture_output(self) -> None: + """Capture stdout/stderr from the server and append to the log file.""" + if not self._process or not self._process.stdout: + return + + try: + with LOG_FILE.open("ab") as fh: + while True: + line = await self._process.stdout.readline() + if not line: + break + # Prefix timestamps for easier debugging + fh.write(line) + fh.flush() + except asyncio.CancelledError: + pass + except Exception: + pass + + async def stop(self, timeout: int = 5) -> None: + """Stop the server process gracefully.""" + proc = self._process + if not proc: + return + + try: + proc.terminate() + await asyncio.wait_for(proc.wait(), timeout=timeout) + except asyncio.TimeoutError: + try: + proc.kill() + except Exception: + pass + except Exception: + pass + + self._process = None + + if self._reader_task and not self._reader_task.done(): + self._reader_task.cancel() + try: + await self._reader_task + except Exception: + pass + + async def health_check(self, timeout: int = 5) -> bool: + """Check the server health endpoint. Returns True if healthy.""" + url = f"http://{self.host}:{self.port}/health" + + if aiohttp: + try: + async with aiohttp.ClientSession() as session: + async with session.get(url, timeout=timeout) as resp: + return resp.status == 200 + except Exception: + return False + + # Fallback: synchronous urllib in thread + import urllib.request + + def _check(): + try: + with urllib.request.urlopen(url, timeout=timeout) as r: + return r.status == 200 + except Exception: + return False + + loop = asyncio.get_running_loop() + return await loop.run_in_executor(None, _check) + + def is_running(self) -> bool: + return self._process is not None and self._process.returncode is None + + +__all__ = ["HexstrikeAdapter"] diff --git a/pentestagent/mcp/manager.py b/pentestagent/mcp/manager.py index c7b96c1..01e5bb7 100644 --- a/pentestagent/mcp/manager.py +++ b/pentestagent/mcp/manager.py @@ -33,6 +33,8 @@ class MCPServerConfig: env: Dict[str, str] = field(default_factory=dict) enabled: bool = True description: str = "" + # Whether to auto-start this server when `connect_all()` is called. + start_on_launch: bool = False @dataclass @@ -96,6 +98,7 @@ class MCPManager: args=config.get("args", []), env=config.get("env", {}), enabled=config.get("enabled", True), + start_on_launch=config.get("start_on_launch", False), description=config.get("description", ""), ) return servers @@ -165,6 +168,24 @@ class MCPManager: for name, config in servers_config.items(): if not config.enabled: continue + # Optionally auto-start vendored servers (e.g., HexStrike subtree) + if getattr(config, "start_on_launch", False): + try: + # Attempt to auto-start vendored HexStrike if present + if "third_party/hexstrike" in " ".join(config.args) or ( + config.command and "third_party/hexstrike" in config.command + ): + try: + from .hexstrike_adapter import HexstrikeAdapter + + adapter = HexstrikeAdapter() + started = await adapter.start() + if started: + print(f"[MCP] Auto-started vendored server for {name}") + except Exception as e: + print(f"[MCP] Failed to auto-start vendored server {name}: {e}") + except Exception: + pass server = await self._connect_server(config) if server: self.servers[name] = server diff --git a/pentestagent/mcp/mcp_servers.json b/pentestagent/mcp/mcp_servers.json index da39e4f..8354adf 100644 --- a/pentestagent/mcp/mcp_servers.json +++ b/pentestagent/mcp/mcp_servers.json @@ -1,3 +1,16 @@ { - "mcpServers": {} + "mcpServers": { + "hexstrike-local": { + "command": "python3", + "args": [ + "third_party/hexstrike/hexstrike_server.py", + "--port", + "8888" + ], + "description": "HexStrike AI (vendored) - local server", + "timeout": 300, + "disabled": false, + "start_on_launch": true + } + } } diff --git a/scripts/add_hexstrike_subtree.sh b/scripts/add_hexstrike_subtree.sh new file mode 100755 index 0000000..1982f9b --- /dev/null +++ b/scripts/add_hexstrike_subtree.sh @@ -0,0 +1,16 @@ +#!/usr/bin/env bash +# Helper script to vendor HexStrike into this repo using git subtree. +# Run from repository root. + +set -euo pipefail + +REPO_URL="https://github.com/0x4m4/hexstrike-ai.git" +PREFIX="third_party/hexstrike" +BRANCH="main" + +echo "This will add HexStrike as a git subtree under ${PREFIX}." +echo "If you already have a subtree, use 'git subtree pull' instead.\n" + +git subtree add --prefix="${PREFIX}" "${REPO_URL}" "${BRANCH}" --squash + +echo "HexStrike subtree added under ${PREFIX}."