ida-pro-mcp-xjoker 1.0.1__py3-none-any.whl

ida_pro_mcp/ida_mcp/port_utils.py

@@ -0,0 +1,104 @@
+"""端口冲突自动递增工具
+
+多个 IDA Pro 实例同时加载 MCP 插件时，提供端口自动递增重试功能。
+"""
+
+import errno
+import logging
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from .zeromcp.mcp import McpServer
+
+logger = logging.getLogger(__name__)
+
+# 端口已占用的 errno 值（跨平台）
+# macOS: 48, Linux: 98, Windows: 10048
+_ADDR_IN_USE_ERRNOS = {errno.EADDRINUSE, 48, 98, 10048}
+
+DEFAULT_MAX_RETRIES = 10
+
+
+def try_serve_with_port_retry(
+    mcp_server: "McpServer",
+    host: str,
+    base_port: int,
+    *,
+    max_retries: int = DEFAULT_MAX_RETRIES,
+    request_handler=None,
+) -> tuple[int, list[int]]:
+    """尝试启动 MCP 服务器，端口冲突时自动递增重试。
+
+    Args:
+        mcp_server: McpServer 实例
+        host: 绑定地址
+        base_port: 起始端口
+        max_retries: 最大重试次数（包含首次尝试共 max_retries 次）
+        request_handler: HTTP 请求处理器类
+
+    Returns:
+        (actual_port, failed_ports) 元组：
+        - actual_port: 实际绑定成功的端口
+        - failed_ports: 绑定失败的端口列表（为空表示首次即成功）
+
+    Raises:
+        OSError: 所有端口尝试均失败，或遇到非端口冲突的错误
+    """
+    kwargs = {}
+    if request_handler is not None:
+        kwargs["request_handler"] = request_handler
+
+    failed_ports: list[int] = []
+    last_error: OSError | None = None
+
+    for i in range(max_retries):
+        port = base_port + i
+        if port > 65535:
+            break
+        try:
+            mcp_server.serve(host, port, **kwargs)
+            return port, failed_ports
+        except OSError as e:
+            if e.errno in _ADDR_IN_USE_ERRNOS:
+                logger.debug(f"Port {port} in use, trying next")
+                failed_ports.append(port)
+                last_error = e
+            else:
+                # 非端口冲突错误，立即抛出
+                raise
+
+    # 所有端口都失败了
+    assert last_error is not None
+    raise last_error
+
+
+def format_port_exhausted_message(
+    host: str, base_port: int, failed_ports: list[int]
+) -> str:
+    """生成端口耗尽时的用户友好提示信息。
+
+    Args:
+        host: 绑定地址
+        base_port: 起始端口
+        failed_ports: 尝试失败的端口列表
+
+    Returns:
+        格式化的错误提示字符串
+    """
+    if not failed_ports:
+        return f"[MCP] Error: Could not bind to port on {host}"
+
+    first = failed_ports[0]
+    last = failed_ports[-1]
+    ports_str = ", ".join(str(p) for p in failed_ports)
+
+    return (
+        f"[MCP] Error: Could not find an available port.\n"
+        f"  Tried ports: {ports_str}\n"
+        f"  All ports in range {first}-{last} are in use.\n"
+        f"\n"
+        f"  To manually set a different port, run in IDA Python console:\n"
+        f"    from ida_mcp.http import set_server_config\n"
+        f"    set_server_config({{'host': '{host}', 'port': 15000}})\n"
+        f"  Then toggle: Edit -> Plugins -> MCP Server"
+    )

ida_pro_mcp/ida_mcp/rpc.py

