pentestagent/ghostcrew/mcp/discovery.py

"""MCP tool discovery for GhostCrew."""

import json
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, List, Optional


@dataclass
class DiscoveredServer:
    """A discovered MCP server."""

    name: str
    description: str
    type: str  # "stdio" or "sse"
    command: Optional[str] = None
    args: Optional[List[str]] = None
    url: Optional[str] = None
    tools: List[dict] = None

    def __post_init__(self):
        if self.tools is None:
            self.tools = []


class MCPDiscovery:
    """Discovers available MCP servers and tools."""

    # Known MCP servers for security tools
    KNOWN_SERVERS = [
        {
            "name": "nmap",
            "description": "Network scanning and host discovery",
            "type": "stdio",
            "command": "npx",
            "args": ["-y", "@modelcontextprotocol/server-nmap"],
        },
        {
            "name": "filesystem",
            "description": "File system operations",
            "type": "stdio",
            "command": "npx",
            "args": ["-y", "@modelcontextprotocol/server-filesystem"],
        },
        {
            "name": "fetch",
            "description": "HTTP requests and web fetching",
            "type": "stdio",
            "command": "npx",
            "args": ["-y", "@modelcontextprotocol/server-fetch"],
        },
    ]

    def __init__(self, config_path: Path = Path("mcp.json")):
        """
        Initialize MCP discovery.

        Args:
            config_path: Path to the MCP configuration file
        """
        self.config_path = config_path

    def discover_local(self) -> List[DiscoveredServer]:
        """
        Discover locally installed MCP servers.

        Returns:
            List of discovered servers
        """
        discovered = []

        # Check for npm global packages
        # Check for Python packages
        # This is a simplified implementation

        for server_info in self.KNOWN_SERVERS:
            discovered.append(DiscoveredServer(**server_info))

        return discovered

    def load_from_config(self) -> List[Dict[str, Any]]:
        """
        Load server configurations from file.

        Returns:
            List of server configurations
        """
        if not self.config_path.exists():
            return []

        try:
            config = json.loads(self.config_path.read_text(encoding="utf-8"))
            return config.get("servers", [])
        except json.JSONDecodeError:
            return []

    def add_server_to_config(
        self,
        name: str,
        server_type: str,
        command: Optional[str] = None,
        args: Optional[List[str]] = None,
        url: Optional[str] = None,
        env: Optional[Dict[str, str]] = None,
    ) -> bool:
        """
        Add a server to the configuration file.

        Args:
            name: Server name
            server_type: "stdio" or "sse"
            command: Command for stdio servers
            args: Arguments for stdio servers
            url: URL for SSE servers
            env: Environment variables

        Returns:
            True if added successfully
        """
        # Load existing config
        if self.config_path.exists():
            try:
                config = json.loads(self.config_path.read_text(encoding="utf-8"))
            except json.JSONDecodeError:
                config = {"servers": []}
        else:
            config = {"servers": []}

        # Check if server already exists
        for existing in config["servers"]:
            if existing.get("name") == name:
                return False

        # Build server config
        server_config = {"name": name, "type": server_type, "enabled": True}

        if server_type == "stdio":
            server_config["command"] = command
            server_config["args"] = args or []
            if env:
                server_config["env"] = env
        elif server_type == "sse":
            server_config["url"] = url

        config["servers"].append(server_config)

        # Save config
        self.config_path.write_text(json.dumps(config, indent=2), encoding="utf-8")

        return True

    def remove_server_from_config(self, name: str) -> bool:
        """
        Remove a server from the configuration file.

        Args:
            name: Server name to remove

        Returns:
            True if removed successfully
        """
        if not self.config_path.exists():
            return False

        try:
            config = json.loads(self.config_path.read_text(encoding="utf-8"))
        except json.JSONDecodeError:
            return False

        original_count = len(config.get("servers", []))
        config["servers"] = [
            s for s in config.get("servers", []) if s.get("name") != name
        ]

        if len(config["servers"]) == original_count:
            return False

        self.config_path.write_text(json.dumps(config, indent=2), encoding="utf-8")

        return True

    def generate_default_config(self) -> Dict[str, Any]:
        """
        Generate a default MCP configuration.

        Returns:
            Default configuration dictionary
        """
        return {
            "servers": [
                {
                    "name": "filesystem",
                    "type": "stdio",
                    "command": "npx",
                    "args": ["-y", "@modelcontextprotocol/server-filesystem", "/tmp"],
                    "enabled": True,
                }
            ]
        }

    def save_default_config(self):
        """Save the default configuration to file."""
        config = self.generate_default_config()
        self.config_path.write_text(json.dumps(config, indent=2), encoding="utf-8")