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

ida_pro_mcp/ida_mcp/api_memory.py

@@ -0,0 +1,304 @@
+"""Memory reading and writing operations for IDA Pro MCP.
+
+This module provides batch operations for reading and writing memory at various
+granularities (bytes, integers, strings) and patching binary data.
+"""
+
+import re
+
+from typing import Annotated
+import ida_bytes
+import idaapi
+
+from .rpc import tool
+from .sync import idasync
+from .utils import (
+    IntRead,
+    IntWrite,
+    MemoryPatch,
+    MemoryRead,
+    normalize_list_input,
+    parse_address,
+)
+
+
+# ============================================================================
+# Memory Reading Operations
+# ============================================================================
+
+
+@tool
+@idasync
+def get_bytes(regions: list[MemoryRead] | MemoryRead) -> list[dict]:
+    """Read bytes from memory addresses"""
+    if isinstance(regions, dict):
+        regions = [regions]
+
+    results = []
+    for item in regions:
+        addr = item.get("addr", "")
+        size = item.get("size", 0)
+
+        try:
+            ea = parse_address(addr)
+            data = " ".join(f"{x:#02x}" for x in ida_bytes.get_bytes(ea, size))
+            results.append({"addr": addr, "data": data})
+        except Exception as e:
+            results.append({"addr": addr, "data": None, "error": str(e)})
+
+    return results
+
+
+_INT_CLASS_RE = re.compile(r"^(?P<sign>[iu])(?P<bits>8|16|32|64)(?P<endian>le|be)?$")
+
+
+def _parse_int_class(text: str) -> tuple[int, bool, str, str]:
+    if not text:
+        raise ValueError("Missing integer class")
+
+    cleaned = text.strip().lower()
+    match = _INT_CLASS_RE.match(cleaned)
+    if not match:
+        raise ValueError(f"Invalid integer class: {text}")
+
+    bits = int(match.group("bits"))
+    signed = match.group("sign") == "i"
+    endian = match.group("endian") or "le"
+    byte_order = "little" if endian == "le" else "big"
+    normalized = f"{'i' if signed else 'u'}{bits}{endian}"
+    return bits, signed, byte_order, normalized
+
+
+def _parse_int_value(text: str, signed: bool, bits: int) -> int:
+    if text is None:
+        raise ValueError("Missing integer value")
+
+    value_text = str(text).strip()
+    try:
+        value = int(value_text, 0)
+    except ValueError:
+        raise ValueError(f"Invalid integer value: {text}")
+
+    if not signed and value < 0:
+        raise ValueError(f"Negative value not allowed for u{bits}")
+
+    return value
+
+
+@tool
+@idasync
+def get_int(
+    queries: Annotated[
+        list[IntRead] | IntRead,
+        "Integer read requests (ty, addr). ty: i8/u64/i16le/i16be/etc",
+    ],
+) -> list[dict]:
+    """Read integer values from memory addresses"""
+    if isinstance(queries, dict):
+        queries = [queries]
+
+    results = []
+    for item in queries:
+        addr = item.get("addr", "")
+        ty = item.get("ty", "")
+
+        try:
+            bits, signed, byte_order, normalized = _parse_int_class(ty)
+            ea = parse_address(addr)
+            size = bits // 8
+            data = ida_bytes.get_bytes(ea, size)
+            if not data or len(data) != size:
+                raise ValueError(f"Failed to read {size} bytes at {addr}")
+
+            value = int.from_bytes(data, byte_order, signed=signed)
+            results.append(
+                {"addr": addr, "ty": normalized, "value": value, "error": None}
+            )
+        except Exception as e:
+            results.append({"addr": addr, "ty": ty, "value": None, "error": str(e)})
+
+    return results
+
+
+@tool
+@idasync
+def get_string(
+    addrs: Annotated[list[str] | str, "Addresses to read strings from"],
+) -> list[dict]:
+    """Read strings from memory addresses"""
+    addrs = normalize_list_input(addrs)
+    results = []
+
+    for addr in addrs:
+        try:
+            ea = parse_address(addr)
+            raw = idaapi.get_strlit_contents(ea, -1, 0)
+            if not raw:
+                results.append(
+                    {"addr": addr, "value": None, "error": "No string at address"}
+                )
+                continue
+            value = raw.decode("utf-8", errors="replace")
+            results.append({"addr": addr, "value": value})
+        except Exception as e:
+            results.append({"addr": addr, "value": None, "error": str(e)})
+
+    return results
+
+
+def get_global_variable_value_internal(ea: int) -> str:
+    import ida_typeinf
+    import ida_nalt
+    import ida_bytes
+    from .sync import IDAError
+
+    tif = ida_typeinf.tinfo_t()
+    if not ida_nalt.get_tinfo(tif, ea):
+        if not ida_bytes.has_any_name(ea):
+            raise IDAError(f"Failed to get type information for variable at {ea:#x}")
+
+        size = ida_bytes.get_item_size(ea)
+        if size == 0:
+            raise IDAError(f"Failed to get type information for variable at {ea:#x}")
+    else:
+        size = tif.get_size()
+
+    if size == 0 and tif.is_array() and tif.get_array_element().is_decl_char():
+        raw = idaapi.get_strlit_contents(ea, -1, 0)
+        if not raw:
+            return '""'
+        return_string = raw.decode("utf-8", errors="replace").strip()
+        return f'"{return_string}"'
+    elif size == 1:
+        return hex(ida_bytes.get_byte(ea))
+    elif size == 2:
+        return hex(ida_bytes.get_word(ea))
+    elif size == 4:
+        return hex(ida_bytes.get_dword(ea))
+    elif size == 8:
+        return hex(ida_bytes.get_qword(ea))
+    else:
+        return " ".join(hex(x) for x in ida_bytes.get_bytes(ea, size))
+
+
+@tool
+@idasync
+def get_global_value(
+    queries: Annotated[
+        list[str] | str, "Global variable addresses or names to read values from"
+    ],
+) -> list[dict]:
+    """Read global variable values by address or name
+    (auto-detects hex addresses vs names)"""
+    from .utils import looks_like_address
+
+    queries = normalize_list_input(queries)
+    results = []
+
+    for query in queries:
+        try:
+            ea = idaapi.BADADDR
+
+            # Try as address first if it looks like one
+            if looks_like_address(query):
+                try:
+                    ea = parse_address(query)
+                except Exception:
+                    ea = idaapi.BADADDR
+
+            # Fall back to name lookup
+            if ea == idaapi.BADADDR:
+                ea = idaapi.get_name_ea(idaapi.BADADDR, query)
+
+            if ea == idaapi.BADADDR:
+                results.append({"query": query, "value": None, "error": "Not found"})
+                continue
+
+            value = get_global_variable_value_internal(ea)
+            results.append({"query": query, "value": value, "error": None})
+        except Exception as e:
+            results.append({"query": query, "value": None, "error": str(e)})
+
+    return results
+
+
+# ============================================================================
+# Batch Data Operations
+# ============================================================================
+
+
+@tool
+@idasync
+def patch(patches: list[MemoryPatch] | MemoryPatch) -> list[dict]:
+    """Patch bytes at memory addresses with hex data"""
+    if isinstance(patches, dict):
+        patches = [patches]
+
+    results = []
+
+    for patch in patches:
+        try:
+            ea = parse_address(patch["addr"])
+            data = bytes.fromhex(patch["data"])
+
+            ida_bytes.patch_bytes(ea, data)
+            results.append(
+                {"addr": patch["addr"], "size": len(data), "ok": True, "error": None}
+            )
+
+        except Exception as e:
+            results.append({"addr": patch.get("addr"), "size": 0, "error": str(e)})
+
+    return results
+
+
+@tool
+@idasync
+def put_int(
+    items: Annotated[
+        list[IntWrite] | IntWrite,
+        "Integer write requests (ty, addr, value). value is a string; supports 0x.. and negatives",
+    ],
+) -> list[dict]:
+    """Write integer values to memory addresses"""
+    if isinstance(items, dict):
+        items = [items]
+
+    results = []
+    for item in items:
+        addr = item.get("addr", "")
+        ty = item.get("ty", "")
+        value_text = item.get("value")
+
+        try:
+            bits, signed, byte_order, normalized = _parse_int_class(ty)
+            value = _parse_int_value(value_text, signed, bits)
+            size = bits // 8
+            try:
+                data = value.to_bytes(size, byte_order, signed=signed)
+            except OverflowError:
+                raise ValueError(f"Value {value_text} does not fit in {normalized}")
+
+            ea = parse_address(addr)
+            ida_bytes.patch_bytes(ea, data)
+            results.append(
+                {
+                    "addr": addr,
+                    "ty": normalized,
+                    "value": str(value_text),
+                    "ok": True,
+                    "error": None,
+                }
+            )
+        except Exception as e:
+            results.append(
+                {
+                    "addr": addr,
+                    "ty": ty,
+                    "value": str(value_text) if value_text is not None else None,
+                    "ok": False,
+                    "error": str(e),
+                }
+            )
+
+    return results