@@ -0,0 +1,187 @@
+import json
+import os
+from typing import Any, Optional
+from .zeromcp import McpRpcRegistry, McpServer, McpToolError, McpHttpRequestHandler
+
+MCP_UNSAFE: set[str] = set()
+MCP_EXTENSIONS: dict[str, set[str]] = {}  # group -> set of function names
+MCP_SERVER = McpServer("ida-pro-mcp", extensions=MCP_EXTENSIONS)
+
+# ============================================================================
+# Output Size Limiting
+# ============================================================================
+
+OUTPUT_LIMIT_MAX_CHARS = 50000
+OUTPUT_CACHE_MAX_SIZE = 100
+_output_cache: dict[str, Any] = {}
+_download_base_url: str = os.environ.get("IDA_MCP_URL", "http://127.0.0.1:13337")
+
+
+def set_download_base_url(url: str) -> None:
+    global _download_base_url
+    _download_base_url = url.rstrip("/")
+
+
+def get_download_base_url() -> str:
+    return _download_base_url
+
+
+def _generate_output_id() -> str:
+    import uuid
+
+    return str(uuid.uuid4())
+
+
+OUTPUT_LIMIT_PREVIEW_ITEMS = 10
+OUTPUT_LIMIT_PREVIEW_STR_LEN = 1000
+
+
+def _truncate_value(value: Any, depth: int = 0) -> Any:
+    if depth > 5:
+        return value
+
+    if isinstance(value, str) and len(value) > OUTPUT_LIMIT_PREVIEW_STR_LEN:
+        return value[:OUTPUT_LIMIT_PREVIEW_STR_LEN] + f"... [{len(value)} chars total]"
+
+    if isinstance(value, list):
+        truncated_list = [
+            _truncate_value(item, depth + 1)
+            for item in value[:OUTPUT_LIMIT_PREVIEW_ITEMS]
+        ]
+        if len(value) > OUTPUT_LIMIT_PREVIEW_ITEMS:
+            truncated_list.append(
+                {
+                    "_truncated": f"... and {len(value) - OUTPUT_LIMIT_PREVIEW_ITEMS} more items"
+                }
+            )
+        return truncated_list
+
+    if isinstance(value, dict):
+        return {k: _truncate_value(v, depth + 1) for k, v in value.items()}
+
+    return value
+
+
+def _add_download_info(result: Any, output_id: str, total_chars: int) -> Any:
+    download_url = f"{_download_base_url}/output/{output_id}.json"
+    info = {
+        "_output_truncated": True,
+        "_total_chars": total_chars,
+        "_output_id": output_id,
+        "_download_url": download_url,
+        "_download_hint": f"Output truncated. Run: curl -o .ida-mcp/{output_id}.json {download_url}",
+    }
+
+    if isinstance(result, dict):
+        return {**result, **info}
+
+    if isinstance(result, list) and result:
+        result = list(result)
+        if isinstance(result[0], dict):
+            result[0] = {**result[0], **info}
+        else:
+            result.insert(0, info)
+        return result
+
+    return {"_preview": result, **info}
+
+
+def get_cached_output(output_id: str) -> Optional[Any]:
+    return _output_cache.get(output_id)
+
+
+def _cache_output(output_id: str, data: Any) -> None:
+    if len(_output_cache) >= OUTPUT_CACHE_MAX_SIZE:
+        oldest_key = next(iter(_output_cache))
+        del _output_cache[oldest_key]
+    _output_cache[output_id] = data
+
+
+def _install_tools_call_patch() -> None:
+    original = MCP_SERVER.registry.methods["tools/call"]
+
+    def patched(
+        name: str, arguments: Optional[dict] = None, _meta: Optional[dict] = None
+    ) -> dict:
+        response = original(name, arguments, _meta)
+
+        if response.get("isError"):
+            return response
+
+        structured = response.get("structuredContent")
+        if structured is None:
+            return response
+
+        serialized = json.dumps(structured)
+        if len(serialized) <= OUTPUT_LIMIT_MAX_CHARS:
+            return response
+
+        output_id = _generate_output_id()
+        _cache_output(output_id, structured)
+
+        preview = _truncate_value(structured)
+        preview = _add_download_info(preview, output_id, len(serialized))
+
+        return {
+            "structuredContent": preview,
+            "content": response.get("content", []),
+            "isError": False,
+        }
+
+    MCP_SERVER.registry.methods["tools/call"] = patched
+
+
+# Install the output limiting patch
+_install_tools_call_patch()
+
+
+# ============================================================================
+# Decorators
+# ============================================================================
+
+
+def tool(func):
+    return MCP_SERVER.tool(func)
+
+
+def resource(uri):
+    return MCP_SERVER.resource(uri)
+
+
+def unsafe(func):
+    MCP_UNSAFE.add(func.__name__)
+    return func
+
+
+def ext(group: str):
+    """Mark a tool as belonging to an extension group.
+
+    Tools in extension groups are hidden by default. Enable via ?ext=group query param.
+    Example: @ext("dbg") marks debugger tools that require ?ext=dbg to be visible.
+    """
+
+    def decorator(func):
+        if group not in MCP_EXTENSIONS:
+            MCP_EXTENSIONS[group] = set()
+        MCP_EXTENSIONS[group].add(func.__name__)
+        return func
+
+    return decorator
+
+
+__all__ = [
+    "McpRpcRegistry",
+    "McpServer",
+    "McpToolError",
+    "McpHttpRequestHandler",
+    "MCP_SERVER",
+    "MCP_UNSAFE",
+    "MCP_EXTENSIONS",
+    "tool",
+    "unsafe",
+    "ext",
+    "resource",
+    "get_cached_output",
+    "set_download_base_url",
+    "get_download_base_url",
+]

