Add safe IDA item definition tool

README.md

@@ -6,7 +6,7 @@ Standalone MCP (Model Context Protocol) server plugin for **IDA Pro 9.x** that e
 
 ## Features
 
-- **41 MCP tools** covering binary analysis, decompilation, cross-references, symbol management, type system, navigation, patching, export, and more
+- **42 MCP tools** covering binary analysis, decompilation, cross-references, symbol management, type system, navigation, patching, export, and more
 - **6 consolidated tools** with `action`/`format`/`direction` parameters for comments, variables, types, xrefs, bookmarks, and code
 - **8 MCP resources** for browsable binary metadata (triage, functions, imports, exports, strings, info, segments, sections)
 - **7 guided prompts** for common reverse engineering workflows (function analysis, vulnerability identification, documentation, data flow tracing, function comparison, struct recovery, network protocol analysis)
@@ -155,7 +155,7 @@ Configure via environment variables with the `IDASSISTMCP_` prefix:
 | `get_segments` | Memory segments with permissions |
 | `get_entry_points` | All binary entry points |
 
-### Data Analysis (6)
+### Data Analysis (7)
 | Tool | Description |
 |------|-------------|
 | `read_memory` | Read raw bytes at address |
@@ -164,6 +164,7 @@ Configure via environment variables with the `IDASSISTMCP_` prefix:
 | `search_strings` | String search with pagination |
 | `create_data_var` | Define data variable at address (byte/word/dword/qword/float/ascii/C type) |
 | `get_data_vars` | List defined data variables (non-code items) |
+| `define` | **IDA `U`/`C`/`P` hotkeys** - actions: `undefine`, `code`, `function` (bytes untouched) |
 
 ### Patching (3)
 | Tool | Description |
@@ -200,7 +201,7 @@ IDAssistMCP/
         ├── __init__.py
         ├── server.py                  # FastMCP server + transport
         ├── context.py                 # Single-binary IDA context
-        ├── tools.py                   # 41 MCP tools (IDA API)
+        ├── tools.py                   # 42 MCP tools (IDA API)
         ├── resources.py               # 8 MCP resources
         ├── prompts.py                 # 7 guided workflow prompts
         ├── config.py                  # Pydantic settings

src/idassist_mcp/tools.py

@@ -1,11 +1,11 @@
 """
 Comprehensive MCP tool implementations for IDAssistMCP
 
-This module provides 41 IDA Pro integration tools registered as
+This module provides 42 IDA Pro integration tools registered as
 FastMCP tools. All tools that call IDA APIs use @_ida_main_thread to dispatch
 onto IDA's main thread (required for both reads and writes).
 
-Consolidated tools (5): get_code, comments, variables, types, xrefs
+Consolidated tools (6): get_code, comments, variables, types, xrefs, bookmarks
 Standalone tools (36): see register_tools() for the full list
 """
 
@@ -47,6 +47,7 @@
     import ida_nalt
     import ida_segment
     import ida_typeinf
+    import ida_ua
     import ida_xref
     _IN_IDA = True
 except ImportError:
@@ -2025,7 +2026,113 @@ def get_data_vars(ctx: Context, segment_name: str = "",
         }
 
     # ================================================================== #
