From c9208de05d8583b677117fe43c9d3567e89eb2ce Mon Sep 17 00:00:00 2001 From: blasty Date: Fri, 7 Aug 2026 12:39:54 +0200 Subject: Rebase MISTER EXO's ida-codemode port onto the current tree Mechanical part of the port: the 27-file patch was cut against a base ~148 commits behind us, so it did not apply. Resolved 11 conflicts (all of them diff drift, not semantic clashes) and the three file deletions: - app.py: the patch re-inserted _do_rename/_do_name_addr/_seek_split etc. as "theirs" because our tree moved them to edit_ctl.py/trace_ctl.py. Kept ours and applied the real intent (WorkerClient->CodeModeClient, .call->.invoke, _open_worker_client->_open_database_client) at their current homes. - domain.py: kept Head as a NamedTuple -- the patch reverted it to a frozen dataclass, which the perf work measured at 2.9us vs 1.9us per row on a quarter-million-row walk. Dropped _fetch_output (no download_url under Code Mode) and its now-dead urllib/json imports. - pane.py: the patch's deletion swallowed our zellij support along with the worker-reaping block it meant to remove. Kept zellij, removed the reaping. - test_scenarios.py: the idb_save->save_database teardown hunk belongs to tests/_fixtures.py in our tree; applied it there and kept our pc_num_format scenario that the drift landed on. Three defects in the patch itself, fixed here: - It made "import idatui" hard-require ida_codemode, so every offline suite died at import -- including the pure ones (graph/index/trace) that are the house rule for "tests/run.py --fast". The import is now deferred and gated on the binding, which is also what lets the port's own contract tests inject a fake DatabaseHandle. - project.stage() inlined an ida_codemode.registry import and treated "library not installed" as "someone owns this database", which broke IDA-free project staging. Ownership lookup moved to codemode_client.database_owner(). - tests/test_codemode_client.py had no NEEDS_IDA marker, which tests/run.py rejects outright. Offline suite: 301 passed, 0 failed. Against master's 344 the whole delta is accounted for: -40 worker_client (module deleted), -18 launch sweep checks (behaviour deliberately removed) +3 guarding that it stays removed, +2 pool (GUI-save semantics), +13 new codemode_client contract tests. NOT yet done, and the port is not functional without it: the adapter is missing five operations our tree grew since the patch's base (flowchart, op_format, pc_nums, pc_num_format, survey_binary) and its "heads" predates back-walking and digest/expect. --- idatui/codemode_client.py | 1164 +++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 1164 insertions(+) create mode 100644 idatui/codemode_client.py (limited to 'idatui/codemode_client.py') diff --git a/idatui/codemode_client.py b/idatui/codemode_client.py new file mode 100644 index 0000000..77a0227 --- /dev/null +++ b/idatui/codemode_client.py @@ -0,0 +1,1164 @@ +"""Client adapter from ida-tui's domain operations to IDA Code Mode. + +``DatabaseHandle`` is the lifecycle boundary: it discovers an already-registered +GUI database, reuses a shared managed idalib worker, or starts one when needed. +The TUI never owns or terminates an IDA process. Closing this client releases +only its lease. + +The Code Mode transport intentionally exposes one broad operation, +``execute_python``. ``CodeModeClient.invoke`` turns the small, address-centric +operations needed by the paging layer into self-contained snippets. The +snippets prefer the public ``ida-domain`` ``db`` object. A handful of features +that ida-domain does not currently expose (IDA-coloured listing rows, creating +instructions, ARM T-state, and detailed Hex-Rays line maps/failures) use the +IDAPython modules that Code Mode deliberately makes importable. +""" +from __future__ import annotations + +import json +import os +import shlex +import threading +import time +from textwrap import dedent +from typing import Any + +from .errors import IDAConnectionError, IDATimeoutError, IDAToolError, Session + +# ida_codemode is imported EAGERLY-IF-PRESENT but never at hard import cost. +# +# The paging/graph/trace layers and their offline test suites must keep importing +# `idatui` on a machine with no IDA and no Code Mode installed -- that is the +# house rule the stdlib-only worker client used to satisfy for free, and +# `tests/run.py --fast` (257 checks, any python3) depends on it. A hard top-level +# import here makes the whole package unimportable, so the failure is deferred to +# the first operation that genuinely needs the library. +_CODEMODE_ERROR: Exception | None = None +try: + from ida_codemode.client import ( + ClientError, + DatabaseHandle, + InstanceDisconnectedError, + RemoteError, + ) + from ida_codemode.registry import ( + REGISTRY_DIR, + FileLock, + RegistryEntry, + canonical_path, + idb_key, + scan_instances, + ) + from ida_codemode.resolver import IdbBusy, expected_idb_path +except ImportError as _exc: # library absent: usable only for offline layers + _CODEMODE_ERROR = _exc + # Bound to None rather than left undefined so the names stay patchable: the + # offline contract tests inject a fake DatabaseHandle here. + ClientError = InstanceDisconnectedError = RemoteError = None # type: ignore[assignment,misc] + DatabaseHandle = RegistryEntry = FileLock = None # type: ignore[assignment,misc] + REGISTRY_DIR = canonical_path = idb_key = scan_instances = None # type: ignore[assignment] + IdbBusy = expected_idb_path = None # type: ignore[assignment] + + +def _require_codemode() -> None: + """Raise an actionable error when the Code Mode library is missing. + + Gated on the binding, not on the original import result, so a test that + injects a fake ``DatabaseHandle`` exercises the real adapter logic. + """ + if DatabaseHandle is None: + raise IDAConnectionError( + "ida-codemode-mcp is not installed in this environment " + f"({_CODEMODE_ERROR}). Install it (e.g. `uv sync`, or " + "`pip install -e ../ida-codemode-mcp`) so ida-tui can lease a " + "database.") from _CODEMODE_ERROR + + +def database_owner(idb_path: str, staged_path: str | None = None): + """The registry entry that owns ``idb_path``/``staged_path``, else None. + + Returns None when the Code Mode library is absent: with no library there is + no client in this environment that could be holding the database, and the + IDA-free layers (project staging) must keep working. Registry errors that + happen WITH the library installed still propagate -- those mean "we could + not determine ownership", which is not the same as "nobody owns it". + """ + if DatabaseHandle is None: + return None + expected_key = idb_key(idb_path) + staged = canonical_path(staged_path) if staged_path else None + for item in scan_instances(timeout=0.5): + entry = item.entry + if entry.idb_key == expected_key: + return entry + if staged and entry.exe_path and canonical_path(entry.exe_path) == staged: + return entry + return None + + +def registered_database(path: str, output_database: str | None = None) -> bool: + """Whether a live/lock-held Code Mode instance owns this target.""" + _require_codemode() + source = canonical_path(path) + expected = canonical_path(output_database) if output_database else expected_idb_path(source) + expected_key = idb_key(expected) + for instance in scan_instances(timeout=0.5): + entry = instance.entry + if entry.idb_key == expected_key: + return True + if not output_database and entry.backend == "gui" and entry.exe_path: + if canonical_path(entry.exe_path) == source: + return True + return False + + +class _NoopKeepAlive: + """Compatibility shim: the DatabaseHandle's SSE lease is the heartbeat.""" + + def __init__(self) -> None: + self.beats = self.failures = 0 + + def start(self) -> "_NoopKeepAlive": + return self + + def stop(self) -> None: + pass + + +def _parse_load_args(value: str) -> tuple[str | None, int | None, str | None]: + """Translate ida-tui's legacy first-open switches to Code Mode options. + + Code Mode has typed options for processor, natural loading address and file + type. It deliberately has no arbitrary command-line escape hatch; reject + switches we cannot represent instead of silently loading a blob wrongly. + """ + processor: str | None = None + loading_address: int | None = None + file_type: str | None = None + unsupported: list[str] = [] + try: + words = shlex.split(value or "", posix=os.name != "nt") + except ValueError as exc: + raise ValueError(f"invalid IDA load options: {exc}") from exc + for word in words: + if word.startswith("-p") and len(word) > 2: + processor = word[2:] + elif word.startswith("-b") and len(word) > 2: + try: + # IDA's -b is in 16-byte paragraphs. DatabaseHandle expects the + # natural address, which is the safer public API. + loading_address = int(word[2:], 16) << 4 + except ValueError as exc: + raise ValueError(f"invalid IDA loading address: {word!r}") from exc + elif word.startswith("-T") and len(word) > 2: + file_type = word[2:] + else: + unsupported.append(word) + if unsupported: + joined = " ".join(unsupported) + raise ValueError( + "ida-codemode cannot represent arbitrary IDA load options: " + f"{joined!r}; use processor/base/file type options instead" + ) + return processor, loading_address, file_type + + +def _script(args: dict[str, Any], body: str) -> str: + """Bind JSON arguments without interpolating user text into Python code.""" + encoded = json.dumps(args, ensure_ascii=False, separators=(",", ":")) + return f"import json\na = json.loads({encoded!r})\n{dedent(body).strip()}\n" + + +# Rich flat-listing generation is the largest ida-domain gap in this port. +# ida-domain can enumerate heads and render plain disassembly, but it does not +# expose undefined runs, IDA colour spans, function banners, or expanded UDT +# members. Keep that IDAPython-only logic isolated in this one operation. +_HEADS = r''' +import ida_bytes, ida_funcs, ida_idaapi, ida_lines, ida_name, ida_nalt, ida_segment, ida_typeinf +start = int(str(a["addr"]), 16) +count = max(1, min(int(a.get("count", 200)), 2000)) +offset = max(0, int(a.get("offset", 0))) +annotate = bool(a.get("annotate", False)) +seg = db.segments.get_at(start) +if seg is None: + result = {"addr": a["addr"], "error": "no segment", "heads": [], "cursor": {"done": True}} +else: + lo, hi = int(seg.start_ea), int(seg.end_ea) + if a.get("end"): + hi = min(hi, int(str(a["end"]), 16)) + + span_names = { + "insn": ("SCOLOR_INSN", "SCOLOR_KEYWORD", "SCOLOR_ASMDIR", "SCOLOR_MACRO"), + "reg": ("SCOLOR_REG",), + "num": ("SCOLOR_NUMBER", "SCOLOR_CHAR", "SCOLOR_BINPREF"), + "str": ("SCOLOR_STRING",), + "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",), + } + tag_kinds = {} + for kind, names in span_names.items(): + for name in names: + value = getattr(ida_lines, name, None) + if isinstance(value, str) and value: + tag_kinds[value[0]] = kind + elif isinstance(value, int): + tag_kinds[chr(value)] = kind + + def spans(tagged): + 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)) + out, stack, buf = [], [], [] + def flush(): + if buf: + out.append([stack[-1] if stack else "text", "".join(buf)]) + buf.clear() + i = 0 + while i < len(tagged): + ch = tagged[i] + if ch == on and i + 1 < len(tagged): + tag = tagged[i + 1] + if tag == addr_tag: + i += 2 + addr_len + continue + flush(); stack.append(tag_kinds.get(tag, "text")); i += 2; continue + if ch == off and i + 1 < len(tagged): + flush() + if stack: stack.pop() + i += 2; continue + if ch == esc and i + 1 < len(tagged): + buf.append(tagged[i + 1]); i += 2; continue + buf.append(ch); i += 1 + flush() + collapsed, previous_space = [], False + for kind, text in out: + acc = [] + for ch in text: + if ch.isspace(): + if previous_space: continue + acc.append(" "); previous_space = True + else: + acc.append(ch); previous_space = False + if acc: collapsed.append([kind, "".join(acc)]) + if collapsed: + collapsed[0][1] = collapsed[0][1].lstrip() + collapsed[-1][1] = collapsed[-1][1].rstrip() + return [[kind, text] for kind, text in collapsed if text] + + def row(ea): + flags = ida_bytes.get_flags(ea) + kind = "code" if ida_bytes.is_code(flags) else ("data" if ida_bytes.is_data(flags) else "unknown") + tagged = ida_lines.generate_disasm_line(ea, 0) or "" + text = " ".join(ida_lines.tag_remove(tagged).split()) if tagged else "" + item = {"ea": hex(ea), "kind": kind, "size": int(ida_bytes.get_item_size(ea)), "text": text} + if tagged: + rich = spans(tagged) + if " ".join("".join(x[1] for x in rich).split()) == text: + item["spans"] = rich + name = ida_name.get_ea_name(ea) + if name: item["name"] = name + return item + + def unknown_row(ea, size): + if size <= 1: return row(ea) + item = {"ea": hex(ea), "kind": "unknown", "size": int(size), "text": f"db {size} dup(?)"} + name = ida_name.get_ea_name(ea) + if name: item["name"] = name + return item + + def members(ea): + tif = db.types.get_at(ea) + if tif is None or not tif.is_udt(): return [] + answer = [] + for member in db.types.get_udt_members(tif): + type_text = member.type.dstr() or "" + text = f"+{member.offset:X} {member.name}" + (f" {type_text}" if type_text else "") + answer.append({"ea": hex(ea + member.offset), "kind": "member", + "size": int(member.size), "text": text}) + return answer + + def is_unknown(ea): + flags = ida_bytes.get_flags(ea) + return not (ida_bytes.is_code(flags) or ida_bytes.is_data(flags)) + def run_end(ea): + nxt = ida_bytes.next_head(ea, hi) + return nxt if nxt != ida_idaapi.BADADDR and ea < nxt <= hi else hi + def advance(ea): + if is_unknown(ea): return run_end(ea) + nxt = ida_bytes.get_item_end(ea) + return nxt if nxt > ea else ea + 1 + def rows_for(ea): + if is_unknown(ea): return [unknown_row(ea, run_end(ea) - ea)] + fn = db.functions.get_at(ea) if annotate else None + at_start = fn is not None and int(fn.start_ea) == ea + answer = [] + if at_start: + name = db.functions.get_name(fn) or f"sub_{ea:X}" + answer += [ + {"ea": hex(ea), "kind": "sep", "size": 0, "text": ""}, + {"ea": hex(ea), "kind": "sep", "size": 0, + "text": "; " + "=" * 15 + " S U B R O U T I N E " + "=" * 15}, + {"ea": hex(ea), "kind": "funchdr", "size": 0, + "text": name + " proc", "name": name}, + ] + item = row(ea) + if at_start: + item["name"] = None + elif annotate and item["kind"] == "code" and item.get("name"): + name = item["name"] + answer.append({"ea": hex(ea), "kind": "label", "size": 0, + "text": name + ":", "name": name}) + item["name"] = None + answer.append(item) + if item["kind"] == "data": answer += members(ea) + if fn is not None and ida_bytes.get_item_end(ea) >= int(fn.end_ea): + name = db.functions.get_name(fn) or f"sub_{int(fn.start_ea):X}" + answer += [ + {"ea": hex(ea), "kind": "funchdr", "size": 0, + "text": name + " endp", "name": name}, + {"ea": hex(ea), "kind": "sep", "size": 0, "text": "; " + "-" * 60}, + ] + return answer + + ea = ida_bytes.get_item_head(start) + if ea == ida_idaapi.BADADDR: ea = start + for _ in range(offset): + if ea >= hi: break + ea = advance(ea) + rows = [] + more = False + while ea != ida_idaapi.BADADDR and ea < hi: + if len(rows) >= count: + more = True; break + rows += rows_for(ea) + ea = advance(ea) + result = {"addr": a["addr"], "heads": rows, + "cursor": {"next": hex(ea)} if more else {"done": True}} +result +''' + + +_DECOMP_MAP_HELPER = r''' +def line_map(cfunc): + import ida_hexrays + answer = [] + for sl in cfunc.get_pseudocode(): + tagged, eas, seen = sl.line, [], set() + for x in range(len(tagged) + 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(tagged, x, False, head, item, tail): continue + text = item.dstr() or "" + try: ea = int(text.split(": ", 1)[0], 16) + except (ValueError, IndexError): continue + if ea not in seen: seen.add(ea); eas.append(ea) + answer.append(eas) + return answer +''' + + +_OPERATIONS: dict[str, str] = { + "list_funcs": r''' +import fnmatch +queries = a.get("queries") or [{}] +q = queries[0] +offset, count = max(0, int(q.get("offset", 0))), max(1, int(q.get("count", 500))) +pattern = str(q.get("filter") or "").lower() +if pattern and not any(ch in pattern for ch in "*?["): pattern = "*" + pattern + "*" +rows = [] +for fn in db.functions.get_all(): + name = db.functions.get_name(fn) or f"sub_{int(fn.start_ea):X}" + if pattern and not fnmatch.fnmatchcase(name.lower(), pattern): continue + rows.append({"addr": hex(int(fn.start_ea)), "name": name, + "size": int(fn.end_ea) - int(fn.start_ea)}) +page = rows[offset:offset + count] +result = {"result": [{"data": page, "next_offset": offset + len(page), "total": len(rows)}]} +result +''', + "disasm": r''' +ea = int(str(a["addr"]), 16) +fn = db.functions.get_at(ea) +if fn is None: + result = {"instructions": [], "total_instructions": 0, "instruction_count": 0} +else: + instructions = list(db.functions.get_instructions(fn)) + limit = max(1, int(a.get("max_instructions", len(instructions) or 1))) + rows = [{"addr": hex(int(insn.ea)), "instruction": db.instructions.get_disassembly(insn)} + for insn in instructions[:limit]] + result = {"instructions": rows, "total_instructions": len(instructions), + "instruction_count": len(instructions)} +result +''', + "file_regions": r''' +import idaapi +rows = [] +for seg in db.segments.get_all(): + try: file_off = int(idaapi.get_fileregion_offset(seg.start_ea)) + except Exception: file_off = -1 + if file_off < 0 or file_off >= (1 << 48): file_off = -1 + rows.append({"start": hex(int(seg.start_ea)), "end": hex(int(seg.end_ea)), + "file_off": file_off, "name": db.segments.get_name(seg) or ""}) +result = {"regions": rows} +result +''', + "read_raw": r''' +import ida_bytes +ea, size = int(str(a["addr"]), 16), max(0, int(a["size"])) +raw = ida_bytes.get_bytes(ea, size) or b"" +raw = raw[:size] + b"\xff" * max(0, size - len(raw)) +data = bytearray(raw) +for index, value in enumerate(data): + if value == 0xFF and not ida_bytes.is_loaded(ea + index): data[index] = 0 +result = {"addr": a["addr"], "hex": bytes(data).hex(), "n": len(data)} +result +''', + "get_bytes": r''' +rows = [] +for region in a.get("regions", []): + ea, size = int(str(region["addr"]), 16), int(region["size"]) + raw = db.bytes.get_bytes_at(ea, size) or b"" + rows.append({"addr": region["addr"], "data": " ".join(f"{b:02x}" for b in raw)}) +result = {"result": rows} +result +''', + "search_structs": r''' +needle = str(a.get("filter") or "").lower() +rows = [] +for tif in db.types.get_all(): + name = tif.get_type_name() or "" + if not name or needle not in name.lower() or not tif.is_udt(): continue + members = list(db.types.get_udt_members(tif)) + rows.append({"name": name, "size": int(tif.get_size()), "is_union": bool(tif.is_union()), + "cardinality": len(members), "ordinal": int(tif.get_ordinal())}) +result = {"result": rows} +result +''', + "type_inspect": r''' +rows = [] +for query in a.get("queries", []): + name = str(query.get("name") or "") + tif = db.types.get_by_name(name) + if tif is None: + rows.append({"name": name, "error": "type not found"}); continue + members = [{"name": m.name, "type": m.type.dstr() or str(m.type), + "offset": int(m.offset), "size": int(m.size)} + for m in db.types.get_udt_members(tif)] if tif.is_udt() else [] + rows.append({"name": name, "size": int(tif.get_size()), "is_union": bool(tif.is_union()), + "members": members}) +result = {"result": rows} +result +''', + "declare_type": r''' +import ida_typeinf +decls = a.get("decls", "") +if isinstance(decls, str): decls = [decls] +rows = [] +for declaration in decls: + try: + errors = int(db.types.parse_declarations(ida_typeinf.get_idati(), declaration)) + rows.append({"ok": errors == 0, **({} if errors == 0 else {"error": f"{errors} parse error(s)"})}) + except Exception as exc: + rows.append({"ok": False, "error": str(exc)}) +result = {"result": rows} +result +''', + "del_type": r''' +import ida_typeinf +name = str(a["name"]) +ok = bool(ida_typeinf.del_named_type(ida_typeinf.get_idati(), name, ida_typeinf.NTF_TYPE)) +result = {"name": name, "deleted": ok, **({} if ok else {"error": f"Type {name!r} not found or could not be deleted"})} +result +''', + "func_types": r''' +import ida_typeinf +ea = int(str(a["addr"]), 16) +fn = db.functions.get_at(ea) +if fn is None: + result = {"addr": a["addr"], "error": "no function at address"} +else: + pseudo = db.pseudocode.decompile(fn) + name = db.functions.get_name(fn) or "" + tif = pseudo.get_func_type() + try: prototype = ida_typeinf.print_tinfo("", 0, 0, ida_typeinf.PRTYPE_1LINE, tif, name, "") if tif else "" + except Exception: prototype = tif.dstr() if tif else "" + lvars = [{"name": var.name, "type": var.type_info.dstr() if var.type_info else "", + "is_arg": bool(var.is_arg)} for var in pseudo.local_variables] + result = {"addr": hex(int(fn.start_ea)), "name": name, + "prototype": (prototype or "").strip(), "lvars": lvars} +result +''', + "set_lvar_type": r''' +import ida_typeinf +ea, variable, declaration = int(str(a["addr"]), 16), str(a["variable"]), str(a["type"]) +fn = db.functions.get_at(ea) +if fn is None: + result = {"error": "no function at address"} +else: + pseudo = db.pseudocode.decompile(fn) + var = pseudo.find_local_variable(variable) + if var is None: + result = {"error": f"local variable {variable!r} not found"} + else: + try: + tif = db.types.parse_one_declaration(ida_typeinf.get_idati(), declaration) + accepted = bool(var.set_type(tif)) + saved = bool(pseudo.save_local_variable_info(var, save_type=True)) if accepted else False + result = {"addr": hex(int(fn.start_ea)), "variable": variable, + "type": declaration, "ok": accepted and saved} + except Exception as exc: + result = {"error": f"bad type {declaration!r}: {exc}"} +result +''', + "set_type": r''' +from ida_domain.types import TypeApplyFlags +rows = [] +for edit in a.get("edits", []): + ea = int(str(edit["addr"]), 16) + declaration = str(edit.get("signature") or edit.get("type") or "") + try: + ok = bool(db.types.apply_declaration_at(ea, declaration, TypeApplyFlags.DEFINITE)) + rows.append({"addr": hex(ea), "ok": ok, **({} if ok else {"error": "IDA rejected the type"})}) + except Exception as exc: + rows.append({"addr": hex(ea), "ok": False, "error": str(exc)}) +result = {"result": rows} +result +''', + "data_type": r''' +ea = int(str(a["addr"]), 16) +try: + tif = db.types.get_at(ea) + fn = db.functions.get_at(ea) + result = {"addr": hex(ea), "name": db.names.get_at(ea) or "", + "type": tif.dstr() if tif else "", "size": int(db.heads.size(ea)) if db.heads.is_head(ea) else 0, + "is_func": bool(fn)} +except Exception as exc: + result = {"addr": hex(ea), "error": str(exc)} +result +''', + "force_recompile": r''' +import ida_hexrays +rows = [] +for item in a.get("items", []): + ea = int(str(item["addr"]), 16) + ida_hexrays.mark_cfunc_dirty(ea, False) + rows.append({"addr": hex(ea), "ok": True}) +result = {"result": rows} +result +''', + "undefine": r''' +import ida_bytes +rows = [] +for item in a.get("items", []): + ea = int(str(item["addr"]), 16) + size = max(1, int(item.get("size") or ida_bytes.get_item_size(ea) or 1)) + ok = bool(ida_bytes.del_items(ea, ida_bytes.DELIT_SIMPLE, size)) + rows.append({"addr": hex(ea), "ok": ok, **({} if ok else {"error": "delete items failed"})}) +result = {"result": rows} +result +''', + "define_code": r''' +import ida_ua +rows = [] +for item in a.get("items", []): + ea = int(str(item["addr"]), 16); size = int(ida_ua.create_insn(ea)) + rows.append({"addr": hex(ea), "ok": size > 0, "size": size, + **({} if size > 0 else {"error": "instruction did not decode"})}) +result = {"result": rows} +result +''', + "define_func": r''' +rows = [] +for item in a.get("items", []): + ea = int(str(item["addr"]), 16); ok = bool(db.functions.create(ea)) + rows.append({"addr": hex(ea), "ok": ok, **({} if ok else {"error": "IDA refused the function"})}) +result = {"result": rows} +result +''', + "make_data": r''' +import ida_bytes, ida_idaapi, ida_typeinf +from ida_domain.types import TypeApplyFlags +rows = [] +for item in a.get("items", []): + ea, declaration = int(str(item["addr"]), 16), str(item["type"]) + try: + tif = db.types.parse_one_declaration(ida_typeinf.get_idati(), declaration) + size = max(1, int(tif.get_size())) + saved_names = [(addr, name) for addr, name in db.names.get_all() + if ea <= int(addr) < ea + size] + ida_bytes.del_items(ea, ida_bytes.DELIT_EXPAND | ida_bytes.DELIT_DELNAMES, + max(size, int(ida_bytes.get_item_size(ea) or 1))) + created = bool(ida_bytes.create_data(ea, ida_bytes.FF_BYTE, size, ida_idaapi.BADADDR)) + ok = created and bool(db.types.apply_at(tif, ea, TypeApplyFlags.DEFINITE)) + for address, name in saved_names: + db.names.set_name(int(address), name) + if ok and item.get("name"): ok = bool(db.names.set_name(ea, str(item["name"]))) + rows.append({"addr": hex(ea), "ok": ok, "size": size, + **({} if ok else {"error": "IDA rejected the data type"})}) + except Exception as exc: + rows.append({"addr": hex(ea), "ok": False, "error": str(exc)}) +result = {"result": rows} +result +''', + "make_string": r''' +from ida_domain.strings import StringType +ea, length = int(str(a["addr"]), 16), max(0, int(a.get("length", 0))) +kind = {"c": StringType.C, "c16": StringType.C_16, "c32": StringType.C_32, + "pascal": StringType.PASCAL}.get(str(a.get("kind", "c")).lower(), StringType.C) +import ida_bytes +try: + ida_bytes.del_items(ea, ida_bytes.DELIT_SIMPLE, length if length > 0 else 1) +except Exception: + pass +try: + ok = bool(db.bytes.create_string_at(ea, length or None, kind)) + text = db.bytes.get_string_at(ea) or "" if ok else "" + result = {"addr": hex(ea), "ok": ok, "size": int(db.heads.size(ea)) if ok else 0, "text": text} +except Exception as exc: + result = {"addr": hex(ea), "ok": False, "error": str(exc)} +result +''', + "list_strings": r''' +from ida_domain.strings import StringListConfig +offset, count, min_len = max(0, int(a.get("offset", 0))), max(1, int(a.get("count", 2000))), max(1, int(a.get("min_len", 4))) +if offset == 0 or a.get("refresh"): + from ida_domain.strings import StringType + db.strings.rebuild(StringListConfig(string_types=list(StringType), min_len=min_len, + only_ascii_7bit=False)) +items = list(db.strings.get_all()) +page = items[offset:offset + count] +rows = [] +for item in page: + try: text = str(item) + except Exception: text = item.contents.decode("utf-8", "replace") if item.contents else "" + rows.append({"addr": hex(int(item.address)), "text": text, "len": int(item.length), "type": item.type.name}) +result = {"strings": rows, "total": len(items), "next_offset": offset + len(rows)} +result +''', + "list_linkage": r''' +imports = [{"addr": hex(int(item.address)), "name": item.name, "module": item.module_name} + for item in db.imports.get_all_imports() if item.name] +exports = [{"addr": hex(int(item.address)), "name": item.name, "ordinal": int(item.ordinal)} + for item in db.entries.get_all() if item.name] +result = {"imports": imports, "exports": exports, + "n_imports": len(imports), "n_exports": len(exports)} +result +''', + "lookup_funcs": r''' +rows = [] +for query in a.get("queries", []): + raw = str(query) + try: ea = int(raw, 16) + except ValueError: + fn = db.functions.get_by_name(raw); ea = int(fn.start_ea) if fn else None + else: fn = db.functions.get_at(ea) + if fn is None: + rows.append({"query": raw, "fn": None}) + else: + rows.append({"query": raw, "fn": {"addr": hex(int(fn.start_ea)), + "name": db.functions.get_name(fn) or f"sub_{int(fn.start_ea):X}", + "size": int(fn.end_ea) - int(fn.start_ea)}}) +result = {"result": rows} +result +''', + "resolve_names": r''' +import ida_idaapi, ida_name +rows = [] +for query in a.get("queries", []): + name = str(query).strip(); ea = ida_name.get_name_ea(ida_idaapi.BADADDR, name) + rows.append({"query": name, "ea": hex(int(ea)) if ea != ida_idaapi.BADADDR else None}) +result = {"result": rows} +result +''', + "xref_types": r''' +queries = a.get("queries") or [] +all_results = [] +for query in queries: + ea, direction = int(str(query["addr"]), 16), str(query.get("direction", "both")) + refs = [] + if direction in ("to", "both"): refs += list(db.xrefs.to_ea(ea)) + if direction in ("from", "both"): refs += list(db.xrefs.from_ea(ea)) + rows, seen = [], set() + for ref in refs: + key = (int(ref.from_ea), int(ref.to_ea), int(ref.type)) + if query.get("dedup") and key in seen: continue + seen.add(key) + fn = db.functions.get_at(int(ref.from_ea)) + kind = ("call" if ref.is_call else "jump" if ref.is_jump else "flow" if ref.is_flow + else "read" if ref.is_read else "write" if ref.is_write else ref.type.name.lower()) + row = {"from": hex(int(ref.from_ea)), "to": hex(int(ref.to_ea)), + "type": "code" if ref.is_code else "data", "kind": kind} + if query.get("include_fn") and fn is not None: + row["fn"] = {"addr": hex(int(fn.start_ea)), "name": db.functions.get_name(fn) or ""} + rows.append(row) + if len(rows) >= int(query.get("count", 2000)): break + all_results.append({"data": rows}) +result = {"result": all_results} +result +''', + "xref_query": r''' +queries = a.get("queries") or [] +all_results = [] +for query in queries: + ea, direction = int(str(query["addr"]), 16), str(query.get("direction", "both")) + refs = [] + if direction in ("to", "both"): refs += list(db.xrefs.to_ea(ea)) + if direction in ("from", "both"): refs += list(db.xrefs.from_ea(ea)) + rows = [] + for ref in refs[:int(query.get("count", 2000))]: + fn = db.functions.get_at(int(ref.from_ea)) + row = {"from": hex(int(ref.from_ea)), "to": hex(int(ref.to_ea)), + "type": "code" if ref.is_code else "data"} + if query.get("include_fn") and fn is not None: + row["fn"] = {"addr": hex(int(fn.start_ea)), "name": db.functions.get_name(fn) or ""} + rows.append(row) + all_results.append({"data": rows}) +result = {"result": all_results} +result +''', + "set_comments": r''' +rows = [] +for item in a.get("items", []): + ea, text = int(str(item["addr"]), 16), str(item.get("comment") or "") + try: + if text: ok = bool(db.comments.set_at(ea, text)) + else: db.comments.delete_at(ea); ok = True + rows.append({"addr": hex(ea), "ok": ok}) + except Exception as exc: + rows.append({"addr": hex(ea), "ok": False, "error": str(exc)}) +result = {"result": rows} +result +''', + "rename": r''' +import ida_idaapi, ida_name, ida_typeinf +batch = a.get("batch") or {} +out = {}; ok_count = failed = 0 +for category, edit in batch.items(): + try: + if category == "func": + ea, new = int(str(edit["addr"]), 16), str(edit["name"]) + fn = db.functions.get_at(ea); ok = bool(fn and db.functions.set_name(fn, new)) + elif category == "data": + new = str(edit.get("new") or "") + if edit.get("addr") is not None: ea = int(str(edit["addr"]), 16) + else: ea = int(ida_name.get_name_ea(ida_idaapi.BADADDR, str(edit.get("old") or ""))) + ok = bool(db.names.set_name(ea, new)) + elif category in ("local", "stack"): + ea, old, new = int(str(edit["func_addr"]), 16), str(edit["old"]), str(edit["new"]) + pseudo = db.pseudocode.decompile(ea); var = pseudo.find_local_variable(old) + if var is None: ok = False + else: + var.set_user_name(new) + ok = bool(pseudo.save_local_variable_info(var, save_name=True)) + else: + raise ValueError(f"unsupported rename category: {category}") + row = {"ok": ok, **({} if ok else {"error": "IDA rejected the name"})} + except Exception as exc: + row = {"ok": False, "error": str(exc)} + out[category] = [row] + if row["ok"]: ok_count += 1 + else: failed += 1 +out["summary"] = {"ok": ok_count, "failed": failed} +result = out +result +''', +} + + +_OPERATIONS["decompile"] = _DECOMP_MAP_HELPER + r''' +ea = int(str(a["addr"]), 16) +fn = db.functions.get_at(ea) +if fn is None: + result = {"error": f"no function at {ea:#x}"} +else: + pseudo = db.pseudocode.decompile(fn) + mapping = line_map(pseudo.raw_cfunc) + plain = pseudo.to_text() + marked = [line + (f" /*0x{eas[0]:X}*/" if eas else "") + for line, eas in zip(plain, mapping)] + import ida_name + refs, seen = [], set() + for expr in pseudo.find_objects(): + target = int(expr.obj_ea) + if target in seen or not (db.is_valid_ea(target) or db.is_private_ea(target)): continue + seen.add(target) + name = expr.obj_name or ida_name.get_name(target) or "" + try: string = db.bytes.get_string_at(target) if db.is_valid_ea(target) else None + except Exception: string = None + refs.append({"addr": hex(target), "name": name, "string": string}) + result = {"addr": hex(int(fn.start_ea)), "code": "\n".join(marked), "refs": refs} +result +''' + +_OPERATIONS["decomp_map"] = _DECOMP_MAP_HELPER + r''' +ea = int(str(a["addr"]), 16) +fn = db.functions.get_at(ea) +if fn is None: + result = {"error": f"no function at {ea:#x}"} +else: + pseudo = db.pseudocode.decompile(fn) + mapping = line_map(pseudo.raw_cfunc) + result = {"addr": hex(int(fn.start_ea)), + "lines": [{"ea": hex(eas[0]) if eas else None, + "eas": [hex(item) for item in eas]} for eas in mapping]} +result +''' + +_OPERATIONS["define_code_run"] = r''' +import ida_bytes, ida_idp, ida_segment, ida_ua, idaapi +ea, limit = int(str(a["addr"]), 16), max(1, min(int(a.get("limit", 20000)), 200000)) +seg = ida_segment.getseg(ea) +if seg is None: + result = {"addr": a["addr"], "error": "no segment", "count": 0} +else: + start, count, stopped, hi = ea, 0, "limit", int(seg.end_ea) + 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): stopped = "defined"; break + size = int(ida_ua.create_insn(ea)) + if size <= 0: stopped = "undecodable"; break + count += 1 + insn = ida_ua.insn_t() + if ida_ua.decode_insn(insn, ea) > 0: + try: is_ret = bool(ida_idp.is_ret_insn(insn)) + except Exception: is_ret = False + if is_ret or (insn.get_canon_feature() & idaapi.CF_STOP): + ea += size; stopped = "flow"; break + ea += size + result = {"start": hex(start), "end": hex(ea), "count": count, "stopped": stopped} +result +''' + +_OPERATIONS["define_func_run"] = r''' +import ida_bytes, ida_funcs, ida_segment +ea = int(str(a["addr"]), 16) +fn = db.functions.get_at(ea) +if fn is not None and int(fn.start_ea) == ea: + result = {"addr": hex(ea), "ok": True, "start": hex(ea), "end": hex(int(fn.end_ea)), "how": "existed"} +else: + automatic = bool(db.functions.create(ea)) + if not automatic: + seg = db.segments.get_at(ea); end = ea; hi = int(seg.end_ea) if seg else ea + while end < hi and ida_bytes.is_code(ida_bytes.get_flags(end)): + nxt = int(ida_bytes.get_item_end(end)) + if nxt <= end: break + end = nxt + ok = bool(end > ea and ida_funcs.add_func(ea, end)) + else: ok = True + fn = db.functions.get_at(ea) + result = ({"addr": hex(ea), "ok": True, "start": hex(int(fn.start_ea)), + "end": hex(int(fn.end_ea)), "how": "auto" if automatic else "explicit-end"} + if ok and fn is not None else + {"addr": hex(ea), "ok": False, "error": f"IDA refused a function at {ea:#x}"}) +result +''' + +_OPERATIONS["set_thumb"] = r''' +import ida_bytes, ida_ida, ida_idp, ida_segment, ida_segregs +ea = int(str(a["addr"]), 16); treg = ida_idp.str2reg("T") +seg = ida_segment.getseg(ea) +if treg is None or treg < 0: + result = {"addr": hex(ea), "error": "no T register (not an ARM database)"} +elif seg is None: + result = {"addr": hex(ea), "error": "no segment"} +else: + current = ida_segregs.get_sreg(ea, treg) + current = 0 if current in (None, 0xFFFFFFFF, -1) else int(current) + want = {"on": 1, "off": 0}.get(str(a.get("mode", "toggle")).lower(), 0 if current else 1) + changed = False + if want and seg.bitness != 1: + ida_segment.set_segm_addressing(seg, 1); changed = True + size = max(int(ida_bytes.get_item_size(ea)), 2) + 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) + result = {"addr": hex(ea), "thumb": bool(now), "was": bool(current), "ok": ok, + "bitness": ida_segment.getseg(ea).bitness, "forced_32bit": changed, + "db_64bit": bool(ida_ida.inf_get_app_bitness() == 64 and want)} +result +''' + +_OPERATIONS["thumb_scan"] = r''' +import ida_bytes, ida_funcs, ida_idp, ida_segment, ida_segregs, ida_ua +lo, hi = int(str(a["start"]), 16), int(str(a["end"]), 16) +apply, limit = bool(a.get("apply", True)), int(a.get("limit", 512)) +treg = ida_idp.str2reg("T"); found = []; applied = 0; cursor = lo +while cursor + 4 <= hi and len(found) < limit: + at = cursor; value = int(ida_bytes.get_dword(cursor)); cursor += 4 + if not value & 1: continue + target = value & ~1; seg = ida_segment.getseg(target) + if seg is None or not (seg.perm & ida_segment.SEGPERM_EXEC or seg.perm == 0): continue + flags = ida_bytes.get_flags(target) + if ida_bytes.is_data(flags): continue + item = {"at": hex(at), "value": hex(value), "target": hex(target), + "was_code": bool(ida_bytes.is_code(flags))}; found.append(item) + if not apply: continue + if treg is not None and treg >= 0: ida_segregs.split_sreg_range(target, treg, 1, ida_segregs.SR_user) + if not ida_bytes.is_code(ida_bytes.get_flags(target)): + ida_bytes.del_items(target, 0, 2) + if ida_ua.create_insn(target) <= 0: item["decoded"] = False; continue + item["decoded"] = True; item["function"] = bool(db.functions.get_at(target) or db.functions.create(target)); applied += 1 +result = {"start": hex(lo), "end": hex(hi), "found": found, "applied": applied, "n": len(found)} +result +''' + +_OPERATIONS["decomp_error"] = r''' +import ida_hexrays, ida_ida +ea = int(str(a["addr"]), 16); fn = db.functions.get_at(ea) +result = {"addr": hex(ea), "bitness": ida_ida.inf_get_app_bitness()} +if fn is None: + result["reason"] = "no function here" +else: + try: + failure = ida_hexrays.hexrays_failure_t(); cfunc = ida_hexrays.decompile_func(fn, failure) + if cfunc is not None: result["reason"] = "" + else: + result.update({"reason": failure.desc() or f"error {failure.code}", + "code": int(failure.code), "errea": hex(int(failure.errea))}) + except Exception as exc: result["reason"] = f"{type(exc).__name__}: {exc}" +result +''' + + +class CodeModeClient: + """A leased GUI/idalib database accessed through ``ida_codemode``.""" + + def __init__( + self, + binary_path: str, + *, + ttl: int = 0, + load_args: str = "", + processor: str | None = None, + loading_address: int | None = None, + file_type: str | None = None, + output_database: str | None = None, + spawn: bool = True, + new_database: bool = False, + ) -> None: + del ttl # managed-worker lifetime is lease-based, not idle-TTL based + self._path = os.path.abspath(os.path.expanduser(binary_path)) + parsed_processor, parsed_address, parsed_file_type = _parse_load_args(load_args) + self._processor = processor or parsed_processor + self._loading_address = loading_address if loading_address is not None else parsed_address + self._file_type = file_type or parsed_file_type + self._output_database = output_database + self._spawn = spawn + self._new_database = new_database + self._handle: DatabaseHandle | None = None + self._last_entry: RegistryEntry | None = None + self._connect_lock = threading.Lock() + + def connect(self, timeout: float = 1800.0, progress=None) -> "CodeModeClient": + _require_codemode() + with self._connect_lock: + if self._handle is not None and self._handle.connected: + return self + if progress: + progress(f"discovering Code Mode database for {os.path.basename(self._path)}…") + try: + # A Ctrl+L reload releases its current managed-worker lease, but + # that worker remains registered during Code Mode's final-lease + # grace period. Retry only that known handoff window. A GUI or + # another long-lived client remains busy and yields a clear + # failure rather than being modified underneath its owner. + deadline = time.monotonic() + min(timeout, 60.0) + while True: + try: + handle = DatabaseHandle.open( + self._path, + spawn=self._spawn, + timeout=max(0.1, timeout), + output_database=self._output_database, + processor=self._processor, + loading_address=self._loading_address, + file_type=self._file_type, + new_database=self._new_database, + ) + break + except IdbBusy: + if not self._new_database or time.monotonic() >= deadline: + raise + if progress: + progress("waiting for the previous Code Mode lease to close…") + # Remember the record before managed shutdown withdraws + # its JSON. The lifetime lock remains held until IDA has + # actually closed the IDB; waiting on it avoids racing a + # replacement worker into the old process's file lock. + expected = canonical_path( + self._output_database or expected_idb_path(self._path) + ) + owners = [item.entry for item in scan_instances(timeout=0.5) + if item.entry.idb_key == idb_key(expected)] + if owners: + self._wait_for_entry_release( + owners[0], max(0.0, deadline - time.monotonic()) + ) + else: + time.sleep(0.2) + if progress: + backend = handle.entry.backend + progress(f"attached to {backend} database; waiting for auto-analysis…") + handle.wait_autoanalysis(timeout=timeout) + except Exception as exc: # normalize the dependency's transport errors + raise self._connection_error(exc) from exc + self._handle = handle + self._last_entry = handle.entry + return self + + @staticmethod + def _connection_error(exc: BaseException) -> IDAConnectionError: + return IDAConnectionError(str(exc) or type(exc).__name__) + + @property + def connected(self) -> bool: + return self._handle is not None and self._handle.connected + + @property + def pid(self) -> int | None: + return self._handle.entry.pid if self._handle is not None else None + + @property + def backend(self) -> str | None: + return self._handle.entry.backend if self._handle is not None else None + + def execute_python(self, code: str, *, timeout: float | None = None) -> Any: + if not self.connected: + self.connect() + handle = self._handle + if handle is None: + raise IDAConnectionError("Code Mode database is not connected") + try: + response = handle.execute_python(code, timeout=timeout) + except RemoteError as exc: + details = exc.details or {} + message = str(exc) + if details.get("traceback"): + message += f"\n{details['traceback']}" + if exc.code == "operation_timeout": + raise IDATimeoutError(message) from exc + raise IDAToolError("execute_python", message) from exc + except (InstanceDisconnectedError, ClientError) as exc: + raise self._connection_error(exc) from exc + if not isinstance(response, dict) or "result" not in response: + raise IDAToolError("execute_python", "Code Mode returned an invalid execution result") + return response["result"] + + def invoke(self, operation: str, *, timeout: float | None = None, **args) -> Any: + """Execute one TUI domain operation through Code Mode.""" + if operation in ("idb_save", "save"): + return self.save_database() + if operation in ("server_health", "ping", "health", "state"): + return self.health() + body = _HEADS if operation == "heads" else _OPERATIONS.get(operation) + if body is None: + raise IDAToolError(operation, f"unknown ida-tui Code Mode operation: {operation}") + try: + return self.execute_python(_script(args, body), timeout=timeout) + except IDAToolError as exc: + if exc.tool == "execute_python": + raise IDAToolError(operation, exc.message) from exc + raise + + # Temporary source compatibility for external drivers/tests that used the + # old WorkerClient. Application code uses the accurately named invoke(). + call = invoke + + def save_database(self) -> dict[str, Any]: + if not self.connected: + self.connect() + handle = self._handle + if handle is None: + raise IDAConnectionError("Code Mode database is not connected") + try: + return handle.save_database() + except RemoteError as exc: + raise IDAToolError("save_database", str(exc)) from exc + except (InstanceDisconnectedError, ClientError) as exc: + raise self._connection_error(exc) from exc + + def health(self) -> dict[str, Any]: + if not self.connected: + self.connect() + assert self._handle is not None + entry = self._handle.entry + module = os.path.basename(entry.exe_path or entry.idb_path or self._path) + return { + "ok": self._handle.connected, + "module": module, + "backend": entry.backend, + "record_id": entry.record_id, + "input_path": entry.exe_path, + "idb_path": entry.idb_path, + } + + def keepalive(self, interval: float = 120.0) -> _NoopKeepAlive: + del interval + return _NoopKeepAlive() + + def resolve_db(self) -> str: + if not self.connected: + self.connect() + assert self._handle is not None + return self._handle.entry.record_id + + def set_db(self, db: str | None) -> None: + del db # one handle is permanently bound to one registered database + + def list_sessions(self) -> list[Session]: + if not self.connected: + self.connect() + assert self._handle is not None + entry = self._handle.entry + path = entry.exe_path or entry.idb_path or self._path + return [Session(session_id=entry.record_id, filename=os.path.basename(path), + input_path=path, is_active=True)] + + def close(self, grace: float = 0.0) -> None: + del grace + with self._connect_lock: + handle, self._handle = self._handle, None + if handle is not None: + self._last_entry = handle.entry + handle.close() # release our lease; never close a GUI/other client's DB + + @staticmethod + def _wait_for_entry_release(entry: "RegistryEntry", timeout: float) -> bool: + _require_codemode() + path = REGISTRY_DIR / f"{entry.record_id}.lock" + deadline = time.monotonic() + max(0.0, timeout) + while True: + lock = FileLock(path) + try: + if lock.try_acquire(): + return True + except OSError: + pass + finally: + lock.close() + if time.monotonic() >= deadline: + return False + time.sleep(min(0.1, deadline - time.monotonic())) + + def wait_released(self, timeout: float = 45.0) -> bool: + """Wait until a managed instance releases its lifetime lock. + + Normal application shutdown must not wait: another client may retain the + worker. This is an explicit test/maintenance helper for deleting a + temporary IDB safely after this client closes. GUI instances return + ``False`` immediately because clients never own their lifetime. + """ + entry = self._last_entry + if entry is None or entry.backend != "idalib": + return False + return self._wait_for_entry_release(entry, timeout) + + def __enter__(self) -> "CodeModeClient": + return self.connect() + + def __exit__(self, *exc) -> None: + self.close() -- cgit v1.3.1-sl0p From 6e58e9e4ea72771619ded48b0ac2390a58d31cfb Mon Sep 17 00:00:00 2001 From: blasty Date: Fri, 7 Aug 2026 12:56:27 +0200 Subject: codemode: fix DatabaseHandle.open kwarg, and check kwargs against the real signature ida-codemode is now cloned at ../ida-codemode (0.3.1) and installed into ~/ida-venv, so the adapter can be checked against the library instead of against assumptions. First thing it found: connect() passed loading_address=, which DatabaseHandle.open() does not have. The real parameter is image_base, and it already wants the natural 16-byte-aligned address we compute, so this is a rename. Every connect would have died with TypeError on the first call. The port's own contract test could not catch it: its fake handle takes **kwargs, so any keyword at all looks accepted. The test now also validates the keywords we send against inspect.signature(DatabaseHandle.open) when the library is importable, and skips that one check when it is not. Offline suite: 302 passed with the library installed, 302 without it. --- idatui/codemode_client.py | 5 ++++- tests/test_codemode_client.py | 24 +++++++++++++++++++++++- 2 files changed, 27 insertions(+), 2 deletions(-) (limited to 'idatui/codemode_client.py') diff --git a/idatui/codemode_client.py b/idatui/codemode_client.py index 77a0227..b43bfca 100644 --- a/idatui/codemode_client.py +++ b/idatui/codemode_client.py @@ -975,7 +975,10 @@ class CodeModeClient: timeout=max(0.1, timeout), output_database=self._output_database, processor=self._processor, - loading_address=self._loading_address, + # DatabaseHandle calls this image_base and wants the + # natural (16-byte aligned) address; it does the + # conversion to IDA's paragraph-based -b itself. + image_base=self._loading_address, file_type=self._file_type, new_database=self._new_database, ) diff --git a/tests/test_codemode_client.py b/tests/test_codemode_client.py index 0954918..2f70ba3 100644 --- a/tests/test_codemode_client.py +++ b/tests/test_codemode_client.py @@ -76,6 +76,21 @@ class FakeDatabaseHandle: return FakeHandle(path) +def _open_kwargs_are_real(sent: dict): + """(ok, detail) for the kwargs the adapter passes to DatabaseHandle.open. + + Skips (passes) when ida_codemode is not installed, so the file stays pure. + """ + try: + import inspect + from ida_codemode.client import DatabaseHandle as Real + except ImportError: + return True, "ida_codemode not installed - signature not checked" + accepted = set(inspect.signature(Real.open).parameters) + unknown = sorted(set(sent) - accepted) + return not unknown, f"open() rejects {unknown}" + + def main() -> int: proc, base, file_type = _parse_load_args("-parm:ARMv7-M -b800000 -TRaw") check("legacy switches map to typed Code Mode options", @@ -103,8 +118,15 @@ def main() -> int: FakeDatabaseHandle.opened == path and handle is not None) check("typed loader options cross the dependency boundary", FakeDatabaseHandle.kwargs["processor"] == "arm:ARMv7-A" - and FakeDatabaseHandle.kwargs["loading_address"] == 0x1000, + and FakeDatabaseHandle.kwargs["image_base"] == 0x1000, FakeDatabaseHandle.kwargs) + # A fake that swallows **kwargs cannot catch a keyword the real + # library does not have -- which is exactly how this port shipped + # `loading_address` (the real name is `image_base`) and would have + # raised TypeError on the very first connect. Check the names we + # send against the real signature whenever it is importable. + check("every open() keyword exists in the real library", + *_open_kwargs_are_real(FakeDatabaseHandle.kwargs)) check("connect waits for Code Mode autoanalysis", handle.waited == 42, getattr(handle, "waited", None)) check("progress distinguishes discovery and backend attachment", -- cgit v1.3.1-sl0p From 5de318682c527976f3a298d38a0dc947c3a8c0dd Mon Sep 17 00:00:00 2001 From: blasty Date: Fri, 7 Aug 2026 13:02:27 +0200 Subject: codemode: restore the graph view and pseudocode comments Verified live now: ida-codemode 0.3.1 spawns a managed idalib worker on this box, so the pilot suite runs against the port. flowchart: the port simply does not have the operation, so domain.get_flowchart returned None and every graph key reported 'no control-flow graph for this function'. Ported ours onto ida_gdl (ida-domain exposes no basic-block or edge-kind surface). Blocks stay address RANGES, never text -- that is what lets graph boxes reuse the listing's own rows. Graph suite: 0 -> 50 passed. set_comments: the port set only the disassembly comment via db.comments.set_at(), so a comment never appeared in the pseudocode. A Hex-Rays comment is anchored to a ctree location and an anchor the ctree does not own is discarded as an orphan, so the itp slot must be searched until one sticks, and the entry ea is a function comment instead. Ported that logic back. survey_binary: added as the (caught) fallback domain.py expects behind file_regions, so the fallback path is real rather than always empty. --- idatui/codemode_client.py | 102 +++++++++++++++++++++++++++++++++++++++++++--- 1 file changed, 97 insertions(+), 5 deletions(-) (limited to 'idatui/codemode_client.py') diff --git a/idatui/codemode_client.py b/idatui/codemode_client.py index b43bfca..0eac29a 100644 --- a/idatui/codemode_client.py +++ b/idatui/codemode_client.py @@ -719,16 +719,60 @@ for query in queries: result = {"result": all_results} result ''', + # A comment must land in BOTH views, and the pseudocode half is not a + # simple set: db.comments.set_at() alone leaves the pseudocode unchanged. + # Hex-Rays comments are anchored to a ctree location (treeloc_t), and an + # anchor the ctree does not actually own is dropped as an "orphan" -- so the + # itp slot has to be searched until one sticks, exactly as IDA's own UI does. + # Without it a comment silently never appears in the decompilation. "set_comments": r''' +import idaapi, idc, ida_hexrays rows = [] for item in a.get("items", []): - ea, text = int(str(item["addr"]), 16), str(item.get("comment") or "") + addr_s = str(item.get("addr", "")) + text = str(item.get("comment") or "") try: - if text: ok = bool(db.comments.set_at(ea, text)) - else: db.comments.delete_at(ea); ok = True - rows.append({"addr": hex(ea), "ok": ok}) + ea = int(addr_s, 16) + if not idaapi.set_cmt(ea, text, False): + rows.append({"addr": addr_s, + "error": f"Failed to set disassembly comment at {hex(ea)}"}) + continue + if not ida_hexrays.init_hexrays_plugin(): + rows.append({"addr": addr_s}); continue + try: + cfunc = ida_hexrays.decompile(ea) + except Exception: + cfunc = None + if cfunc is None: + rows.append({"addr": addr_s}); continue + if ea == cfunc.entry_ea: + # The signature line carries no ctree item: it is a function comment. + idc.set_func_cmt(ea, text, True) + cfunc.refresh_func_ctext() + rows.append({"addr": addr_s}); continue + eamap = cfunc.get_eamap() + if ea not in eamap: + rows.append({"addr": addr_s, + "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 + placed = False + for itp in range(idaapi.ITP_SEMI, idaapi.ITP_COLON): + tl.itp = itp + cfunc.set_user_cmt(tl, text) + cfunc.save_user_cmts() + cfunc.refresh_func_ctext() + if not cfunc.has_orphan_cmts(): + placed = True; break + cfunc.del_orphan_cmts(); cfunc.save_user_cmts() + rows.append({"addr": addr_s} if placed else + {"addr": addr_s, + "error": f"Failed to set decompiler comment at {hex(ea)}"}) except Exception as exc: - rows.append({"addr": hex(ea), "ok": False, "error": str(exc)}) + rows.append({"addr": addr_s, "error": str(exc)}) result = {"result": rows} result ''', @@ -923,6 +967,54 @@ else: result ''' +# The graph view's only backend call. Blocks are address RANGES, never text: +# the client re-renders them with `heads`, so boxes reuse the exact listing rows +# (colours, operand marks, trail painting) instead of growing a second renderer. +# +# ida-domain exposes no basic-block/edge-kind surface, so this stays on ida_gdl. +_OPERATIONS["flowchart"] = r''' +import ida_funcs, ida_gdl +ea = int(str(a["addr"]), 16) +fn = ida_funcs.get_func(ea) +if fn is None: + result = {"addr": hex(ea), "error": "no function at that address", "blocks": []} +else: + fc = ida_gdl.FlowChart(fn, flags=ida_gdl.FC_PREDS) + index, order = {}, [] + for bb in fc: + index[bb.start_ea] = len(order) + order.append(bb) + blocks = [] + for bb in order: + sl = [s for s in bb.succs() if s.start_ea in index] + succs = [] + for s in sl: + # Edge kind is what the graph view colours by: an n-way dispatch is + # "switch", a successor that is literally the next address falls + # through, anything else is a taken branch. + if len(sl) > 2: kind = "switch" + elif s.start_ea == bb.end_ea: kind = "fall" + else: kind = "jump" + succs.append([index[s.start_ea], kind]) + blocks.append({"id": index[bb.start_ea], "start": hex(int(bb.start_ea)), + "end": hex(int(bb.end_ea)), "succs": succs}) + result = {"addr": hex(ea), + "func": {"addr": hex(int(fn.start_ea)), "end": hex(int(fn.end_ea)), + "name": ida_funcs.get_func_name(fn.start_ea) or ""}, + "entry": index.get(fn.start_ea, 0), "blocks": blocks} +result +''' + +# Only ever reached as domain.py's fallback when file_regions yields nothing. +_OPERATIONS["survey_binary"] = r''' +segments = [] +for seg in db.segments.get_all(): + segments.append({"start": hex(int(seg.start_ea)), "end": hex(int(seg.end_ea)), + "name": db.segments.get_name(seg) or ""}) +result = {"segments": segments} +result +''' + class CodeModeClient: """A leased GUI/idalib database accessed through ``ida_codemode``.""" -- cgit v1.3.1-sl0p From 3edb58cc0b93fdd90b68c6aacc8af91874142099 Mon Sep 17 00:00:00 2001 From: blasty Date: Fri, 7 Aug 2026 13:10:53 +0200 Subject: codemode: port the xref tools' real contract, order included The pseudocode follow's address fallback broke: following a call landed on the NEXT LINE instead of the callee (decomp_nav's stale-name check, cur=0x20dd want=0x2060). The port's xref_query returned rows in raw IDA order, and at a call site IDA yields the ordinary-flow xref (fl_F, the next instruction) before the call xref (fl_CN), so 'first code xref' picked the fall-through. The tool ida-tui was written against sorts rows by the far-end address and dedups by default; sorted, 0x2060 precedes 0x210e and the follow is correct. That ordering is load-bearing, so it is now part of the port rather than an accident of the old implementation. Also fixed: the port attached 'fn' to ref.from_ea for both directions, where a from-xref must describe its TARGET (the xref dialog shows the wrong function otherwise), and the envelope was missing direction/addr/total/next_offset/resolved_addr. xref_types (ours, the kind badges in the xref dialog) is ported verbatim and deliberately stays UNsorted -- that dialog lists xrefs in IDA's own order. decomp_nav, follow_xrefs, xref_labels, decomp_follow_self: 15 passed, 0 failed. --- idatui/codemode_client.py | 143 +++++++++++++++++++++++++++++++++++----------- 1 file changed, 111 insertions(+), 32 deletions(-) (limited to 'idatui/codemode_client.py') diff --git a/idatui/codemode_client.py b/idatui/codemode_client.py index 0eac29a..18504a5 100644 --- a/idatui/codemode_client.py +++ b/idatui/codemode_client.py @@ -673,49 +673,128 @@ for query in a.get("queries", []): result = {"result": rows} result ''', + # Ours: the coarse code/data type plus a fine `kind` (call/jump/flow, + # read/write/offset/text/info) that the xref dialog draws its badges from. + # Deliberately NOT sorted -- the dialog lists xrefs in IDA's own order. "xref_types": r''' +import idaapi, idautils, ida_bytes, ida_funcs, 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): + return (code_kind if xr.iscode else data_kind).get(xr.type, "code" if xr.iscode else "data") +def _fn(ea): + f = ida_funcs.get_func(ea) + return {"addr": hex(int(f.start_ea)), "name": ida_funcs.get_func_name(f.start_ea) or ""} if f else None queries = a.get("queries") or [] all_results = [] for query in queries: - ea, direction = int(str(query["addr"]), 16), str(query.get("direction", "both")) - refs = [] - if direction in ("to", "both"): refs += list(db.xrefs.to_ea(ea)) - if direction in ("from", "both"): refs += list(db.xrefs.from_ea(ea)) - rows, seen = [], set() - for ref in refs: - key = (int(ref.from_ea), int(ref.to_ea), int(ref.type)) - if query.get("dedup") and key in seen: continue - seen.add(key) - fn = db.functions.get_at(int(ref.from_ea)) - kind = ("call" if ref.is_call else "jump" if ref.is_jump else "flow" if ref.is_flow - else "read" if ref.is_read else "write" if ref.is_write else ref.type.name.lower()) - row = {"from": hex(int(ref.from_ea)), "to": hex(int(ref.to_ea)), - "type": "code" if ref.is_code else "data", "kind": kind} - if query.get("include_fn") and fn is not None: - row["fn"] = {"addr": hex(int(fn.start_ea)), "name": db.functions.get_name(fn) or ""} - rows.append(row) - if len(rows) >= int(query.get("count", 2000)): break - all_results.append({"data": rows}) + query = query if isinstance(query, dict) else {"addr": query} + raw = str(query.get("addr", "")).strip() + direction = str(query.get("direction", "to") or "to").lower() + include_fn = bool(query.get("include_fn", True)) + dedup = bool(query.get("dedup", True)) + try: count = int(query.get("count", 2000) or 2000) + except (TypeError, ValueError): count = 2000 + try: target = int(raw, 16) + except ValueError: target = idaapi.get_name_ea(idaapi.BADADDR, 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(int(xr.frm)), "from": hex(int(xr.frm)), + "to": hex(int(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(int(xr.to)), "from": hex(int(target)), + "to": hex(int(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, deduped = set(), [] + 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] + all_results.append({"query": raw, "data": rows, "next_offset": None}) result = {"result": all_results} result ''', + # Mirrors the tool ida-tui was written against, ORDER INCLUDED. The rows are + # sorted by the far-end address and deduped by default, and the pseudocode + # follow's address fallback silently depends on it: at a call site the raw + # IDA order yields the ordinary-flow xref (the next instruction) first, so an + # unsorted result makes "follow the call" land on the following line instead. "xref_query": r''' +import idaapi, idautils, ida_bytes, ida_funcs +def _fn(ea): + f = ida_funcs.get_func(ea) + return {"addr": hex(int(f.start_ea)), "name": ida_funcs.get_func_name(f.start_ea) or ""} if f else None queries = a.get("queries") or [] all_results = [] for query in queries: - ea, direction = int(str(query["addr"]), 16), str(query.get("direction", "both")) - refs = [] - if direction in ("to", "both"): refs += list(db.xrefs.to_ea(ea)) - if direction in ("from", "both"): refs += list(db.xrefs.from_ea(ea)) - rows = [] - for ref in refs[:int(query.get("count", 2000))]: - fn = db.functions.get_at(int(ref.from_ea)) - row = {"from": hex(int(ref.from_ea)), "to": hex(int(ref.to_ea)), - "type": "code" if ref.is_code else "data"} - if query.get("include_fn") and fn is not None: - row["fn"] = {"addr": hex(int(fn.start_ea)), "name": db.functions.get_name(fn) or ""} - rows.append(row) - all_results.append({"data": rows}) + raw = str(query.get("addr", "")).strip() + direction = str(query.get("direction", "both") or "both").lower() + if direction not in ("to", "from", "both"): direction = "both" + xref_type = str(query.get("xref_type", "any") or "any").lower() + if xref_type not in ("any", "code", "data"): xref_type = "any" + include_fn = bool(query.get("include_fn", True)) + dedup = bool(query.get("dedup", True)) + sort_by = str(query.get("sort_by", "addr") or "addr") + descending = bool(query.get("descending", False)) + try: offset = max(0, int(query.get("offset", 0) or 0)) + except (TypeError, ValueError): offset = 0 + try: count = max(0, min(int(query.get("count", 200) or 200), 5000)) + except (TypeError, ValueError): count = 200 + try: + try: target = int(raw, 16) + except ValueError: + target = idaapi.get_name_ea(idaapi.BADADDR, raw) + if target == idaapi.BADADDR: raise ValueError(f"Failed to resolve address/name: {raw}") + if not ida_bytes.is_mapped(target): raise ValueError(f"Address not mapped: {raw}") + rows = [] + if direction in ("to", "both"): + for xr in idautils.XrefsTo(target, 0): + kind = "code" if xr.iscode else "data" + if xref_type != "any" and kind != xref_type: continue + row = {"direction": "to", "addr": hex(int(xr.frm)), "from": hex(int(xr.frm)), + "to": hex(int(target)), "type": kind} + if include_fn: row["fn"] = _fn(xr.frm) + rows.append(row) + if direction in ("from", "both"): + for xr in idautils.XrefsFrom(target, 0): + kind = "code" if xr.iscode else "data" + if xref_type != "any" and kind != xref_type: continue + row = {"direction": "from", "addr": hex(int(xr.to)), "from": hex(int(target)), + "to": hex(int(xr.to)), "type": kind} + if include_fn: row["fn"] = _fn(xr.to) + rows.append(row) + if dedup: + seen, deduped = set(), [] + for row in rows: + key = (row["direction"], row["from"], row["to"], row["type"]) + if key in seen: continue + seen.add(key); deduped.append(row) + rows = deduped + if sort_by == "type": + rows.sort(key=lambda r: (str(r.get("type", "")), int(str(r["addr"]), 16)), reverse=descending) + else: + rows.sort(key=lambda r: int(str(r["addr"]), 16), reverse=descending) + page = rows[offset:offset + count] if count else rows[offset:] + nxt = offset + len(page) + all_results.append({"target": raw, "resolved_addr": hex(int(target)), "direction": direction, + "xref_type": xref_type, "data": page, + "next_offset": nxt if nxt < len(rows) else None, + "total": len(rows), "error": None}) + except Exception as exc: + all_results.append({"target": raw, "resolved_addr": None, "direction": direction, + "xref_type": xref_type, "data": [], "next_offset": None, + "total": 0, "error": str(exc)}) result = {"result": all_results} result ''', -- cgit v1.3.1-sl0p From ddeb7afea4553e73a7c1decd78fde968e40d186e Mon Sep 17 00:00:00 2001 From: blasty Date: Fri, 7 Aug 2026 13:34:03 +0200 Subject: codemode: carry over the listing + operand-format tools, and stop reshipping them This closes the five operations the port was missing and restores the listing's own tooling instead of a re-implementation of it. idatui/remote_tools.py is the port's IDAPython island: `heads` (the continuous listing) and `op_format`/`pc_nums`/`pc_num_format` (`o`/`O`), extracted verbatim from the BODY that server/patch_server.py used to inject. They are real, diffable source shipped to the database process as text, not string literals, because this is the most performance-tuned and behaviour-sensitive code in the project. Why carry `heads` over rather than keep the port's version: the port's rewrite emitted no per-operand extents ("ops"), so no keypress could show which literal it would reformat (opfmt_highlight had no two-operand row to find); it had no digest/`expect` support, so every page was re-sent after any edit; and its span walk was the per-character loop ours had already been rewritten out of. It also dropped struct-member expansion sizing and the func banner/label rows' exact shapes. The library is installed ONCE per database process (sys.modules, keyed by a hash of the source) and then called by name. Code Mode's execute_python builds a fresh namespace per call, so a library exec'd inline is rebuilt every time and its module-level caches thrown away -- the per-line render lru_cache in particular, which the perf work sized to 65536 entries. Installing it once took `heads` count=200 from 181ms to 92ms; the cache reports 211 hits on a second call where it previously reported none. (Extraction footgun recorded: ast FunctionDef.lineno points at `def`, not at the decorators, so a naive slice silently drops @lru_cache.) Also ported: flowchart, survey_binary, and the xref contract. Live pilot suite on targets/echo: 301 passed, 0 failed -- identical to master. Known, quantified, and NOT fixed here: Code Mode's transport is much slower than the unix-socket worker for the listing's paging. heads count=200 is 2.6ms on master vs 92ms here, count=500 is 6.3ms vs 214ms. Roughly half of that is to_jsonable + HTTP framing per call and is inherent to the architecture; the empty round trip alone is 2ms. The digest/`expect` path (unchanged pages) is the main mitigation and is restored. --- idatui/codemode_client.py | 250 +++------ idatui/remote_tools.py | 1367 +++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 1442 insertions(+), 175 deletions(-) create mode 100644 idatui/remote_tools.py (limited to 'idatui/codemode_client.py') diff --git a/idatui/codemode_client.py b/idatui/codemode_client.py index 18504a5..069cf98 100644 --- a/idatui/codemode_client.py +++ b/idatui/codemode_client.py @@ -15,11 +15,13 @@ IDAPython modules that Code Mode deliberately makes importable. """ from __future__ import annotations +import hashlib import json import os import shlex import threading import time +from pathlib import Path from textwrap import dedent from typing import Any @@ -169,180 +171,6 @@ def _script(args: dict[str, Any], body: str) -> str: return f"import json\na = json.loads({encoded!r})\n{dedent(body).strip()}\n" -# Rich flat-listing generation is the largest ida-domain gap in this port. -# ida-domain can enumerate heads and render plain disassembly, but it does not -# expose undefined runs, IDA colour spans, function banners, or expanded UDT -# members. Keep that IDAPython-only logic isolated in this one operation. -_HEADS = r''' -import ida_bytes, ida_funcs, ida_idaapi, ida_lines, ida_name, ida_nalt, ida_segment, ida_typeinf -start = int(str(a["addr"]), 16) -count = max(1, min(int(a.get("count", 200)), 2000)) -offset = max(0, int(a.get("offset", 0))) -annotate = bool(a.get("annotate", False)) -seg = db.segments.get_at(start) -if seg is None: - result = {"addr": a["addr"], "error": "no segment", "heads": [], "cursor": {"done": True}} -else: - lo, hi = int(seg.start_ea), int(seg.end_ea) - if a.get("end"): - hi = min(hi, int(str(a["end"]), 16)) - - span_names = { - "insn": ("SCOLOR_INSN", "SCOLOR_KEYWORD", "SCOLOR_ASMDIR", "SCOLOR_MACRO"), - "reg": ("SCOLOR_REG",), - "num": ("SCOLOR_NUMBER", "SCOLOR_CHAR", "SCOLOR_BINPREF"), - "str": ("SCOLOR_STRING",), - "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",), - } - tag_kinds = {} - for kind, names in span_names.items(): - for name in names: - value = getattr(ida_lines, name, None) - if isinstance(value, str) and value: - tag_kinds[value[0]] = kind - elif isinstance(value, int): - tag_kinds[chr(value)] = kind - - def spans(tagged): - 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)) - out, stack, buf = [], [], [] - def flush(): - if buf: - out.append([stack[-1] if stack else "text", "".join(buf)]) - buf.clear() - i = 0 - while i < len(tagged): - ch = tagged[i] - if ch == on and i + 1 < len(tagged): - tag = tagged[i + 1] - if tag == addr_tag: - i += 2 + addr_len - continue - flush(); stack.append(tag_kinds.get(tag, "text")); i += 2; continue - if ch == off and i + 1 < len(tagged): - flush() - if stack: stack.pop() - i += 2; continue - if ch == esc and i + 1 < len(tagged): - buf.append(tagged[i + 1]); i += 2; continue - buf.append(ch); i += 1 - flush() - collapsed, previous_space = [], False - for kind, text in out: - acc = [] - for ch in text: - if ch.isspace(): - if previous_space: continue - acc.append(" "); previous_space = True - else: - acc.append(ch); previous_space = False - if acc: collapsed.append([kind, "".join(acc)]) - if collapsed: - collapsed[0][1] = collapsed[0][1].lstrip() - collapsed[-1][1] = collapsed[-1][1].rstrip() - return [[kind, text] for kind, text in collapsed if text] - - def row(ea): - flags = ida_bytes.get_flags(ea) - kind = "code" if ida_bytes.is_code(flags) else ("data" if ida_bytes.is_data(flags) else "unknown") - tagged = ida_lines.generate_disasm_line(ea, 0) or "" - text = " ".join(ida_lines.tag_remove(tagged).split()) if tagged else "" - item = {"ea": hex(ea), "kind": kind, "size": int(ida_bytes.get_item_size(ea)), "text": text} - if tagged: - rich = spans(tagged) - if " ".join("".join(x[1] for x in rich).split()) == text: - item["spans"] = rich - name = ida_name.get_ea_name(ea) - if name: item["name"] = name - return item - - def unknown_row(ea, size): - if size <= 1: return row(ea) - item = {"ea": hex(ea), "kind": "unknown", "size": int(size), "text": f"db {size} dup(?)"} - name = ida_name.get_ea_name(ea) - if name: item["name"] = name - return item - - def members(ea): - tif = db.types.get_at(ea) - if tif is None or not tif.is_udt(): return [] - answer = [] - for member in db.types.get_udt_members(tif): - type_text = member.type.dstr() or "" - text = f"+{member.offset:X} {member.name}" + (f" {type_text}" if type_text else "") - answer.append({"ea": hex(ea + member.offset), "kind": "member", - "size": int(member.size), "text": text}) - return answer - - def is_unknown(ea): - flags = ida_bytes.get_flags(ea) - return not (ida_bytes.is_code(flags) or ida_bytes.is_data(flags)) - def run_end(ea): - nxt = ida_bytes.next_head(ea, hi) - return nxt if nxt != ida_idaapi.BADADDR and ea < nxt <= hi else hi - def advance(ea): - if is_unknown(ea): return run_end(ea) - nxt = ida_bytes.get_item_end(ea) - return nxt if nxt > ea else ea + 1 - def rows_for(ea): - if is_unknown(ea): return [unknown_row(ea, run_end(ea) - ea)] - fn = db.functions.get_at(ea) if annotate else None - at_start = fn is not None and int(fn.start_ea) == ea - answer = [] - if at_start: - name = db.functions.get_name(fn) or f"sub_{ea:X}" - answer += [ - {"ea": hex(ea), "kind": "sep", "size": 0, "text": ""}, - {"ea": hex(ea), "kind": "sep", "size": 0, - "text": "; " + "=" * 15 + " S U B R O U T I N E " + "=" * 15}, - {"ea": hex(ea), "kind": "funchdr", "size": 0, - "text": name + " proc", "name": name}, - ] - item = row(ea) - if at_start: - item["name"] = None - elif annotate and item["kind"] == "code" and item.get("name"): - name = item["name"] - answer.append({"ea": hex(ea), "kind": "label", "size": 0, - "text": name + ":", "name": name}) - item["name"] = None - answer.append(item) - if item["kind"] == "data": answer += members(ea) - if fn is not None and ida_bytes.get_item_end(ea) >= int(fn.end_ea): - name = db.functions.get_name(fn) or f"sub_{int(fn.start_ea):X}" - answer += [ - {"ea": hex(ea), "kind": "funchdr", "size": 0, - "text": name + " endp", "name": name}, - {"ea": hex(ea), "kind": "sep", "size": 0, "text": "; " + "-" * 60}, - ] - return answer - - ea = ida_bytes.get_item_head(start) - if ea == ida_idaapi.BADADDR: ea = start - for _ in range(offset): - if ea >= hi: break - ea = advance(ea) - rows = [] - more = False - while ea != ida_idaapi.BADADDR and ea < hi: - if len(rows) >= count: - more = True; break - rows += rows_for(ea) - ea = advance(ea) - result = {"addr": a["addr"], "heads": rows, - "cursor": {"next": hex(ea)} if more else {"done": True}} -result -''' - - _DECOMP_MAP_HELPER = r''' def line_map(cfunc): import ida_hexrays @@ -1046,6 +874,72 @@ else: result ''' +# `heads` and the operand-format tools are the port's IDAPython island: the +# continuous listing's presentation model (undefined runs, colour spans, operand +# extents, banners, struct members, the digest protocol) and IDA/Hex-Rays number +# formats have no ida-domain surface. Rather than paraphrase ~1100 lines of +# performance-tuned, behaviour-sensitive code into string literals, they stay +# real, diffable source in idatui/remote_tools.py and are shipped to the database +# process as text. Read once at import; the file ships beside this module. +_REMOTE_LIB = (Path(__file__).with_name("remote_tools.py")).read_text(encoding="utf-8") + +#: Versioned by content, so editing remote_tools.py re-installs it instead of +#: silently running the copy a long-lived worker already has. +_REMOTE_MODULE = "_idatui_remote_" + hashlib.sha1( + _REMOTE_LIB.encode("utf-8")).hexdigest()[:12] + +#: Sent back when the database process has not got the library yet; the client +#: installs it and retries once. Amortised, a worker receives it exactly once. +_NEED_LIB = "__idatui_needs_remote_lib__" + +#: Installs the library as a real module in the database process. Persisting it +#: in sys.modules is what makes the module-level caches (the tag maps, and the +#: line-render lru_cache the listing's throughput depends on) survive between +#: calls -- execute_python builds a fresh namespace every time, so a library +#: exec'd inline is rebuilt, and its caches thrown away, on every single call. +_INSTALL_LIB = f''' +import sys, types +_m = types.ModuleType({_REMOTE_MODULE!r}) +exec(compile(a["source"], {_REMOTE_MODULE!r}, "exec"), _m.__dict__) +sys.modules[{_REMOTE_MODULE!r}] = _m +result = True +result +''' + + +def _remote_op(call: str) -> str: + """A snippet that calls one of the carried-over tools by its real signature. + + Costs one short request: the library is imported from the database process's + own sys.modules, not shipped again. + """ + return (f"import sys\n" + f"_m = sys.modules.get({_REMOTE_MODULE!r})\n" + f"result = {{{_NEED_LIB!r}: True}} if _m is None else _m.{call}\n" + f"result\n") + + +_OPERATIONS["op_format"] = _remote_op( + 'op_format(addr=a["addr"], mode=a.get("mode", "cycle"),' + ' col=int(a.get("col", -1)), n=int(a.get("n", -1)))') +_OPERATIONS["pc_nums"] = _remote_op('pc_nums(addr=a["addr"])') +_OPERATIONS["pc_num_format"] = _remote_op( + 'pc_num_format(addr=a["addr"], mode=a.get("mode", "cycle"),' + ' line=int(a.get("line", -1)), col=int(a.get("col", -1)),' + ' ea=a.get("ea", ""), opnum=int(a.get("opnum", -1)))') + +# The listing walker itself. Replaces the port's re-implementation, which +# rendered no per-operand extents (so no keypress could say which literal it +# would reformat) and had no digest/expect support (so every page was re-sent +# after any edit), and whose span walk was the per-character loop our own +# version had already been rewritten to avoid. +_HEADS = _remote_op( + 'heads(addr=a["addr"], count=int(a.get("count", 200)),' + ' offset=int(a.get("offset", 0)), end=a.get("end", ""),' + ' back=bool(a.get("back", False)), annotate=bool(a.get("annotate", False)),' + ' expect=a.get("expect", ""))') + + # The graph view's only backend call. Blocks are address RANGES, never text: # the client re-renders them with `heads`, so boxes reuse the exact listing rows # (colours, operand marks, trail painting) instead of growing a second renderer. @@ -1232,7 +1126,13 @@ class CodeModeClient: if body is None: raise IDAToolError(operation, f"unknown ida-tui Code Mode operation: {operation}") try: - return self.execute_python(_script(args, body), timeout=timeout) + answer = self.execute_python(_script(args, body), timeout=timeout) + if isinstance(answer, dict) and answer.get(_NEED_LIB): + # First call against this database process (or a restarted one). + self.execute_python(_script({"source": _REMOTE_LIB}, _INSTALL_LIB), + timeout=timeout) + answer = self.execute_python(_script(args, body), timeout=timeout) + return answer except IDAToolError as exc: if exc.tool == "execute_python": raise IDAToolError(operation, exc.message) from exc diff --git a/idatui/remote_tools.py b/idatui/remote_tools.py new file mode 100644 index 0000000..41de0de --- /dev/null +++ b/idatui/remote_tools.py @@ -0,0 +1,1367 @@ +"""The IDAPython ida-tui runs inside the Code Mode sandbox. + +Two features have no ida-domain surface at all and are carried over VERBATIM +from the tools ida-tui was developed against (`server/patch_server.py`'s +injected BODY, which the Code Mode port deletes): + +* `heads` -- the continuous listing. ida-domain enumerates defined heads and + renders plain disassembly; the listing also needs coalesced undefined runs, + IDA colour-tag spans, PER-OPERAND EXTENTS, function banners, code labels and + expanded struct members, plus the digest/`expect` protocol the paging layer + uses to skip re-sending a page that has not changed. +* `op_format` / `pc_nums` / `pc_num_format` -- `o`/`O`. IDA's operand types and + Hex-Rays' per-(ea, opnum) numforms are separate sets, and neither is exposed. + +Keeping the originals rather than paraphrasing them is deliberate: this is the +most performance-tuned and most behaviour-sensitive code in the project (the +span walker is a single regex pass because a per-character loop was the most +expensive thing the listing did, and the cycle only offers stops that change +what you see). A re-implementation drifts from it silently. + +This file is SOURCE SHIPPED AS TEXT to the database process; it is never +imported here, because the ida_* modules do not exist in the TUI's interpreter. +`codemode_client` reads it and prepends it to the relevant snippets. Keep it +self-contained: no relative imports, nothing beyond what Code Mode provides. +""" +# ruff: noqa +import re as _re + +from typing import Annotated # the extracted tool signatures still carry these + + +class IDAError(Exception): + """The MCP host's error type; the tools raise/catch it by name.""" + + +def parse_address(addr): + """ida_pro_mcp.utils.parse_address: hex/decimal string, int, or a symbol.""" + if isinstance(addr, int): + return addr + try: + return int(addr, 0) + except ValueError: + import idaapi + ea = idaapi.get_name_ea(idaapi.BADADDR, str(addr).strip()) + if ea != idaapi.BADADDR: + return ea + raise IDAError(f"Not found: {addr!r}") + + +#: Byte-identical to ida_pro_mcp.utils._STRING_OR_SPACES_RE: the pseudocode +#: column coordinates the client holds depend on collapsing exactly the same way. +_IDATUI_STRING_OR_SPACES_RE = _re.compile( + r'"(?:[^"\\]|\\.)*"' # double-quoted string + r"|'(?:[^'\\]|\\.)*'" # single-quoted string / char + r"|[ \t]{2,}" # run of 2+ whitespace (outside strings) +) + + +def compact_whitespace(line: str) -> str: + """ida_pro_mcp.utils.compact_whitespace: collapse runs of 2+ spaces/tabs to + one, preserving string literals.""" + stripped = line.lstrip(" \t") + if not stripped: + return line + lead = line[: len(line) - len(stripped)] + + def _repl(m): + s = m.group() + if s[0] in ('"', "'"): + return s # preserve string content + return " " + + return lead + _IDATUI_STRING_OR_SPACES_RE.sub(_repl, stripped) + + +def _idatui_head_row(ea, flags=None): + """One flat-listing row for the head at ``ea``: kind (code/data/unknown), + byte size, rendered text, and any symbol name. + + ``flags`` lets a caller that already asked for them say so -- the walk in + ``heads`` used to fetch them three times per head (here, in _is_unknown from + _advance, and again from _rows_for). + """ + import ida_bytes + import ida_lines + import ida_name + + f = ida_bytes.get_flags(ea) if flags is None else flags + 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, spans, ops = _idatui_line_parts(line) if line else ("", None, None) + row = { + "ea": hex(ea), + "kind": kind, + "size": int(ida_bytes.get_item_size(ea)), + "text": text, + } + if spans is not None: + 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 + + +import functools as _idatui_functools + + +import os as _idatui_os + + +_IDATUI_LINE_CACHE = int(_idatui_os.environ.get("IDATUI_LINE_CACHE") or 65536) + + +def _idatui_line_parts(line): + """``(text, spans, ops)`` for one tagged disassembly line -- memoised. + + A function of the tagged line and nothing else, so the same line always + gives the same answer: a rename changes the line, which changes the key. + And listings repeat themselves hard -- 196k lines of bash are 53k distinct + ones, so a 16k-entry cache serves ~70% of them and takes the per-line cost + from 10.4us to 3.9us. This is the most expensive thing the backend does per + listing row, and a jump to an address near the end of a big binary walks + hundreds of thousands of them. + + ``spans`` is None when the tag walk and the plain text disagree about what + the line says (then the text wins and the row renders unhighlighted). + + The returned lists are SHARED between every row that has the same line; + treat them as read-only. Pickle notices the sharing too, so a page of + repetitive disassembly also serialises smaller. + """ + import ida_lines + text = " ".join(ida_lines.tag_remove(line).split()) # collapse the padding + spans, ops = _idatui_spans(line) + # Built from the SAME line as `text`, then whitespace-collapsed identically, + # so the two can never disagree about what the row says. + joined = "".join([t for _k, t in spans]) + if " ".join(joined.split()) != text: + return (text, None, None) + return (text, spans, ops) + + +def _idatui_head_row(ea, flags=None): + """One flat-listing row for the head at ``ea``: kind (code/data/unknown), + byte size, rendered text, and any symbol name. + + ``flags`` lets a caller that already asked for them say so -- the walk in + ``heads`` used to fetch them three times per head (here, in _is_unknown from + _advance, and again from _rows_for). + """ + import ida_bytes + import ida_lines + import ida_name + + f = ida_bytes.get_flags(ea) if flags is None else flags + 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, spans, ops = _idatui_line_parts(line) if line else ("", None, None) + row = { + "ea": hex(ea), + "kind": kind, + "size": int(ida_bytes.get_item_size(ea)), + "text": text, + } + if spans is not None: + 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 + + +import functools as _idatui_functools + + +import os as _idatui_os + + +_IDATUI_LINE_CACHE = int(_idatui_os.environ.get("IDATUI_LINE_CACHE") or 65536) + + +@_idatui_functools.lru_cache(maxsize=_IDATUI_LINE_CACHE) +def _idatui_line_parts(line): + """``(text, spans, ops)`` for one tagged disassembly line -- memoised. + + A function of the tagged line and nothing else, so the same line always + gives the same answer: a rename changes the line, which changes the key. + And listings repeat themselves hard -- 196k lines of bash are 53k distinct + ones, so a 16k-entry cache serves ~70% of them and takes the per-line cost + from 10.4us to 3.9us. This is the most expensive thing the backend does per + listing row, and a jump to an address near the end of a big binary walks + hundreds of thousands of them. + + ``spans`` is None when the tag walk and the plain text disagree about what + the line says (then the text wins and the row renders unhighlighted). + + The returned lists are SHARED between every row that has the same line; + treat them as read-only. Pickle notices the sharing too, so a page of + repetitive disassembly also serialises smaller. + """ + import ida_lines + text = " ".join(ida_lines.tag_remove(line).split()) # collapse the padding + spans, ops = _idatui_spans(line) + # Built from the SAME line as `text`, then whitespace-collapsed identically, + # so the two can never disagree about what the row says. + joined = "".join([t for _k, t in spans]) + if " ".join(joined.split()) != text: + return (text, None, None) + return (text, spans, ops) + + +_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 + + +_IDATUI_CTL = None # re: a tag = one of three control chars plus its argument + + +_IDATUI_TAGINFO = 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, _IDATUI_CTL, _IDATUI_TAGINFO + 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() + if _IDATUI_CTL is None: + import re as _re + # One capturing split gives [text, tag, text, tag, ..., text] in a + # single C pass. A per-character python loop over the line used to be + # the most expensive thing the `heads` tool did, and a line is ~54 + # characters but only ~13 tags -- everything between two tags is already + # exactly one span's worth of text. + _IDATUI_CTL = _re.compile("([\\x01\\x02\\x03](?s:.))") + if _IDATUI_TAGINFO is None: + _IDATUI_TAGINFO = { + tag: (_IDATUI_TAGS.get(tag, "text"), _IDATUI_OPND_TAGS.get(tag)) + for tag in set(_IDATUI_TAGS) | set(_IDATUI_OPND_TAGS)} + taginfo = _IDATUI_TAGINFO + plain_tag = ("text", None) + 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)) + parts = _IDATUI_CTL.split(line) + spans, stack = [], [] # stack entries: (kind, operand index|None) + kind, opnd = "text", None # state the current run of text belongs to + pend = "" + skip = 0 # characters of an address payload still due + i, n = 0, len(parts) + while i < n: + txt = parts[i] + i += 1 + if skip: + if len(txt) <= skip: + skip -= len(txt) + txt = "" + else: + txt = txt[skip:] + skip = 0 + if txt: + pend += txt + if i >= n: + break + pair = parts[i] + i += 1 + if skip: # a tag INSIDE an address payload: 2 chars + skip = skip - 2 if skip > 2 else 0 + continue + ch = pair[0] + if ch == esc: # escaped literal: keep the char it guards + pend += pair[1] + continue + tag = pair[1] + if ch == on and tag == addr_tag: + # An embedded target address, not display text: 16 hex digits that + # must not reach the screen. Deliberately NOT a span boundary. + skip = addr_len + continue + if pend: + spans.append([kind, pend, opnd]) + pend = "" + if ch == on: + stack.append((kind, opnd)) + kind, o = taginfo.get(tag, plain_tag) + if o is not None: + opnd = o # operands nest: an inner colour keeps the operand + elif stack: + kind, opnd = stack.pop() + else: + kind, opnd = "text", None + if pend: + spans.append([kind, pend, opnd]) + # Collapse IDA's column padding EXACTLY as the plain text does. A run of + # spaces can straddle two spans, so the leading space of a span is dropped + # when the previous one ended in space — otherwise the spans and `text` + # disagree about the line and the row silently loses its highlighting. + # ``" ".join(txt.split())`` splits on exactly what str.isspace() calls + # whitespace, which is what the character walk this replaces tested. + out, prev_space = [], False + for kind, txt, opnd in spans: + core = " ".join(txt.split()) + if core == txt: + # Nothing to collapse and no edge whitespace -- which is the common + # case ("mov", "rax", ", ") and skips both isspace() probes below. + prev_space = False + out.append([kind, txt, opnd]) + continue + if not core: # the span is nothing but padding + if not prev_space: + prev_space = True + out.append([kind, " ", opnd]) + continue + acc = core + if txt[0].isspace() and not prev_space: + acc = " " + acc + if txt[-1].isspace(): + acc += " " + prev_space = acc[-1] == " " + out.append([kind, 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_rows_digest(rows): + """A value that changes whenever any of ``rows`` would render differently. + + Covers everything a client keeps off a row: address, kind, size, the plain + text, the symbol name and the colour spans (which is what makes it exact + rather than a heuristic -- two lines can collapse to the same text and still + be coloured differently). + + Uses the interpreter's own ``hash``, deliberately. It never has to mean + anything outside this process: the client stores what a page hashed to when + it loaded it and hands the same number back to ask whether the page still + hashes to that. One worker, one process, one hash seed. + """ + acc = 0 + # The per-line render is memoised, so one spans list is shared by every row + # that says the same thing -- about 45% of them within a page. Hash each + # distinct list once and key that by identity, rather than rebuilding a + # tuple of tuples per row (which is the exact cost that was measured and + # removed from the client side for the same reason). + seen = {} + for r in rows: + sp = r.get("spans") + if sp is None: + sh = None + else: + key = id(sp) + sh = seen.get(key) + if sh is None: + sh = seen[key] = hash(tuple(map(tuple, sp))) + acc = hash((acc, r.get("ea"), r.get("kind"), r.get("size"), + r.get("text"), r.get("name"), sh)) + return acc + + +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}, + ] + + +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, + expect: Annotated[str, "Digest a caller already holds: the rows are omitted when they still hash to it"] = "", +) -> 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_f(f): + 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, f): + if _is_unknown_f(f): + return _run_end(e) + nxt = ida_bytes.get_item_end(e) + return nxt if nxt > e else e + 1 + + # The function the walk is currently inside, reused while it stays inside. + # get_func is ~0.5us and the walk asks per head; a head is nearly always in + # the same function as the one before it. Only ever consulted when ``e`` + # falls in [start_ea, end_ea), so a tail chunk elsewhere cannot be + # misattributed -- checked against get_func over 437k heads of + # bash/ls_ttl/echo with zero disagreements. + fn_cache = [None] + + def _func_at(e): + cur = fn_cache[0] + if cur is not None and cur.start_ea <= e < cur.end_ea: + return cur + cur = idaapi.get_func(e) + fn_cache[0] = cur + return cur + + def _rows_for(e, f): + if _is_unknown_f(f): + return [_idatui_unknown_row(e, _run_end(e) - e)] + func = _func_at(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, f) + 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) + get_flags = ida_bytes.get_flags + for _ in range(offset): + if ea >= hi or ea == idaapi.BADADDR: + break + ea = _advance(ea, get_flags(ea)) + more = False + while ea != idaapi.BADADDR and ea < hi: + if len(rows) >= count: + more = True + break + f = get_flags(ea) # once per head, not once per consumer + rows.extend(_rows_for(ea, f)) # a struct head expands into member rows + ea = _advance(ea, f) + cursor = {"next": hex(ea)} if more else {"done": True} + dig = _idatui_rows_digest(rows) + out = {"addr": str(addr), "cursor": cursor, "digest": dig, "count": len(rows)} + # ``expect`` says "I already hold a page that hashed to this". The rows are + # built either way -- generate_disasm_line is the floor and there is no way + # to know a line is unchanged without rendering it -- but pickling several + # hundred rows with their colour spans, unpickling them and rebuilding Heads + # is about 40% of what a page costs, and after a rename almost every page + # comes back identical. + # + # It carries the expected value rather than being a yes/no "digest mode" so + # that a page which HAS changed still costs one round trip: asking first and + # fetching afterwards made every changed page two. + if not (expect and str(dig) == expect): + out["heads"] = rows + return out + + +_IDATUI_FMT_CYCLE = ("hex", "dec", "bin", "char", "offset", "default") + + +_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)), "" + + +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 + + +_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)) + + +_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 + + +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)} + + +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 -- cgit v1.3.1-sl0p From df4a42e25881531469a0a8399f457ecf58cab099 Mon Sep 17 00:00:00 2001 From: blasty Date: Fri, 7 Aug 2026 13:50:42 +0200 Subject: codemode: rename takes a LIST of edits per category, not just one Found by tests/test_rawimage_rpc.py, which the earlier runs had not covered: every rename_many check failed with {"ok": 0, "failed": 2, "errors": [{"addr": null, "error": "list indices must be integers or slices, not str"}]} The port's rename read each category as a single edit (edit["addr"]), but the batch shape is {func: [{addr,name}, ...], data: [...], local/stack: [...]} -- a list per category, with a single dict accepted as shorthand. Indexing the list with "addr" raised, and because the whole category was one try block the error came back attached to addr=null, naming nothing. That is the entire point of the rename_many RPC verb: a firmware image arrives with hundreds of names from a loader map or an emulator's symbols.json, and applying them one at a time costs a navigation plus two prompt round trips each. Only the single-rename UI path worked. Now mirrors the real tool: one row per EDIT (addr/old/name plus a per-row error), a summary counting edits rather than categories, conflict detection before the write, and dry_run/allow_overwrite/stop_on_error. Renaming a function refreshes Hex-Rays' ctext, whose cache is per function and persisted in the .i64 -- without it the pseudocode keeps calling the old name forever while every other readback reports the new one. Clearing a label with an empty new name is kept as a real request (the scenarios revert with it) rather than being rejected as a missing argument. tests/test_rawimage_rpc.py: 14 passed/7 failed -> 21 passed, 0 failed. --- .gitignore | 1 + idatui/codemode_client.py | 145 +++++++++++++++++++++++++++++++++++++--------- 2 files changed, 120 insertions(+), 26 deletions(-) (limited to 'idatui/codemode_client.py') diff --git a/.gitignore b/.gitignore index 65cc5f3..ff80117 100644 --- a/.gitignore +++ b/.gitignore @@ -19,3 +19,4 @@ bin/ # core dumps (idalib/SWIG can segfault under differential probes) core core.* +.fastfeedback/ diff --git a/idatui/codemode_client.py b/idatui/codemode_client.py index 069cf98..5e8089b 100644 --- a/idatui/codemode_client.py +++ b/idatui/codemode_client.py @@ -683,36 +683,129 @@ for item in a.get("items", []): result = {"result": rows} result ''', + # Every category takes EITHER one edit or a LIST of them, and the answer is + # one row per edit. The port accepted only a single dict, so any batch path + # (rpc rename_many applying a whole symbol file, which is the entire point of + # that verb) died with "list indices must be integers or slices, not str" and + # reported the failure against addr=null. Mirrors the real tool: conflict + # detection before the write, dry_run/allow_overwrite/stop_on_error, per-row + # addr/old/name, and a summary counting EDITS rather than categories. "rename": r''' -import ida_idaapi, ida_name, ida_typeinf +import idaapi, ida_hexrays, ida_name batch = a.get("batch") or {} -out = {}; ok_count = failed = 0 -for category, edit in batch.items(): - try: - if category == "func": - ea, new = int(str(edit["addr"]), 16), str(edit["name"]) - fn = db.functions.get_at(ea); ok = bool(fn and db.functions.set_name(fn, new)) - elif category == "data": - new = str(edit.get("new") or "") - if edit.get("addr") is not None: ea = int(str(edit["addr"]), 16) - else: ea = int(ida_name.get_name_ea(ida_idaapi.BADADDR, str(edit.get("old") or ""))) - ok = bool(db.names.set_name(ea, new)) - elif category in ("local", "stack"): - ea, old, new = int(str(edit["func_addr"]), 16), str(edit["old"]), str(edit["new"]) - pseudo = db.pseudocode.decompile(ea); var = pseudo.find_local_variable(old) - if var is None: ok = False +dry_run = bool(batch.get("dry_run", False)) +allow_overwrite = bool(batch.get("allow_overwrite", False)) +stop_on_error = bool(batch.get("stop_on_error", False)) + +def _items(value): + if value is None: return [] + if isinstance(value, dict): return [value] + if isinstance(value, list): return [i for i in value if isinstance(i, dict)] + return [] + +def _set_name_checked(ea, new): + conflict = idaapi.get_name_ea(idaapi.BADADDR, new) + if conflict != idaapi.BADADDR and conflict != ea and not allow_overwrite: + return False, f"can't rename at {hex(ea)} as {new!r}: name already used at {hex(conflict)}" + if dry_run: + return True, None + flags = idaapi.SN_CHECK + if allow_overwrite: flags |= int(getattr(idaapi, "SN_FORCE", 0)) + if not idaapi.set_name(ea, new, flags): + return False, (f"Rename failed at {hex(ea)}: IDA rejected name {new!r} " + "(invalid identifier or internal conflict)") + return True, None + +def _refresh_ctext(fn_addr): + # A renamed function must invalidate Hex-Rays' cache, which is per function + # and persisted in the .i64: without this the pseudocode keeps calling the + # old name forever while every other readback reports the new one. + if not ida_hexrays.init_hexrays_plugin(): return + failure = ida_hexrays.hexrays_failure_t() + cfunc = ida_hexrays.decompile_func(fn_addr, failure, ida_hexrays.DECOMP_WARNINGS) + if cfunc: cfunc.refresh_func_ctext() + +out = {}; ok_count = failed = 0; halted = False +for category in ("func", "data", "local", "stack"): + if category not in batch: continue + rows = [] + for edit in _items(batch.get(category)): + try: + if category == "func": + addr_text = edit.get("addr") or edit.get("func_addr") or edit.get("func") + new = edit.get("name") or edit.get("new") or edit.get("new_name") + if not addr_text or not new: + row = {"addr": addr_text, "name": new, + "error": "Function rename requires addr + name"} + else: + ea = int(str(addr_text), 16) + fn = idaapi.get_func(ea) + if fn is None: + row = {"addr": addr_text, "name": new, "error": "Function not found"} + else: + old = idaapi.get_name(fn.start_ea) or None + ok, err = _set_name_checked(fn.start_ea, str(new)) + row = {"addr": addr_text, "old": old, "name": str(new)} + if err: row["error"] = err + if dry_run: row["dry_run"] = True + if ok and not dry_run: _refresh_ctext(fn.start_ea) + elif category == "data": + addr_text = edit.get("addr") + old = edit.get("old") or edit.get("old_name") + new = edit.get("new") or edit.get("new_name") or edit.get("name") + if not new and new != "": + row = {"old": old, "new": None, + "error": "Global rename requires target and new name"} + else: + if addr_text is not None: + ea = int(str(addr_text), 16) + old = old or (idaapi.get_name(ea) or None) + else: + ea = idaapi.get_name_ea(idaapi.BADADDR, str(old or "")) + if ea == idaapi.BADADDR: + row = {"old": old, "new": str(new), "error": f"Global {old!r} not found"} + else: + # An empty new name CLEARS the label; that is a real + # request (tests revert with it), not a missing argument. + if str(new) == "": + ok = bool(ida_name.set_name(ea, "", idaapi.SN_CHECK)) + err = None if ok else f"Failed to clear the name at {hex(ea)}" + else: + ok, err = _set_name_checked(ea, str(new)) + row = {"addr": hex(ea), "old": old, "new": str(new)} + if err: row["error"] = err + if dry_run: row["dry_run"] = True else: - var.set_user_name(new) - ok = bool(pseudo.save_local_variable_info(var, save_name=True)) - else: - raise ValueError(f"unsupported rename category: {category}") - row = {"ok": ok, **({} if ok else {"error": "IDA rejected the name"})} - except Exception as exc: - row = {"ok": False, "error": str(exc)} - out[category] = [row] - if row["ok"]: ok_count += 1 - else: failed += 1 + fa, old, new = edit.get("func_addr"), edit.get("old"), edit.get("new") + if not fa or not old or not new: + row = {"old": old, "new": new, + "error": f"{category} rename requires func_addr + old + new"} + else: + ea = int(str(fa), 16) + pseudo = db.pseudocode.decompile(ea) + var = pseudo.find_local_variable(str(old)) + if var is None: + row = {"func_addr": fa, "old": old, "new": new, + "error": f"no local {old!r} in that function"} + elif dry_run: + row = {"func_addr": fa, "old": old, "new": new, "dry_run": True} + else: + var.set_user_name(str(new)) + ok = bool(pseudo.save_local_variable_info(var, save_name=True)) + row = {"func_addr": fa, "old": old, "new": new} + if not ok: row["error"] = "IDA rejected the local variable name" + except Exception as exc: + row = {"addr": edit.get("addr"), "error": str(exc)} + rows.append(row) + if row.get("error"): failed += 1 + else: ok_count += 1 + if row.get("error") and stop_on_error: + halted = True; break + out[category] = rows + if halted: break out["summary"] = {"ok": ok_count, "failed": failed} +if dry_run: out["summary"]["dry_run"] = True +if halted: out["summary"]["halted"] = True result = out result ''', -- cgit v1.3.1-sl0p From ecc58d7725db6d4929ae3e299dea4f0202a81946 Mon Sep 17 00:00:00 2001 From: blasty Date: Fri, 7 Aug 2026 14:08:12 +0200 Subject: codemode: three defects the A/B benchmark found in the decompiler path Benchmarking the port against master op-by-op (rather than only asking whether tests pass) turned up three real bugs, all in the most user-visible path: opening pseudocode. 1. decompile was doing decomp_map's job. It called the full per-column line map purely to fill in each line's /*0xEA*/ anchor. The tool ida-tui was written against takes ONE get_line_item at column 0 per line; the port took one per COLUMN, i.e. thousands of get_line_item+dstr() calls per function instead of one per line. Every pseudocode open cost the same as opening the split view. Carried the real implementation over: 1888ms -> 53ms. 2. decomp_map used the pre-optimisation line map. Ours memoises obj_id -> ea for the whole function (commit 853d90c: dstr() was 79% of the tool, and consecutive columns report the same ctree item), the port's did not. 1925ms -> 287ms. 3. _idatui_compact imported ida_pro_mcp on every call. Under Code Mode that package is not installed in the database process, so the import failed every time -- and a FAILED import is never cached, so each one re-searched the whole of sys.path: 422 failed imports per pc_nums call, which was most of its runtime. 1428ms -> 257ms. The same bug was a correctness bug hiding behind the perf bug: the fallback path collapsed whitespace INSIDE string literals, where the real function preserves it. Pseudocode columns are served in those coordinates, so on any line containing a string with two spaces, every literal's mark and every reformat would have been placed on the wrong column. It never fired on master because ida_pro_mcp is installed there. Now calls the byte-identical module-level shim directly, with the deviation from the extracted original documented in place. Narrow verification: decomp/split_view/opfmt/follow/comment/structs scenarios, 72 passed, 0 failed. Full gate running separately. --- idatui/codemode_client.py | 65 ++------------- idatui/remote_tools.py | 199 +++++++++++++++++++++++++++++++++++++++++++--- 2 files changed, 198 insertions(+), 66 deletions(-) (limited to 'idatui/codemode_client.py') diff --git a/idatui/codemode_client.py b/idatui/codemode_client.py index 5e8089b..3bf01cc 100644 --- a/idatui/codemode_client.py +++ b/idatui/codemode_client.py @@ -171,24 +171,6 @@ def _script(args: dict[str, Any], body: str) -> str: return f"import json\na = json.loads({encoded!r})\n{dedent(body).strip()}\n" -_DECOMP_MAP_HELPER = r''' -def line_map(cfunc): - import ida_hexrays - answer = [] - for sl in cfunc.get_pseudocode(): - tagged, eas, seen = sl.line, [], set() - for x in range(len(tagged) + 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(tagged, x, False, head, item, tail): continue - text = item.dstr() or "" - try: ea = int(text.split(": ", 1)[0], 16) - except (ValueError, IndexError): continue - if ea not in seen: seen.add(ea); eas.append(ea) - answer.append(eas) - return answer -''' - - _OPERATIONS: dict[str, str] = { "list_funcs": r''' import fnmatch @@ -812,45 +794,6 @@ result } -_OPERATIONS["decompile"] = _DECOMP_MAP_HELPER + r''' -ea = int(str(a["addr"]), 16) -fn = db.functions.get_at(ea) -if fn is None: - result = {"error": f"no function at {ea:#x}"} -else: - pseudo = db.pseudocode.decompile(fn) - mapping = line_map(pseudo.raw_cfunc) - plain = pseudo.to_text() - marked = [line + (f" /*0x{eas[0]:X}*/" if eas else "") - for line, eas in zip(plain, mapping)] - import ida_name - refs, seen = [], set() - for expr in pseudo.find_objects(): - target = int(expr.obj_ea) - if target in seen or not (db.is_valid_ea(target) or db.is_private_ea(target)): continue - seen.add(target) - name = expr.obj_name or ida_name.get_name(target) or "" - try: string = db.bytes.get_string_at(target) if db.is_valid_ea(target) else None - except Exception: string = None - refs.append({"addr": hex(target), "name": name, "string": string}) - result = {"addr": hex(int(fn.start_ea)), "code": "\n".join(marked), "refs": refs} -result -''' - -_OPERATIONS["decomp_map"] = _DECOMP_MAP_HELPER + r''' -ea = int(str(a["addr"]), 16) -fn = db.functions.get_at(ea) -if fn is None: - result = {"error": f"no function at {ea:#x}"} -else: - pseudo = db.pseudocode.decompile(fn) - mapping = line_map(pseudo.raw_cfunc) - result = {"addr": hex(int(fn.start_ea)), - "lines": [{"ea": hex(eas[0]) if eas else None, - "eas": [hex(item) for item in eas]} for eas in mapping]} -result -''' - _OPERATIONS["define_code_run"] = r''' import ida_bytes, ida_idp, ida_segment, ida_ua, idaapi ea, limit = int(str(a["addr"]), 16), max(1, min(int(a.get("limit", 20000)), 200000)) @@ -877,6 +820,7 @@ else: result ''' + _OPERATIONS["define_func_run"] = r''' import ida_bytes, ida_funcs, ida_segment ea = int(str(a["addr"]), 16) @@ -901,6 +845,7 @@ else: result ''' + _OPERATIONS["set_thumb"] = r''' import ida_bytes, ida_ida, ida_idp, ida_segment, ida_segregs ea = int(str(a["addr"]), 16); treg = ida_idp.str2reg("T") @@ -926,6 +871,7 @@ else: result ''' + _OPERATIONS["thumb_scan"] = r''' import ida_bytes, ida_funcs, ida_idp, ida_segment, ida_segregs, ida_ua lo, hi = int(str(a["start"]), 16), int(str(a["end"]), 16) @@ -950,6 +896,7 @@ result = {"start": hex(lo), "end": hex(hi), "found": found, "applied": applied, result ''' + _OPERATIONS["decomp_error"] = r''' import ida_hexrays, ida_ida ea = int(str(a["addr"]), 16); fn = db.functions.get_at(ea) @@ -1016,6 +963,10 @@ _OPERATIONS["op_format"] = _remote_op( 'op_format(addr=a["addr"], mode=a.get("mode", "cycle"),' ' col=int(a.get("col", -1)), n=int(a.get("n", -1)))') _OPERATIONS["pc_nums"] = _remote_op('pc_nums(addr=a["addr"])') +_OPERATIONS["decompile"] = _remote_op( + 'decompile(addr=a["addr"],' + ' include_addresses=bool(a.get("include_addresses", True)))') +_OPERATIONS["decomp_map"] = _remote_op('decomp_map(addr=a["addr"])') _OPERATIONS["pc_num_format"] = _remote_op( 'pc_num_format(addr=a["addr"], mode=a.get("mode", "cycle"),' ' line=int(a.get("line", -1)), col=int(a.get("col", -1)),' diff --git a/idatui/remote_tools.py b/idatui/remote_tools.py index 41de0de..6fb6436 100644 --- a/idatui/remote_tools.py +++ b/idatui/remote_tools.py @@ -1050,15 +1050,26 @@ _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) + a column in what the client SHOWS can be mapped back to Hex-Rays' line. + + DEVIATION FROM THE EXTRACTED ORIGINAL, deliberately: this used to be + ``from ida_pro_mcp.ida_mcp.utils import compact_whitespace`` inside a + try/except, with a plain ``[ \\t]{2,}`` regex as the fallback. Under Code + Mode ida_pro_mcp is not installed in the database process, so BOTH halves + of that were wrong: + + * the import failed on every call, and a failed import is never cached, so + each one re-searched the whole of sys.path -- 422 failures per pc_nums + call, which was the majority of its runtime; + * the fallback collapses runs of spaces INSIDE STRING LITERALS, which the + real function preserves. Pseudocode columns are served in these + coordinates, so a line containing a string with two spaces would have put + every literal's mark, and every reformat, on the wrong column. + + The module-level shim above is byte-identical to the original regex, so + call it directly. + """ + return compact_whitespace(line) def _idatui_compact_col(plain, compact, col): @@ -1365,3 +1376,173 @@ def pc_num_format( out["text"] = out["before"] out["warn"] = f"re-render failed: {e}" return out + + +def decompile(addr, include_addresses=True): + """Pseudocode for the function at ``addr``, plus the objects it references. + + Faithful to the tool ida-tui was written against, and in particular to its + COST: the per-line address anchor comes from ONE ``get_line_item`` at column + 0 per line. The Code Mode port asked for the full per-column line map (what + ``decomp_map`` is for) purely to fill in that anchor, which is thousands of + ``get_line_item``+``dstr()`` calls per function instead of one per line, and + made every pseudocode open cost the same as opening the split view. + + Text is whitespace-collapsed exactly as the client displays it, because + ``pc_nums`` reports literal columns in those coordinates. + """ + import ida_bytes + import ida_hexrays + import ida_lines + import ida_name + import idaapi + + try: + ea = parse_address(addr) + except Exception as e: + return {"addr": str(addr), "code": None, "error": str(e)} + fn = idaapi.get_func(ea) + if fn is None: + return {"addr": str(addr), "code": None, "error": f"no function at {ea:#x}"} + if not ida_hexrays.init_hexrays_plugin(): + return {"addr": hex(int(fn.start_ea)), "code": None, "error": "no decompiler"} + failure = ida_hexrays.hexrays_failure_t() + try: + cfunc = ida_hexrays.decompile_func(fn, failure) + except Exception as e: + return {"addr": hex(int(fn.start_ea)), "code": None, + "error": f"Decompilation failed at {ea:#x}: {e}"} + if cfunc is None: + return {"addr": hex(int(fn.start_ea)), "code": None, + "error": failure.desc() or f"Decompilation failed at {ea:#x}"} + + lines = [] + for sl in cfunc.get_pseudocode(): + head = ida_hexrays.ctree_item_t() + item = ida_hexrays.ctree_item_t() + tail = ida_hexrays.ctree_item_t() + line_ea = None + if include_addresses and cfunc.get_line_item(sl.line, 0, False, head, item, tail): + parts = (item.dstr() or "").split(": ") + if len(parts) == 2: + try: + line_ea = int(parts[0], 16) + except ValueError: + line_ea = None + text = compact_whitespace(ida_lines.tag_remove(sl.line)) + lines.append(f"{text} /*{line_ea:#x}*/" if line_ea is not None else text) + + refs, seen = [], set() + + class _RefVisitor(ida_hexrays.ctree_visitor_t): + def __init__(self): + ida_hexrays.ctree_visitor_t.__init__(self, ida_hexrays.CV_FAST) + + def visit_expr(self, e): + if e.op == ida_hexrays.cot_obj: + target = int(e.obj_ea) + if target != idaapi.BADADDR and target not in seen: + seen.add(target) + try: + raw = ida_bytes.get_strlit_contents(target, -1, 0) + text = raw.decode("utf-8", "replace") if raw else None + except Exception: + text = None + refs.append({"addr": hex(target), + "name": ida_name.get_name(target) or "", + "string": text}) + return 0 + + try: + _RefVisitor().apply_to(cfunc.body, None) + except Exception: + pass + return {"addr": hex(int(fn.start_ea)), "code": "\n".join(lines), "refs": refs} + + +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"} + import ida_lines + # Three things this loop must not do, each measured on real functions (the 25 + # largest of bash went 68.3s -> 6.5s; echo's 60 largest 5.4s -> 0.6s, with + # byte-identical output): + # + # * allocate ctree_item_t's per COLUMN. They are SWIG objects and this is + # the innermost loop; one per call is enough, and head/tail are never + # read, so don't ask for them at all. + # * sweep the TAGGED length. ``x`` is a screen column but ``sl.line`` still + # carries IDA's colour tags, so a 23-column line was swept 124 times. + # * call dstr() per column. It formats a whole 'EA: description' string -- + # 24us a call, which is 79% of this tool. Comparing against the PREVIOUS + # column's item id is not enough: items interleave, so `foo(a, b)` flips + # call -> arg -> call -> arg and every flip re-formats an item already + # seen (106 594 calls for 15 417 lines of bash). Memoise id -> ea for the + # whole function instead: obj_id is unique within a cfunc, so the same id + # always yields the same string, and the result is deduped by ``seen`` + # anyway. Items with no ctree node (it is None) have no id to key on and + # still pay per occurrence. + item = ida_hexrays.ctree_item_t() + tag_remove = ida_lines.tag_remove + get_line_item = cfunc.get_line_item + ea_of_id = {} + lines = [] + for sl in cfunc.get_pseudocode(): + line = sl.line + eas, seen = [], set() + prev_id = None + for x in range(len(tag_remove(line)) + 1): + if not get_line_item(line, x, False, None, item, None): + continue + it = item.it + if it is not None: + oid = it.obj_id + if oid == prev_id: + continue + prev_id = oid + if oid in ea_of_id: + e = ea_of_id[oid] + if e is not None and e not in seen: + seen.add(e) + eas.append(hex(e)) + continue + else: + oid = None + prev_id = None + # Match the /*ea*/ marker's source (decompile_function_safe): the + # item's dstr() is 'EA: description'; get_ea() reports a different ea. + e = None + dstr = item.dstr() + if dstr: + parts = dstr.split(": ", 1) + if len(parts) == 2: + try: + e = int(parts[0], 16) + except ValueError: + e = None + if oid is not None: + ea_of_id[oid] = e + if e is not None and 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} -- cgit v1.3.1-sl0p From 21971d0a9ce9b8a2fb1db2ec67303696cea9cbbb Mon Sep 17 00:00:00 2001 From: blasty Date: Fri, 7 Aug 2026 14:32:56 +0200 Subject: tests: blob_ui 39.8s -> 4.1s, and a fatal it was hiding Three separate wastes, all of the same family: waiting on a guess instead of a signal, and paying for work that never had to be repeated. 1. The 64KB blob was built with os.urandom into a fresh TemporaryDirectory on every run. New bytes at a new path means the pristine-database cache can never apply, so full auto-analysis of 64KB of AArch64-decoded noise was paid every single run. It is now built from a seeded PRNG at a stable path (tests/.synthetic/, gitignored) and staged through the existing cache. Determinism is also a correctness fix: whether 64KB of chance bytes contains something IDA reads as a function is luck, and this suite asserts "and really has no functions". 2. `wait(lambda: lst.model is not old, ..., 30)` after commenting. The perf work made an item edit KEEP the listing's walk and re-render in place, so the model object is never replaced and this waited out its full 30s timeout on every run -- and then "commenting leaves the view where it was" passed vacuously, because nothing had happened at all. A test that burns 30s to check nothing is worse than no test. 3. Two `pause(2.0)`/`pause(2.5)` after a carve, replaced with settle() on a real condition. The second one deliberately has NO predicate: that spot is random data, so the carve may legitimately produce nothing, and "the row became code" would never hold -- gating on it cost another 30s timeout. What that check is about is the VIEW not moving, so the gate is "the app finished reacting". Fixing (1) exposed a real bug in the client, fixed here too: reopening a database that already exists while passing loader switches is FATAL in IDA -- FATAL ERROR: Switch '-b400' can be used only when loading a new file which kills the worker before it can report anything. Loader switches describe an IMPORT and are recorded in the database they produce, so they are now sent only when there is an import to describe. This was never reachable from the old suite (a fresh random blob never had a database to reopen), but it is reachable by any user who opens a raw blob with --ida-args twice. 30 passed, 0 failed. --- .gitignore | 1 + idatui/codemode_client.py | 25 +++++++++-- tests/_fixtures.py | 30 +++++++++++++ tests/test_blob_ui.py | 109 +++++++++++++++++++++++++++++++++++----------- 4 files changed, 136 insertions(+), 29 deletions(-) (limited to 'idatui/codemode_client.py') diff --git a/.gitignore b/.gitignore index ff80117..9a7b5b2 100644 --- a/.gitignore +++ b/.gitignore @@ -20,3 +20,4 @@ bin/ core core.* .fastfeedback/ +tests/.synthetic/ diff --git a/idatui/codemode_client.py b/idatui/codemode_client.py index 3bf01cc..9d91335 100644 --- a/idatui/codemode_client.py +++ b/idatui/codemode_client.py @@ -1062,6 +1062,17 @@ class CodeModeClient: self._last_entry: RegistryEntry | None = None self._connect_lock = threading.Lock() + def _database_exists(self) -> bool: + """Whether the IDB this open would target is already on disk. + + Its loader switches are baked in, so they must not be sent again. + """ + try: + target = self._output_database or expected_idb_path(self._path) + except Exception: # noqa: BLE001 -- resolver unavailable: assume fresh + return False + return bool(target) and os.path.exists(target) + def connect(self, timeout: float = 1800.0, progress=None) -> "CodeModeClient": _require_codemode() with self._connect_lock: @@ -1078,17 +1089,25 @@ class CodeModeClient: deadline = time.monotonic() + min(timeout, 60.0) while True: try: + # Loader switches describe how to IMPORT a raw file and + # are recorded in the database it produces. Sending them + # again for a database that already exists is a FATAL + # error in IDA itself ("Switch '-b400' can be used only + # when loading a new file"), which kills the worker + # before it can report anything useful. So: describe the + # import only when there is an import to describe. + fresh = self._new_database or not self._database_exists() handle = DatabaseHandle.open( self._path, spawn=self._spawn, timeout=max(0.1, timeout), output_database=self._output_database, - processor=self._processor, + processor=self._processor if fresh else None, # DatabaseHandle calls this image_base and wants the # natural (16-byte aligned) address; it does the # conversion to IDA's paragraph-based -b itself. - image_base=self._loading_address, - file_type=self._file_type, + image_base=self._loading_address if fresh else None, + file_type=self._file_type if fresh else None, new_database=self._new_database, ) break diff --git a/tests/_fixtures.py b/tests/_fixtures.py index c7a45d3..0b72856 100644 --- a/tests/_fixtures.py +++ b/tests/_fixtures.py @@ -37,6 +37,36 @@ def cache_is_fresh(binary: str) -> bool: return os.path.exists(c) and os.path.getmtime(c) >= os.path.getmtime(binary) +#: Generated targets live here so their pristine caches survive between runs. +#: Gitignored; safe to delete (the next run rebuilds both). +SYNTHETIC_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), ".synthetic") + + +def synthetic(name: str, build) -> str: + """A generated binary at a STABLE path, rebuilt only when its bytes change. + + Generated targets used to be written into a fresh TemporaryDirectory on + every run, which quietly defeated the whole pristine-cache scheme: a new + path with new bytes every time means auto-analysis is paid in full, every + run, forever. `test_blob_ui`'s 64KB blob cost ~40s a run that way. + + ``build()`` must be DETERMINISTIC and return bytes. That is also what makes + the suites reproducible: a blob built from os.urandom can, by luck, contain + something IDA reads as a function, and then a test asserting "no functions" + fails for reasons no one can reproduce. + """ + os.makedirs(SYNTHETIC_DIR, exist_ok=True) + path = os.path.join(SYNTHETIC_DIR, name) + data = build() + if not os.path.exists(path) or open(path, "rb").read() != data: + with open(path, "wb") as fh: # content changed -> cache is stale + fh.write(data) + for stale in (cache_path(path), path + ".i64"): + if os.path.exists(stale): + os.remove(stale) + return path + + async def build_pristine(binary: str, cache: str, app_factory) -> None: """Analyse ``binary`` once and keep the database as a golden copy. diff --git a/tests/test_blob_ui.py b/tests/test_blob_ui.py index f60e914..1055e5a 100644 --- a/tests/test_blob_ui.py +++ b/tests/test_blob_ui.py @@ -17,10 +17,13 @@ import sys import tempfile sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) from textual.widgets import Input, Static # noqa: E402 from idatui.app import ConfirmScreen, IdaTui, ListingView # noqa: E402 +from _fixtures import staged, synthetic # noqa: E402 +from idatui._sync import settle # noqa: E402 PASS = FAIL = 0 @@ -46,28 +49,63 @@ async def wait(pred, pilot, t=240.0): return False -async def run() -> int: - with tempfile.TemporaryDirectory() as tmp: - # Random bytes so IDA finds no functions... but with REAL AArch64 - # instructions planted at a known offset. Whether arbitrary random bytes - # happen to decode is chance, and a test that depends on chance tells you - # nothing on the run where it fails. - data = bytearray(os.urandom(64 * 1024)) +#: File offset of the planted instruction run -> ea 0x4000 + PLANTED. +PLANTED = 0x40 + + +def _blob_bytes() -> bytes: + """A DETERMINISTIC pseudo-random blob with real AArch64 instructions planted. + + Seeded, not os.urandom: the bytes must be identical every run or the + pristine-database cache can never apply (this suite used to pay ~40s of + auto-analysis per run because the content, and the path, changed each time). + Determinism also removes a genuine flake -- whether 64KB of chance bytes + contains something IDA reads as a function is luck, and "and really has no + functions" is asserted below. + """ + import random + data = bytearray(random.Random(0xB10BCAFE).randbytes(64 * 1024)) # -parm puts IDA in AArch64 mode, so these are A64 encodings; the ARM32 # spelling of a nop (0xE1A00000) is NOT decodable there and made this # test fail for a reason that had nothing to do with what it checks. - planted = 0x40 # file offset -> ea 0x4040 - for k, insn in enumerate((0xD503201F, # nop - 0xD503201F, # nop - 0xD65F03C0)): # ret <- the run must stop here - data[planted + k * 4:planted + k * 4 + 4] = insn.to_bytes(4, "little") - blob = os.path.join(tmp, "rnd.bin") - with open(blob, "wb") as f: - f.write(bytes(data)) + for k, insn in enumerate((0xD503201F, # nop + 0xD503201F, # nop + 0xD65F03C0)): # ret <- the run must stop here + data[PLANTED + k * 4:PLANTED + k * 4 + 4] = insn.to_bytes(4, "little") + return bytes(data) + + +#: -parm puts IDA in AArch64 mode (the ARM32 spelling of a nop is not decodable +#: there); -b400 sets the image base. The cached database must be built with the +#: SAME switches, so both go through one factory. +BLOB_ARGS = "-parm -b400" + +def _blob_app(path): + return IdaTui(open_path=path, keepalive=False, load_args=BLOB_ARGS) + + +def head_at(lst, ea): + """The listing row for ``ea`` off the LIVE model, or None. + + Always re-reads ``lst.model``: an edit may rebuild the model, and holding + the old object shows pre-edit rows -- which looks exactly like the edit + silently failing. + """ + m = lst.model + if m is None: + return None + i = m.index_of_ea(ea) + return m.get(i) if i is not None and i >= 0 else None + + +async def run() -> int: + blob_src = synthetic("rnd.bin", _blob_bytes) + # staged() analyses once ever and copies the result in on later runs. + async with staged(blob_src, _blob_app) as blob: # Skip the dialog by answering up front; this test is about what # happens AFTER a described blob turns out to contain nothing. - app = IdaTui(open_path=blob, keepalive=False, load_args="-parm -b400") + app = _blob_app(blob) async with app.run_test(size=(140, 44)) as pilot: ok = await wait(lambda: app._func_index is not None and app._func_index.complete, pilot) @@ -126,7 +164,7 @@ async def run() -> int: i >= 0 and m.get(i).ea == 0x4021, f"row={i} ea={m.get(i).ea if i >= 0 else None}") - target = 0x4000 + planted # a NOP we put there ourselves + target = 0x4000 + PLANTED # a NOP we put there ourselves lst.cursor = m.index_of_ea(target) lst._scroll_cursor_into_view() await pilot.pause(0.1) @@ -134,12 +172,18 @@ async def run() -> int: lst._cursor_ea() == target, f"{lst._cursor_ea():#x} want {target:#x}") await pilot.press("c") - await pilot.pause(2.0) - # Defining an item REBUILDS the listing model, so re-read it from the - # view: holding the old object shows the pre-edit rows and looks - # exactly like the edit silently failing. + # settle(), not a fixed sleep AND not a bare predicate: an edit can + # look done for a moment and then be replaced when a queued listing + # rebuild lands, so the gate has to be "the row is code AND the app + # has stopped working". settle() is the same helper the app's own + # RPC layer uses, so tests and driver agree on what "done" means. + await settle(app, lambda: (lambda h: h is not None and h.kind == "code")( + head_at(lst, target)), timeout=30) + # Re-read the model: defining an item rebuilds it, and holding the + # old object shows pre-edit rows -- which looks exactly like the + # edit silently failing. m = lst.model - h = m.get(m.index_of_ea(target)) + h = head_at(lst, target) check("`c` on a chosen byte carves an instruction there", h is not None and h.kind == "code", f"kind={h.kind if h else None} text={h.text if h else None!r}") @@ -184,9 +228,17 @@ async def run() -> int: for ch in "note": await pilot.press(ch) await pilot.press("enter") - await wait(lambda: lst.model is not old and lst.model is not None, - pilot, 30) - await pilot.pause(0.4) + # Wait for the COMMENT ITSELF to show up, not for the model object to + # be replaced: a comment now re-renders the listing in place (the + # walk is kept), so `model is not old` never becomes true and this + # burned its full 30s timeout on every run -- after which the check + # below passed vacuously, because nothing had happened at all. + # The prompt closing plus quiescence is the real end of the edit. + # (The listing re-renders its text lazily, so the comment is not + # necessarily visible in model rows the moment the worker returns -- + # which is why this waits for the app, not for the text.) + await settle(app, lambda: not app.query_one("#comment", Input).display, + timeout=30) check("commenting leaves the view where it was", lst.model.get(round(lst.scroll_offset.y)).ea == ctop and lst._cursor_ea() == ccur, @@ -208,7 +260,12 @@ async def run() -> int: check("scrolled somewhere with rows above us", round(lst.scroll_offset.y) > 0, f"top={lst.scroll_offset.y}") await pilot.press("c") - await pilot.pause(2.5) + # No predicate here on purpose: this spot is random data, so the + # carve may legitimately produce nothing and "the row became code" + # would never hold (it timed out for 30s and then passed anyway). + # What is being checked is that the VIEW did not move, so the gate + # is simply "the app has finished reacting". + await settle(app, timeout=30) m2 = lst.model top_after = m2.get(round(lst.scroll_offset.y)).ea check("carving leaves the scroll position where it was", -- cgit v1.3.1-sl0p From 55e9d9fdca4baafce728659d43614e203bfbea9a Mon Sep 17 00:00:00 2001 From: blasty Date: Fri, 7 Aug 2026 15:02:05 +0200 Subject: codemode: close the performance gap with the old worker (heads 35x -> 2.2x) Two changes, both about work that was never ours to do, found by profiling the A/B benchmark rather than guessing. 1. Serialise inside the database process. Code Mode runs to_jsonable() over whatever a snippet returns, walking the entire structure to make it JSON-safe. Our answers are already JSON-safe and they are large: a 200-row listing page is ~10k small objects, and walking them cost 66ms of the page's 92ms -- 114x what json.dumps of the very same data costs (0.58ms). Snippets now return one pre-serialised string, so that walk is O(1) and the client parses a payload it was going to parse anyway. heads(200): 92ms -> 24.7ms. 2. Detach the runtime's trace hook while our snippet runs. ida_codemode.runtime wraps every execute_python in sys.settrace(timeout_trace) to enforce deadlines, and timeout_trace RETURNS ITSELF -- which switches on LINE tracing in every frame it sees. Every line of every function we call pays a Python-level callback. Measured here: ida_bytes.get_flags 0.106us untraced 5.49us traced 52x (plain idalib, no Code Mode: 0.119us -- i.e. untraced == native) heads(200 rows) 2.0ms untraced 20.2ms traced 10x That one hook was the entire residual gap against the old unix-socket worker. The snippet now detaches it and restores it in a finally. What that gives up, stated plainly: the deadline is no longer enforced for a pure-Python loop inside our snippet. The runtime's other cancellation path -- a threading.Timer calling ida_kernwin.set_cancelled() -- does not go through the trace and still fires, so a long IDA operation remains interruptible, and every operation here is bounded by its own count/limit argument. Set IDATUI_CODEMODE_TRACE=1 to keep the stock behaviour. Against the worker backend, same box, targets/echo (worker -> codemode): heads_200 2.65ms -> 5.85ms 2.2x (was 35x) heads_500 6.21ms -> 10.67ms 1.7x heads expect-hit 2.16ms -> 4.61ms 2.1x disasm_200 9.03ms -> 5.25ms 0.6x faster decompile_cold 162.62ms -> 30.66ms 0.2x faster decompile_warm 30.02ms -> 25.39ms 0.8x faster decomp_map 45.82ms -> 47.85ms 1.0x parity pc_nums 19.90ms -> 22.56ms 1.1x parity rename_func 254.08ms -> 255.54ms 1.0x parity connect 550.0ms -> 410.0ms 0.7x faster What is left is the transport floor: an empty execute_python round trip is 2.0ms, so trivial calls (data_type 0.07ms -> 2.63ms, force_recompile, a single xref query) look like 40x while being 2.5ms of wall clock. Reducing those needs fewer calls, not faster ones -- the digest/expect path already does that for the listing, which is where call volume actually is. Full suite: 788 passed, 0 failed, 115.3s (was 146.3s; the pilot alone went 80.9s -> 62.2s). --- idatui/codemode_client.py | 71 +++++++++++++++++++++++++++++++++++++++++++---- 1 file changed, 66 insertions(+), 5 deletions(-) (limited to 'idatui/codemode_client.py') diff --git a/idatui/codemode_client.py b/idatui/codemode_client.py index 9d91335..88d7b31 100644 --- a/idatui/codemode_client.py +++ b/idatui/codemode_client.py @@ -22,7 +22,7 @@ import shlex import threading import time from pathlib import Path -from textwrap import dedent +from textwrap import dedent, indent from typing import Any from .errors import IDAConnectionError, IDATimeoutError, IDAToolError, Session @@ -165,10 +165,63 @@ def _parse_load_args(value: str) -> tuple[str | None, int | None, str | None]: return processor, loading_address, file_type +#: Key of the pre-serialised payload envelope. See _script(). +_PACKED = "__idatui_json__" + +#: Serialise the answer INSIDE the database process and hand back one string. +#: +#: Code Mode runs to_jsonable() over whatever a snippet returns, walking the +#: whole structure to make it JSON-safe. Our answers are already JSON-safe, and +#: they are big: a 200-row listing page is ~10k small objects, which costs 66ms +#: to walk -- 72% of the page's total cost, and 114x what json.dumps of the very +#: same data costs (0.58ms). Returning a STRING makes that walk O(1); the client +#: parses it, which it was going to do at the transport layer anyway. +_PACK_EPILOGUE = ( + '\n{"' + _PACKED + '": json.dumps(result, separators=(",", ":"), default=str)}\n' +) + + +#: Keep Code Mode's per-line trace hook installed while our snippet runs. +#: Set IDATUI_CODEMODE_TRACE=1 to restore the stock behaviour. +_KEEP_TRACE = os.environ.get("IDATUI_CODEMODE_TRACE", "") not in ("", "0") + + def _script(args: dict[str, Any], body: str) -> str: - """Bind JSON arguments without interpolating user text into Python code.""" + """Bind JSON arguments without interpolating user text into Python code. + + Also runs the body with Code Mode's trace hook detached, which is worth an + order of magnitude. The runtime wraps every execute_python in + sys.settrace(timeout_trace), and that trace function RETURNS ITSELF, which + turns on line tracing in every frame it sees -- so every line of every + function we call pays a Python-level callback. Measured on this box: + ida_bytes.get_flags is 0.106us untraced (0.119us in a plain idalib process) + and 5.49us traced, 52x; a 200-row listing page is 2.0ms untraced and 20.2ms + traced. That single hook was the whole residual gap against the old worker. + + What this gives up: the deadline is no longer enforced for a pure-Python + loop inside our snippet. The runtime's OTHER cancellation path -- a + threading.Timer that calls ida_kernwin.set_cancelled() -- is independent of + the trace and still fires, so a long IDA operation is still interruptible; + and every operation here is bounded by its own count/limit argument. The + trace is restored in a finally, so a raising snippet cannot leak the change. + """ encoded = json.dumps(args, ensure_ascii=False, separators=(",", ":")) - return f"import json\na = json.loads({encoded!r})\n{dedent(body).strip()}\n" + head = f"import json\na = json.loads({encoded!r})\n" + if _KEEP_TRACE: + return f"{head}{dedent(body).strip()}\n{_PACK_EPILOGUE}" + return ( + f"{head}" + "import sys\n" + "_idatui_trace = sys.gettrace()\n" + "sys.settrace(None)\n" + "try:\n" + f"{indent(dedent(body).strip(), ' ')}\n" + ' _idatui_packed = {"' + _PACKED + '": json.dumps(' + 'result, separators=(",", ":"), default=str)}\n' + "finally:\n" + " sys.settrace(_idatui_trace)\n" + "_idatui_packed\n" + ) _OPERATIONS: dict[str, str] = { @@ -1179,6 +1232,13 @@ class CodeModeClient: raise IDAToolError("execute_python", "Code Mode returned an invalid execution result") return response["result"] + @staticmethod + def _unpack(answer: Any) -> Any: + """Undo _PACK_EPILOGUE. Anything else passes through untouched.""" + if isinstance(answer, dict) and _PACKED in answer: + return json.loads(answer[_PACKED]) + return answer + def invoke(self, operation: str, *, timeout: float | None = None, **args) -> Any: """Execute one TUI domain operation through Code Mode.""" if operation in ("idb_save", "save"): @@ -1189,12 +1249,13 @@ class CodeModeClient: if body is None: raise IDAToolError(operation, f"unknown ida-tui Code Mode operation: {operation}") try: - answer = self.execute_python(_script(args, body), timeout=timeout) + answer = self._unpack(self.execute_python(_script(args, body), timeout=timeout)) if isinstance(answer, dict) and answer.get(_NEED_LIB): # First call against this database process (or a restarted one). self.execute_python(_script({"source": _REMOTE_LIB}, _INSTALL_LIB), timeout=timeout) - answer = self.execute_python(_script(args, body), timeout=timeout) + answer = self._unpack( + self.execute_python(_script(args, body), timeout=timeout)) return answer except IDAToolError as exc: if exc.tool == "execute_python": -- cgit v1.3.1-sl0p