ida_pro_mcp/ida_mcp/api_modify.py

@@ -0,0 +1,406 @@
+import idaapi
+import idautils
+import idc
+import ida_hexrays
+import ida_bytes
+import ida_typeinf
+import ida_frame
+import ida_dirtree
+
+from .rpc import tool
+from .sync import idasync, IDAError
+from .cache import invalidate_function_caches, decompile_cache
+from .utils import (
+    parse_address,
+    decompile_checked,
+    refresh_decompiler_ctext,
+    CommentOp,
+    AsmPatchOp,
+    FunctionRename,
+    GlobalRename,
+    LocalRename,
+    StackRename,
+    RenameBatch,
+)
+
+
+# ============================================================================
+# Modification Operations
+# ============================================================================
+
+
+@tool
+@idasync
+def set_comments(items: list[CommentOp] | CommentOp):
+    """Set comments at addresses (both disassembly and decompiler views)"""
+    if isinstance(items, dict):
+        items = [items]
+
+    results = []
+    for item in items:
+        addr_str = item.get("addr", "")
+        comment = item.get("comment", "")
+
+        try:
+            ea = parse_address(addr_str)
+
+            # Invalidate decompile cache for this function
+            func = idaapi.get_func(ea)
+            if func:
+                decompile_cache.invalidate(hex(func.start_ea))
+
+            if not idaapi.set_cmt(ea, comment, False):
+                results.append(
+                    {
+                        "addr": addr_str,
+                        "error": f"Failed to set disassembly comment at {hex(ea)}",
+                    }
+                )
+                continue
+
+            if not ida_hexrays.init_hexrays_plugin():
+                results.append({"addr": addr_str, "ok": True})
+                continue
+
+            try:
+                cfunc = decompile_checked(ea)
+            except IDAError:
+                results.append({"addr": addr_str, "ok": True})
+                continue
+
+            if ea == cfunc.entry_ea:
+                idc.set_func_cmt(ea, comment, True)
+                cfunc.refresh_func_ctext()
+                results.append({"addr": addr_str, "ok": True})
+                continue
+
+            eamap = cfunc.get_eamap()
+            if ea not in eamap:
+                results.append(
+                    {
+                        "addr": addr_str,
+                        "ok": True,
+                        "error": f"Failed to set decompiler comment at {hex(ea)}",
+                    }
+                )
+                continue
+            nearest_ea = eamap[ea][0].ea
+
+            if cfunc.has_orphan_cmts():
+                cfunc.del_orphan_cmts()
+                cfunc.save_user_cmts()
+
+            tl = idaapi.treeloc_t()
+            tl.ea = nearest_ea
+            for itp in range(idaapi.ITP_SEMI, idaapi.ITP_COLON):
+                tl.itp = itp
+                cfunc.set_user_cmt(tl, comment)
+                cfunc.save_user_cmts()
+                cfunc.refresh_func_ctext()
+                if not cfunc.has_orphan_cmts():
+                    results.append({"addr": addr_str, "ok": True})
+                    break
+                cfunc.del_orphan_cmts()
+                cfunc.save_user_cmts()
+            else:
+                results.append(
+                    {
+                        "addr": addr_str,
+                        "ok": True,
+                        "error": f"Failed to set decompiler comment at {hex(ea)}",
+                    }
+                )
+        except Exception as e:
+            results.append({"addr": addr_str, "error": str(e)})
+
+    return results
+
+
+@tool
+@idasync
+def patch_asm(items: list[AsmPatchOp] | AsmPatchOp) -> list[dict]:
+    """Patch assembly instructions at addresses"""
+    if isinstance(items, dict):
+        items = [items]
+
+    results = []
+    for item in items:
+        addr_str = item.get("addr", "")
+        instructions = item.get("asm", "")
+
+        try:
+            ea = parse_address(addr_str)
+            assembles = instructions.split(";")
+            for assemble in assembles:
+                assemble = assemble.strip()
+                try:
+                    (check_assemble, bytes_to_patch) = idautils.Assemble(ea, assemble)
+                    if not check_assemble:
+                        results.append(
+                            {
+                                "addr": addr_str,
+                                "error": f"Failed to assemble: {assemble}",
+                            }
+                        )
+                        break
+                    ida_bytes.patch_bytes(ea, bytes_to_patch)
+                    ea += len(bytes_to_patch)
+                except Exception as e:
+                    results.append(
+                        {"addr": addr_str, "error": f"Failed at {hex(ea)}: {e}"}
+                    )
+                    break
+            else:
+                results.append({"addr": addr_str, "ok": True})
+        except Exception as e:
+            results.append({"addr": addr_str, "error": str(e)})
+
+    return results
+
+
+@tool
+@idasync
+def rename(batch: RenameBatch) -> dict:
+    """Unified rename operation for functions, globals, locals, and stack variables"""
+
+    def _normalize_items(items):
+        """Convert single item or None to list"""
+        if items is None:
+            return []
+        return [items] if isinstance(items, dict) else items
+
+    def _has_user_name(ea: int) -> bool:
+        flags = idaapi.get_flags(ea)
+        checker = getattr(idaapi, "has_user_name", None)
+        if checker is not None:
+            return checker(flags)
+        try:
+            import ida_name
+
+            checker = getattr(ida_name, "has_user_name", None)
+            if checker is not None:
+                return checker(flags)
+        except Exception:
+            pass
+        return False
+
+    def _place_func_in_vibe_dir(ea: int) -> tuple[bool, str | None]:
+        tree = ida_dirtree.get_std_dirtree(ida_dirtree.DIRTREE_FUNCS)
+        if tree is None:
+            return False, "Function dirtree not available"
+
+        if not tree.load():
+            return False, "Failed to load function dirtree"
+
+        vibe_path = "/vibe/"
+        if not tree.isdir(vibe_path):
+            err = tree.mkdir(vibe_path)
+            if err not in (ida_dirtree.DTE_OK, ida_dirtree.DTE_ALREADY_EXISTS):
+                return False, f"mkdir failed: {err}"
+
+        old_cwd = tree.getcwd()
+        try:
+            if tree.chdir(vibe_path) != ida_dirtree.DTE_OK:
+                return False, "Failed to chdir to vibe"
+            err = tree.link(ea)
+            if err not in (ida_dirtree.DTE_OK, ida_dirtree.DTE_ALREADY_EXISTS):
+                return False, f"link failed: {err}"
+            if not tree.save():
+                return False, "Failed to save function dirtree"
+        finally:
+            if old_cwd:
+                tree.chdir(old_cwd)
+
+        return True, None
+
+    def _rename_funcs(items: list[FunctionRename]) -> list[dict]:
+        results = []
+        for item in items:
+            try:
+                ea = parse_address(item["addr"])
+                had_user_name = _has_user_name(ea)
+                success = idaapi.set_name(ea, item["name"], idaapi.SN_CHECK)
+                if success:
+                    # Invalidate caches for this function
+                    invalidate_function_caches(ea)
+                    func = idaapi.get_func(ea)
+                    if func:
+                        refresh_decompiler_ctext(func.start_ea)
+                    if not had_user_name and func:
+                        placed, place_error = _place_func_in_vibe_dir(func.start_ea)
+                    else:
+                        placed, place_error = None, None
+                results.append(
+                    {
+                        "addr": item["addr"],
+                        "name": item["name"],
+                        "ok": success,
+                        "error": None if success else "Rename failed",
+                        "dir": "vibe" if success and placed else None,
+                        "dir_error": place_error if success else None,
+                    }
+                )
+            except Exception as e:
+                results.append({"addr": item.get("addr"), "error": str(e)})
+        return results
+
+    def _rename_globals(items: list[GlobalRename]) -> list[dict]:
+        results = []
+        for item in items:
+            try:
+                ea = idaapi.get_name_ea(idaapi.BADADDR, item["old"])
+                if ea == idaapi.BADADDR:
+                    results.append(
+                        {
+                            "old": item["old"],
+                            "new": item["new"],
+                            "ok": False,
+                            "error": f"Global '{item['old']}' not found",
+                        }
+                    )
+                    continue
+                success = idaapi.set_name(ea, item["new"], idaapi.SN_CHECK)
+                results.append(
+                    {
+                        "old": item["old"],
+                        "new": item["new"],
+                        "ok": success,
+                        "error": None if success else "Rename failed",
+                    }
+                )
+            except Exception as e:
+                results.append({"old": item.get("old"), "error": str(e)})
+        return results
+
+    def _rename_locals(items: list[LocalRename]) -> list[dict]:
+        results = []
+        for item in items:
+            try:
+                func = idaapi.get_func(parse_address(item["func_addr"]))
+                if not func:
+                    results.append(
+                        {
+                            "func_addr": item["func_addr"],
+                            "old": item["old"],
+                            "new": item["new"],
+                            "ok": False,
+                            "error": "No function found",
+                        }
+                    )
+                    continue
+                success = ida_hexrays.rename_lvar(
+                    func.start_ea, item["old"], item["new"]
+                )
+                if success:
+                    refresh_decompiler_ctext(func.start_ea)
+                results.append(
+                    {
+                        "func_addr": item["func_addr"],
+                        "old": item["old"],
+                        "new": item["new"],
+                        "ok": success,
+                        "error": None if success else "Rename failed",
+                    }
+                )
+            except Exception as e:
+                results.append({"func_addr": item.get("func_addr"), "error": str(e)})
+        return results
+
+    def _rename_stack(items: list[StackRename]) -> list[dict]:
+        results = []
+        for item in items:
+            try:
+                func = idaapi.get_func(parse_address(item["func_addr"]))
+                if not func:
+                    results.append(
+                        {
+                            "func_addr": item["func_addr"],
+                            "old": item["old"],
+                            "new": item["new"],
+                            "ok": False,
+                            "error": "No function found",
+                        }
+                    )
+                    continue
+
+                frame_tif = ida_typeinf.tinfo_t()
+                if not ida_frame.get_func_frame(frame_tif, func):
+                    results.append(
+                        {
+                            "func_addr": item["func_addr"],
+                            "old": item["old"],
+                            "new": item["new"],
+                            "ok": False,
+                            "error": "No frame",
+                        }
+                    )
+                    continue
+
+                idx, udm = frame_tif.get_udm(item["old"])
+                if not udm:
+                    results.append(
+                        {
+                            "func_addr": item["func_addr"],
+                            "old": item["old"],
+                            "new": item["new"],
+                            "ok": False,
+                            "error": f"'{item['old']}' not found",
+                        }
+                    )
+                    continue
+
+                tid = frame_tif.get_udm_tid(idx)
+                if ida_frame.is_special_frame_member(tid):
+                    results.append(
+                        {
+                            "func_addr": item["func_addr"],
+                            "old": item["old"],
+                            "new": item["new"],
+                            "ok": False,
+                            "error": "Special frame member",
+                        }
+                    )
+                    continue
+
+                udm = ida_typeinf.udm_t()
+                frame_tif.get_udm_by_tid(udm, tid)
+                offset = udm.offset // 8
+                if ida_frame.is_funcarg_off(func, offset):
+                    results.append(
+                        {
+                            "func_addr": item["func_addr"],
+                            "old": item["old"],
+                            "new": item["new"],
+                            "ok": False,
+                            "error": "Argument member",
+                        }
+                    )
+                    continue
+
+                sval = ida_frame.soff_to_fpoff(func, offset)
+                success = ida_frame.define_stkvar(func, item["new"], sval, udm.type)
+                results.append(
+                    {
+                        "func_addr": item["func_addr"],
+                        "old": item["old"],
+                        "new": item["new"],
+                        "ok": success,
+                        "error": None if success else "Rename failed",
+                    }
+                )
+            except Exception as e:
+                results.append({"func_addr": item.get("func_addr"), "error": str(e)})
+        return results
+
+    # Process each category
+    result = {}
+    if "func" in batch:
+        result["func"] = _rename_funcs(_normalize_items(batch["func"]))
+    if "data" in batch:
+        result["data"] = _rename_globals(_normalize_items(batch["data"]))
+    if "local" in batch:
+        result["local"] = _rename_locals(_normalize_items(batch["local"]))
+    if "stack" in batch:
+        result["stack"] = _rename_stack(_normalize_items(batch["stack"]))
+
+    return result