-    #  38-41. Task Management
+    #  38. define (IDA U / C / P hotkeys)
+    # ================================================================== #
+
+    @_tool("define", annotations=NON_IDEMPOTENT)
+    @_ida_main_thread
+    def define(action: str, address: str, ctx: Context,
+               end_address: str = "", size: int = 0) -> dict:
+        """Undefine items, force code, or create a function (IDA U/C/P hotkeys).
+
+        Bytes are never modified — only the IDB interpretation of them.
+        Typical order when reclaiming misparsed data: 'undefine', then 'code',
+        then 'function'.
+
+        Args:
+            action: 'undefine' (U), 'code' (C), or 'function' (P)
+            address: Hex start address
+            end_address: Optional exclusive end address for range operations
+            size: Optional byte count, alternative to end_address
+
+        Returns:
+            Dictionary with operation status, or an error.
+        """
+        ea = parse_address(address)
+        if ea is None:
+            return {"error": f"Invalid address: {address}"}
+
+        if end_address:
+            end_ea = parse_address(end_address)
+            if end_ea is None:
+                return {"error": f"Invalid end_address: {end_address}"}
+            if end_ea <= ea:
+                return {"error": "end_address must be greater than address"}
+        elif size:
+            if size < 0:
+                return {"error": "size must be positive"}
+            end_ea = ea + size
+        else:
+            end_ea = idaapi.BADADDR
+
+        if action == "undefine":
+            count = (end_ea - ea) if end_ea != idaapi.BADADDR else (ida_bytes.get_item_size(ea) or 1)
+            flags = ida_bytes.DELIT_SIMPLE if end_ea != idaapi.BADADDR else ida_bytes.DELIT_EXPAND
+            if not ida_bytes.del_items(ea, flags, count):
+                return {"error": f"del_items failed at {hex(ea)}"}
+            return {
+                "status": "ok",
+                "address": hex(ea),
+                "end_address": hex(ea + count),
+                "size": count,
+            }
+
+        elif action == "code":
+            stop = end_ea if end_ea != idaapi.BADADDR else ea + 1
+            if end_ea != idaapi.BADADDR:
+                cur = ea
+                while cur < stop:
+                    length = ida_ua.decode_insn(ida_ua.insn_t(), cur)
+                    if not length:
+                        return {"error": f"Cannot decode instruction at {hex(cur)}"}
+                    if cur + length > stop:
+                        return {
+                            "error": f"Instruction at {hex(cur)} crosses end_address {hex(stop)}",
+                            "address": hex(ea),
+                            "instructions": 0,
+                        }
+                    cur += length
+
+            cur = ea
+            made = 0
+            while cur < stop:
+                length = ida_ua.create_insn(cur)
+                if not length:
+                    return {
+                        "error": f"Cannot decode instruction at {hex(cur)} "
+                                 "(undefine the range first?)",
+                        "address": hex(ea),
+                        "instructions": made,
+                    }
+                cur += length
+                made += 1
+            return {
+                "status": "ok",
+                "address": hex(ea),
+                "end_address": hex(cur),
+                "instructions": made,
+            }
+
+        elif action == "function":
+            # BADADDR end lets IDA determine the function end from control flow.
+            if not ida_funcs.add_func(ea, end_ea):
+                return {"error": f"add_func failed at {hex(ea)}"}
+            func = ida_funcs.get_func(ea)
+            if not func:
+                return {"error": f"Function created but not found at {hex(ea)}"}
+            return {
+                "status": "ok",
+                "name": ida_funcs.get_func_name(func.start_ea) or f"sub_{func.start_ea:x}",
+                "address": hex(func.start_ea),
+                "end": hex(func.end_ea),
+                "size": func.end_ea - func.start_ea,
+            }
+
+        else:
+            return {"error": f"Unknown action '{action}'. Use 'undefine', 'code', or 'function'."}
+
+    # ================================================================== #
+    #  39-42. Task Management
     # ================================================================== #
 
     @_tool("start_task", annotations=NON_IDEMPOTENT)

src/idassist_mcp/ui/tool_registry.py

@@ -18,7 +18,7 @@ class ToolInfo:
     annotation: str  # "read_only", "modify", or "non_idempotent"
 
 