ida_pro_mcp/ida_mcp/server_manager.py

@@ -0,0 +1,339 @@
+"""IDA MCP Server Manager
+
+Manages multiple MCP server instances with support for:
+- Starting/stopping individual servers
+- Status monitoring
+- Configuration-based server management
+"""
+
+import logging
+import threading
+from dataclasses import dataclass, field
+from typing import Optional, Callable, TYPE_CHECKING
+
+from .config import ServerInstanceConfig, McpConfig, get_config, save_config
+from .auth import AuthMiddleware
+from .port_utils import try_serve_with_port_retry
+
+if TYPE_CHECKING:
+    from .zeromcp.mcp import McpServer
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class ServerInstance:
+    """Represents a running server instance"""
+
+    config: ServerInstanceConfig
+    server: Optional["McpServer"] = None
+    auth: AuthMiddleware = field(default_factory=AuthMiddleware)
+    error: Optional[str] = None
+    _actual_port: Optional[int] = field(default=None, repr=False)
+
+    @property
+    def actual_port(self) -> int:
+        """实际绑定的端口（可能因端口冲突自动递增）。"""
+        if self._actual_port is not None:
+            return self._actual_port
+        return self.config.port
+
+    @property
+    def is_running(self) -> bool:
+        return self.server is not None and self.server._running
+
+    @property
+    def status(self) -> str:
+        if self.error:
+            return f"error: {self.error}"
+        if self.is_running:
+            return "running"
+        return "stopped"
+
+    def to_status_dict(self) -> dict:
+        result = {
+            "instance_id": self.config.instance_id,
+            "host": self.config.host,
+            "port": self.config.port,
+            "status": self.status,
+            "auth_enabled": self.config.auth_enabled,
+            "address": self.config.address,
+        }
+        if self._actual_port is not None and self._actual_port != self.config.port:
+            result["actual_port"] = self._actual_port
+            result["address"] = f"{self.config.host}:{self._actual_port}"
+        return result
+
+
+class ServerManager:
+    """Manages multiple MCP server instances.
+
+    Thread-safe manager for starting, stopping, and monitoring
+    multiple MCP server instances.
+    """
+
+    def __init__(self, create_server_func: Optional[Callable[[], "McpServer"]] = None):
+        """Initialize the server manager.
+
+        Args:
+            create_server_func: Factory function to create new McpServer instances.
+                               If None, uses default McpServer constructor.
+        """
+        self._instances: dict[str, ServerInstance] = {}
+        self._lock = threading.RLock()
+        self._create_server = create_server_func
+        self._request_handler_class: Optional[type] = None
+
+    def set_server_factory(self, factory: Callable[[], "McpServer"]) -> None:
+        """Set the server factory function."""
+        self._create_server = factory
+
+    def set_request_handler(self, handler_class: type) -> None:
+        """Set the request handler class for new servers."""
+        self._request_handler_class = handler_class
+
+    def _create_server_instance(self, config: ServerInstanceConfig) -> ServerInstance:
+        """Create a new server instance from configuration."""
+        instance = ServerInstance(
+            config=config,
+            auth=AuthMiddleware(config.api_key, config.auth_enabled),
+        )
+        return instance
+
+    def add_server(self, config: ServerInstanceConfig) -> str:
+        """Add a new server configuration.
+
+        Args:
+            config: Server instance configuration
+
+        Returns:
+            The instance ID
+
+        Raises:
+            ValueError: If instance ID already exists
+        """
+        with self._lock:
+            if config.instance_id in self._instances:
+                raise ValueError(f"Server instance '{config.instance_id}' already exists")
+
+            instance = self._create_server_instance(config)
+            self._instances[config.instance_id] = instance
+            return config.instance_id
+
+    def remove_server(self, instance_id: str) -> bool:
+        """Remove a server instance (stops it first if running).
+
+        Args:
+            instance_id: The instance ID to remove
+
+        Returns:
+            True if removed, False if not found
+        """
+        with self._lock:
+            if instance_id not in self._instances:
+                return False
+
+            # Stop if running
+            self.stop_server(instance_id)
+
+            del self._instances[instance_id]
+            return True
+
+    def start_server(self, instance_id: str) -> bool:
+        """Start a server instance.
+
+        Args:
+            instance_id: The instance ID to start
+
+        Returns:
+            True if started successfully, False otherwise
+        """
+        with self._lock:
+            instance = self._instances.get(instance_id)
+            if not instance:
+                logger.error(f"Server instance '{instance_id}' not found")
+                return False
+
+            if instance.is_running:
+                logger.info(f"Server '{instance_id}' is already running")
+                return True
+
+            if not self._create_server:
+                logger.error("No server factory configured")
+                return False
+
+            try:
+                # Create new server
+                server = self._create_server()
+
+                # Configure authentication
+                if hasattr(server, '_auth'):
+                    server._auth = instance.auth
+
+                # Start the server with port retry
+                kwargs = {}
+                if self._request_handler_class:
+                    kwargs["request_handler"] = self._request_handler_class
+
+                actual_port, failed_ports = try_serve_with_port_retry(
+                    server,
+                    instance.config.host,
+                    instance.config.port,
+                    **kwargs,
+                )
+
+                instance.server = server
+                instance.error = None
+                instance._actual_port = actual_port if actual_port != instance.config.port else None
+
+                if failed_ports:
+                    logger.info(
+                        f"Server '{instance_id}': port {instance.config.port} in use, "
+                        f"auto-selected port {actual_port}"
+                    )
+                logger.info(
+                    f"Started server '{instance_id}' on {instance.config.host}:{actual_port}"
+                )
+                return True
+
+            except OSError as e:
+                if e.errno in (48, 98, 10048):  # Address already in use
+                    instance.error = (
+                        f"All ports {instance.config.port}-{instance.config.port + 9} are in use"
+                    )
+                else:
+                    instance.error = str(e)
+                logger.error(f"Failed to start server '{instance_id}': {instance.error}")
+                return False
+            except Exception as e:
+                instance.error = str(e)
+                logger.error(f"Failed to start server '{instance_id}': {e}")
+                return False
+
+    def stop_server(self, instance_id: str) -> bool:
+        """Stop a server instance.
+
+        Args:
+            instance_id: The instance ID to stop
+
+        Returns:
+            True if stopped, False if not found or not running
+        """
+        with self._lock:
+            instance = self._instances.get(instance_id)
+            if not instance:
+                return False
+
+            if not instance.is_running:
+                return True
+
+            try:
+                instance.server.stop()
+                instance.server = None
+                instance.error = None
+                logger.info(f"Stopped server '{instance_id}'")
+                return True
+            except Exception as e:
+                instance.error = str(e)
+                logger.error(f"Error stopping server '{instance_id}': {e}")
+                return False
+
+    def restart_server(self, instance_id: str) -> bool:
+        """Restart a server instance.
+
+        Args:
+            instance_id: The instance ID to restart
+
+        Returns:
+            True if restarted successfully
+        """
+        self.stop_server(instance_id)
+        return self.start_server(instance_id)
+
+    def get_status(self) -> dict[str, dict]:
+        """Get status of all server instances.
+
+        Returns:
+            Dictionary mapping instance IDs to their status
+        """
+        with self._lock:
+            return {
+                instance_id: instance.to_status_dict()
+                for instance_id, instance in self._instances.items()
+            }
+
+    def get_instance(self, instance_id: str) -> Optional[ServerInstance]:
+        """Get a server instance by ID.
+
+        Args:
+            instance_id: The instance ID
+
+        Returns:
+            The server instance or None
+        """
+        with self._lock:
+            return self._instances.get(instance_id)
+
+    def load_from_config(self, config: Optional[McpConfig] = None) -> None:
+        """Load server instances from configuration.
+
+        Args:
+            config: Configuration to load from (uses global config if None)
+        """
+        if config is None:
+            config = get_config()
+
+        with self._lock:
+            # Stop and remove existing instances
+            for instance_id in list(self._instances.keys()):
+                self.remove_server(instance_id)
+
+            # Add instances from config
+            for server_config in config.servers:
+                if server_config.enabled:
+                    self.add_server(server_config)
+
+    def start_auto_servers(self) -> int:
+        """Start all servers marked with auto_start.
+
+        Returns:
+            Number of servers started
+        """
+        count = 0
+        with self._lock:
+            for instance_id, instance in self._instances.items():
+                if instance.config.auto_start and not instance.is_running:
+                    if self.start_server(instance_id):
+                        count += 1
+        return count
+
+    def stop_all(self) -> None:
+        """Stop all running server instances."""
+        with self._lock:
+            for instance_id in list(self._instances.keys()):
+                self.stop_server(instance_id)
+
+    def __len__(self) -> int:
+        return len(self._instances)
+
+    def __contains__(self, instance_id: str) -> bool:
+        return instance_id in self._instances
+
+
+# Global server manager instance
+_manager: Optional[ServerManager] = None
+
+
+def get_server_manager() -> ServerManager:
+    """Get the global server manager instance."""
+    global _manager
+    if _manager is None:
+        _manager = ServerManager()
+    return _manager
+
+
+__all__ = [
+    "ServerInstance",
+    "ServerManager",
+    "get_server_manager",
+]