#!/usr/bin/env python3 """Inject idatui's extra ida-pro-mcp tools into the installed server package. DEPRECATED along with the ida-pro-mcp transport: the default backend is now the idalib worker (idatui/worker.py), which registers these same tools in-process and needs no patching. Kept only for `--backend mcp`; slated for removal. ida-pro-mcp lacks a few tools idatui needs. Rather than vendor/fork the server, we keep the tool source here and inject it (idempotently) into the installed ``api_types.py``. That module is imported by every worker (``python -m ida_pro_mcp.idalib_server``), so the tools register themselves via ``@tool`` on the shared ``MCP_SERVER`` — no server code is forked, and re-running this (spawn.sh does, on every start) re-applies it after a reinstall/upgrade. Injected tools: * ``del_type`` — delete a named local type (struct editor CRUD). * ``func_types`` — structured decompiler types for a function (prototype + local variables), so clients don't parse pseudocode text. * ``set_lvar_type`` — set a decompiler local variable's type; works on auto/ register vars too (the stock set_type only updates lvars that already have user-saved info). The block between the BEGIN/END markers is *replaced* on each run, so editing BODY here and restarting the supervisor updates the tools. Run with the *same* interpreter the server uses (the idalib-mcp entry point's ``/usr/bin/python``), so it patches the file the workers actually import. Changing a tool needs a supervisor restart so workers respawn. """ from __future__ import annotations import importlib.util import pathlib import sys BEGIN = "# >>> idatui-ext: begin (auto-injected by server/patch_server.py) >>>" END = "# <<< idatui-ext: end <<<" # Appended to ida_pro_mcp/ida_mcp/api_types.py, which already imports # ``Annotated``, ``tool``, ``idasync``, ``ida_typeinf``, ``parse_address`` and # ``_parse_type_tinfo``. BODY = ''' def _idatui_lv_get(x): return x() if callable(x) else x @tool @idasync def resolve_names( queries: Annotated[list, "Symbol name(s) to resolve to their OWN address"], ) -> list: """Resolve named locations (functions, labels like loc_/locret_, data) to the exact address the NAME denotes, via get_name_ea. Unlike lookup_funcs, a mid-function label resolves to the label's address, not the containing function's entry.""" import idaapi qs = queries if isinstance(queries, list) else [queries] out = [] for q in qs: q = str(q).strip() ea = idaapi.get_name_ea(idaapi.BADADDR, q) out.append({"query": q, "ea": (hex(ea) if ea != idaapi.BADADDR else None)}) return out @tool @idasync def del_type( name: Annotated[str, "Local type name to delete (struct/union/enum/typedef)"], ) -> dict: """Delete a named local type from the local type library.""" til = ida_typeinf.get_idati() ok = ida_typeinf.del_named_type(til, name, ida_typeinf.NTF_TYPE) if not ok: return {"name": name, "error": f"Type '{name}' not found or could not be deleted"} return {"name": name, "deleted": True} @tool @idasync def func_types( addr: Annotated[str, "Function address or name"], ) -> dict: """Structured decompiler types for a function: its prototype plus each local variable (name/type/is_arg). Lets clients read/edit types without parsing pseudocode text.""" import ida_hexrays import idaapi def _tstr(tif): try: s = tif.dstr() if s: return s except Exception: pass return str(tif) ea = parse_address(addr) f = idaapi.get_func(ea) if not f: return {"addr": str(addr), "error": "no function at address"} try: cf = ida_hexrays.decompile(f.start_ea) except Exception as e: return {"addr": hex(f.start_ea), "error": f"decompile failed: {e}"} if cf is None: return {"addr": hex(f.start_ea), "error": "decompilation failed"} name = idaapi.get_func_name(f.start_ea) or "" try: proto = ida_typeinf.print_tinfo( "", 0, 0, ida_typeinf.PRTYPE_1LINE, cf.type, name, "") except Exception: proto = "" lvars = [] for lv in cf.get_lvars(): try: ty = _tstr(_idatui_lv_get(lv.type)) except Exception: ty = "" lvars.append({ "name": _idatui_lv_get(lv.name), "type": ty, "is_arg": bool(_idatui_lv_get(lv.is_arg_var)), }) return { "addr": hex(f.start_ea), "name": name, "prototype": (proto or "").strip(), "lvars": lvars, } @tool @idasync def set_lvar_type( addr: Annotated[str, "Function address or name"], variable: Annotated[str, "Local variable name"], type: Annotated[str, "New C type for the variable"], ) -> dict: """Set a decompiler local variable's type. Handles auto/register vars (unlike set_type, which only updates lvars that already have user-saved info).""" import ida_hexrays import idaapi ea = parse_address(addr) f = idaapi.get_func(ea) if not f: return {"error": "no function at address"} try: cf = ida_hexrays.decompile(f.start_ea) except Exception as e: return {"error": f"decompile failed: {e}"} if cf is None: return {"error": "decompilation failed"} target = None for lv in cf.get_lvars(): if _idatui_lv_get(lv.name) == variable: target = lv break if target is None: return {"error": f"local variable {variable!r} not found"} try: tif = _parse_type_tinfo(type) except Exception as e: return {"error": f"bad type {type!r}: {e}"} lsi = ida_hexrays.lvar_saved_info_t() try: lsi.ll = target except Exception: try: lsi.ll.location = _idatui_lv_get(target.location) lsi.ll.defea = target.defea except Exception as e: return {"error": f"could not locate variable: {e}"} lsi.type = tif ok = bool(ida_hexrays.modify_user_lvar_info( f.start_ea, ida_hexrays.MLI_TYPE, lsi)) return {"addr": hex(f.start_ea), "variable": variable, "type": type, "ok": ok} @tool @idasync def file_regions() -> dict: """Loaded segments mapped to their raw file offsets (get_fileregion_offset), so clients can convert a virtual address to an on-disk file offset without a format-specific header parser. file_off is -1 for non-file-backed segments (e.g. .bss).""" import ida_segment import idaapi out = [] seg = ida_segment.get_first_seg() while seg is not None: try: fo = int(idaapi.get_fileregion_offset(seg.start_ea)) except Exception: fo = -1 if fo < 0 or fo >= (1 << 48): fo = -1 try: nm = ida_segment.get_segm_name(seg) or "" except Exception: nm = "" out.append({"start": hex(seg.start_ea), "end": hex(seg.end_ea), "file_off": fo, "name": nm}) seg = ida_segment.get_next_seg(seg.start_ea) return {"regions": out} @tool @idasync def make_string( addr: Annotated[str, "Address of the string start"], length: Annotated[int, "Length in bytes (0 = auto-detect to the terminator)"] = 0, kind: Annotated[str, "String kind: c | c16 | c32 | pascal"] = "c", ) -> dict: """Create a string literal at ``addr`` (IDA's 'A'). ``length`` 0 auto-detects to the terminator. Undefines any items in the way first, like the UI does. Returns the created byte size and the decoded contents.""" import ida_bytes import ida_nalt ea = parse_address(addr) strtype = { "c": ida_nalt.STRTYPE_C, "c16": ida_nalt.STRTYPE_C_16, "c32": ida_nalt.STRTYPE_C_32, "pascal": ida_nalt.STRTYPE_PASCAL, }.get(str(kind).lower(), ida_nalt.STRTYPE_C) n = max(int(length), 0) # Free any existing item(s) so create_strlit can carve the literal. ida_bytes.del_items(ea, ida_bytes.DELIT_SIMPLE, n if n > 0 else 1) ok = bool(ida_bytes.create_strlit(ea, n, strtype)) if not ok: return {"addr": addr, "ok": False, "error": "create_strlit failed"} size = int(ida_bytes.get_item_size(ea)) try: raw = ida_bytes.get_strlit_contents(ea, -1, strtype) text = raw.decode("utf-8", "replace") if raw else "" except Exception: text = "" return {"addr": addr, "ok": True, "size": size, "text": text} @tool @idasync def read_raw( addr: Annotated[str, "Start address (hex or name)"], size: Annotated[int, "Number of bytes to read"], ) -> dict: """Read ``size`` bytes at ``addr`` as ONE contiguous lowercase hex string (no per-byte '0x'/spaces). The hot path for the hex view and disasm opcode bytes. Fast: does a single bulk ``ida_bytes.get_bytes`` (C-speed) instead of the per-byte read_bytes_bss_safe loop (2 IDA calls/byte). Unloaded bytes come back from IDA as the 0xFF sentinel, so we only re-check is_loaded for the (usually sparse) 0xFF bytes and zero the genuinely-unloaded ones — matching get_bytes' bss semantics without paying per-byte for the whole range. Encoding is compact hex (~2.5x smaller than get_bytes' '0x..'-with-spaces) and, unlike get_bytes, does not truncate on large reads.""" import ida_bytes ea = parse_address(addr) n = max(int(size), 0) if n == 0: return {"addr": addr, "hex": "", "n": 0} raw = ida_bytes.get_bytes(ea, n) if raw is None or len(raw) < n: # nothing (or not all) mapped base = bytearray(raw or b"") base.extend(b"\\xff" * (n - len(base))) raw = bytes(base) ba = bytearray(raw) # Only unloaded bytes read as 0xFF; correct just those to 0 (bss => zero). i = ba.find(0xFF) while i != -1: if not ida_bytes.is_loaded(ea + i): ba[i] = 0 i = ba.find(0xFF, i + 1) return {"addr": addr, "hex": bytes(ba).hex(), "n": len(ba)} def _idatui_head_row(ea): """One flat-listing row for the head at ``ea``: kind (code/data/unknown), byte size, rendered text, and any symbol name.""" import ida_bytes import ida_lines import ida_name f = ida_bytes.get_flags(ea) if ida_bytes.is_code(f): kind = "code" elif ida_bytes.is_data(f): kind = "data" else: kind = "unknown" line = ida_lines.generate_disasm_line(ea, 0) text = ida_lines.tag_remove(line) if line else "" text = " ".join(text.split()) # collapse IDA's column padding row = { "ea": hex(ea), "kind": kind, "size": int(ida_bytes.get_item_size(ea)), "text": text, } if line: # Keep IDA's own token classification for syntax highlighting. Built from # the SAME line as `text`, then whitespace-collapsed identically so the # two never disagree about what the row says. spans, ops = _idatui_spans(line) joined = "".join(t for _k, t in spans) if " ".join(joined.split()) == text: row["spans"] = spans # Where each operand sits in `text`. Comes out of the same tag walk # (free), and is what lets the client show WHICH literal a keypress # would reformat before you press it. if ops: row["ops"] = ops nm = ida_name.get_ea_name(ea) if nm: row["name"] = nm return row #: IDA colour tag -> the semantic kind the TUI styles. IDA already classifies #: every token in a disassembly line, for every processor it supports, so there #: is nothing to lex: generate_disasm_line emits \x01text\x02 and the #: tag says what the text IS. A pygments assembly lexer would be a worse guess at #: this and would need one dialect per architecture. _IDATUI_SPAN_KINDS = { "insn": ("SCOLOR_INSN", "SCOLOR_KEYWORD", "SCOLOR_ASMDIR", "SCOLOR_MACRO"), "reg": ("SCOLOR_REG",), "num": ("SCOLOR_NUMBER", "SCOLOR_CHAR", "SCOLOR_BINPREF"), "str": ("SCOLOR_STRING",), # NB the real constant names: DATNAME/CODNAME, not "DNAME". Guessing here # fails silently — an unmapped tag renders as plain body text, so symbols # just quietly aren't blue and nothing tells you why. "name": ("SCOLOR_DATNAME", "SCOLOR_CODNAME", "SCOLOR_LOCNAME", "SCOLOR_IMPNAME", "SCOLOR_DEMNAME", "SCOLOR_LIBNAME", "SCOLOR_CNAME", "SCOLOR_DNAME", "SCOLOR_CREF", "SCOLOR_DREF", "SCOLOR_CREFTAIL", "SCOLOR_DREFTAIL"), "seg": ("SCOLOR_SEGNAME",), "cmt": ("SCOLOR_AUTOCMT", "SCOLOR_REGCMT", "SCOLOR_RPTCMT", "SCOLOR_VOIDOP"), "punct": ("SCOLOR_SYMBOL", "SCOLOR_ALTOP", "SCOLOR_HIDNAME"), "err": ("SCOLOR_ERROR",), } def _idatui_tag_map(): """{tag character: kind}, built once from whatever this IDA actually has.""" import ida_lines out = {} for kind, names in _IDATUI_SPAN_KINDS.items(): for n in names: v = getattr(ida_lines, n, None) if isinstance(v, str) and v: out[v[0]] = kind elif isinstance(v, int): out[chr(v)] = kind return out _IDATUI_TAGS = None _IDATUI_OPND_TAGS = None def _idatui_opnd_tag_map(): """{tag character: operand index}. IDA wraps each operand of a disassembly line in COLOR_OPND1..8, so the line already says where operand N starts and ends -- no need to re-render operands with print_operand to find out (and the two agree exactly; checked over thousands of instructions).""" import ida_lines out = {} for i in range(1, 9): v = getattr(ida_lines, "COLOR_OPND%d" % i, None) if isinstance(v, int): out[chr(v)] = i - 1 elif isinstance(v, str) and v: out[v[0]] = i - 1 return out def _idatui_spans(line): """(spans, ops) for a tagged disasm line. ``spans`` is [[kind, text], ...] with colour tags resolved; ``ops`` is [[start, end, n], ...], the extent of each operand in the SAME (collapsed) coordinates the row's ``text`` uses -- which is what lets a cursor column name the operand it is standing on. Unknown tags become 'text' rather than being dropped: a processor module can emit a colour we don't classify, and losing the characters would corrupt the line.""" global _IDATUI_TAGS, _IDATUI_OPND_TAGS import ida_lines if _IDATUI_TAGS is None: _IDATUI_TAGS = _idatui_tag_map() if _IDATUI_OPND_TAGS is None: _IDATUI_OPND_TAGS = _idatui_opnd_tag_map() on, off, esc = "\x01", "\x02", "\x03" addr_tag = chr(getattr(ida_lines, "COLOR_ADDR", 0x28)) addr_len = int(getattr(ida_lines, "COLOR_ADDR_SIZE", 16)) spans, stack, buf = [], [], [] # stack entries: (kind, operand index|None) i, n = 0, len(line) def _opnd(): for _k, o in reversed(stack): if o is not None: return o return None def flush(): if buf: spans.append([stack[-1][0] if stack else "text", "".join(buf), _opnd()]) del buf[:] while i < n: ch = line[i] if ch == on and i + 1 < n: tag = line[i + 1] if tag == addr_tag: # An embedded target address, not display text: 16 hex digits # that must not reach the screen. i += 2 + addr_len continue flush() stack.append((_IDATUI_TAGS.get(tag, "text"), _IDATUI_OPND_TAGS.get(tag))) i += 2 continue if ch == off and i + 1 < n: flush() if stack: stack.pop() i += 2 continue if ch == esc and i + 1 < n: # escaped literal buf.append(line[i + 1]) i += 2 continue buf.append(ch) i += 1 flush() # Collapse IDA's column padding EXACTLY as the plain text does. A run of # spaces can straddle two spans, so this walks characters rather than # collapsing each span on its own — otherwise the spans and `text` disagree # about the line and the row silently loses its highlighting. out, prev_space = [], False for kind, txt, opnd in spans: acc = [] for ch in txt: if ch.isspace(): if prev_space: continue acc.append(" ") prev_space = True else: acc.append(ch) prev_space = False if acc: out.append([kind, "".join(acc), opnd]) while out and out[0][1] == " ": out.pop(0) while out and out[-1][1] == " ": out.pop() if out and out[0][1].startswith(" "): out[0][1] = out[0][1].lstrip() if out and out[-1][1].endswith(" "): out[-1][1] = out[-1][1].rstrip() out = [s for s in out if s[1]] # Operand extents, in the coordinates of the collapsed text these spans # spell out. Adjacent spans of the same operand merge, so an operand like # ``[rbp+var_40]`` (five differently-coloured tokens) comes back as ONE # range -- which is the thing a cursor is inside of, and the thing a format # change applies to. ops, pos, cur, start = [], 0, None, 0 for _kind, txt, opnd in out: if opnd != cur: if cur is not None and pos > start: ops.append([start, pos, cur]) cur, start = opnd, pos pos += len(txt) if cur is not None and pos > start: ops.append([start, pos, cur]) text = "".join(t for _k, t, _o in out) trimmed = [] for lo, hi, k in ops: # don't let a range own trailing space while hi > lo and text[hi - 1].isspace(): hi -= 1 while lo < hi and text[lo].isspace(): lo += 1 if hi > lo: trimmed.append([lo, hi, k]) return [[k, t] for k, t, _o in out], trimmed def _idatui_unknown_row(ea, size): """One collapsed row for a run of ``size`` undefined bytes starting at ``ea``. A single byte is rendered normally (shows its value); a longer run collapses to ``db N dup(?)`` so a big .bss/gap doesn't explode into millions of one-byte rows.""" import ida_name if size <= 1: return _idatui_head_row(ea) row = {"ea": hex(ea), "kind": "unknown", "size": int(size), "text": f"db {size} dup(?)"} nm = ida_name.get_ea_name(ea) if nm: row["name"] = nm return row def _idatui_struct_member_rows(ea): """Indented member rows for a struct-typed data item at ``ea`` (expansion), or [] if it isn't a struct. Top-level fields only.""" import ida_nalt import ida_typeinf import idaapi tif = ida_typeinf.tinfo_t() if not (ida_nalt.get_tinfo(tif, ea) and tif.is_udt()): return [] udt = ida_typeinf.udt_type_data_t() if not tif.get_udt_details(udt): return [] rows = [] for m in udt: off = m.begin() // 8 try: mtype = m.type._print() or "" except Exception: mtype = "" try: sz = int(m.type.get_size()) if sz == idaapi.BADSIZE: sz = 0 except Exception: sz = 0 name = m.name or "" text = f"+{off:X} {name}" + (f" {mtype}" if mtype else "") rows.append({"ea": hex(ea + off), "kind": "member", "size": sz, "text": text}) return rows def _idatui_func_header_rows(ea): """IDA-style subroutine banner rows shown just before a function's entry.""" import ida_funcs name = ida_funcs.get_func_name(ea) or "sub_%X" % ea bar = "=" * 15 + " S U B R O U T I N E " + "=" * 15 return [ {"ea": hex(ea), "kind": "sep", "size": 0, "text": ""}, {"ea": hex(ea), "kind": "sep", "size": 0, "text": "; " + bar}, {"ea": hex(ea), "kind": "funchdr", "size": 0, "text": name + " proc", "name": name}, ] def _idatui_func_footer_rows(ea, func): """End-of-function marker shown just after a function's last item.""" import ida_funcs name = ida_funcs.get_func_name(func.start_ea) or "sub_%X" % func.start_ea return [ {"ea": hex(ea), "kind": "funchdr", "size": 0, "text": name + " endp", "name": name}, {"ea": hex(ea), "kind": "sep", "size": 0, "text": "; " + "-" * 60}, ] @tool @idasync def heads( addr: Annotated[str, "Start address or name to walk from"], count: Annotated[int, "Max heads to return (default 200, max 2000)"] = 200, offset: Annotated[int, "Skip first N heads from addr (default 0)"] = 0, end: Annotated[str, "Optional exclusive end address; default = segment end"] = "", back: Annotated[bool, "Walk backwards: return the count heads ENDING just before addr, in forward order"] = False, annotate: Annotated[bool, "Emit IDA-style function boundary banner rows (kind sep/funchdr)"] = False, ) -> dict: """Walk item heads from ``addr`` as a flat listing: every head is rendered (code OR data OR undefined) via generate_disasm_line and stepped with next_head/prev_head. Unlike ``disasm`` (code-only, bails at the first data byte) this shows db/dw/dd/... lines for data and undefined regions — IDA's real disassembly view. Address-paged: page forward by re-calling with ``addr`` = the returned cursor.next; page up with ``back=true``.""" import ida_bytes import ida_segment import idaapi count = 2000 if count > 2000 else (1 if count < 1 else count) offset = max(int(offset), 0) try: start = parse_address(addr) except Exception as e: return {"addr": str(addr), "error": str(e), "heads": [], "cursor": {"done": True}} seg = ida_segment.getseg(start) if not seg: return {"addr": str(addr), "error": "no segment", "heads": [], "cursor": {"done": True}} lo, hi = seg.start_ea, seg.end_ea if end: try: hi = min(hi, parse_address(end)) except Exception: pass rows = [] if back: # Collect up to (count+offset) heads strictly before `start`, then take # the window closest to `start`, returned in forward order. walk = [] cur = ida_bytes.prev_head(start, lo) while cur != idaapi.BADADDR and cur >= lo and len(walk) < count + offset: walk.append(cur) cur = ida_bytes.prev_head(cur, lo) walk.reverse() chosen = walk[: len(walk) - offset] if offset else walk chosen = chosen[-count:] rows = [_idatui_head_row(e) for e in chosen] first = chosen[0] if chosen else start pea = ida_bytes.prev_head(first, lo) cursor = {"done": True} if pea == idaapi.BADADDR or pea < lo else {"prev": hex(pea)} return {"addr": str(addr), "heads": rows, "cursor": cursor} # Walk by item END (not next_head): next_head SKIPS undefined bytes, but a # flat listing must show them (IDA renders undefined as `db ?` lines, and # navigating to an unmarked address must land ON it). Defined items advance # by get_item_end; a run of undefined bytes is COLLAPSED into one row (its # end found in O(1) via next_head, which skips undefined) so a large .bss or # gap doesn't explode into millions of one-byte rows. def _is_unknown(e): f = ida_bytes.get_flags(e) return not (ida_bytes.is_code(f) or ida_bytes.is_data(f)) def _run_end(e): """End (exclusive) of the undefined run starting at ``e``.""" nh = ida_bytes.next_head(e, hi) return nh if (nh != idaapi.BADADDR and e < nh <= hi) else hi def _advance(e): if _is_unknown(e): return _run_end(e) nxt = ida_bytes.get_item_end(e) return nxt if nxt > e else e + 1 def _rows_for(e): if _is_unknown(e): return [_idatui_unknown_row(e, _run_end(e) - e)] func = idaapi.get_func(e) if annotate else None at_start = func is not None and func.start_ea == e out = [] if at_start: out.extend(_idatui_func_header_rows(e)) row = _idatui_head_row(e) if at_start: row = dict(row) row["name"] = None # the name is shown on the proc header line elif annotate and row.get("kind") == "code" and row.get("name"): # A code label (loc_XXX/jump target) gets its OWN line at depth 0, # like IDA; strip it from the instruction row below. nm = row["name"] out.append({"ea": hex(e), "kind": "label", "size": 0, "text": nm + ":", "name": nm}) row = dict(row) row["name"] = None out.append(row) if row.get("kind") == "data": out.extend(_idatui_struct_member_rows(e)) # expand struct fields if func is not None and ida_bytes.get_item_end(e) >= func.end_ea: out.extend(_idatui_func_footer_rows(e, func)) return out ea = ida_bytes.get_item_head(start) for _ in range(offset): if ea >= hi or ea == idaapi.BADADDR: break ea = _advance(ea) more = False while ea != idaapi.BADADDR and ea < hi: if len(rows) >= count: more = True break rows.extend(_rows_for(ea)) # a struct head expands into member rows ea = _advance(ea) cursor = {"next": hex(ea)} if more else {"done": True} return {"addr": str(addr), "heads": rows, "cursor": cursor} @tool @idasync def xref_types( queries: Annotated[list, "[{addr, direction:'to'|'from'|'both', include_fn, dedup, count}]"], ) -> dict: """Like xref_query, but every row carries a fine-grained ``kind`` derived from the IDA xref type \u2014 call/jump/flow for code, read/write/offset/text/info for data \u2014 alongside the coarse ``type`` (code/data). Feeds the xref dialog's r/w/call badges. Same query/envelope shape as xref_query.""" import idaapi, idautils, ida_funcs, ida_bytes, ida_xref code_kind = {ida_xref.fl_CF: "call", ida_xref.fl_CN: "call", ida_xref.fl_JF: "jump", ida_xref.fl_JN: "jump", ida_xref.fl_F: "flow"} data_kind = {ida_xref.dr_O: "offset", ida_xref.dr_W: "write", ida_xref.dr_R: "read", ida_xref.dr_T: "text", ida_xref.dr_I: "info"} def _kind(xr): table = code_kind if xr.iscode else data_kind return table.get(xr.type, "code" if xr.iscode else "data") def _fn(ea): f = ida_funcs.get_func(ea) if not f: return None return {"addr": hex(f.start_ea), "name": ida_funcs.get_func_name(f.start_ea)} def _resolve(raw): raw = str(raw).strip() try: return int(raw, 16) # handles '0x2490' and '2490' except ValueError: return idaapi.get_name_ea(idaapi.BADADDR, raw) qs = queries if isinstance(queries, list) else [queries] result = [] for q in qs: q = q if isinstance(q, dict) else {"addr": q} raw = str(q.get("addr", "")).strip() direction = str(q.get("direction", "to") or "to").lower() include_fn = bool(q.get("include_fn", True)) dedup = bool(q.get("dedup", True)) try: count = int(q.get("count", 2000) or 2000) except (TypeError, ValueError): count = 2000 target = _resolve(raw) rows = [] if target is not None and target != idaapi.BADADDR and ida_bytes.is_mapped(target): if direction in ("to", "both"): for xr in idautils.XrefsTo(target, 0): row = {"direction": "to", "addr": hex(xr.frm), "from": hex(xr.frm), "to": hex(target), "type": "code" if xr.iscode else "data", "kind": _kind(xr)} if include_fn: row["fn"] = _fn(xr.frm) rows.append(row) if direction in ("from", "both"): for xr in idautils.XrefsFrom(target, 0): row = {"direction": "from", "addr": hex(xr.to), "from": hex(target), "to": hex(xr.to), "type": "code" if xr.iscode else "data", "kind": _kind(xr)} if include_fn: row["fn"] = _fn(xr.to) rows.append(row) if dedup: seen = set() deduped = [] for r in rows: k = (r["direction"], r["from"], r["to"], r["kind"]) if k in seen: continue seen.add(k) deduped.append(r) rows = deduped rows = rows[:count] result.append({"query": raw, "data": rows, "next_offset": None}) return {"result": result} @tool @idasync def data_type( addr: Annotated[str, "Address or name of a data item / global"], ) -> dict: """The current C type of a data item, for prefilling a retype prompt: {addr, name, type, size, is_func}. ``type`` is empty when the item is untyped; ``is_func`` distinguishes a global from a function so the caller knows which flavour of set_type to use.""" import idaapi import ida_bytes import ida_name import idc raw = str(addr).strip() try: ea = int(raw, 16) except ValueError: ea = idaapi.get_name_ea(idaapi.BADADDR, raw) if ea == idaapi.BADADDR or not ida_bytes.is_mapped(ea): return {"addr": raw, "error": f"not a mapped address: {raw}"} return { "addr": hex(ea), "name": ida_name.get_name(ea) or "", "type": idc.get_type(ea) or "", "size": int(ida_bytes.get_item_size(ea) or 0), "is_func": bool(idaapi.get_func(ea)), } @tool @idasync def decomp_map( addr: Annotated[str, "Function address or name"], ) -> dict: """Per-pseudocode-line instruction coverage for the split view's region highlight: for each line, the set of EAs the decompiler attributes to it, swept across the line's columns via get_line_item. Shape: {addr, lines:[{ea: primary|None, eas:[hex,...]}, ...]}.""" import ida_hexrays import idaapi try: ea = int(str(addr), 16) except ValueError: ea = idaapi.get_name_ea(idaapi.BADADDR, str(addr).strip()) func = idaapi.get_func(ea) if not func: return {"error": f"no function at {addr}"} try: cfunc = ida_hexrays.decompile(func.start_ea) except Exception as e: # noqa: BLE001 return {"error": f"decompile failed: {e}"} if cfunc is None: return {"error": "decompile failed"} lines = [] for sl in cfunc.get_pseudocode(): line = sl.line eas, seen = [], set() for x in range(len(line) + 1): head = ida_hexrays.ctree_item_t() item = ida_hexrays.ctree_item_t() tail = ida_hexrays.ctree_item_t() if not cfunc.get_line_item(line, x, False, head, item, tail): continue # Match the /*ea*/ marker's source (decompile_function_safe): the # item's dstr() is 'EA: description'; get_ea() reports a different ea. dstr = item.dstr() if not dstr: continue parts = dstr.split(": ", 1) if len(parts) != 2: continue try: e = int(parts[0], 16) except ValueError: continue if e not in seen: seen.add(e) eas.append(hex(e)) lines.append({"ea": eas[0] if eas else None, "eas": eas}) return {"addr": hex(func.start_ea), "lines": lines} _idatui_strings_cache = {} def _idatui_build_strings(min_len): """[(ea, text, length, typename)] for every string IDA found, cached by min_len (rebuilding the list is O(n) and the browser pages through it).""" import idautils import ida_nalt hit = _idatui_strings_cache.get(min_len) if hit is not None: return hit tnames = {} for nm, lbl in (("STRTYPE_C", "C"), ("STRTYPE_C_16", "utf16"), ("STRTYPE_C_32", "utf32"), ("STRTYPE_PASCAL", "pascal")): v = getattr(ida_nalt, nm, None) if v is not None: tnames[v & 0xFF] = lbl items = [] for s in idautils.Strings(): if s is None: continue try: text = str(s) except Exception: # noqa: BLE001 -- undecodable literal continue if len(text) < min_len: continue st = getattr(s, "strtype", 0) & 0xFF items.append((s.ea, text, getattr(s, "length", len(text)), tnames.get(st, "t%d" % st))) _idatui_strings_cache[min_len] = items return items @tool @idasync def list_strings( offset: Annotated[int, "Start index into the strings list"] = 0, count: Annotated[int, "Max strings to return (page size)"] = 2000, min_len: Annotated[int, "Minimum string length to include"] = 4, refresh: Annotated[bool, "Rebuild the cached strings list"] = False, ) -> dict: """Every string literal IDA found in the binary (IDA's Shift+F12 window), paginated: {strings:[{addr,text,len,type}], total, next_offset}. Feeds the TUI's strings browser.""" try: min_len = max(int(min_len), 1) except (TypeError, ValueError): min_len = 4 try: offset = max(int(offset), 0) except (TypeError, ValueError): offset = 0 try: count = max(int(count), 1) except (TypeError, ValueError): count = 2000 if refresh: _idatui_strings_cache.pop(min_len, None) items = _idatui_build_strings(min_len) page = items[offset:offset + count] return { "strings": [{"addr": hex(ea), "text": text, "len": ln, "type": ty} for (ea, text, ln, ty) in page], "total": len(items), "next_offset": offset + len(page), } @tool @idasync def list_linkage( kind: Annotated[str, "'import', 'export' or 'both'"] = "both", ) -> dict: """What this binary imports from, and exports to, other modules: {imports:[{addr,name,module}], exports:[{addr,name,ordinal}]}. Feeds the project-wide import/export join, which resolves a PLT stub in one binary to the real implementation in another.""" import idaapi import idautils import ida_nalt want = str(kind or "both").lower() imports = [] exports = [] if want in ("import", "both"): n = ida_nalt.get_import_module_qty() for i in range(n): mod = ida_nalt.get_import_module_name(i) or "" def _cb(ea, name, ordinal, _mod=mod): # An ordinal-only import has no name; skip rather than invent one. if name: imports.append({"addr": hex(ea), "name": name, "module": _mod}) return True ida_nalt.enum_import_names(i, _cb) if want in ("export", "both"): for index, ordinal, ea, name in idautils.Entries(): if name: exports.append({"addr": hex(ea), "name": name, "ordinal": int(ordinal)}) return {"imports": imports, "exports": exports, "n_imports": len(imports), "n_exports": len(exports)} @tool @idasync def define_code_run( addr: Annotated[str, "Address to start disassembling from"], limit: Annotated[int, "Max instructions to create (safety stop)"] = 20000, ) -> dict: """Disassemble CONSECUTIVELY from ``addr`` until something stops it, the way IDA's 'c' does — one instruction is rarely what you want when carving a raw image. Returns {start,end,count,stopped} where ``stopped`` says why: 'undecodable' (bytes aren't an instruction), 'flow' (the last instruction doesn't fall through, e.g. RET/B), 'defined' (ran into existing code/data), 'segment' (hit the end) or 'limit'. Runs in-process: doing this from the client would be one round trip per instruction, which is minutes on a real firmware image.""" import ida_bytes import ida_idp import ida_segment import ida_ua import idaapi try: ea = parse_address(addr) except Exception as e: return {"addr": str(addr), "error": str(e), "count": 0} seg = ida_segment.getseg(ea) if not seg: return {"addr": str(addr), "error": "no segment", "count": 0} hi = seg.end_ea try: limit = max(1, min(int(limit), 200000)) except (TypeError, ValueError): limit = 20000 start, count, stopped = ea, 0, "limit" while count < limit: if ea >= hi: stopped = "segment" break flags = ida_bytes.get_flags(ea) if ida_bytes.is_code(flags) or ida_bytes.is_data(flags): # Already defined: stop rather than clobber. Undefining someone's # existing work to keep a speculative run going is not a trade the # user asked for. stopped = "defined" break n = ida_ua.create_insn(ea) if n <= 0: stopped = "undecodable" break count += 1 # Stop where control flow stops. Past a RET the next bytes are usually # padding or a new function's data, and running on turns a clean carve # into a mess that has to be undone by hand. # # Ask ida_idp.is_ret_insn, NOT the canonical feature bits: on AArch64 # get_canon_feature() returns 0 for RET, so a CF_STOP test silently never # fires and the run walks straight through the end of the routine. insn = ida_ua.insn_t() if ida_ua.decode_insn(insn, ea) > 0: try: is_ret = ida_idp.is_ret_insn(insn) except Exception: is_ret = False if is_ret or (insn.get_canon_feature() & idaapi.CF_STOP): ea += n stopped = "flow" break ea += n return {"start": hex(start), "end": hex(ea), "count": count, "stopped": stopped} @tool @idasync def set_thumb( addr: Annotated[str, "Address to change the ARM decoding mode at"], mode: Annotated[str, "'toggle', 'on' (Thumb) or 'off' (ARM)"] = "toggle", end: Annotated[str, "Optional exclusive end address (default: this item)"] = "", ) -> dict: """Switch ARM/Thumb decoding at ``addr`` (IDA's T segment register). Thumb is not a property of the bytes, it's a mode the CPU is in, so a raw image gives IDA no way to know: at a Thumb entry point it decodes 16-bit instructions as 32-bit ARM and produces confident nonsense (``push {r3,lr}`` reads as ``SVCLT 0xBF00``). Also forces the segment to 32-bit when turning Thumb ON. Thumb does not exist in AArch64, and a headerless blob loaded with -parm defaults to 64-bit — so setting T alone changes nothing and looks broken. Asking for Thumb IS asking for ARM32.""" import ida_bytes import ida_idp import ida_segment import ida_segregs try: ea = parse_address(addr) except Exception as e: return {"addr": str(addr), "error": str(e)} treg = ida_idp.str2reg("T") if treg is None or treg < 0: return {"addr": hex(ea), "error": "no T register (not an ARM database)"} seg = ida_segment.getseg(ea) if not seg: return {"addr": hex(ea), "error": "no segment"} import ida_ida db64 = ida_ida.inf_get_app_bitness() == 64 cur = ida_segregs.get_sreg(ea, treg) cur = 0 if cur in (None, 0xFFFFFFFF, -1) else int(cur) want = {"on": 1, "off": 0}.get(str(mode).lower(), 0 if cur else 1) changed_bits = False if want and seg.bitness != 1: ida_segment.set_segm_addressing(seg, 1) changed_bits = True try: stop = parse_address(end) if end else 0 except Exception: stop = 0 size = max(int(stop) - ea, 0) or max(ida_bytes.get_item_size(ea), 2) # The bytes are currently decoded in the OLD mode; leaving that item defined # pins the wrong instruction length and the new mode has nothing to apply to. ida_bytes.del_items(ea, 0, size) ok = bool(ida_segregs.split_sreg_range(ea, treg, want, ida_segregs.SR_user)) now = ida_segregs.get_sreg(ea, treg) return {"addr": hex(ea), "thumb": bool(now), "was": bool(cur), "ok": ok, "bitness": ida_segment.getseg(ea).bitness, "forced_32bit": changed_bits, # The DATABASE's bitness is fixed at load and can't be corrected # here (setting it post-hoc makes the decompiler INTERR). In a # 64-bit database a 32-bit function disassembles but Hex-Rays # refuses it outright, so say so instead of leaving the user to # discover that F5 does nothing. "db_64bit": bool(db64 and want)} def _idatui_add_func(ea): """add_func at ``ea``, falling back to an explicit end. ida_funcs.add_func(ea) asks IDA to find the end and on carved or freshly-marked code it often can't, failing with no reason given.""" import ida_bytes import ida_funcs import ida_segment import idaapi if idaapi.get_func(ea) is not None: return True if ida_funcs.add_func(ea): return True seg = ida_segment.getseg(ea) hi = seg.end_ea if seg else ea end = ea while end < hi and ida_bytes.is_code(ida_bytes.get_flags(end)): nxt = ida_bytes.get_item_end(end) if nxt <= end: break end = nxt return bool(end > ea and ida_funcs.add_func(ea, end)) @tool @idasync def define_func_run( addr: Annotated[str, "Entry point of the function to create"], ) -> dict: """Create a function at ``addr``, working out its end if IDA can't. ida_funcs.add_func(ea) asks IDA to find the end itself, and on hand-carved code it often can't — a run that ends in a tail call, or whose last instruction isn't recognised as a return, simply fails with no reason given. You then have a disassembled routine that refuses to become a function, and F5 has nothing to work with. So: try IDA's way, and if that fails, use the end of the contiguous instruction run starting at ``addr``.""" import ida_bytes import ida_funcs import ida_segment import idaapi try: ea = parse_address(addr) except Exception as e: return {"addr": str(addr), "error": str(e), "ok": False} fn = idaapi.get_func(ea) if fn is not None and fn.start_ea == ea: return {"addr": hex(ea), "ok": True, "start": hex(fn.start_ea), "end": hex(fn.end_ea), "how": "existed"} auto = ida_funcs.add_func(ea) if not auto and not _idatui_add_func(ea): return {"addr": hex(ea), "ok": False, "error": f"IDA refused a function at {ea:#x}"} f = idaapi.get_func(ea) if f is None: return {"addr": hex(ea), "ok": False, "error": "function did not stick"} return {"addr": hex(ea), "ok": True, "start": hex(f.start_ea), "end": hex(f.end_ea), "how": "auto" if auto else "explicit-end"} @tool @idasync def decomp_error( addr: Annotated[str, "Address of the function that failed to decompile"], ) -> dict: """Why Hex-Rays refused this function, in its own words. The plain decompile tool reports "Decompilation failed at 0x0" and drops the reason, which is the only useful part. Hex-Rays fills in a hexrays_failure_t saying things like "only 64-bit functions can be decompiled in the current database" — that one is unfixable in place (the database's bitness is set at load), so a user who can't see it has no way to know they must reload.""" import ida_funcs import ida_hexrays import ida_ida try: ea = parse_address(addr) except Exception as e: return {"addr": str(addr), "error": str(e)} out = {"addr": hex(ea), "bitness": ida_ida.inf_get_app_bitness()} fn = ida_funcs.get_func(ea) if fn is None: out["reason"] = "no function here" return out try: if not ida_hexrays.init_hexrays_plugin(): out["reason"] = "the decompiler is not available for this processor" return out hf = ida_hexrays.hexrays_failure_t() cf = ida_hexrays.decompile_func(fn, hf) if cf is not None: out["reason"] = "" # it decompiles now return out out["reason"] = hf.desc() or f"error {hf.code}" out["code"] = int(hf.code) out["errea"] = hex(hf.errea) except Exception as e: # noqa: BLE001 out["reason"] = f"{type(e).__name__}: {e}" return out @tool @idasync def thumb_scan( start: Annotated[str, "Start of the range to scan for entry pointers"] = "", end: Annotated[str, "Exclusive end of the range (default: 1KB from start)"] = "", apply: Annotated[bool, "Mark the targets as Thumb and disassemble them"] = True, limit: Annotated[int, "Max entries to act on"] = 512, ) -> dict: """Find Thumb entry points from ODD pointers, e.g. a Cortex-M vector table. An ARM function pointer carries the mode in bit 0: odd means Thumb. A vector table is therefore a list of Thumb entry points that IDA won't follow on a headerless image, because nothing tells it those words are pointers at all. Being wrong here is expensive — marking a data word as code corrupts the listing — so a word only counts when it is odd, lands inside a loaded segment, and its target is EXECUTABLE and not already defined as data. The even words in a vector table (the initial stack pointer) fail the first test, which is the point.""" import ida_bytes import ida_funcs import ida_idp import ida_segment import ida_segregs import ida_ua seg0 = ida_segment.getseg(parse_address(start)) if start else None if seg0 is None: seg0 = ida_segment.getnseg(0) if seg0 is None: return {"error": "no segments", "found": [], "applied": 0} try: lo = parse_address(start) if start else seg0.start_ea hi = parse_address(end) if end else min(lo + 0x400, seg0.end_ea) except Exception as e: return {"error": str(e), "found": [], "applied": 0} treg = ida_idp.str2reg("T") found, applied = [], 0 ea = lo while ea + 4 <= hi and len(found) < limit: w = ida_bytes.get_dword(ea) ea += 4 if not (w & 1): continue # even: not a Thumb pointer tgt = w & ~1 seg = ida_segment.getseg(tgt) if seg is None or not (seg.perm & ida_segment.SEGPERM_EXEC or seg.perm == 0): continue # points outside the image, or at data f = ida_bytes.get_flags(tgt) if ida_bytes.is_data(f): continue # already something else; don't fight it rec = {"at": hex(ea - 4), "value": hex(w), "target": hex(tgt), "was_code": bool(ida_bytes.is_code(f))} found.append(rec) if not apply: continue if treg is not None and treg >= 0: ida_segregs.split_sreg_range(tgt, treg, 1, ida_segregs.SR_user) if not ida_bytes.is_code(ida_bytes.get_flags(tgt)): ida_bytes.del_items(tgt, 0, 2) if ida_ua.create_insn(tgt) <= 0: rec["decoded"] = False continue rec["decoded"] = True rec["function"] = _idatui_add_func(tgt) applied += 1 return {"start": hex(lo), "end": hex(hi), "found": found, "applied": applied, "n": len(found)} # --------------------------------------------------------------------------- # # operand display formats (IDA's 'o' family: hex / dec / char / offset / ...) # --------------------------------------------------------------------------- # #: The stops a cycle walks, in order, before filtering to the ones that make #: sense for the operand in hand. Octal is deliberately NOT one of them -- every #: extra stop is another keypress and nobody reads octal -- but it is still #: reachable by name. "default" hands the operand back to IDA's own choice, #: which for data is how you get an auto-detected offset/string back. _IDATUI_FMT_CYCLE = ("hex", "dec", "bin", "char", "offset", "default") #: Formats we can re-apply from a name alone. enum/stroff/custom carry an id #: (which enum, which struct) that a nibble doesn't record, so they are never #: cycled INTO -- and cycling out of one is called out in ``warn``. _IDATUI_FMT_SETTABLE = ("hex", "dec", "oct", "bin", "char", "offset", "seg", "float", "stack", "default") def _idatui_fmt_nibbles(): """{format name: IDA operand-type nibble}. Built on call, not at import: this module is injected into a file that is imported before a database is open.""" import ida_bytes return { "default": ida_bytes.FF_N_VOID, "hex": ida_bytes.FF_N_NUMH, "dec": ida_bytes.FF_N_NUMD, "char": ida_bytes.FF_N_CHAR, "seg": ida_bytes.FF_N_SEG, "offset": ida_bytes.FF_N_OFF, "bin": ida_bytes.FF_N_NUMB, "oct": ida_bytes.FF_N_NUMO, "enum": ida_bytes.FF_N_ENUM, "forced": ida_bytes.FF_N_FOP, "stroff": ida_bytes.FF_N_STRO, "stack": ida_bytes.FF_N_STK, "float": ida_bytes.FF_N_FLT, "custom": ida_bytes.FF_N_CUST, } def _idatui_fmt_name(nib): for name, v in _idatui_fmt_nibbles().items(): if v == nib: return name return "default" def _idatui_op_fmt(ea, n): """The format operand ``n`` of the item at ``ea`` is currently displayed in. Reads the nibble IDA keeps per operand rather than guessing from the text -- ``1`` renders identically in hex and decimal, so the rendered line cannot answer this.""" import ida_bytes F = ida_bytes.get_flags(ea) nib = (F >> ida_bytes.get_operand_type_shift(int(n))) & 0xF return _idatui_fmt_name(nib) def _idatui_op_value(ea, n): """(value, byte width) of operand ``n``, or (None, 0) if it hasn't got one. The value is what decides which formats are OFFERED: a character constant for 0x38A9 or an offset to an unmapped address are stops worth skipping.""" import ida_bytes import ida_ua F = ida_bytes.get_flags(ea) if ida_bytes.is_code(F): insn = ida_ua.insn_t() if ida_ua.decode_insn(insn, ea) <= 0: return None, 0 try: op = insn.ops[int(n)] except Exception: return None, 0 if op.type == ida_ua.o_void: return None, 0 v = op.value if op.type == ida_ua.o_imm else op.addr try: size = int(ida_ua.get_dtype_size(op.dtype)) except Exception: size = 0 return int(v), size size = int(ida_bytes.get_item_size(ea)) read = {1: ida_bytes.get_byte, 2: ida_bytes.get_word, 4: ida_bytes.get_dword, 8: ida_bytes.get_qword}.get(size) if read is None: return None, size try: return int(read(ea)), size except Exception: return None, size def _idatui_printable(v): """Whether ``v`` would actually render as a character constant. IDA accepts op_chr on anything and then prints the number anyway, so a cycle that offers 'char' for 0x18 has a stop where nothing visibly happens.""" if v is None or v < 0 or v > 0xFFFFFFFF: return False bs, x = [], int(v) while True: bs.append(x & 0xFF) x >>= 8 if not x: break return all(0x20 <= b <= 0x7E or b in (9, 10, 13) for b in bs) def _idatui_offset_worth(v): """Whether 'offset' is worth OFFERING as a cycle stop for value ``v``. Making an offset is not free: IDA invents a dummy name at the target (``off_18``) and that name STAYS once you cycle past it. So the ring only stops there when the target is already something you could name -- a symbol, a function, or an item something else references. In a PIE at base 0 half the small constants in a function are 'mapped' (they land in the ELF header); ``sub rsp, 18h`` is not a reference and must not offer to become one on the way past. An explicit request still converts anything mapped: that's a decision, not a keypress that happened to land here. After it, the target HAS a name, so the ring includes the stop from then on.""" import ida_bytes import ida_name return bool(v and ida_bytes.is_mapped(v) and ida_name.get_ea_name(v)) def _idatui_op_candidates(ea): """Operand indices at ``ea`` whose display format is worth changing. Immediates and displacements -- the literals. Deliberately NOT: * branch targets (o_near/o_far), or every jump on the listing would offer to become a bare number, on a view you navigate by label; * memory references (o_mem), e.g. x86-64's RIP-relative ``lea rdi, name``. IDA prints those from the reference, not from the operand's number format, so setting one is accepted and changes nothing on screen -- a keypress that appears to do nothing is worse than one that says it can't. An explicit ``n`` still reaches them; this is what a bare cursor picks.""" import ida_bytes import ida_ua F = ida_bytes.get_flags(ea) if ida_bytes.is_data(F): return [0] # a data item's value is operand 0 if not ida_bytes.is_code(F): return [] # undefined bytes: IDA refuses a format outright insn = ida_ua.insn_t() if ida_ua.decode_insn(insn, ea) <= 0: return [] want = (ida_ua.o_imm, ida_ua.o_displ) out = [] for i in range(len(insn.ops)): op = insn.ops[i] if op.type == ida_ua.o_void: break if op.type in want: out.append(i) return out def _idatui_op_spans(ea, text): """[(start, end, n)] -- where each operand sits inside ``text`` (the whitespace-collapsed line the TUI shows), so a cursor column can name the operand it is standing on. Read out of IDA's own COLOR_OPND markers on the line, which is both free (the line is generated anyway) and exact. print_operand is kept as a fallback for a processor module that emits no operand markers -- it agrees with the tags where both exist, but it re-renders every operand to say so. """ import ida_lines import ida_ua line = ida_lines.generate_disasm_line(ea, 0) if line: _spans, ops = _idatui_spans(line) if ops: return [tuple(o) for o in ops] out, pos = [], 0 for n in range(8): try: raw = ida_ua.print_operand(ea, n) except Exception: raw = None if not raw: continue op = " ".join(ida_lines.tag_remove(raw).split()) if not op: continue i = text.find(op, pos) if i < 0: # duplicated operand text (mov eax, eax) i = text.find(op) if i < 0: continue out.append((i, i + len(op), n)) pos = i + len(op) return out def _idatui_line_text(ea): import ida_lines line = ida_lines.generate_disasm_line(ea, 0) return " ".join(ida_lines.tag_remove(line).split()) if line else "" def _idatui_op_text(ea, text, n): """How operand ``n`` reads on the line, for a message that names it.""" for lo, hi, i in _idatui_op_spans(ea, text): if i == int(n): return text[lo:hi].strip() return "" def _idatui_apply_fmt(ea, n, fmt): """Set operand ``n``'s display format. Returns (ok, error).""" import ida_bytes import ida_offset import idaapi n = int(n) if fmt == "default": return bool(ida_bytes.clr_op_type(ea, n)), "" if fmt == "offset": base = ida_offset.calc_offset_base(ea, n) if base in (idaapi.BADADDR, None) or base < 0: base = 0 return bool(ida_offset.op_plain_offset(ea, n, base)), "" fn = {"hex": ida_bytes.op_hex, "dec": ida_bytes.op_dec, "oct": ida_bytes.op_oct, "bin": ida_bytes.op_bin, "char": ida_bytes.op_chr, "seg": ida_bytes.op_seg, "float": ida_bytes.op_flt, "stack": ida_bytes.op_stkvar}.get(fmt) if fn is None: return False, (f"can't set {fmt!r} from a name alone" if fmt in _idatui_fmt_nibbles() else f"unknown format {fmt!r}") return bool(fn(ea, n)), "" @tool @idasync def op_format( addr: Annotated[str, "Address of the instruction or data item"], mode: Annotated[str, "cycle | back | show | hex | dec | oct | bin | char | offset | stack | default"] = "cycle", col: Annotated[int, "Cursor column inside the rendered line (-1: first literal)"] = -1, n: Annotated[int, "Operand index; -1 derives it from ``col``"] = -1, ) -> dict: """Change how a literal is DISPLAYED (IDA's 'o' family): hex, decimal, binary, character, or an offset to the address it names. The value in the bytes never changes -- only the representation IDA renders and remembers. ``cycle``/``back`` step the stops that make sense for THIS operand: 'char' is skipped unless the value prints as one, 'offset' unless the target is already named, so no press is ever a no-op you have to press again. ``show`` reports without changing anything. A format the ring can't hold (a stack variable, an enum) is reported in ``warn`` on the way out, with what to do about it -- ``mode`` takes any of the names above outright, which is also how you put one back. Which operand: ``n`` if given, else the one under ``col`` (a column in the whitespace-collapsed line, as ``heads`` renders it), else the first literal on the line.""" import ida_bytes try: ea = ida_bytes.get_item_head(parse_address(addr)) except Exception as e: return {"addr": str(addr), "error": str(e)} before = _idatui_line_text(ea) cands = _idatui_op_candidates(ea) n = int(n) if n < 0: n = -1 if int(col) >= 0: for lo, hi, i in _idatui_op_spans(ea, before): if not (lo <= int(col) < hi): continue if i in cands: n = i break # The cursor IS on an operand, just not one with a format. The # client highlights what the cursor is on, so quietly moving to # a different operand would make that highlight a lie -- say # which one can be changed instead. where = before[lo:hi].strip() alt = (f"; the literal on this line is operand {cands[0]} " f"({_idatui_op_text(ea, before, cands[0])})" if cands else "") return {"addr": hex(ea), "n": i, "text": before, "error": f"operand {i} ({where}) has no format to " f"change{alt}"} if n < 0: if not cands: F = ida_bytes.get_flags(ea) why = ("no literal on this line to reformat" if ida_bytes.is_code(F) or ida_bytes.is_data(F) else "undefined bytes have no format to change -- define " "them first ('d' makes data, 'c' makes code)") return {"addr": hex(ea), "text": before, "error": why} n = cands[0] cur = _idatui_op_fmt(ea, n) value, width = _idatui_op_value(ea, n) mapped = value is not None and value != 0 and ida_bytes.is_mapped(value) # The ring is a property of the OPERAND, not of what you last pressed: every # stop is one that changes what you see for this value, and it is the same # ring at every step, so a lap always comes home. choices = [f for f in _IDATUI_FMT_CYCLE if (f != "char" or _idatui_printable(value)) and (f != "offset" or _idatui_offset_worth(value))] # A stack variable is deliberately NOT a stop: ``[rbp+var_40]`` is a frame # member, not a way of writing a number, and IDA's own "is this a stack # variable" test isn't exposed to Python here (calc_stkvar_struc_offset # happily answers for ``[r14+8]`` too, which would put a bogus stop in the # ring). Leaving one is reported instead, with the command that undoes it. lossy = cur not in choices and cur != "default" mode = str(mode or "cycle").lower() if mode == "show": return {"addr": hex(ea), "n": n, "format": cur, "prev": cur, "choices": choices, "text": before, "before": before, "value": None if value is None else hex(value), "width": width, "applied": False} if mode in ("cycle", "back"): step = 1 if mode == "cycle" else -1 if cur in choices: want = choices[(choices.index(cur) + step) % len(choices)] else: # Standing on a format the ring can't hold (an enum names a type a # nibble doesn't record): enter the ring at its end, don't skip a # stop working out where we "would have" been. want = choices[0] if step > 0 else choices[-1] else: want = mode if want not in _idatui_fmt_nibbles(): return {"addr": hex(ea), "n": n, "text": before, "error": f"unknown format {mode!r}; one of " + ", ".join(_IDATUI_FMT_SETTABLE)} if want == "offset" and not mapped: return {"addr": hex(ea), "n": n, "text": before, "format": cur, "error": (f"{'0x%x' % value if value is not None else 'this operand'}" " isn't a mapped address -- an offset to it would" " invent a name for nothing")} ok, err = _idatui_apply_fmt(ea, n, want) if err: return {"addr": hex(ea), "n": n, "text": before, "format": cur, "error": err} got = _idatui_op_fmt(ea, n) out = {"addr": hex(ea), "n": n, "prev": cur, "format": got, "requested": want, "applied": bool(ok), "choices": choices, "before": before, "text": _idatui_line_text(ea), "value": None if value is None else hex(value), "width": width} if not ok: out["error"] = f"IDA refused {want} on operand {n}" elif lossy: out["warn"] = ( f"operand {n} was {cur} and the ring has no stop there -- " + (f"'{cur}' sets it again" if cur in _IDATUI_FMT_SETTABLE else f"{cur} names a type this can't put back, reassign it by hand")) return out # --------------------------------------------------------------------------- # # the same thing in the decompiler (Hex-Rays keeps its own number formats) # --------------------------------------------------------------------------- # #: Hex-Rays prints C, so two of the listing's stops are missing here: binary #: (C has no binary literal -- the format takes and then renders decimal, which #: would be a lie on screen) and offset (it makes the function fail to #: decompile outright). _IDATUI_PC_FMT_CYCLE = ("hex", "dec", "oct", "char", "default") def _idatui_compact(line): """The ida-pro-mcp whitespace collapse the pseudocode is served through, so a column in what the client SHOWS can be mapped back to Hex-Rays' line.""" try: from ida_pro_mcp.ida_mcp.utils import compact_whitespace return compact_whitespace(line) except Exception: import re as _re stripped = line.lstrip(" \t") lead = line[: len(line) - len(stripped)] return lead + _re.sub(r"[ \t]{2,}", " ", stripped) def _idatui_compact_col(plain, compact, col): """The inverse of ``_idatui_uncompact_col``: a column in Hex-Rays' own line, expressed in the collapsed line the client shows.""" j = 0 for i in range(min(int(col), len(plain))): if j < len(compact) and plain[i] == compact[j]: j += 1 return j def _idatui_uncompact_col(plain, compact, col): """Map a column in the collapsed line back to the same character in the original. The transform only ever DELETES spaces, so walking both in step and skipping what vanished is exact.""" i = 0 for j in range(min(int(col), len(compact))): c = compact[j] while i < len(plain) and plain[i] != c: i += 1 i += 1 return min(i, max(len(plain) - 1, 0)) #: Characters that can be part of a C number literal as Hex-Rays prints one #: (digits, hex letters, the 0x prefix, u/L suffixes). _IDATUI_LIT_CHARS = frozenset("0123456789abcdefABCDEFxXuUlL") def _idatui_lit_extent(plain, x): """The [start, end) of the literal token containing column ``x``. Hex-Rays says WHICH item a column belongs to, but not how wide the printed literal is -- and it attributes neighbouring punctuation to the same item, so ``if ( a1 > 1 )`` reports the closing paren as part of the number. The identity comes from the ctree; the extent is the run of literal characters around the column, which cannot reach a ``)`` or a space.""" if x >= len(plain): return None if plain[x] == "'": # a character constant: '-' end = plain.find("'", x + 1) return (x, end + 1) if end > x else None lo = plain.rfind("'", 0, x) if lo >= 0 and plain.find("'", x) > x and "'" in plain[lo:x] and \ plain[lo:x].count("'") == 1 and " " not in plain[lo:x]: return (lo, plain.find("'", x) + 1) # inside 'c' if plain[x] not in _IDATUI_LIT_CHARS: return None lo = x while lo > 0 and plain[lo - 1] in _IDATUI_LIT_CHARS: lo -= 1 hi = x while hi < len(plain) and plain[hi] in _IDATUI_LIT_CHARS: hi += 1 if lo > 0 and plain[lo - 1] == "-": # a unary minus is part of it lo -= 1 return (lo, hi) def _idatui_pc_nums(cf, sl): """Every number literal on one pseudocode line, as [{x0, x1, ea, opnum, value, nbytes, fmt}]. Asks Hex-Rays what each column belongs to rather than pattern-matching the text: a regex over ``v6 = a1 - 1;`` has to guess which of those characters are a literal, and ``v11`` looks like one.""" import ida_bytes import ida_hexrays import ida_lines import idaapi plain = ida_lines.tag_remove(sl.line) out = [] x = 0 while x < len(plain): ch = plain[x] if ch not in _IDATUI_LIT_CHARS and ch != "'": x += 1 continue head, item, tail = (ida_hexrays.ctree_item_t() for _ in range(3)) if not cf.get_line_item(sl.line, x, True, head, item, tail): x += 1 continue if item.citype != ida_hexrays.VDI_EXPR: x += 1 continue e = item.e if e.op != ida_hexrays.cot_num: x += 1 continue extent = _idatui_lit_extent(plain, x) if extent is None: x += 1 continue nf = e.n.nf opnum = ord(nf.opnum) if isinstance(nf.opnum, str) else int(nf.opnum) nbytes = (ord(nf.org_nbytes) if isinstance(nf.org_nbytes, str) else int(nf.org_nbytes)) ea = int(e.ea) if ea == idaapi.BADADDR: x = extent[1] continue # synthesised: nothing to key on nib = (nf.flags >> ida_bytes.get_operand_type_shift(opnum)) & 0xF # Whether this format is the USER's or Hex-Rays' own guess. The nibble # can't say: an untouched number reads back as whatever it happens to # be printed as, and cycling from there would skip that stop forever # (default already looks like it) and never come back to it. loc = ida_hexrays.operand_locator_t(ea, opnum) user = (ida_hexrays.user_numforms_find(cf.numforms, loc) != ida_hexrays.user_numforms_end(cf.numforms)) out.append({"x0": extent[0], "x1": extent[1], "ea": ea, "opnum": opnum, "value": int(e.n._value), "nbytes": nbytes, "user": user, "fmt": _idatui_fmt_name(nib) if user else "default", "shown": _idatui_fmt_name(nib)}) x = extent[1] # past this literal, not into it return out @tool @idasync def pc_nums( addr: Annotated[str, "Function address (or any address inside it)"], ) -> dict: """Every number literal in a function's pseudocode, as [{line, x0, x1, ea, opnum, value, fmt, user}]. One call per decompilation, so a client can show WHICH literal the cursor is on (and reformat exactly that one) without a round trip per cursor move. Columns are in the same collapsed coordinates the decompile tool serves its text in, i.e. what the client actually displays.""" import ida_hexrays import ida_lines import idaapi if not ida_hexrays.init_hexrays_plugin(): return {"addr": str(addr), "error": "no decompiler", "nums": []} try: f = idaapi.get_func(parse_address(addr)) except Exception as e: return {"addr": str(addr), "error": str(e), "nums": []} if f is None: return {"addr": str(addr), "error": "no function here", "nums": []} try: cf = ida_hexrays.decompile(f.start_ea) except Exception as e: return {"addr": hex(f.start_ea), "error": f"decompile failed: {e}", "nums": []} if cf is None: return {"addr": hex(f.start_ea), "error": "decompilation failed", "nums": []} sv = cf.get_pseudocode() out = [] for i in range(len(sv)): plain = ida_lines.tag_remove(sv[i].line) compact = _idatui_compact(plain) for rec in _idatui_pc_nums(cf, sv[i]): out.append({ "line": i, "x0": _idatui_compact_col(plain, compact, rec["x0"]), "x1": _idatui_compact_col(plain, compact, rec["x1"]), "ea": hex(rec["ea"]), "opnum": rec["opnum"], "value": hex(rec["value"]), "fmt": rec["fmt"], "shown": rec["shown"], "user": bool(rec["user"]), }) return {"addr": hex(f.start_ea), "nums": out, "lines": len(sv)} @tool @idasync def pc_num_format( addr: Annotated[str, "Function address (or any address inside it)"], mode: Annotated[str, "cycle | back | show | hex | dec | oct | char | default"] = "cycle", line: Annotated[int, "0-based pseudocode line index"] = -1, col: Annotated[int, "Cursor column in the DISPLAYED line (-1: first literal)"] = -1, ea: Annotated[str, "Address of the number instead of line/col"] = "", opnum: Annotated[int, "Operand number, with ``ea``"] = -1, ) -> dict: """Change how a number is displayed in the DECOMPILATION (Hex-Rays keeps its own number formats, per (address, operand), independent of the listing). Same stops as ``op_format`` minus the two C can't express: binary (no such literal -- IDA takes the format and prints decimal anyway) and offset (it makes the function stop decompiling). Returns the re-rendered line, and marks the function dirty so the next decompile is the new text.""" import ida_hexrays import ida_lines import idaapi if not ida_hexrays.init_hexrays_plugin(): return {"addr": str(addr), "error": "no decompiler"} try: f = idaapi.get_func(parse_address(addr)) except Exception as e: return {"addr": str(addr), "error": str(e)} if f is None: return {"addr": str(addr), "error": "no function here"} try: cf = ida_hexrays.decompile(f.start_ea) except Exception as e: return {"addr": hex(f.start_ea), "error": f"decompile failed: {e}"} if cf is None: return {"addr": hex(f.start_ea), "error": "decompilation failed"} sv = cf.get_pseudocode() line = int(line) target = None if ea: try: want_ea = parse_address(ea) except Exception as e: return {"addr": hex(f.start_ea), "error": str(e)} for i in range(len(sv)): for rec in _idatui_pc_nums(cf, sv[i]): if rec["ea"] == want_ea and (int(opnum) < 0 or rec["opnum"] == int(opnum)): target, line = rec, i break if target: break elif 0 <= line < len(sv): nums = _idatui_pc_nums(cf, sv[line]) if nums: if int(col) >= 0: plain = ida_lines.tag_remove(sv[line].line) x = _idatui_uncompact_col(plain, _idatui_compact(plain), int(col)) target = next((r for r in nums if r["x0"] <= x < r["x1"]), None) target = target or nums[0] else: return {"addr": hex(f.start_ea), "error": f"line {line} is outside the {len(sv)}-line decompilation"} if target is None: return {"addr": hex(f.start_ea), "line": line, "text": (ida_lines.tag_remove(sv[line].line).strip() if 0 <= line < len(sv) else ""), "error": "no number literal on this line"} cur, value = target["fmt"], target["value"] choices = [c for c in _IDATUI_PC_FMT_CYCLE if c != "char" or _idatui_printable(value)] # Same rule as the listing: one ring per literal, every step. A format the # ring can't hold (an enum set in the GUI) is reported on the way out # instead of being kept for one lap and then lost. lossy = cur not in choices and cur != "default" out = {"addr": hex(f.start_ea), "ea": hex(target["ea"]), "opnum": target["opnum"], "line": line, "prev": cur, "format": cur, "shown": target["shown"], "choices": choices, "value": hex(value), "before": ida_lines.tag_remove(sv[line].line).strip()} mode = str(mode or "cycle").lower() if mode == "show": out["text"] = out["before"] out["applied"] = False return out if mode in ("cycle", "back"): step = 1 if mode == "cycle" else -1 if cur in choices: want = choices[(choices.index(cur) + step) % len(choices)] else: want = choices[0] if step > 0 else choices[-1] else: want = mode if want in ("bin", "offset", "stack", "seg", "float"): out["error"] = (f"Hex-Rays has no {want} format for a number " f"-- set it on the listing instead") out["text"] = out["before"] return out if want not in ("hex", "dec", "oct", "char", "default"): out["error"] = (f"unknown format {mode!r}; one of hex, dec, oct, " f"char, default") out["text"] = out["before"] return out loc = ida_hexrays.operand_locator_t(target["ea"], target["opnum"]) it = ida_hexrays.user_numforms_find(cf.numforms, loc) if it != ida_hexrays.user_numforms_end(cf.numforms): # std::map::insert is a no-op on an existing key, so a format already # set here would silently win over the new one. ida_hexrays.user_numforms_erase(cf.numforms, it) if want != "default": import ida_bytes nf = ida_hexrays.number_format_t(target["opnum"]) nf.flags = ida_bytes.get_operand_flag(_idatui_fmt_nibbles()[want], target["opnum"]) try: nf.org_nbytes = target["nbytes"] except Exception: pass ida_hexrays.user_numforms_insert(cf.numforms, loc, nf) cf.save_user_numforms() try: ida_hexrays.mark_cfunc_dirty(f.start_ea) except Exception: pass out["format"] = want out["applied"] = True if lossy: out["warn"] = (f"this number was {cur}, which names a type a radix " f"can't put back -- reassign it in IDA") try: cf2 = ida_hexrays.decompile(f.start_ea, flags=ida_hexrays.DECOMP_NO_CACHE) sv2 = cf2.get_pseudocode() if cf2 is not None else None out["text"] = (ida_lines.tag_remove(sv2[line].line).strip() if sv2 is not None and line < len(sv2) else out["before"]) except Exception as e: out["text"] = out["before"] out["warn"] = f"re-render failed: {e}" return out ''' SNIPPET = f"{BEGIN}\n{BODY.strip()}\n{END}\n" def api_types_path() -> pathlib.Path | None: """Locate ida_pro_mcp/ida_mcp/api_types.py without importing it (importing the submodule would pull in IDA, which isn't available outside a worker).""" spec = importlib.util.find_spec("ida_pro_mcp") # top-level pkg is IDA-free if spec is None or not spec.submodule_search_locations: return None p = pathlib.Path(spec.submodule_search_locations[0]) / "ida_mcp" / "api_types.py" return p if p.exists() else None def main() -> int: path = api_types_path() if path is None: print("idatui: ida_pro_mcp not found; skipping tool injection", file=sys.stderr) return 0 text = path.read_text() if BEGIN in text and END in text: # replace the existing block in place pre = text[: text.index(BEGIN)].rstrip() post = text[text.index(END) + len(END):].lstrip("\n") new = pre + "\n\n" + SNIPPET + ("\n" + post if post else "") else: new = text.rstrip() + "\n\n" + SNIPPET if new == text: return 0 try: path.write_text(new) except OSError as e: print(f"idatui: could not patch {path}: {e}", file=sys.stderr) return 1 print(f"idatui: injected/updated idatui-ext tools in {path}", file=sys.stderr) return 0 if __name__ == "__main__": raise SystemExit(main())