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

ida_pro_mcp/ida_mcp/cache.py

@@ -0,0 +1,232 @@
+"""LRU Cache implementation for IDA Pro MCP Server.
+
+Provides thread-safe caching with TTL expiration for frequently accessed data
+such as function lookups, decompilation results, and cross-references.
+"""
+
+import threading
+import time
+from collections import OrderedDict
+from typing import Callable, Generic, Optional, TypeVar
+
+T = TypeVar("T")
+
+
+class LRUCache(Generic[T]):
+    """Thread-safe LRU cache with TTL expiration.
+
+    Features:
+    - Automatic eviction of least recently used items when max_size is reached
+    - Time-based expiration (TTL) for cached entries
+    - Thread-safe operations using RLock
+    - Optional key generation function for complex keys
+    """
+
+    def __init__(self, max_size: int = 1000, ttl_seconds: float = 300.0):
+        """Initialize the LRU cache.
+
+        Args:
+            max_size: Maximum number of entries to store
+            ttl_seconds: Time-to-live for entries in seconds (0 = no expiration)
+        """
+        self.max_size = max_size
+        self.ttl = ttl_seconds
+        self._cache: OrderedDict[str, tuple[T, float]] = OrderedDict()
+        self._lock = threading.RLock()
+        self._hits = 0
+        self._misses = 0
+
+    def get(self, key: str) -> Optional[T]:
+        """Get a value from the cache.
+
+        Args:
+            key: Cache key
+
+        Returns:
+            Cached value or None if not found/expired
+        """
+        with self._lock:
+            entry = self._cache.get(key)
+            if entry is None:
+                self._misses += 1
+                return None
+
+            value, timestamp = entry
+
+            # Check TTL expiration
+            if self.ttl > 0 and time.monotonic() - timestamp > self.ttl:
+                del self._cache[key]
+                self._misses += 1
+                return None
+
+            # Move to end (most recently used)
+            self._cache.move_to_end(key)
+            self._hits += 1
+            return value
+
+    def set(self, key: str, value: T) -> None:
+        """Set a value in the cache.
+
+        Args:
+            key: Cache key
+            value: Value to cache
+        """
+        with self._lock:
+            # Remove old entry if exists
+            if key in self._cache:
+                del self._cache[key]
+
+            # Evict LRU entries if at capacity
+            while len(self._cache) >= self.max_size:
+                self._cache.popitem(last=False)
+
+            # Add new entry
+            self._cache[key] = (value, time.monotonic())
+
+    def invalidate(self, key: Optional[str] = None) -> None:
+        """Invalidate cache entries.
+
+        Args:
+            key: Specific key to invalidate, or None to clear all
+        """
+        with self._lock:
+            if key is None:
+                self._cache.clear()
+            elif key in self._cache:
+                del self._cache[key]
+
+    def invalidate_prefix(self, prefix: str) -> int:
+        """Invalidate all entries with keys starting with prefix.
+
+        Args:
+            prefix: Key prefix to match
+
+        Returns:
+            Number of entries invalidated
+        """
+        with self._lock:
+            keys_to_remove = [k for k in self._cache if k.startswith(prefix)]
+            for key in keys_to_remove:
+                del self._cache[key]
+            return len(keys_to_remove)
+
+    def get_or_compute(
+        self,
+        key: str,
+        compute_fn: Callable[[], T],
+        ttl_override: Optional[float] = None,
+    ) -> T:
+        """Get from cache or compute and cache the value.
+
+        Args:
+            key: Cache key
+            compute_fn: Function to compute value if not cached
+            ttl_override: Optional TTL override for this entry
+
+        Returns:
+            Cached or computed value
+        """
+        # Try to get from cache first
+        value = self.get(key)
+        if value is not None:
+            return value
+
+        # Compute value
+        value = compute_fn()
+
+        # Store with optional TTL override
+        if ttl_override is not None:
+            original_ttl = self.ttl
+            self.ttl = ttl_override
+            self.set(key, value)
+            self.ttl = original_ttl
+        else:
+            self.set(key, value)
+
+        return value
+
+    def stats(self) -> dict:
+        """Get cache statistics.
+
+        Returns:
+            Dict with hits, misses, size, and hit_rate
+        """
+        with self._lock:
+            total = self._hits + self._misses
+            return {
+                "hits": self._hits,
+                "misses": self._misses,
+                "size": len(self._cache),
+                "max_size": self.max_size,
+                "hit_rate": self._hits / total if total > 0 else 0.0,
+            }
+
+    def clear_stats(self) -> None:
+        """Reset cache statistics."""
+        with self._lock:
+            self._hits = 0
+            self._misses = 0
+
+
+# ============================================================================
+# Global Cache Instances
+# ============================================================================
+
+# Function lookup cache: name/address -> function info
+# High TTL since function metadata rarely changes
+function_cache = LRUCache(max_size=5000, ttl_seconds=300.0)
+
+# Decompilation cache: address -> pseudocode
+# Moderate size, longer TTL since decompilation is expensive
+decompile_cache = LRUCache(max_size=200, ttl_seconds=600.0)
+
+# Cross-reference cache: address -> xrefs list
+# Moderate size and TTL
+xrefs_cache = LRUCache(max_size=2000, ttl_seconds=300.0)
+
+# String cache: regex pattern -> matches
+# Smaller size, moderate TTL
+string_cache = LRUCache(max_size=500, ttl_seconds=180.0)
+
+
+def invalidate_all_caches() -> None:
+    """Invalidate all global caches.
+
+    Call this when IDB changes significantly (e.g., new analysis, database reload).
+    """
+    function_cache.invalidate()
+    decompile_cache.invalidate()
+    xrefs_cache.invalidate()
+    string_cache.invalidate()
+
+
+def invalidate_function_caches(addr: Optional[int] = None) -> None:
+    """Invalidate caches related to a specific function.
+
+    Args:
+        addr: Function address to invalidate, or None for all functions
+    """
+    if addr is None:
+        function_cache.invalidate()
+        decompile_cache.invalidate()
+        xrefs_cache.invalidate()
+    else:
+        # Invalidate specific function entries
+        addr_hex = hex(addr)
+        function_cache.invalidate(addr_hex)
+        decompile_cache.invalidate(addr_hex)
+        xrefs_cache.invalidate_prefix(addr_hex)
+
+
+def get_cache_stats() -> dict:
+    """Get statistics for all caches.
+
+    Returns:
+        Dict with stats for each cache type
+    """
+    return {
+        "function_cache": function_cache.stats(),
+        "decompile_cache": decompile_cache.stats(),
+        "xrefs_cache": xrefs_cache.stats(),
+        "string_cache": string_cache.stats(),
+    }

