mcghidra/src/ghydramcp/server.py

"""GhydraMCP Server - FastMCP server composing all mixins.

This module creates and configures the FastMCP server by composing
all domain-specific mixins into a single MCP server.
"""

import signal
import sys
import threading
import time
from pathlib import Path
from typing import Optional

from fastmcp import FastMCP

from .config import GhydraConfig, get_config, set_config
from .mixins import (
    AnalysisMixin,
    BookmarksMixin,
    CursorsMixin,
    DataMixin,
    DataTypesMixin,
    DockerMixin,
    FunctionsMixin,
    InstancesMixin,
    MemoryMixin,
    NamespacesMixin,
    SegmentsMixin,
    StructsMixin,
    SymbolsMixin,
    VariablesMixin,
    XrefsMixin,
)


def create_server(
    name: str = "GhydraMCP",
    config: Optional[GhydraConfig] = None,
) -> FastMCP:
    """Create and configure the GhydraMCP server.

    Args:
        name: Server name
        config: Optional configuration override

    Returns:
        Configured FastMCP server instance
    """
    if config:
        set_config(config)

    # Create the FastMCP server
    mcp = FastMCP(name)

    # Instantiate all mixins
    instances_mixin = InstancesMixin()
    functions_mixin = FunctionsMixin()
    data_mixin = DataMixin()
    structs_mixin = StructsMixin()
    analysis_mixin = AnalysisMixin()
    memory_mixin = MemoryMixin()
    xrefs_mixin = XrefsMixin()
    cursors_mixin = CursorsMixin()
    docker_mixin = DockerMixin()
    symbols_mixin = SymbolsMixin()
    segments_mixin = SegmentsMixin()
    variables_mixin = VariablesMixin()
    namespaces_mixin = NamespacesMixin()
    bookmarks_mixin = BookmarksMixin()
    datatypes_mixin = DataTypesMixin()

    # Register all mixins with the server
    # Each mixin registers its tools, resources, and prompts
    instances_mixin.register_all(mcp)
    functions_mixin.register_all(mcp)
    data_mixin.register_all(mcp)
    structs_mixin.register_all(mcp)
    analysis_mixin.register_all(mcp)
    memory_mixin.register_all(mcp)
    xrefs_mixin.register_all(mcp)
    cursors_mixin.register_all(mcp)
    docker_mixin.register_all(mcp)
    symbols_mixin.register_all(mcp)
    segments_mixin.register_all(mcp)
    variables_mixin.register_all(mcp)
    namespaces_mixin.register_all(mcp)
    bookmarks_mixin.register_all(mcp)
    datatypes_mixin.register_all(mcp)

    # Optional feedback collection
    cfg = get_config()
    if cfg.feedback_enabled:
        try:
            from fastmcp_feedback import add_feedback_tools

            db_path = Path(cfg.feedback_db_path)
            db_path.parent.mkdir(parents=True, exist_ok=True)
            add_feedback_tools(mcp, database_url=f"sqlite:///{db_path}")
        except ImportError:
            pass  # fastmcp-feedback not installed — skip silently

    return mcp


def _periodic_discovery(interval: int = 30):
    """Background thread for periodic instance discovery.

    Uses a short timeout per port so a full scan completes quickly
    even when most ports are unreachable.

    Args:
        interval: Seconds between discovery attempts
    """
    import requests as _requests

    from .mixins.base import GhydraMixinBase

    config = get_config()

    while True:
        time.sleep(interval)
        try:
            # Quick scan — use discovery_timeout (0.5s), NOT request_timeout (30s)
            for port in config.quick_discovery_range:
                try:
                    url = f"http://{config.ghidra_host}:{port}/"
                    resp = _requests.get(
                        url,
                        timeout=config.discovery_timeout,
                        headers={"Accept": "application/json"},
                    )
                    if resp.ok:
                        response = resp.json()
                        if response.get("success", False):
                            with GhydraMixinBase._instances_lock:
                                if port not in GhydraMixinBase._instances:
                                    GhydraMixinBase._instances[port] = {
                                        "url": url.rstrip("/"),
                                        "project": response.get("project", ""),
                                        "file": response.get("file", ""),
                                        "discovered_at": time.time(),
                                    }
                except Exception:
                    pass
        except Exception:
            pass


def _handle_sigint(signum, frame):
    """Handle SIGINT gracefully."""
    print("\nShutting down GhydraMCP...", file=sys.stderr)
    sys.exit(0)


def main():
    """Main entry point for the GhydraMCP server."""
    import shutil

    try:
        from importlib.metadata import version
        package_version = version("ghydramcp")
    except Exception:
        package_version = "2025.12.1"

    print(f"🔬 GhydraMCP v{package_version}", file=sys.stderr)
    print("   AI-assisted reverse engineering bridge for Ghidra", file=sys.stderr)

    # Check Docker availability
    docker_available = shutil.which("docker") is not None
    if docker_available:
        print("   🐳 Docker available (use docker_* tools for container management)", file=sys.stderr)
    else:
        print("   ⚠ Docker not found (container management disabled)", file=sys.stderr)

    config = get_config()

    if config.feedback_enabled:
        print(f"   📋 Feedback collection: {config.feedback_db_path}", file=sys.stderr)

    # Create and configure the server
    mcp = create_server()

    # Initial instance discovery
    print(f"   Discovering Ghidra instances on {config.ghidra_host}...", file=sys.stderr)

    from .core.http_client import safe_get
    from .mixins.base import GhydraMixinBase

    found = 0
    for port in config.quick_discovery_range:
        try:
            response = safe_get(port, "")
            if response.get("success", False):
                GhydraMixinBase._instances[port] = {
                    "url": f"http://{config.ghidra_host}:{port}",
                    "project": response.get("project", ""),
                    "file": response.get("file", ""),
                    "discovered_at": time.time(),
                }
                found += 1
                print(f"   ✓ Found instance on port {port}", file=sys.stderr)
        except Exception:
            pass

    if found == 0:
        print("   ⚠ No Ghidra instances found (they can be discovered later)", file=sys.stderr)
    else:
        print(f"   Found {found} Ghidra instance(s)", file=sys.stderr)

    # Start background discovery thread
    discovery_thread = threading.Thread(
        target=_periodic_discovery,
        daemon=True,
        name="GhydraMCP-Discovery",
    )
    discovery_thread.start()

    # Set up signal handler
    signal.signal(signal.SIGINT, _handle_sigint)

    print("   Starting MCP server...", file=sys.stderr)

    # Run the server
    mcp.run(transport="stdio")


if __name__ == "__main__":
    main()