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

ida_pro_mcp/ida_mcp/sync.py

@@ -0,0 +1,233 @@
+import logging
+import queue
+import functools
+import os
+import sys
+import time
+import idaapi
+import idc
+from .rpc import McpToolError
+from .zeromcp.jsonrpc import get_current_cancel_event, RequestCancelledError
+
+# ============================================================================
+# IDA Synchronization & Error Handling
+# ============================================================================
+
+ida_major, ida_minor = map(int, idaapi.get_kernel_version().split("."))
+
+
+class IDAError(McpToolError):
+    def __init__(self, message: str):
+        super().__init__(message)
+
+    @property
+    def message(self) -> str:
+        return self.args[0]
+
+
+class IDASyncError(Exception):
+    pass
+
+
+class CancelledError(RequestCancelledError):
+    """Raised when a request is cancelled via notifications/cancelled."""
+
+    pass
+
+
+logger = logging.getLogger(__name__)
+_TOOL_TIMEOUT_ENV = "IDA_MCP_TOOL_TIMEOUT_SEC"
+_DEFAULT_TOOL_TIMEOUT_SEC = 15.0
+
+
+def _get_tool_timeout_seconds() -> float:
+    value = os.getenv(_TOOL_TIMEOUT_ENV, "").strip()
+    if value == "":
+        return _DEFAULT_TOOL_TIMEOUT_SEC
+    try:
+        return float(value)
+    except ValueError:
+        return _DEFAULT_TOOL_TIMEOUT_SEC
+
+
+# ============================================================================
+# Queue Object Pool - Reduces GC pressure and allocation overhead
+# ============================================================================
+
+
+class QueuePool:
+    """Object pool for Queue instances to reduce allocation overhead.
+
+    Reuses Queue objects instead of creating new ones for each IDA sync call.
+    Thread-safe implementation with automatic cleanup.
+    """
+
+    def __init__(self, max_size: int = 50):
+        self._pool: queue.Queue = queue.Queue(maxsize=max_size)
+        self._max_size = max_size
+
+    def acquire(self) -> queue.Queue:
+        """Get a Queue from the pool or create a new one."""
+        try:
+            return self._pool.get_nowait()
+        except queue.Empty:
+            return queue.Queue()
+
+    def release(self, q: queue.Queue):
+        """Return a Queue to the pool after clearing it."""
+        # Clear the queue before returning to pool
+        while not q.empty():
+            try:
+                q.get_nowait()
+            except queue.Empty:
+                break
+        # Try to return to pool
+        try:
+            self._pool.put_nowait(q)
+        except queue.Full:
+            # Pool full, let GC handle it
+            pass
+
+
+_queue_pool = QueuePool()
+
+
+call_stack = queue.LifoQueue()
+
+
+def _sync_wrapper(ff):
+    """Call a function ff with a specific IDA safety_mode."""
+
+    res_container = _queue_pool.acquire()
+
+    def runned():
+        if not call_stack.empty():
+            last_func_name = call_stack.get()
+            error_str = f"Call stack is not empty while calling the function {ff.__name__} from {last_func_name}"
+            raise IDASyncError(error_str)
+
+        call_stack.put((ff.__name__))
+        try:
+            res_container.put(ff())
+        except Exception as x:
+            res_container.put(x)
+        finally:
+            call_stack.get()
+
+    idaapi.execute_sync(runned, idaapi.MFF_WRITE)
+    res = res_container.get()
+    _queue_pool.release(res_container)
+    if isinstance(res, Exception):
+        raise res
+    return res
+
+
+def _normalize_timeout(value: object) -> float | None:
+    if value is None:
+        return None
+    try:
+        return float(value)
+    except (TypeError, ValueError):
+        return None
+
+
+def sync_wrapper(ff, timeout_override: float | None = None):
+    """Wrapper to enable batch mode during IDA synchronization."""
+    # Capture cancel event from thread-local before execute_sync
+    cancel_event = get_current_cancel_event()
+
+    def _run_with_batch(inner_ff):
+        def _wrapped():
+            old_batch = idc.batch(1)
+            try:
+                return inner_ff()
+            finally:
+                idc.batch(old_batch)
+
+        _wrapped.__name__ = inner_ff.__name__
+        return _wrapped
+
+    timeout = timeout_override
+    if timeout is None:
+        timeout = _get_tool_timeout_seconds()
+    if timeout > 0 or cancel_event is not None:
+
+        def timed_ff():
+            # Calculate deadline when execution starts on IDA main thread,
+            # not when the request was queued (avoids stale deadlines)
+            deadline = time.monotonic() + timeout if timeout > 0 else None
+
+            def profilefunc(frame, event, arg):
+                # Check cancellation first (higher priority)
+                if cancel_event is not None and cancel_event.is_set():
+                    raise CancelledError("Request was cancelled")
+                if deadline is not None and time.monotonic() >= deadline:
+                    raise IDASyncError(f"Tool timed out after {timeout:.2f}s")
+
+            old_profile = sys.getprofile()
+            sys.setprofile(profilefunc)
+            try:
+                return ff()
+            finally:
+                sys.setprofile(old_profile)
+
+        timed_ff.__name__ = ff.__name__
+        return _sync_wrapper(_run_with_batch(timed_ff))
+    return _sync_wrapper(_run_with_batch(ff))
+
+
+def idasync(f):
+    """Run the function on the IDA main thread in write mode.
+
+    This is the unified decorator for all IDA synchronization.
+    Previously there were separate @idaread and @idawrite decorators,
+    but since read-only operations in IDA might actually require write
+    access (e.g., decompilation), we now use a single decorator.
+    """
+
+    @functools.wraps(f)
+    def wrapper(*args, **kwargs):
+        ff = functools.partial(f, *args, **kwargs)
+        ff.__name__ = f.__name__
+        timeout_override = _normalize_timeout(
+            getattr(f, "__ida_mcp_timeout_sec__", None)
+        )
+        return sync_wrapper(ff, timeout_override)
+
+    return wrapper
+
+
+def tool_timeout(seconds: float):
+    """Decorator to override per-tool timeout (seconds).
+
+    IMPORTANT: Must be applied BEFORE @idasync (i.e., listed AFTER it)
+    so the attribute exists when it captures the function in closure.
+
+    Correct order:
+        @tool
+        @idasync
+        @tool_timeout(90.0)  # innermost
+        def my_func(...):
+    """
+
+    def decorator(func):
+        setattr(func, "__ida_mcp_timeout_sec__", seconds)
+        return func
+
+    return decorator
+
+
+def is_window_active():
+    """Returns whether IDA is currently active."""
+    # Source: https://github.com/OALabs/hexcopy-ida/blob/8b0b2a3021d7dc9010c01821b65a80c47d491b61/hexcopy.py#L30
+    using_pyside6 = (ida_major > 9) or (ida_major == 9 and ida_minor >= 2)
+
+    if using_pyside6:
+        from PySide6 import QtWidgets
+    else:
+        from PyQt5 import QtWidgets
+
+    app = QtWidgets.QApplication.instance()
+    if app is None:
+        return False
+    return app.activeWindow() is not None