-# Complete catalog of all 41 MCP tools
+# Complete catalog of all 42 MCP tools
 TOOL_CATALOG: List[ToolInfo] = [
     # Binary Management (2)
     ToolInfo("list_binaries", "List Binaries", "Binary Management", "List the currently loaded binary", "read_only"),
@@ -31,10 +31,10 @@ class ToolInfo:
     ToolInfo("get_basic_blocks", "Get Basic Blocks", "Code Analysis", "Get basic blocks (CFG) for a function", "read_only"),
 
     # Consolidated Tools (4)
-    ToolInfo("xrefs_tool", "Cross-References", "Cross-References", "Get xrefs to/from address, optionally include callers/callees", "read_only"),
-    ToolInfo("comments_tool", "Comments", "Comments & Variables", "Get, set, list, or remove comments (action parameter)", "modify"),
-    ToolInfo("variables_tool", "Variables", "Comments & Variables", "List variables or rename local/global variables (action parameter)", "modify"),
-    ToolInfo("types_tool", "Types", "Types", "List, set, create_struct, or create_enum (action parameter)", "modify"),
+    ToolInfo("xrefs", "Cross-References", "Cross-References", "Get xrefs to/from address, optionally include callers/callees", "read_only"),
+    ToolInfo("comments", "Comments", "Comments & Variables", "Get, set, list, or remove comments (action parameter)", "modify"),
+    ToolInfo("variables", "Variables", "Comments & Variables", "List variables or rename local/global variables (action parameter)", "modify"),
+    ToolInfo("types", "Types", "Types", "List, set, create_struct, or create_enum (action parameter)", "modify"),
 
     # Function Discovery (5)
     ToolInfo("get_functions", "Get Functions", "Function Discovery", "List all functions with filtering and pagination", "read_only"),
@@ -67,15 +67,16 @@ class ToolInfo:
 
     # Navigation (4)
     ToolInfo("navigate_to", "Navigate To", "Navigation", "Move IDA cursor to address", "modify"),
-    ToolInfo("set_bookmark", "Set Bookmark", "Navigation", "Create a position bookmark", "modify"),
+    ToolInfo("bookmarks", "Bookmarks", "Navigation", "List, set, or remove position bookmarks (action parameter)", "modify"),
     ToolInfo("get_current_address", "Get Current Address", "Navigation", "Get address at cursor position", "read_only"),
     ToolInfo("get_current_function", "Get Current Function", "Navigation", "Get function at cursor position", "read_only"),
 
-    # New Feature Parity Tools (3)
+    # New Feature Parity Tools (5)
     ToolInfo("get_function_stack_layout", "Stack Layout", "Code Analysis", "Get stack frame layout for a function", "read_only"),
     ToolInfo("get_classes", "Get Classes", "Types", "Get struct/class types from type library", "read_only"),
     ToolInfo("create_data_var", "Create Data Var", "Data Analysis", "Define a data variable at address", "modify"),
     ToolInfo("get_data_vars", "Get Data Vars", "Data Analysis", "Get defined data variables (non-code items)", "read_only"),
+    ToolInfo("define", "Define", "Data Analysis", "Undefine items, force code, or create a function (IDA U/C/P)", "non_idempotent"),
 
     # Task Management (4)
     ToolInfo("start_task", "Start Task", "Task Management", "Start an async background task", "non_idempotent"),

tests/test_define_tool.py

@@ -0,0 +1,92 @@
+import unittest
+from types import SimpleNamespace
+
+from idassist_mcp import tools
+
+
+class _Mcp:
+    def __init__(self):
+        self.tools = {}
+
+    def tool(self, **_kwargs):
+        def register(fn):
+            self.tools[fn.__name__] = fn
+            return fn
+
+        return register
+
+
+class _Bytes:
+    DELIT_SIMPLE = 0
+    DELIT_EXPAND = 1
+
+    def __init__(self):
+        self.calls = []
+
+    def get_item_size(self, _ea):
+        return 1
+
+    def del_items(self, ea, flags, count):
+        self.calls.append((ea, flags, count))
+        return True
+
+
+class _Funcs:
+    def __init__(self):
+        self.deleted = []
+
+    def del_func(self, ea):
+        self.deleted.append(ea)
+        return True
+
+
+class _Ua:
+    def __init__(self, lengths):
+        self.lengths = lengths
+        self.created = []
+
+    @staticmethod
+    def insn_t():
+        return object()
+
+    def decode_insn(self, _insn, ea):
+        return self.lengths.get(ea, 0)
+
+    def create_insn(self, ea):
+        self.created.append(ea)
+        return self.lengths.get(ea, 0)
+
+
+class DefineToolTests(unittest.TestCase):
+    def setUp(self):
+        self.mcp = _Mcp()
+        self.bytes = _Bytes()
+        self.funcs = _Funcs()
+        tools.idaapi = SimpleNamespace(BADADDR=-1)
+        tools.ida_bytes = self.bytes
+        tools.ida_funcs = self.funcs
+        tools.ida_ua = _Ua({0x3000: 2})
+        tools.register_tools(self.mcp)
+        self.define = self.mcp.tools["define"]
+
+    def test_undefine_preserves_functions(self):
+        result = self.define("undefine", "0x1000", None)
+
+        self.assertEqual(result["status"], "ok")
+        self.assertEqual(self.funcs.deleted, [])
+
+    def test_explicit_undefine_range_does_not_expand(self):
+        result = self.define("undefine", "0x2000", None, size=4)
+
+        self.assertEqual(result["status"], "ok")
+        self.assertEqual(self.bytes.calls, [(0x2000, _Bytes.DELIT_SIMPLE, 4)])
+
+    def test_code_range_is_checked_before_modification(self):
+        result = self.define("code", "0x3000", None, size=1)
+
+        self.assertIn("error", result)
+        self.assertEqual(tools.ida_ua.created, [])
+
+
+if __name__ == "__main__":
+    unittest.main()