pentestagent/pentestagent/mcp/manager.py

"""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 variable: LAUNCH_HEXTRIKE or LAUNCH_HEXSTRIKE
            # If set to a truthy value (1,true,y), force-enable auto-start for vendored HexStrike.
            # If set to a falsy value (0,false,n), force-disable auto-start for vendored HexStrike.
            launch_env = os.environ.get("LAUNCH_HEXTRIKE") or os.environ.get("LAUNCH_HEXSTRIKE")
            if launch_env is not None:
                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():
                    lowered = name.lower() if name else ""
                    is_hex = (
                        "hexstrike" in lowered
                        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 []))
                    )
                    if not is_hex:
                        continue
                    if enable:
                        cfg.start_on_launch = True
                    elif disable:
                        cfg.start_on_launch = False

            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_HEXTRIKE/LAUNCH_HEXSTRIKE env override.
        # If set to a falsy value (0/false/no/n) we will skip connecting to vendored HexStrike servers.
        launch_env = os.environ.get("LAUNCH_HEXTRIKE") or os.environ.get("LAUNCH_HEXSTRIKE")
        launch_disabled = False
        if launch_env is not None:
            v = str(launch_env).strip().lower()
            if v in ("0", "false", "no", "n"):
                launch_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)
            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:
                                # remember adapter so we can stop it later
                                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}")
                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