ida_pro_mcp/ida_mcp/tests/__init__.py

@@ -0,0 +1,14 @@
+"""IDA Pro MCP Test Package.
+
+This package contains test modules for each API module.
+Tests are registered via the @test decorator from the framework module.
+"""
+
+# Import all test modules to register tests when the package is imported
+from . import test_api_core
+from . import test_api_analysis
+from . import test_api_memory
+from . import test_api_modify
+from . import test_api_types
+from . import test_api_stack
+from . import test_api_resources

ida_pro_mcp/ida_mcp/tests/test_api_analysis.py

@@ -0,0 +1,336 @@
+"""Tests for api_analysis API functions."""
+
+# Import test framework from parent
+from ..framework import (
+    test,
+    assert_has_keys,
+    assert_is_list,
+    get_any_function,
+    get_n_functions,
+    get_data_address,
+    get_unmapped_address,
+    get_functions_with_calls,
+    get_functions_with_callers,
+)
+
+# Import functions under test
+from ..api_analysis import (
+    decompile,
+    disasm,
+    xrefs_to,
+    xrefs_to_field,
+    callees,
+    find_bytes,
+    basic_blocks,
+    find,
+    export_funcs,
+    callgraph,
+)
+
+# Import sync module for IDAError
+
+
+# ============================================================================
+# Tests for decompile
+# ============================================================================
+
+
+@test()
+def test_decompile_valid_function():
+    """decompile returns code for a valid function"""
+    fn_addr = get_any_function()
+    if not fn_addr:
+        return
+
+    result = decompile(fn_addr)
+    assert_is_list(result, min_length=1)
+    # Should have code or error
+    r = result[0]
+    assert_has_keys(r, "addr")
+    # Either has code or has an error
+    assert r.get("code") is not None or r.get("error") is not None
+
+
+@test()
+def test_decompile_invalid_address():
+    """decompile handles invalid address gracefully"""
+    result = decompile(get_unmapped_address())
+    assert_is_list(result, min_length=1)
+    # Should have an error
+    assert result[0].get("error") is not None or result[0].get("code") is None
+
+
+@test()
+def test_decompile_batch():
+    """decompile can handle multiple addresses"""
+    addrs = get_n_functions(3)
+    if len(addrs) < 2:
+        return
+
+    result = decompile(addrs)
+    assert len(result) == len(addrs)
+
+
+# ============================================================================
+# Tests for disasm
+# ============================================================================
+
+
+@test()
+def test_disasm_valid_function():
+    """disasm returns assembly for a valid function"""
+    fn_addr = get_any_function()
+    if not fn_addr:
+        return
+
+    result = disasm(fn_addr)
+    assert_is_list(result, min_length=1)
+    r = result[0]
+    assert_has_keys(r, "addr")
+    # Should have asm output or error
+    assert r.get("asm") is not None or r.get("error") is not None
+
+
+@test()
+def test_disasm_pagination():
+    """disasm respects count parameter"""
+    fn_addr = get_any_function()
+    if not fn_addr:
+        return
+
+    result = disasm(fn_addr, count=10)
+    assert_is_list(result, min_length=1)
+
+
+@test()
+def test_disasm_unmapped_address():
+    """disasm handles unmapped address"""
+    result = disasm(get_unmapped_address())
+    assert_is_list(result, min_length=1)
+    # Should have error or empty asm
+    r = result[0]
+    assert r.get("error") is not None or r.get("asm") == "" or r.get("asm") is None
+
+
+@test()
+def test_disasm_data_segment():
+    """disasm handles data segment addresses"""
+    data_addr = get_data_address()
+    if not data_addr:
+        return
+
+    result = disasm(data_addr)
+    assert_is_list(result, min_length=1)
+
+
+# ============================================================================
+# Tests for xrefs_to
+# ============================================================================
+
+
+@test()
+def test_xrefs_to():
+    """xrefs_to returns cross-references for a function"""
+    fn_addrs = get_functions_with_callers()
+    if not fn_addrs:
+        # Fallback to any function
+        fn_addr = get_any_function()
+        if not fn_addr:
+            return
+    else:
+        fn_addr = fn_addrs[0]
+
+    result = xrefs_to(fn_addr)
+    assert_is_list(result, min_length=1)
+    r = result[0]
+    assert_has_keys(r, "addr", "xrefs", "error")
+
+
+@test()
+def test_xrefs_to_invalid():
+    """xrefs_to handles invalid address"""
+    result = xrefs_to(get_unmapped_address())
+    assert_is_list(result, min_length=1)
+    # Should return empty xrefs or error
+    r = result[0]
+    assert_has_keys(r, "addr")
+
+
+# ============================================================================
+# Tests for xrefs_to_field
+# ============================================================================
+
+
+@test()
+def test_xrefs_to_field_nonexistent_struct():
+    """xrefs_to_field handles non-existent struct"""
+    result = xrefs_to_field({"struct": "NonExistentStruct", "field": "nonexistent"})
+    assert_is_list(result, min_length=1)
+    r = result[0]
+    assert r.get("error") is not None
+
+
+@test()
+def test_xrefs_to_field_batch():
+    """xrefs_to_field handles batch queries"""
+    result = xrefs_to_field(
+        [
+            {"struct": "Struct1", "field": "field1"},
+            {"struct": "Struct2", "field": "field2"},
+        ]
+    )
+    assert_is_list(result, min_length=2)
+
+
+# ============================================================================
+# Tests for callees
+# ============================================================================
+
+
+@test()
+def test_callees():
+    """callees returns functions called by a function"""
+    fn_addrs = get_functions_with_calls()
+    if not fn_addrs:
+        fn_addr = get_any_function()
+        if not fn_addr:
+            return
+    else:
+        fn_addr = fn_addrs[0]
+
+    result = callees(fn_addr)
+    assert_is_list(result, min_length=1)
+    r = result[0]
+    assert_has_keys(r, "addr", "callees", "error")
+
+
+@test()
+def test_callees_multiple():
+    """callees handles multiple addresses"""
+    addrs = get_n_functions(3)
+    if len(addrs) < 2:
+        return
+
+    result = callees(addrs)
+    assert len(result) == len(addrs)
+
+
+@test()
+def test_callees_invalid_address():
+    """callees handles invalid address"""
+    result = callees(get_unmapped_address())
+    assert_is_list(result, min_length=1)
+    r = result[0]
+    assert_has_keys(r, "addr")
+
+
+# ============================================================================
+# Tests for find_bytes
+# ============================================================================
+
+
+@test()
+def test_find_bytes():
+    """find_bytes can search for byte patterns"""
+    # Search for common bytes that should exist
+    result = find_bytes("00 00")
+    assert_is_list(result, min_length=1)
+    r = result[0]
+    assert_has_keys(r, "query", "matches", "error")
+
+
+# ============================================================================
+# Tests for basic_blocks
+# ============================================================================
+
+
+@test()
+def test_basic_blocks():
+    """basic_blocks returns blocks for a function"""
+    fn_addr = get_any_function()
+    if not fn_addr:
+        return
+
+    result = basic_blocks(fn_addr)
+    assert_is_list(result, min_length=1)
+    r = result[0]
+    assert_has_keys(r, "addr", "blocks", "error")
+
+
+# ============================================================================
+# Tests for find
+# ============================================================================
+
+
+@test()
+def test_find_string():
+    """find can search for strings"""
+    # Most binaries have some strings
+    result = find("string", query="*")
+    assert_is_list(result, min_length=1)
+    r = result[0]
+    assert_has_keys(r, "query", "matches", "error")
+
+
+@test()
+def test_find_invalid_type():
+    """find handles invalid search type"""
+    result = find("invalid_type", query="test")
+    assert_is_list(result, min_length=1)
+    r = result[0]
+    # Should have error for invalid type
+    assert r.get("error") is not None
+
+
+# ============================================================================
+# Tests for export_funcs
+# ============================================================================
+
+
+@test()
+def test_export_funcs_json():
+    """export_funcs returns JSON format"""
+    fn_addr = get_any_function()
+    if not fn_addr:
+        return
+
+    result = export_funcs(fn_addr, fmt="json")
+    assert_is_list(result, min_length=1)
+    r = result[0]
+    assert_has_keys(r, "addr")
+
+
+@test()
+def test_export_funcs_c_header():
+    """export_funcs returns C header format"""
+    fn_addr = get_any_function()
+    if not fn_addr:
+        return
+
+    result = export_funcs(fn_addr, fmt="c_header")
+    assert_is_list(result, min_length=1)
+
+
+@test()
+def test_export_funcs_invalid_address():
+    """export_funcs handles invalid address"""
+    result = export_funcs(get_unmapped_address())
+    assert_is_list(result, min_length=1)
+
+
+# ============================================================================
+# Tests for callgraph
+# ============================================================================
+
+
+@test()
+def test_callgraph():
+    """callgraph returns call graph data"""
+    fn_addr = get_any_function()
+    if not fn_addr:
+        return
+
+    result = callgraph(fn_addr)
+    assert_is_list(result, min_length=1)
+    r = result[0]
+    assert_has_keys(r, "addr")