mirror of
https://github.com/GH05TCREW/pentestagent.git
synced 2026-03-07 22:33:38 +00:00
511 lines
20 KiB
Python
511 lines
20 KiB
Python
"""MCP server connection manager for PentestAgent.
|
|
|
|
Uses standard MCP configuration format:
|
|
{
|
|
"mcpServers": {
|
|
"server-name": {
|
|
"command": "npx",
|
|
"args": ["-y", "package-name"],
|
|
"env": {"VAR": "value"}
|
|
}
|
|
}
|
|
}
|
|
"""
|
|
|
|
import asyncio
|
|
import json
|
|
import os
|
|
import atexit
|
|
import signal
|
|
from dataclasses import dataclass, field
|
|
from pathlib import Path
|
|
from typing import Any, Dict, List, Optional
|
|
|
|
from .tools import create_mcp_tool
|
|
from .transport import MCPTransport, StdioTransport
|
|
|
|
|
|
@dataclass
|
|
class MCPServerConfig:
|
|
"""Configuration for an MCP server."""
|
|
|
|
name: str
|
|
command: str
|
|
args: List[str] = field(default_factory=list)
|
|
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
|
|
class MCPServer:
|
|
"""Represents a connected MCP server."""
|
|
|
|
name: str
|
|
config: MCPServerConfig
|
|
transport: MCPTransport
|
|
tools: List[dict] = field(default_factory=list)
|
|
connected: bool = False
|
|
# Lock for serializing all communication with this server
|
|
# Prevents message ID collisions and transport interleaving
|
|
_lock: asyncio.Lock = field(default_factory=asyncio.Lock)
|
|
|
|
async def disconnect(self):
|
|
"""Disconnect from the server."""
|
|
if self.connected:
|
|
await self.transport.disconnect()
|
|
self.connected = False
|
|
|
|
|
|
class MCPManager:
|
|
"""Manages MCP server connections and exposes tools to agents."""
|
|
|
|
DEFAULT_CONFIG_PATHS = [
|
|
Path.cwd() / "mcp_servers.json",
|
|
Path.cwd() / "mcp.json",
|
|
Path(__file__).parent / "mcp_servers.json",
|
|
Path.home() / ".pentestagent" / "mcp_servers.json",
|
|
]
|
|
|
|
def __init__(self, config_path: Optional[Path] = None):
|
|
self.config_path = config_path or self._find_config()
|
|
self.servers: Dict[str, MCPServer] = {}
|
|
# Track adapters we auto-started so we can stop them later
|
|
self._started_adapters: Dict[str, object] = {}
|
|
self._message_id = 0
|
|
# Ensure we attempt to clean up vendored servers on process exit
|
|
try:
|
|
atexit.register(self._atexit_cleanup)
|
|
except Exception:
|
|
pass
|
|
|
|
def _find_config(self) -> Path:
|
|
for path in self.DEFAULT_CONFIG_PATHS:
|
|
if path.exists():
|
|
return path
|
|
return self.DEFAULT_CONFIG_PATHS[0]
|
|
|
|
def _get_next_id(self) -> int:
|
|
self._message_id += 1
|
|
return self._message_id
|
|
|
|
def _load_config(self) -> Dict[str, MCPServerConfig]:
|
|
if not self.config_path.exists():
|
|
return {}
|
|
try:
|
|
raw = json.loads(self.config_path.read_text(encoding="utf-8"))
|
|
servers = {}
|
|
mcp_servers = raw.get("mcpServers", {})
|
|
for name, config in mcp_servers.items():
|
|
if not config.get("command"):
|
|
continue
|
|
servers[name] = MCPServerConfig(
|
|
name=name,
|
|
command=config["command"],
|
|
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", ""),
|
|
)
|
|
# Allow override via environment variables for vendored MCP servers.
|
|
# Per-adapter overrides supported:
|
|
# - Hexstrike: LAUNCH_HEXTRIKE or LAUNCH_HEXSTRIKE
|
|
# - Metasploit: LAUNCH_METASPLOIT_MCP
|
|
# If set to a truthy value (1,true,y), force-enable auto-start for matching vendored server.
|
|
# If set to a falsy value (0,false,n), force-disable auto-start for matching vendored server.
|
|
def _apply_launch_override(env_names, match_fn):
|
|
launch_env = None
|
|
for e in env_names:
|
|
launch_env = os.environ.get(e)
|
|
if launch_env is not None:
|
|
break
|
|
if launch_env is None:
|
|
return
|
|
v = str(launch_env).strip().lower()
|
|
enable = v in ("1", "true", "yes", "y")
|
|
disable = v in ("0", "false", "no", "n")
|
|
|
|
for name, cfg in servers.items():
|
|
try:
|
|
if not match_fn(name, cfg):
|
|
continue
|
|
if enable:
|
|
cfg.start_on_launch = True
|
|
elif disable:
|
|
cfg.start_on_launch = False
|
|
except Exception:
|
|
continue
|
|
|
|
# Hexstrike override
|
|
_apply_launch_override(["LAUNCH_HEXTRIKE", "LAUNCH_HEXSTRIKE"],
|
|
lambda name, cfg: (
|
|
(name or "").lower().find("hexstrike") != -1
|
|
or (cfg.command and "third_party/hexstrike" in str(cfg.command))
|
|
or any("third_party/hexstrike" in str(a) for a in (cfg.args or []))
|
|
))
|
|
|
|
# Metasploit override
|
|
_apply_launch_override(["LAUNCH_METASPLOIT_MCP"],
|
|
lambda name, cfg: (
|
|
(name or "").lower().find("metasploit") != -1
|
|
or (cfg.command and "third_party/MetasploitMCP" in str(cfg.command))
|
|
or any("third_party/MetasploitMCP" in str(a) for a in (cfg.args or []))
|
|
))
|
|
|
|
return servers
|
|
except json.JSONDecodeError as e:
|
|
print(f"[MCP] Error loading config: {e}")
|
|
return {}
|
|
|
|
def _save_config(self, servers: Dict[str, MCPServerConfig]):
|
|
config = {"mcpServers": {}}
|
|
for name, server in servers.items():
|
|
server_config = {"command": server.command, "args": server.args}
|
|
if server.env:
|
|
server_config["env"] = server.env
|
|
if server.description:
|
|
server_config["description"] = server.description
|
|
if not server.enabled:
|
|
server_config["enabled"] = False
|
|
config["mcpServers"][name] = server_config
|
|
self.config_path.parent.mkdir(parents=True, exist_ok=True)
|
|
self.config_path.write_text(json.dumps(config, indent=2), encoding="utf-8")
|
|
|
|
def _atexit_cleanup(self):
|
|
"""Synchronous atexit cleanup that attempts to stop adapters and disconnect servers."""
|
|
try:
|
|
# Try to run async shutdown; if an event loop is already running this may fail,
|
|
# but it's best-effort to avoid orphaned vendored servers.
|
|
asyncio.run(self._stop_started_adapters_and_disconnect())
|
|
except Exception:
|
|
# Last-ditch: attempt to stop adapters synchronously.
|
|
# If the adapter exposes a blocking `stop()` call, call it. Otherwise, try
|
|
# to kill the underlying process by PID to avoid asyncio subprocess
|
|
# destructors running after the loop is closed.
|
|
for adapter in list(self._started_adapters.values()):
|
|
try:
|
|
# Prefer adapter-provided synchronous stop hook
|
|
stop_sync = getattr(adapter, "stop_sync", None)
|
|
if stop_sync:
|
|
try:
|
|
stop_sync()
|
|
continue
|
|
except Exception:
|
|
pass
|
|
|
|
# Fallback: try blocking stop() if present
|
|
stop = getattr(adapter, "stop", None)
|
|
if stop and not asyncio.iscoroutinefunction(stop):
|
|
try:
|
|
stop()
|
|
continue
|
|
except Exception:
|
|
pass
|
|
|
|
# Final fallback: kill underlying PID if available
|
|
pid = None
|
|
proc = getattr(adapter, "_process", None)
|
|
if proc is not None:
|
|
pid = getattr(proc, "pid", None)
|
|
if pid:
|
|
try:
|
|
os.kill(pid, signal.SIGTERM)
|
|
except Exception:
|
|
try:
|
|
os.kill(pid, signal.SIGKILL)
|
|
except Exception:
|
|
pass
|
|
except Exception:
|
|
pass
|
|
|
|
async def _stop_started_adapters_and_disconnect(self) -> None:
|
|
# Stop any adapters we started
|
|
for name, adapter in list(self._started_adapters.items()):
|
|
try:
|
|
stop = getattr(adapter, "stop", None)
|
|
if stop:
|
|
if asyncio.iscoroutinefunction(stop):
|
|
await stop()
|
|
else:
|
|
# run blocking stop in executor
|
|
loop = asyncio.get_running_loop()
|
|
await loop.run_in_executor(None, stop)
|
|
except Exception:
|
|
pass
|
|
self._started_adapters.clear()
|
|
|
|
# Disconnect any active MCP server connections
|
|
try:
|
|
await self.disconnect_all()
|
|
except Exception:
|
|
pass
|
|
|
|
def add_server(
|
|
self,
|
|
name: str,
|
|
command: str,
|
|
args: List[str] = None,
|
|
env: Dict[str, str] = None,
|
|
description: str = "",
|
|
):
|
|
servers = self._load_config()
|
|
servers[name] = MCPServerConfig(
|
|
name=name,
|
|
command=command,
|
|
args=args or [],
|
|
env=env or {},
|
|
description=description,
|
|
)
|
|
self._save_config(servers)
|
|
print(f"[MCP] Added server: {name}")
|
|
|
|
def remove_server(self, name: str) -> bool:
|
|
servers = self._load_config()
|
|
if name in servers:
|
|
del servers[name]
|
|
self._save_config(servers)
|
|
return True
|
|
return False
|
|
|
|
def set_enabled(self, name: str, enabled: bool) -> bool:
|
|
"""Enable or disable a configured MCP server in the config file.
|
|
|
|
Returns True if the server existed and was updated, False otherwise.
|
|
"""
|
|
servers = self._load_config()
|
|
if name not in servers:
|
|
return False
|
|
servers[name].enabled = bool(enabled)
|
|
self._save_config(servers)
|
|
return True
|
|
|
|
def list_configured_servers(self) -> List[dict]:
|
|
servers = self._load_config()
|
|
return [
|
|
{
|
|
"name": n,
|
|
"command": s.command,
|
|
"args": s.args,
|
|
"env": s.env,
|
|
"enabled": s.enabled,
|
|
"description": s.description,
|
|
"connected": n in self.servers and self.servers[n].connected,
|
|
}
|
|
for n, s in servers.items()
|
|
]
|
|
|
|
async def connect_all(self) -> List[Any]:
|
|
servers_config = self._load_config()
|
|
# Respect explicit LAUNCH_* env overrides for vendored MCP servers.
|
|
# If set to a falsy value (0/false/no/n) we will skip connecting to matching vendored servers.
|
|
launch_hex_env = os.environ.get("LAUNCH_HEXTRIKE") or os.environ.get("LAUNCH_HEXSTRIKE")
|
|
launch_hex_disabled = False
|
|
if launch_hex_env is not None:
|
|
v = str(launch_hex_env).strip().lower()
|
|
if v in ("0", "false", "no", "n"):
|
|
launch_hex_disabled = True
|
|
|
|
launch_msf_env = os.environ.get("LAUNCH_METASPLOIT_MCP")
|
|
launch_msf_disabled = False
|
|
if launch_msf_env is not None:
|
|
v = str(launch_msf_env).strip().lower()
|
|
if v in ("0", "false", "no", "n"):
|
|
launch_msf_disabled = True
|
|
|
|
all_tools = []
|
|
for name, config in servers_config.items():
|
|
if not config.enabled:
|
|
continue
|
|
# If the user explicitly disabled launching HexStrike, skip hexstrike entries entirely
|
|
lowered = name.lower() if name else ""
|
|
is_hex = (
|
|
"hexstrike" in lowered
|
|
or (config.command and "third_party/hexstrike" in str(config.command))
|
|
or any("third_party/hexstrike" in str(a) for a in (config.args or []))
|
|
)
|
|
if launch_disabled and is_hex:
|
|
print(f"[MCP] Skipping auto-connection for {name} due to LAUNCH_HEXTRIKE={launch_env}")
|
|
continue
|
|
# Optionally auto-start vendored servers (e.g., HexStrike subtree or MetasploitMCP)
|
|
if getattr(config, "start_on_launch", False):
|
|
try:
|
|
args_joined = " ".join(config.args or [])
|
|
cmd_str = config.command or ""
|
|
|
|
# Hexstrike auto-start
|
|
if "third_party/hexstrike" in args_joined or (cmd_str and "third_party/hexstrike" in cmd_str):
|
|
if not launch_hex_disabled:
|
|
try:
|
|
from .hexstrike_adapter import HexstrikeAdapter
|
|
|
|
adapter = HexstrikeAdapter()
|
|
started = await adapter.start()
|
|
if started:
|
|
try:
|
|
self._started_adapters[name] = adapter
|
|
except Exception:
|
|
pass
|
|
print(f"[MCP] Auto-started vendored server for {name}")
|
|
except Exception as e:
|
|
print(f"[MCP] Failed to auto-start vendored server {name}: {e}")
|
|
else:
|
|
print(f"[MCP] Skipping auto-start for {name} due to LAUNCH_HEXTRIKE override")
|
|
|
|
# Metasploit auto-start
|
|
if "third_party/MetasploitMCP" in args_joined or (cmd_str and "third_party/MetasploitMCP" in cmd_str) or (name and "metasploit" in name.lower()):
|
|
if not launch_msf_disabled:
|
|
try:
|
|
from .metasploit_adapter import MetasploitAdapter
|
|
|
|
adapter = MetasploitAdapter()
|
|
started = await adapter.start()
|
|
if started:
|
|
try:
|
|
self._started_adapters[name] = adapter
|
|
except Exception:
|
|
pass
|
|
print(f"[MCP] Auto-started vendored server for {name}")
|
|
except Exception as e:
|
|
print(f"[MCP] Failed to auto-start vendored server {name}: {e}")
|
|
else:
|
|
print(f"[MCP] Skipping auto-start for {name} due to LAUNCH_METASPLOIT_MCP override")
|
|
except Exception:
|
|
pass
|
|
server = await self._connect_server(config)
|
|
if server:
|
|
self.servers[name] = server
|
|
for tool_def in server.tools:
|
|
tool = create_mcp_tool(tool_def, server, self)
|
|
all_tools.append(tool)
|
|
print(f"[MCP] Connected to {name} with {len(server.tools)} tools")
|
|
return all_tools
|
|
|
|
async def connect_server(self, name: str) -> Optional[MCPServer]:
|
|
servers_config = self._load_config()
|
|
if name not in servers_config:
|
|
return None
|
|
config = servers_config[name]
|
|
server = await self._connect_server(config)
|
|
if server:
|
|
self.servers[name] = server
|
|
return server
|
|
|
|
async def _connect_server(self, config: MCPServerConfig) -> Optional[MCPServer]:
|
|
transport = None
|
|
try:
|
|
env = {**os.environ, **config.env}
|
|
transport = StdioTransport(
|
|
command=config.command, args=config.args, env=env
|
|
)
|
|
await transport.connect()
|
|
|
|
await transport.send(
|
|
{
|
|
"jsonrpc": "2.0",
|
|
"method": "initialize",
|
|
"params": {
|
|
"protocolVersion": "2024-11-05",
|
|
"capabilities": {},
|
|
"clientInfo": {"name": "pentestagent", "version": "0.2.0"},
|
|
},
|
|
"id": self._get_next_id(),
|
|
}
|
|
)
|
|
await transport.send(
|
|
{"jsonrpc": "2.0", "method": "notifications/initialized"}
|
|
)
|
|
|
|
tools_response = await transport.send(
|
|
{"jsonrpc": "2.0", "method": "tools/list", "id": self._get_next_id()}
|
|
)
|
|
tools = tools_response.get("result", {}).get("tools", [])
|
|
|
|
return MCPServer(
|
|
name=config.name,
|
|
config=config,
|
|
transport=transport,
|
|
tools=tools,
|
|
connected=True,
|
|
)
|
|
except Exception as e:
|
|
# Clean up transport on failure
|
|
if transport:
|
|
try:
|
|
await transport.disconnect()
|
|
except Exception:
|
|
pass
|
|
print(f"[MCP] Failed to connect to {config.name}: {e}")
|
|
return None
|
|
|
|
async def call_tool(self, server_name: str, tool_name: str, arguments: dict) -> Any:
|
|
server = self.servers.get(server_name)
|
|
if not server or not server.connected:
|
|
raise ValueError(f"Server '{server_name}' not connected")
|
|
|
|
# Serialize all communication with this server to prevent:
|
|
# - Message ID collisions
|
|
# - Transport write interleaving
|
|
# - Response routing issues
|
|
async with server._lock:
|
|
# Use 5 minute timeout for tool calls (scans can take a while)
|
|
response = await server.transport.send(
|
|
{
|
|
"jsonrpc": "2.0",
|
|
"method": "tools/call",
|
|
"params": {"name": tool_name, "arguments": arguments},
|
|
"id": self._get_next_id(),
|
|
},
|
|
timeout=300.0,
|
|
)
|
|
if "error" in response:
|
|
raise RuntimeError(f"MCP error: {response['error'].get('message')}")
|
|
return response.get("result", {}).get("content", [])
|
|
|
|
async def disconnect_server(self, name: str):
|
|
server = self.servers.get(name)
|
|
if server:
|
|
await server.disconnect()
|
|
del self.servers[name]
|
|
# If we started an adapter for this server, stop it as well
|
|
adapter = self._started_adapters.pop(name, None)
|
|
if adapter:
|
|
try:
|
|
stop = getattr(adapter, "stop", None)
|
|
if stop:
|
|
if asyncio.iscoroutinefunction(stop):
|
|
await stop()
|
|
else:
|
|
loop = asyncio.get_running_loop()
|
|
await loop.run_in_executor(None, stop)
|
|
except Exception:
|
|
pass
|
|
|
|
async def disconnect_all(self):
|
|
for server in list(self.servers.values()):
|
|
await server.disconnect()
|
|
self.servers.clear()
|
|
|
|
async def reconnect_all(self) -> List[Any]:
|
|
"""Disconnect all servers and reconnect them.
|
|
|
|
Useful after cancellation leaves servers in a bad state.
|
|
"""
|
|
# Disconnect all
|
|
await self.disconnect_all()
|
|
|
|
# Reconnect all configured servers
|
|
return await self.connect_all()
|
|
|
|
def get_server(self, name: str) -> Optional[MCPServer]:
|
|
return self.servers.get(name)
|
|
|
|
def get_all_servers(self) -> List[MCPServer]:
|
|
return list(self.servers.values())
|
|
|
|
def is_connected(self, name: str) -> bool:
|
|
server = self.servers.get(name)
|
|
return server is not None and server.connected
|