ida_pro_mcp/ida_mcp/config.py

@@ -0,0 +1,228 @@
+"""IDA MCP Configuration System
+
+Provides configuration management for MCP server instances with support for:
+- Multiple server instances (different hosts/ports)
+- API Key authentication
+- Persistent configuration storage
+"""
+
+import os
+import logging
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Optional
+
+logger = logging.getLogger(__name__)
+
+# Configuration directory
+CONFIG_DIR = Path.home() / ".ida_mcp"
+CONFIG_FILE = CONFIG_DIR / "config.toml"
+
+
+@dataclass
+class ServerInstanceConfig:
+    """Configuration for a single MCP server instance"""
+
+    instance_id: str
+    host: str = "127.0.0.1"
+    port: int = 13337
+    enabled: bool = True
+    auth_enabled: bool = False
+    api_key: Optional[str] = None
+    auto_start: bool = False
+
+    def __post_init__(self):
+        # Resolve environment variable references in api_key
+        if self.api_key and self.api_key.startswith("${") and self.api_key.endswith("}"):
+            env_var = self.api_key[2:-1]
+            self.api_key = os.environ.get(env_var)
+
+    @property
+    def address(self) -> str:
+        return f"{self.host}:{self.port}"
+
+    def to_dict(self) -> dict:
+        return {
+            "instance_id": self.instance_id,
+            "host": self.host,
+            "port": self.port,
+            "enabled": self.enabled,
+            "auth_enabled": self.auth_enabled,
+            "api_key": self.api_key,
+            "auto_start": self.auto_start,
+        }
+
+    @classmethod
+    def from_dict(cls, data: dict) -> "ServerInstanceConfig":
+        return cls(
+            instance_id=data.get("instance_id", "default"),
+            host=data.get("host", "127.0.0.1"),
+            port=data.get("port", 13337),
+            enabled=data.get("enabled", True),
+            auth_enabled=data.get("auth_enabled", False),
+            api_key=data.get("api_key"),
+            auto_start=data.get("auto_start", False),
+        )
+
+
+@dataclass
+class McpConfig:
+    """Main configuration container"""
+
+    version: int = 1
+    servers: list[ServerInstanceConfig] = field(default_factory=list)
+    tool_timeout_sec: float = 15.0
+    debug: bool = False
+
+    def __post_init__(self):
+        # Ensure at least one default server exists
+        if not self.servers:
+            self.servers.append(ServerInstanceConfig(instance_id="local"))
+
+    def get_server(self, instance_id: str) -> Optional[ServerInstanceConfig]:
+        for server in self.servers:
+            if server.instance_id == instance_id:
+                return server
+        return None
+
+    def add_server(self, config: ServerInstanceConfig) -> bool:
+        if self.get_server(config.instance_id):
+            return False
+        self.servers.append(config)
+        return True
+
+    def remove_server(self, instance_id: str) -> bool:
+        for i, server in enumerate(self.servers):
+            if server.instance_id == instance_id:
+                self.servers.pop(i)
+                return True
+        return False
+
+    def to_dict(self) -> dict:
+        return {
+            "version": self.version,
+            "tool_timeout_sec": self.tool_timeout_sec,
+            "debug": self.debug,
+            "servers": [s.to_dict() for s in self.servers],
+        }
+
+    @classmethod
+    def from_dict(cls, data: dict) -> "McpConfig":
+        servers = [
+            ServerInstanceConfig.from_dict(s) for s in data.get("servers", [])
+        ]
+        return cls(
+            version=data.get("version", 1),
+            servers=servers,
+            tool_timeout_sec=data.get("tool_timeout_sec", 15.0),
+            debug=data.get("debug", False),
+        )
+
+
+def _ensure_config_dir() -> None:
+    """Create configuration directory if it doesn't exist"""
+    CONFIG_DIR.mkdir(parents=True, exist_ok=True)
+
+
+def load_config() -> McpConfig:
+    """Load configuration from TOML file"""
+    if not CONFIG_FILE.exists():
+        return McpConfig()
+
+    try:
+        # Python 3.11+ has tomllib in stdlib
+        import tomllib
+
+        with open(CONFIG_FILE, "rb") as f:
+            data = tomllib.load(f)
+        return McpConfig.from_dict(data)
+    except ImportError:
+        # Fallback: try toml package
+        try:
+            import toml
+
+            data = toml.load(CONFIG_FILE)
+            return McpConfig.from_dict(data)
+        except ImportError:
+            logger.warning("No TOML parser available, using defaults")
+            return McpConfig()
+    except Exception as e:
+        logger.error(f"Failed to load config: {e}")
+        return McpConfig()
+
+
+def save_config(config: McpConfig) -> bool:
+    """Save configuration to TOML file"""
+    _ensure_config_dir()
+
+    try:
+        # Try toml package for writing (tomllib is read-only)
+        try:
+            import toml
+
+            with open(CONFIG_FILE, "w") as f:
+                toml.dump(config.to_dict(), f)
+            return True
+        except ImportError:
+            # Manual TOML generation for simple config
+            return _write_toml_manual(config)
+    except Exception as e:
+        logger.error(f"Failed to save config: {e}")
+        return False
+
+
+def _write_toml_manual(config: McpConfig) -> bool:
+    """Write config as TOML without external dependencies"""
+    lines = [
+        f"version = {config.version}",
+        f"tool_timeout_sec = {config.tool_timeout_sec}",
+        f"debug = {'true' if config.debug else 'false'}",
+        "",
+    ]
+
+    for server in config.servers:
+        lines.append("[[servers]]")
+        lines.append(f'instance_id = "{server.instance_id}"')
+        lines.append(f'host = "{server.host}"')
+        lines.append(f"port = {server.port}")
+        lines.append(f"enabled = {'true' if server.enabled else 'false'}")
+        lines.append(f"auth_enabled = {'true' if server.auth_enabled else 'false'}")
+        if server.api_key:
+            lines.append(f'api_key = "{server.api_key}"')
+        lines.append(f"auto_start = {'true' if server.auto_start else 'false'}")
+        lines.append("")
+
+    with open(CONFIG_FILE, "w") as f:
+        f.write("\n".join(lines))
+    return True
+
+
+# Global config instance (lazy loaded)
+_config: Optional[McpConfig] = None
+
+
+def get_config() -> McpConfig:
+    """Get the global configuration instance"""
+    global _config
+    if _config is None:
+        _config = load_config()
+    return _config
+
+
+def reload_config() -> McpConfig:
+    """Force reload configuration from disk"""
+    global _config
+    _config = load_config()
+    return _config
+
+
+__all__ = [
+    "ServerInstanceConfig",
+    "McpConfig",
+    "load_config",
+    "save_config",
+    "get_config",
+    "reload_config",
+    "CONFIG_DIR",
+    "CONFIG_FILE",
+]