diff options
Diffstat (limited to 'idatui')
| -rw-r--r-- | idatui/app.py | 238 | ||||
| -rw-r--r-- | idatui/codemode_client.py | 1365 | ||||
| -rw-r--r-- | idatui/domain.py | 649 | ||||
| -rw-r--r-- | idatui/edit_ctl.py | 207 | ||||
| -rw-r--r-- | idatui/pool.py | 22 | ||||
| -rw-r--r-- | idatui/remote_ops.py | 1683 | ||||
| -rw-r--r-- | idatui/remote_tools.py | 610 |
7 files changed, 3108 insertions, 1666 deletions
diff --git a/idatui/app.py b/idatui/app.py index d7212b1..e4c2936 100644 --- a/idatui/app.py +++ b/idatui/app.py @@ -46,8 +46,7 @@ from textual.widgets import ( ) from textual.widgets.option_list import Option -from . import graph -from . import kittygfx +from . import diag, graph, kittygfx from .edit_ctl import EditController from .prompt import PromptBar from .trace_ctl import TraceController @@ -55,7 +54,7 @@ from . import findings, search from .highlight import CTextArea, highlight_c from .journal import Journal -from .errors import IDAToolError, IDAConnectionError +from .errors import IDAConnectionError from .codemode_client import CodeModeClient, registered_database from .domain import Func, Head, ListingModel, Program, Struct @@ -3749,14 +3748,14 @@ _HELP = ( class QuitScreen(ModalScreen): """Asked before exiting with unsaved database changes. - Code Mode clients cannot roll a shared database back. The ``d`` choice means - "do not explicitly save": a GUI keeps the changes dirty, while a managed - idalib worker may persist them when its final lease closes. + A final managed-worker lease can discard the whole session. Shared workers + and GUI databases keep their state: releasing this lease transfers the final + save/discard decision to the remaining client or GUI owner. """ BINDINGS = [ Binding("s", "save", "Save & quit"), - Binding("d", "discard", "Leave & quit"), + Binding("d", "discard", "Discard / leave"), Binding("escape,c", "cancel", "Cancel"), ] @@ -3772,8 +3771,11 @@ class QuitScreen(ModalScreen): body = Text() for label in self._labels: body.append(f" \u2022 {label}\n", _S_LABEL) + body.append( + "\nFinal managed leases discard; shared/GUI sessions stay open.", + _S_DIM) yield Static(body, id="quit-list") - yield Static("s save & quit d leave as-is & quit Esc cancel", + yield Static("s save & quit d discard / leave & quit Esc cancel", id="quit-help") def action_save(self) -> None: @@ -5286,6 +5288,11 @@ class IdaTui(App): self.journal = Journal() self._xref_focus_name: str | None = None self._dirty = False + # One subscription for the active database. CodeModeClient debounces + # bursts off the Textual worker pool; the callback re-enters here on the + # UI thread to invalidate and reload the visible models. + self._idb_event_watch = None + self._idb_refresh_seq = 0 # -- layout ------------------------------------------------------------ # def compose(self) -> ComposeResult: @@ -5447,6 +5454,7 @@ class IdaTui(App): self._ask_load_options(path, label=label) def _release_database(self) -> None: + self._stop_idb_event_watch() if self.program is not None: self.program.close() if self._pool is not None and self._binary is not None: @@ -5625,6 +5633,139 @@ class IdaTui(App): pass return len(text) + # -- live refresh from shared IDB changes ----------------------------- # + def _start_idb_event_watch(self, client: CodeModeClient) -> None: + self._stop_idb_event_watch() + watch = getattr(client, "watch_idb_events", None) + if watch is None: # IDA-free test doubles and pre-event adapters + return + + def changed(events) -> None: # listener thread + try: + self.call_from_thread(self._refresh_idb_events, client, events) + except Exception: # noqa: BLE001 -- app teardown can win this race + pass + + def failed(error: BaseException) -> None: # listener thread + try: + self.call_from_thread(self._idb_event_watch_failed, client, error) + except Exception: # noqa: BLE001 -- app teardown can win this race + pass + + self._idb_event_watch = watch( + changed, on_error=failed, debounce=0.2) + + def _stop_idb_event_watch(self) -> None: + watcher, self._idb_event_watch = self._idb_event_watch, None + if watcher is not None: + watcher.close() + + def _idb_event_watch_failed( + self, client: CodeModeClient, error: BaseException + ) -> None: + if client is not self.client: + return + if isinstance(error, IDAConnectionError): + self._on_connection_lost() + else: + self._status(f"live database refresh stopped: {error}") + + def _listing_event_anchor(self) -> ViewAnchor: + """Capture the listing position even when the split's decompiler has focus.""" + anchor = ViewAnchor(view=self._active) + listing = self.query_one(ListingView) + model = listing.model + if model is None: + return anchor + anchor.cursor_x = listing.cursor_x + anchor.ea = listing._cursor_ea() + top = round(listing.scroll_offset.y) + head = model.cached_line(top) or model.get(top) + anchor.top_ea = getattr(head, "ea", None) + return anchor + + def _refresh_idb_events( + self, client: CodeModeClient, events: tuple[dict, ...] + ) -> None: + """Invalidate once per external edit burst and reload the active surface.""" + program = self.program + if not events or client is not self.client or program is None \ + or program.client is not client: + return + self._idb_refresh_seq += 1 + seq = self._idb_refresh_seq + entry = self._cur + anchor = self._listing_event_anchor() + hex_ea = self.query_one(HexView).cursor_va() if self.is_hex else None + graph_ea = self.query_one(GraphView)._cursor_ea() if self.is_graph else None + decomp = self.query_one(DecompView) + if entry is not None and decomp.loaded_ea == entry.ea: + entry.dec_cursor = decomp.cursor + entry.dec_cursor_x = decomp.cursor_x + entry.dec_scroll_y = round(decomp.scroll_offset.y) + entry.dec_scroll_x = round(decomp.scroll_offset.x) + + program.invalidate_external() + self._status( + f"{len(events)} external database " + f"change{'s' if len(events) != 1 else ''} — refreshing…") + self._reindex_functions() + + if self.is_hex: + self.query_one(HexView).model = None + self._load_hex_model(hex_ea) + return + if self.is_graph and entry is not None: + self._load_graph(entry.ea, graph_ea or entry.ea) + return + if entry is None: + return + + # A split needs both halves rebuilt; a decompiler-only view still keeps + # the hidden listing fresh so Tab does not reveal pre-event rows. + decomp.loaded_ea = None + self._reload_idb_listing(program, seq, entry, anchor) + if self.is_decomp or self._split: + self._show_active() + + @work(thread=True, exclusive=True, group="idb-refresh") + def _reload_idb_listing( + self, program: Program, seq: int, entry: NavEntry, anchor: ViewAnchor + ) -> None: + target = anchor.ea if anchor.ea is not None else entry.ea + model = program.listing(target) + cursor = top = -1 + if model is not None: + model.ensure_ea(target) + cursor, top = self._anchor_rows(anchor, model, target) + fn = program.function_of(entry.ea) + name = fn.name if fn is not None else program.region_label(entry.ea) + self.app.call_from_thread( + self._apply_idb_listing, program, seq, entry, model, + cursor, top, anchor.cursor_x, name, fn is None) + + def _apply_idb_listing( + self, program: Program, seq: int, entry: NavEntry, model, + cursor: int, top: int, cursor_x: int, name: str, is_region: bool, + ) -> None: + if program is not self.program or seq != self._idb_refresh_seq \ + or entry is not self._cur: + return + if model is None: + self._status(f"{entry.ea:#x} is no longer in a loaded segment") + return + entry.name = name + entry.is_region = is_region + entry.cursor = max(cursor, 0) + entry.cursor_x = cursor_x + if top >= 0: + entry.scroll_y = top + self.query_one(ListingView).load( + model, name, cursor=entry.cursor, cursor_x=entry.cursor_x, + scroll_y=top if top >= 0 else None) + if self.is_listing: + self._show_active() + # -- connection loss / recovery --------------------------------------- # def _handle_exception(self, error: BaseException) -> None: """Intercept a lost Code Mode lease so the app can rediscover the DB. @@ -5662,15 +5803,22 @@ class IdaTui(App): @work(thread=True, exclusive=True, group="reconnect") def _reconnect(self) -> None: - # The registered instance disappeared. Rediscover it; Code Mode may find - # a GUI/replacement worker, then we rebuild caches against the new handle. + # Rediscovery is attach-only. If a GUI owner closes its database, a TUI + # must not silently reopen it by spawning a headless worker. try: if self._open_path is None: self.app.call_from_thread(self._reconnect_failed, "no binary to reopen") return - client = CodeModeClient(self._open_path, ttl=self._ttl, - load_args=self._load_args) + if self._project is not None and self._binary is not None: + ref = self._project.by_label(self._binary) + client = CodeModeClient( + ref.staged, ttl=self._ttl, load_args=ref.load_args, + output_database=ref.db, spawn=False) + else: + client = CodeModeClient( + self._open_path, ttl=self._ttl, + load_args=self._load_args, spawn=False) client.connect(progress=lambda m: self.app.call_from_thread( self._conn_note, m)) except Exception as e: # noqa: BLE001 @@ -5679,20 +5827,32 @@ class IdaTui(App): self.app.call_from_thread(self._after_reconnect, client, Program(client)) def _after_reconnect(self, client: "CodeModeClient", program: "Program") -> None: + old_client, old_program = self.client, self.program + self._stop_idb_event_watch() + if old_program is not None: + old_program.close() + if self._pool is not None and self._binary is not None: + self._pool.replace_client(self._binary, old_client, client) + if old_client is not None and old_client is not client: + old_client.close() self.client = client self.program = program + self._start_idb_event_watch(client) self._reconnecting = False + self._dirty = False self._dismiss_conn() - self._status("reconnected \u2014 reloading\u2026") - self._load_functions() # rebuild the function index against the new client + self._status("reattached — reloading persisted state…") + self._load_functions() cur = self._cur - if cur is not None: # refresh the current view with the new program + if cur is not None: self._open_entry(cur, push=False) def _reconnect_failed(self, why: str) -> None: self._reconnecting = False - self._conn_note(f"reconnect failed: {why} \u2014 retry on next action, or 'q'") - self._status(f"reconnect failed: {why}") + note = (f"database owner closed: {why} — reopen it in IDA, then " + "Esc and retry an action; or q to quit") + self._conn_note(note) + self._status(note) # -- connection + initial load ---------------------------------------- # @work(thread=True, exclusive=True, group="connect") @@ -5719,6 +5879,7 @@ class IdaTui(App): return self.client = client self.program = program + self._start_idb_event_watch(client) self._new_database = False self.app.call_from_thread( self._status, f"{module} [{client.backend}] — loading functions…") @@ -6115,11 +6276,13 @@ class IdaTui(App): def _on_quit_choice(self, choice: str | None) -> None: if choice == "discard": - # Code Mode has no rollback/close-without-save operation. For GUI - # sessions this leaves changes dirty in IDA; a managed worker owns - # its final save policy and may persist them on final lease release. - self._save_on_exit = False - self.exit() + # Whole-session discard is legal only for the final managed lease. + # GUI/shared sessions retain state and inherit finalization. + dirty = self._dirty_labels() + self._loading_screen = LoadingScreen( + "discarding", note="finalizing database leases…") + self.push_screen(self._loading_screen) + self._discard_then_exit(dirty) elif choice == "save": # Save with the overlay up: writing a big .i64 takes seconds, and # doing it during teardown would look like a hang with no UI left. @@ -6129,6 +6292,31 @@ class IdaTui(App): # None: cancel, stay put @work(thread=True, exclusive=True, group="save-exit") + def _discard_then_exit(self, dirty: list[str]) -> None: + try: + if self._pool is not None: + transferred = self._pool.discard_changes(dirty) + elif self.client is not None: + transferred = [] if self.client.discard_database() else dirty + else: + transferred = dirty + except Exception as exc: # noqa: BLE001 -- keep the app open on failure + self.app.call_from_thread(self._discard_failed, str(exc)) + return + self.app.call_from_thread(self._finish_discard, transferred) + + def _discard_failed(self, why: str) -> None: + self._dismiss_loading() + self._status(f"discard failed: {why}", priority=True) + + def _finish_discard(self, transferred: list[str]) -> None: + if transferred and self._loading_screen is not None: + labels = ", ".join(transferred) + self._loading_screen.update_note( + f"finalization transferred: {labels}") + self._finish_exit() + + @work(thread=True, exclusive=True, group="save-exit") def _save_then_exit(self) -> None: try: if self._pool is not None: @@ -6140,7 +6328,9 @@ class IdaTui(App): self.app.call_from_thread(self._finish_exit) def _finish_exit(self) -> None: - self._save_on_exit = False # already written above + # Teardown must not save again: save, discard, or ownership transfer was + # already decided by the quit path. + self._save_on_exit = False self._dirty = False self.exit() @@ -6207,6 +6397,7 @@ class IdaTui(App): def _after_switch(self, label, client, program, st, reuse) -> None: # type: ignore[no-untyped-def] self.client = client self.program = program + self._start_idb_event_watch(client) self._binary = label self._pool.set_active(label) self._open_path = self._project.by_label(label).staged @@ -8374,6 +8565,7 @@ class IdaTui(App): # -- teardown ---------------------------------------------------------- # async def on_unmount(self) -> None: + self._stop_idb_event_watch() if self._rpc is not None: await self._rpc.stop() if self._ka is not None: diff --git a/idatui/codemode_client.py b/idatui/codemode_client.py index 8eb66b5..e95238f 100644 --- a/idatui/codemode_client.py +++ b/idatui/codemode_client.py @@ -5,24 +5,20 @@ 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. +Remote operations are ordinary typed Python functions declared in +``idatui.remote_ops``. Code Mode installs their content-addressed modules once +per handle; subsequent calls send only encoded arguments. The optimized +IDAPython listing/decompiler implementation remains real source in +``idatui.remote_tools`` and is installed through the same module interface. """ + 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 collections.abc import Callable from typing import Any from .errors import IDAConnectionError, IDATimeoutError, IDAToolError, Session @@ -68,7 +64,8 @@ def _require_codemode() -> None: "ida-codemode is not installed in this environment " f"({_CODEMODE_ERROR}). Install it (e.g. `uv sync`, or " "`pip install ida-codemode`) so ida-tui can lease a " - "database.") from _CODEMODE_ERROR + "database." + ) from _CODEMODE_ERROR def database_owner(idb_path: str, staged_path: str | None = None): @@ -155,1106 +152,119 @@ 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. -#: -#: Written when Code Mode ran to_jsonable() over every snippet result, walking -#: the whole structure in Python to make it JSON-safe: a 200-row listing page is -#: ~10k small objects, which cost 66ms to walk -- 72% of the page's total cost, -#: and 114x what json.dumps of the same data cost (0.58ms). -#: -#: ida-codemode 0.3.2 removed that reason: serialization.dumps_json now hands -#: the structure straight to the C encoder and only falls back to the walker for -#: values json.dumps rejects. Re-measured against 0.3.2, packing buys 0.97x on -#: that same page (experiments/bench_pack_trace.py) -- i.e. nothing, because the -#: dodged walk is replaced by a double encode. -#: -#: It is kept anyway, on correctness rather than speed: packing pins OUR encoder -#: settings (compact separators, default=str) inside the database process, so an -#: un-encodable IDA object degrades to repr() at a point we control instead of -#: depending on the runtime's fallback. Delete it if that stops being worth a -#: protocol step -- it is no longer load-bearing for performance. -_PACK_EPILOGUE = ( - '\n{"' + _PACKED + '": json.dumps(result, separators=(",", ":"), default=str)}\n' -) - +class IDBEventListener: + """Debounced, closeable delivery of another client's IDB changes. -def _script(args: dict[str, Any], body: str) -> str: - """Bind JSON arguments without interpolating user text into Python code. - - This used to also run the body with Code Mode's trace hook detached - (sys.settrace(None) + restore), because the runtime wrapped every - execute_python in a trace function that returned ITSELF -- enabling line - tracing in every frame it saw, so every line of every function we called - paid a Python-level callback (ida_bytes.get_flags: 0.106us -> 5.49us, 52x). - - ida-codemode 0.3.2 deleted that hook; cancellation is now a C-level thread - interrupt (runtime._interrupt_thread) that costs nothing while idle. The - workaround measured 0.99x on a 200-row listing page against 0.3.2 -- pure - noise -- so it is gone, and with it the caveat that a pure-Python loop in a - snippet escaped its deadline. See experiments/bench_pack_trace.py. + Code Mode's subscription is a blocking iterator, so one daemon thread reads + it and a second waits for a quiet period before handing a batch to the UI. + Keeping the debounce here avoids a permanent Textual worker (which would + make the app's worker-idle contract impossible) and bounds refresh work to + one pass per edit burst. """ - encoded = json.dumps(args, ensure_ascii=False, separators=(",", ":")) - head = f"import json\na = json.loads({encoded!r})\n" - return f"{head}{dedent(body).strip()}\n{_PACK_EPILOGUE}" - - -_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 -''', - # Everything a person ADDED to the database: comments, non-dummy names, and - # the prototypes they set. - # - # Names come from IDA's name list, which is already an index -- no scan at - # all. Comments have no index, so they need a walk, and the walk is over - # HEADS: `next_that`'s predicate is a *Python* callback (SWIG calls it with - # one argument, so `f_has_cmt` does not even fit), which would be one call - # per BYTE -- 400 million of them on a big image. `max_scan` bounds it and - # reports `truncated` rather than sitting there. - "list_annotations": r''' -import ida_bytes, ida_funcs, ida_lines, ida_nalt, ida_name -import ida_segment, ida_typeinf, idautils -limit = max(1, int(a.get("limit", 4000))) -max_scan = max(1000, int(a.get("max_scan", 2000000))) -comments, names = [], [] -scanned = 0 - -def _line(ea): - try: - txt = ida_lines.generate_disasm_line(ea, ida_lines.GENDSM_REMOVE_TAGS) - except Exception: - txt = "" - return " ".join((txt or "").split()) -for ea, nm in idautils.Names(): - if len(names) >= limit: - break - if not nm or not ida_bytes.has_user_name(ida_bytes.get_flags(ea)): - continue - fn = ida_funcs.get_func(ea) - is_fn = fn is not None and int(fn.start_ea) == int(ea) - proto = None - if is_fn: - try: - ti = ida_typeinf.tinfo_t() - if ida_nalt.get_tinfo(ti, ea): - proto = str(ti) - except Exception: - proto = None - seg = ida_segment.getseg(ea) - names.append({"addr": hex(int(ea)), "name": nm, "func": is_fn, - "size": (int(fn.end_ea - fn.start_ea) if is_fn else 0), - "proto": proto, - "seg": (ida_segment.get_segm_name(seg) if seg else "")}) + def __init__( + self, + client: "CodeModeClient", + callback: Callable[[tuple[dict[str, Any], ...]], None], + *, + on_error: Callable[[BaseException], None] | None = None, + debounce: float = 0.2, + ) -> None: + self._client = client + self._callback = callback + self._on_error = on_error + self._debounce = max(float(debounce), 0.0) + self._condition = threading.Condition() + self._closed = False + self._subscription = None + self._pending: list[dict[str, Any]] = [] + self._deadline = 0.0 + self._reader = threading.Thread( + target=self._read, name="idatui-idb-events", daemon=True + ) + self._deliverer = threading.Thread( + target=self._deliver, name="idatui-idb-refresh", daemon=True + ) + self._deliverer.start() + self._reader.start() -for i in range(ida_segment.get_segm_qty()): - seg = ida_segment.getnseg(i) - if seg is None or len(comments) >= limit or scanned >= max_scan: - continue - for ea in idautils.Heads(seg.start_ea, seg.end_ea): - scanned += 1 - if len(comments) >= limit or scanned >= max_scan: - break - if not ida_bytes.has_cmt(ida_bytes.get_flags(ea)): - continue - for rep in (False, True): - text = ida_bytes.get_cmt(ea, rep) - if text: - fn = ida_funcs.get_func(ea) - comments.append({ - "addr": hex(int(ea)), "text": text, "repeatable": rep, - "line": _line(ea), "seg": ida_segment.get_segm_name(seg), - "func": (ida_funcs.get_func_name(fn.start_ea) if fn else None), - "func_addr": (hex(int(fn.start_ea)) if fn else None)}) + def _report(self, error: BaseException) -> None: + disconnected = DatabaseDisconnectedError + if isinstance(disconnected, type) and isinstance(error, disconnected): + error = self._client._connection_error(error) + with self._condition: + closed = self._closed + if not closed and self._on_error is not None: + self._on_error(error) -# Whole-function comments are not on the byte flags, so the scan cannot see them. -for fn_ea in idautils.Functions(): - fn = ida_funcs.get_func(fn_ea) - if fn is None or len(comments) >= limit: - continue - for rep in (False, True): - text = ida_funcs.get_func_cmt(fn, rep) - if text: - seg = ida_segment.getseg(fn_ea) - comments.append({"addr": hex(int(fn_ea)), "text": text, - "repeatable": rep, "line": "", "whole_func": True, - "seg": (ida_segment.get_segm_name(seg) if seg else ""), - "func": ida_funcs.get_func_name(fn_ea), - "func_addr": hex(int(fn_ea))}) -result = {"comments": comments, "names": names, "scanned": scanned, - "truncated": (len(comments) >= limit or len(names) >= limit - or scanned >= max_scan)} -result -''', - # The findings journal (idatui/journal.py). A netnode blob rides along in - # the .i64, so "what did I work out here" survives closing the database. - "journal_get": r''' -import ida_netnode -n = ida_netnode.netnode(a.get("node", "$ idatui.journal")) -blob = n.getblob(0, "I") if ida_netnode.exist(n) else None -result = {"data": blob.decode("utf-8", "replace") if blob else ""} -result -''', - "journal_put": r''' -import ida_netnode -n = ida_netnode.netnode(a.get("node", "$ idatui.journal"), 0, True) -payload = (a.get("data") or "").encode("utf-8") -n.setblob(payload, 0, "I") -result = {"ok": True, "bytes": len(payload)} -result -''', - # Database-wide search (Ctrl+F), two kinds. - # - # BYTES uses IDA's own `find_bytes`, which already understands the pattern - # language people expect -- "B8 ? ? ? ? 90", nibble wildcards ("48 8? ??") - # and quoted literals -- so we neither parse nor match anything ourselves. - # Iterating is match+1, per its documented contract. - "search_bytes": r''' -import ida_bytes, ida_funcs, ida_idaapi, ida_lines, ida_segment -pat = str(a.get("pattern", "")).strip() -limit = max(1, int(a.get("limit", 500))) -lo = int(a.get("start", 0)) -hi = int(a.get("end", 0)) or ida_idaapi.BADADDR -flags = ida_bytes.BIN_SEARCH_FORWARD | ida_bytes.BIN_SEARCH_NOSHOW -if a.get("case"): - flags |= ida_bytes.BIN_SEARCH_CASE -rows, err, ea = [], None, lo -while len(rows) < limit: - try: - hit = ida_bytes.find_bytes(pat, range_start=ea, range_end=hi, flags=flags) - except Exception as exc: - err = str(exc) or exc.__class__.__name__ - break - if hit is None or hit == ida_idaapi.BADADDR: - break - head = ida_bytes.get_item_head(hit) - fn = ida_funcs.get_func(hit) - seg = ida_segment.getseg(hit) - try: - line = ida_lines.generate_disasm_line(head, ida_lines.GENDSM_REMOVE_TAGS) or "" - except Exception: - line = "" - rows.append({"addr": hex(int(hit)), "head": hex(int(head)), - "line": " ".join(line.split()), - "func": (ida_funcs.get_func_name(fn.start_ea) if fn else None), - "func_addr": (hex(int(fn.start_ea)) if fn else None), - "seg": (ida_segment.get_segm_name(seg) if seg else "")}) - ea = int(hit) + 1 -result = {"hits": rows, "error": err, "truncated": len(rows) >= limit} -result -''', - # TEXT walks the listing the way a person reads it: every head's rendered - # disassembly line, which is why it finds "call cs:__isoc99_scanf" and - # "0deadbeefh" alike. Bounded by max_scan, so a 400MB image reports partial - # results instead of stalling. - "search_text": r''' -import ida_lines, ida_funcs, ida_segment, idautils -import re as _re -q = str(a.get("query", "")) -limit = max(1, int(a.get("limit", 500))) -max_scan = max(1000, int(a.get("max_scan", 3000000))) -ci = (not a.get("case")) and q.islower() # smartcase, like the in-view search -rx, err = None, None -if a.get("regex"): - try: - rx = _re.compile(q, _re.I if ci else 0) - except Exception as exc: - err = "bad regex: " + str(exc) -needle = q.lower() if ci else q -rows, scanned = [], 0 -if err is None and q: - for i in range(ida_segment.get_segm_qty()): - seg = ida_segment.getnseg(i) - if seg is None or len(rows) >= limit or scanned >= max_scan: - continue - for ea in idautils.Heads(seg.start_ea, seg.end_ea): - scanned += 1 - if len(rows) >= limit or scanned >= max_scan: - break - try: - line = ida_lines.generate_disasm_line(ea, ida_lines.GENDSM_REMOVE_TAGS) or "" - except Exception: - continue - # Match what the user SEES, not IDA's column padding: nobody types - # "call" + four spaces + "cs:getenv_ptr". - line = " ".join(line.split()) - hay = line.lower() if ci else line - if (rx.search(line) if rx is not None else (needle in hay)): - fn = ida_funcs.get_func(ea) - rows.append({"addr": hex(int(ea)), "head": hex(int(ea)), - "line": line, - "func": (ida_funcs.get_func_name(fn.start_ea) if fn else None), - "func_addr": (hex(int(fn.start_ea)) if fn else None), - "seg": ida_segment.get_segm_name(seg)}) -result = {"hits": rows, "error": err, "scanned": scanned, - "truncated": len(rows) >= limit or scanned >= max_scan} -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 -''', - # 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: - 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: - 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 -''', - # 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", []): - addr_s = str(item.get("addr", "")) - text = str(item.get("comment") or "") - try: - 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 + def _read(self) -> None: 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": addr_s, "error": str(exc)}) -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 idaapi, ida_hexrays, ida_name -batch = a.get("batch") or {} -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)): + subscription = self._client.subscribe_idb_events() + except Exception as exc: # noqa: BLE001 -- surfaced through on_error + self._report(exc) + with self._condition: + self._closed = True + self._pending.clear() + self._condition.notify_all() + return + with self._condition: + if self._closed: + subscription.close() + return + self._subscription = subscription 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: - 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 -''', -} - - -_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 -''' - -# `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["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)),' - ' 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. -#: Row count + seek anchors for a whole segment, in ONE call. See -#: remote_tools.segment_index: the alternative is fetching every row. -_OPERATIONS["segment_index"] = _remote_op( - 'segment_index(addr=a["addr"], end=a.get("end", ""),' - ' page_rows=int(a.get("page_rows", 500)), detail=bool(a.get("detail", False)))') - -_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", ""), text=bool(a.get("text", True)))') - + for event in subscription: + with self._condition: + if self._closed: + break + if self._client.owns_event(event): + continue + with self._condition: + if self._closed: + break + self._pending.append(event) + self._deadline = time.monotonic() + self._debounce + self._condition.notify_all() + except Exception as exc: # noqa: BLE001 -- stream failures are recoverable + self._report(exc) + finally: + subscription.close() + with self._condition: + if self._subscription is subscription: + self._subscription = None + self._closed = True + self._pending.clear() + self._condition.notify_all() -# 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 -''' + def _deliver(self) -> None: + while True: + with self._condition: + while not self._closed and not self._pending: + self._condition.wait() + if self._closed: + return + remaining = self._deadline - time.monotonic() + if remaining > 0: + self._condition.wait(remaining) + continue + batch = tuple(self._pending) + self._pending.clear() + try: + self._callback(batch) + except Exception as exc: # noqa: BLE001 -- keep the stream alive + self._report(exc) -# 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 -''' + def close(self) -> None: + """Stop delivery and unblock the subscription reader.""" + with self._condition: + if self._closed: + return + self._closed = True + self._pending.clear() + subscription = self._subscription + self._condition.notify_all() + if subscription is not None: + subscription.close() class CodeModeClient: @@ -1277,7 +287,9 @@ class CodeModeClient: 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._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 @@ -1289,10 +301,17 @@ class CodeModeClient: 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 + handle = self._handle + if handle is not None: + if handle.connected: + return self + raise IDAConnectionError( + "Code Mode database disconnected; explicit rediscovery required" + ) if progress: - progress(f"discovering Code Mode database for {os.path.basename(self._path)}…") + 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 @@ -1321,7 +340,9 @@ class CodeModeClient: if not self._new_database or time.monotonic() >= deadline: raise if progress: - progress("waiting for the previous Code Mode lease to close…") + progress( + "waiting for the previous Code Mode lease to close…" + ) owner = find_database_owner( self._path, output_database=self._output_database, @@ -1336,7 +357,9 @@ class CodeModeClient: time.sleep(0.2) if progress: backend = handle.instance.backend - progress(f"attached to {backend} database; waiting for auto-analysis…") + 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 @@ -1360,61 +383,60 @@ class CodeModeClient: def backend(self) -> str | None: return self._handle.instance.backend if self._handle is not None else None - def execute_python(self, code: str, *, timeout: float | None = None) -> Any: + def owns_event(self, event: dict[str, Any]) -> bool: + """Whether ``event`` was produced through this client's handle.""" + handle = self._handle + return handle is not None and handle.owns_event(event) + + def subscribe_idb_events(self): + """Open Code Mode's closeable IDB-change iterator.""" 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 + return handle.subscribe_idb_events() except (DatabaseDisconnectedError, CodeModeConnectionError) 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"] - @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 watch_idb_events( + self, + callback: Callable[[tuple[dict[str, Any], ...]], None], + *, + on_error: Callable[[BaseException], None] | None = None, + debounce: float = 0.2, + ) -> IDBEventListener: + """Deliver external IDB changes in debounced batches.""" + return IDBEventListener(self, callback, on_error=on_error, debounce=debounce) - 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}") + def call(self, operation: Callable[..., Any], /, **args) -> Any: + """Execute one source-backed remote declaration through this client.""" + name = getattr(operation, "__name__", "remote operation") try: - 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._unpack( - 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 - raise + from .remote_ops import bind - # Temporary source compatibility for external drivers/tests that used the - # old WorkerClient. Application code uses the accurately named invoke(). - call = invoke + remote = bind(operation) + except KeyError as exc: + raise IDAToolError( + name, f"remote operation {name!r} is not registered" + ) from exc + if not self.connected: + self.connect() + handle = self._handle + if handle is None: + raise IDAConnectionError("Code Mode database is not connected") + try: + return remote(handle, **args) + except RemoteError as exc: + message = str(exc) + if exc.details.get("traceback"): + message += f"\n{exc.details['traceback']}" + if exc.code == "operation_timeout": + raise IDATimeoutError(message) from exc + raise IDAToolError(name, message) from exc + except (DatabaseDisconnectedError, CodeModeConnectionError) as exc: + raise self._connection_error(exc) from exc def save_database(self) -> dict[str, Any]: if not self.connected: @@ -1429,6 +451,35 @@ class CodeModeClient: except (DatabaseDisconnectedError, CodeModeConnectionError) as exc: raise self._connection_error(exc) from exc + def discard_database(self, timeout: float = 5.0) -> bool: + """Discard a final managed-worker lease; otherwise transfer finalization. + + ``False`` is an expected ownership result: a GUI owns its session, or + another lease still shares the managed worker. A busy final worker is + retried briefly so background reads finishing during quit do not turn a + real discard into an implicit save. + """ + handle = self._handle + if handle is None or not handle.connected: + return False + entry = handle.instance + if entry.backend != "idalib" or not getattr(entry, "managed", False): + return False + deadline = time.monotonic() + max(float(timeout), 0.0) + while True: + try: + handle.shutdown_database(save=False) + return True + except RemoteError as exc: + if exc.code in ("instance_shared", "shutdown_not_supported"): + return False + if exc.code == "instance_busy" and time.monotonic() < deadline: + time.sleep(0.05) + continue + raise IDAToolError("shutdown_database", str(exc)) from exc + except (DatabaseDisconnectedError, CodeModeConnectionError) as exc: + raise self._connection_error(exc) from exc + def health(self) -> dict[str, Any]: if not self.connected: self.connect() @@ -1463,8 +514,14 @@ class CodeModeClient: assert self._handle is not None entry = self._handle.instance 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)] + 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 diff --git a/idatui/domain.py b/idatui/domain.py index 699f33d..b042332 100644 --- a/idatui/domain.py +++ b/idatui/domain.py @@ -22,13 +22,12 @@ import bisect import re import threading from base64 import b64decode +from collections.abc import Callable, Sequence from concurrent.futures import ThreadPoolExecutor -from collections.abc import Sequence from dataclasses import dataclass, field, replace -from typing import NamedTuple -from typing import Callable, TYPE_CHECKING +from typing import TYPE_CHECKING, NamedTuple -from . import diag +from . import remote_ops from .errors import IDAToolError if TYPE_CHECKING: # type hint only @@ -37,8 +36,7 @@ if TYPE_CHECKING: # type hint only # Clamps derived from measured caps (list ~700, disasm ~500). Margin included. LIST_PAGE = 500 DISASM_BLOCK = 256 # instructions per cached/fetched block (<= disasm cap) -HEX_BLOCK = 16384 # bytes per cached/fetched hex block (compact read_raw -> cheap) -DECOMPILE_TIMEOUT = 15.0 # s; cap per decompile so a failing one can't hang the CLI +HEX_BLOCK = 16384 # bytes per cached/fetched hex block (compact read_raw -> cheap) _TRUNC_RE = re.compile(r"\[(\d+) chars total\]\s*$") @@ -107,7 +105,7 @@ class Head(NamedTuple): """ ea: int - kind: str # 'code' | 'data' | 'unknown' | 'member' + kind: str # 'code' | 'data' | 'unknown' | 'member' size: int text: str name: str | None = None @@ -199,10 +197,10 @@ class Ref: @dataclass class Xref: - frm: int # the referencing address - to: int | None # the referenced address - type: str # coarse: "code" | "data" - fn_name: str | None # function containing `frm` + frm: int # the referencing address + to: int | None # the referenced address + type: str # coarse: "code" | "data" + fn_name: str | None # function containing `frm` fn_addr: int | None kind: str | None = None # fine: call/jump/flow/read/write/offset/text/info @@ -218,7 +216,7 @@ class LVar: class FuncTypes: addr: int name: str - prototype: str # e.g. 'int __fastcall foo(int a, char *b)' + prototype: str # e.g. 'int __fastcall foo(int a, char *b)' lvars: list[LVar] @@ -227,7 +225,7 @@ class Struct: name: str size: int is_union: bool - members: int # field count + members: int # field count ordinal: int @classmethod @@ -244,6 +242,7 @@ class Struct: @dataclass(frozen=True) class StrLit: """A string literal IDA found in the binary (the Shift+F12 list).""" + addr: int text: str length: int @@ -271,6 +270,7 @@ class SearchHit: an instruction, so ``head`` is the item to navigate to and ``line`` is what that item renders as. """ + addr: int head: int line: str = "" @@ -287,6 +287,7 @@ class Comment: report can show what was being commented ON without a second round trip. ``whole_func`` marks a function comment rather than an instruction one. """ + addr: int text: str repeatable: bool = False @@ -302,6 +303,7 @@ class NamedItem: """An address carrying a real name -- one you typed, or one the file's own symbols supplied. IDA records both as "user" names and does not remember which was which, so a report must say so rather than claim authorship.""" + addr: int name: str is_func: bool = False @@ -319,6 +321,7 @@ class Linkage: ``name`` is the joinable name; ``raw`` keeps the spelling IDA reported, which is what the user sees in the listing. """ + addr: int name: str module: str = "" @@ -374,7 +377,9 @@ class FunctionIndex: query: dict = {"offset": offset, "count": LIST_PAGE} if self.filter: query["filter"] = self.filter - data = _query_data(self._prog.client.invoke("list_funcs", queries=[query])) + data = _query_data( + self._prog.client.call(remote_ops.list_funcs, queries=[query]) + ) added = 0 with self._lock: for d in data: @@ -468,7 +473,7 @@ class DisasmModel: self._blocks: dict[int, list[Line]] = {} self._total: int | None = None self._ea_list: list[int] | None = None - self._max_raw = 0 # widest opcode length seen (bytes) + self._max_raw = 0 # widest opcode length seen (bytes) self._func_end: int | None = None self._func_end_done = False self._lock = threading.Lock() @@ -484,8 +489,8 @@ class DisasmModel: code function this equals the heads row count that backs the lines.""" if self._total is not None: return self._total - payload = self._prog.client.invoke( - "disasm", addr=hex(self.ea), max_instructions=1, include_total=True + payload = self._prog.client.call( + remote_ops.disasm, addr=hex(self.ea), max_instructions=1, include_total=True ) total = payload.get("total_instructions") if total is None: @@ -527,7 +532,7 @@ class DisasmModel: nxt = lines[i + 1].ea if i + 1 < len(lines) else last_end length = max(nxt - ln.ea, 0) off = ln.ea - start - b = bytes(data[off:off + length]) + b = bytes(data[off : off + length]) biggest = max(biggest, len(b)) out.append(replace(ln, raw=b)) with self._lock: @@ -538,20 +543,22 @@ class DisasmModel: @staticmethod def _line_from_head(r: dict) -> Line: """Adapt a ``heads`` row to a disasm Line (label = the head's name).""" - return Line(ea=_as_int(r["ea"]), text=r.get("text", ""), - label=r.get("name")) + return Line(ea=_as_int(r["ea"]), text=r.get("text", ""), label=r.get("name")) def _fetch_block(self, b: int) -> list[Line]: # The function disasm view is a listing filtered to the function: fetch a # block of heads (one per instruction for code). Over-fetch one row so # the block knows where its last instruction ends (opcode-byte sizing). - payload = self._prog.client.invoke( - "heads", addr=hex(self.ea), offset=b * self.BLOCK, - count=self.BLOCK + 1, **self._end_kw(), + payload = self._prog.client.call( + remote_ops.heads, + addr=hex(self.ea), + offset=b * self.BLOCK, + count=self.BLOCK + 1, + **self._end_kw(), ) rows = payload.get("heads", []) if isinstance(payload, dict) else [] fetched = [self._line_from_head(r) for r in rows] - lines = fetched[:self.BLOCK] + lines = fetched[: self.BLOCK] if len(fetched) > self.BLOCK: end_ea: int | None = fetched[self.BLOCK].ea else: # this block ends the function @@ -607,7 +614,7 @@ class DisasmModel: block = self._get_block(b) lo = start - b * self.BLOCK if b == b0 else 0 hi = end - b * self.BLOCK if b == b1 else self.BLOCK - out.extend(block[max(lo, 0):hi]) + out.extend(block[max(lo, 0) : hi]) if prefetch: self._prefetch_block(b1 + 1) # forward scroll self._prefetch_block(b0 - 1) # backward scroll @@ -699,8 +706,9 @@ class ListingModel: #: _text_gen, which counts up from 0, so such a page always reads as stale. _SKELETON_GEN = -1 - def __init__(self, program: "Program", seg_start: int, seg_end: int, - name: str | None = None): + def __init__( + self, program: "Program", seg_start: int, seg_end: int, name: str | None = None + ): self._prog = program self.seg_start = seg_start self.seg_end = seg_end @@ -714,7 +722,7 @@ class ListingModel: # N bytes PRESENTS as N rows and the text for each is synthesised on # demand. _row_at[i] is the logical row where physical head i starts. self._row_at: list[int] = [] - self._head_eas: list[int] = [] # parallel to _heads, for bisect + self._head_eas: list[int] = [] # parallel to _heads, for bisect #: Which name generation each head's TEXT was rendered at, parallel to #: _heads. A rename bumps :attr:`_text_gen`; the rows themselves stay #: (their addresses and row numbers are unchanged) and are re-rendered a @@ -740,7 +748,7 @@ class ListingModel: #: means something DID move the walk. Program.listing() throws the model #: away when it sees this, so the next read rebuilds from scratch. self.stale_structure = False - self._rows = 0 # total logical rows loaded + self._rows = 0 # total logical rows loaded self._ubytes: dict[int, bytes] = {} # lazily-read bytes for those rows self._next: int | None = seg_start # next address to fetch from self._done = False @@ -788,7 +796,7 @@ class ListingModel: size = int(r.get("size") or 0) if size > 0: off = _as_int(r["ea"]) - lo - raw = bytes(data[off:off + size]) + raw = bytes(data[off : off + size]) if len(raw) > biggest: biggest = len(raw) try: @@ -827,19 +835,26 @@ class ListingModel: """ with self._lock: if self._done and self._heads: - return True # already indexed; re-priming is a no-op + return True # already indexed; re-priming is a no-op try: - idx = self._prog.client.invoke( - "segment_index", addr=hex(self.seg_start), end=hex(self.seg_end), - page_rows=self.PAGE, detail=True) + idx = self._prog.client.call( + remote_ops.segment_index, + addr=hex(self.seg_start), + end=hex(self.seg_end), + page_rows=self.PAGE, + detail=True, + ) except Exception: # noqa: BLE001 -- fall back to streaming return False if not isinstance(idx, dict) or idx.get("error") or "eas" not in idx: return False try: - eas = array.array("Q"); eas.frombytes(b64decode(idx["eas"])) - kinds = array.array("B"); kinds.frombytes(b64decode(idx["kinds"])) - sizes = array.array("I"); sizes.frombytes(b64decode(idx["sizes"])) + eas = array.array("Q") + eas.frombytes(b64decode(idx["eas"])) + kinds = array.array("B") + kinds.frombytes(b64decode(idx["kinds"])) + sizes = array.array("I") + sizes.frombytes(b64decode(idx["sizes"])) except Exception: # noqa: BLE001 return False names = idx.get("kind_names") or [] @@ -881,7 +896,8 @@ class ListingModel: self._page_digest = [None] * len(anchors) self._page_rows = [ (anchors[k + 1][2] if k + 1 < len(anchors) else n) - anchors[k][2] - for k in range(len(anchors))] + for k in range(len(anchors)) + ] self._skeleton = True self._done = True self._next = None @@ -912,8 +928,9 @@ class ListingModel: if self._done or self._next is None: return 0 frm = self._next - payload = self._prog.client.invoke( - "heads", addr=hex(frm), count=self.PAGE, annotate=True, text=text) + payload = self._prog.client.call( + remote_ops.heads, addr=hex(frm), count=self.PAGE, annotate=True, text=text + ) rows = payload.get("heads", []) if isinstance(payload, dict) else [] cur = payload.get("cursor", {}) if isinstance(payload, dict) else {} page = self._build_page(rows, raw=text) @@ -925,8 +942,9 @@ class ListingModel: self._skeleton = True self._page_head.append(len(self._heads)) self._page_addr.append(frm) - self._page_digest.append(payload.get("digest") - if isinstance(payload, dict) else None) + self._page_digest.append( + payload.get("digest") if isinstance(payload, dict) else None + ) self._page_rows.append(len(rows)) for h in page: # Banner/label rows (function headers, separators, code labels) @@ -983,7 +1001,7 @@ class ListingModel: self._ubytes[b0] = blk off = a - b0 take = min(BLK - off, n - len(out)) - chunk = blk[off:off + take] if blk else b"" + chunk = blk[off : off + take] if blk else b"" if not chunk: break out += chunk @@ -1004,8 +1022,9 @@ class ListingModel: ea = h.ea + off b = self._unknown_bytes(ea, 1) text = f"db {b[0]:02X}h" if b else "db ?" - return Head(ea=ea, kind="unknown", size=1, text=text, - name=h.name if off == 0 else None) + return Head( + ea=ea, kind="unknown", size=1, text=text, name=h.name if off == 0 else None + ) def ensure(self, n: int) -> None: """Ensure at least ``n`` logical rows are loaded (or all, if fewer).""" @@ -1022,8 +1041,11 @@ class ListingModel: return idx with self._lock: have = self._rows - last_ea = (self._heads[-1].ea + max(self._heads[-1].size, 1) - 1 - if self._heads else -1) + last_ea = ( + self._heads[-1].ea + max(self._heads[-1].size, 1) - 1 + if self._heads + else -1 + ) done = self._done if done or (have and last_ea >= ea): # Loaded past ea without an exact head hit: return the first head @@ -1086,13 +1108,13 @@ class ListingModel: """ with self._lock: if not (self.seg_start <= ea < self.seg_end): - return True # another segment; nothing moved here + return True # another segment; nothing moved here if len(self._page_head) < 3: - return False # barely walked; a rebuild is cheaper + return False # barely walked; a rebuild is cheaper p = bisect.bisect_right(self._page_addr, ea) - 1 p = max(p - 1, 0) if p <= 0: - return False # the edit is in the first pages + return False # the edit is in the first pages keep = self._page_head[p] if keep <= 0: return False @@ -1110,7 +1132,7 @@ class ListingModel: last = self._heads[-1] self._rows = self._row_at[-1] + self._span(last) self._done = False - self._ubytes.clear() # undefined-run bytes behind the drop point + self._ubytes.clear() # undefined-run bytes behind the drop point return True def invalidate_text(self) -> None: @@ -1156,8 +1178,9 @@ class ListingModel: def _page_bounds(self, p: int) -> tuple[int, int]: """[first, last) head index of page ``p`` (caller holds the lock).""" lo = self._page_head[p] - hi = (self._page_head[p + 1] if p + 1 < len(self._page_head) - else len(self._heads)) + hi = ( + self._page_head[p + 1] if p + 1 < len(self._page_head) else len(self._heads) + ) return lo, hi def _ensure_page(self, p: int) -> int: @@ -1181,13 +1204,20 @@ class ListingModel: # the expectation rather than asking first means a page that HAS changed # still costs one round trip. try: - payload = self._prog.client.invoke( - "heads", addr=hex(addr), count=self.PAGE, annotate=True, - expect="" if want_digest is None else str(want_digest)) + payload = self._prog.client.call( + remote_ops.heads, + addr=hex(addr), + count=self.PAGE, + annotate=True, + expect="" if want_digest is None else str(want_digest), + ) except Exception: # noqa: BLE001 -- keep the old text rather than blank return p + 1 - if (isinstance(payload, dict) and "heads" not in payload - and payload.get("count") == want_rows): + if ( + isinstance(payload, dict) + and "heads" not in payload + and payload.get("count") == want_rows + ): with self._lock: if self._text_gen == gen and len(self._heads) >= hi: for k in range(lo, hi): @@ -1212,8 +1242,9 @@ class ListingModel: # what it once loaded. Leaving it stale is how a literal cycling # hex -> dec -> hex ends up declared "unchanged" while the row still # shows the decimal it was refetched with in between. - self._page_digest[p] = (payload.get("digest") - if isinstance(payload, dict) else None) + self._page_digest[p] = ( + payload.get("digest") if isinstance(payload, dict) else None + ) for k in range(lo, hi): self._head_gen[k] = gen return p + 1 @@ -1225,8 +1256,9 @@ class ListingModel: j, off = self._phys(i) if j < 0: return None - stale = ((self._renamed or self._skeleton) - and self._head_gen[j] != self._text_gen) + stale = (self._renamed or self._skeleton) and self._head_gen[ + j + ] != self._text_gen if not stale: span = self._span(self._heads[j]) h = self._heads[j] @@ -1265,8 +1297,9 @@ class ListingModel: spans = [self._phys(i) for i in range(max(start, 0), max(rows, 0))] heads = self._heads plain = [(j, off, heads[j]) for j, off in spans if j >= 0] - return [self._row_head(j, off) if self._span(h) > 1 else h - for j, off, h in plain] + return [ + self._row_head(j, off) if self._span(h) > 1 else h for j, off, h in plain + ] def index_of_ea(self, ea: int) -> int: with self._lock: @@ -1361,7 +1394,7 @@ class HexModel: if block is None: return (va, None) bo = off - b * self.BLOCK - return (va, block[bo:bo + 16]) + return (va, block[bo : bo + 16]) def ensure(self, r0: int, count: int) -> None: """Blocking: fetch the blocks covering rows [r0, r0+count) if missing.""" @@ -1459,20 +1492,32 @@ class Program: return self._segments_cache segs: list[tuple[int, int, int, str]] = [] try: - r = self.client.invoke("file_regions") - for d in (r.get("regions", []) if isinstance(r, dict) else []): + r = self.client.call(remote_ops.file_regions) + for d in r.get("regions", []) if isinstance(r, dict) else []: if isinstance(d, dict) and "start" in d: - segs.append((_as_int(d["start"]), _as_int(d["end"]), - int(d.get("file_off", -1)), d.get("name", "") or "")) + segs.append( + ( + _as_int(d["start"]), + _as_int(d["end"]), + int(d.get("file_off", -1)), + d.get("name", "") or "", + ) + ) except IDAToolError: segs = [] if not segs: # older server without file_regions -> survey_binary (slow) try: - sb = self.client.invoke("survey_binary") - for s in (sb.get("segments", []) if isinstance(sb, dict) else []): + sb = self.client.call(remote_ops.survey_binary) + for s in sb.get("segments", []) if isinstance(sb, dict) else []: try: - segs.append((_as_int(s["start"]), _as_int(s["end"]), -1, - s.get("name", "") or "")) + segs.append( + ( + _as_int(s["start"]), + _as_int(s["end"]), + -1, + s.get("name", "") or "", + ) + ) except (KeyError, ValueError, TypeError): continue except Exception: # noqa: BLE001 -- best-effort; callers handle empty @@ -1535,21 +1580,27 @@ class Program: return b"" if not self._no_read_raw: try: - r = self.client.invoke("read_raw", addr=hex(ea), size=int(n)) + r = self.client.call(remote_ops.read_raw, addr=hex(ea), size=int(n)) h = r.get("hex") if isinstance(r, dict) else None if isinstance(h, str): out = bytes.fromhex(h) return out[:n] if len(out) >= n else out + b"\x00" * (n - len(out)) except IDAToolError as e: # Tool missing on this server: stop trying it, use get_bytes. - if "read_raw" in str(e) or "Unknown tool" in str(e) or "not found" in str(e): + if ( + "read_raw" in str(e) + or "Unknown tool" in str(e) + or "not found" in str(e) + ): self._no_read_raw = True else: return b"\x00" * n except (ValueError, KeyError): pass # malformed hex -> fall through to the legacy decoder try: - r = self.client.invoke("get_bytes", regions=[{"addr": hex(ea), "size": int(n)}]) + r = self.client.call( + remote_ops.get_bytes, regions=[{"addr": hex(ea), "size": int(n)}] + ) except IDAToolError: return b"\x00" * n res = r.get("result", []) if isinstance(r, dict) else [] @@ -1590,7 +1641,7 @@ class Program: with self._lock: m = self._listings.get(start) if m is not None and m.stale_structure: - m = None # a refresh found the walk had moved; start over + m = None # a refresh found the walk had moved; start over if m is None: m = ListingModel(self, start, end, name) self._listings[start] = m @@ -1600,11 +1651,15 @@ class Program: def list_structs(self, filter: str = "") -> list[Struct]: """All local structs/unions (optionally name-substring filtered), sorted by name.""" - payload = self.client.invoke("search_structs", filter=filter) + payload = self.client.call(remote_ops.search_structs, filter=filter) res = payload.get("result", []) if isinstance(payload, dict) else [] - out = [Struct.from_raw(d) for d in res - if isinstance(d, dict) and d.get("name") - and not str(d["name"]).startswith("$")] # skip anonymous UDTs + out = [ + Struct.from_raw(d) + for d in res + if isinstance(d, dict) + and d.get("name") + and not str(d["name"]).startswith("$") + ] # skip anonymous UDTs out.sort(key=lambda s: s.name.lower()) return out @@ -1612,8 +1667,9 @@ class Program: """A C definition for ``name`` reconstructed from its member layout (the remote operation exposes members, not printable source). Faithful to IDA's field names/types; array dims are moved after the field name.""" - payload = self.client.invoke( - "type_inspect", queries=[{"name": name, "include_members": True}]) + payload = self.client.call( + remote_ops.type_inspect, queries=[{"name": name, "include_members": True}] + ) res = payload.get("result", []) if isinstance(payload, dict) else [] info = res[0] if res and isinstance(res[0], dict) else {} kw = "union" if info.get("is_union") else "struct" @@ -1635,7 +1691,7 @@ class Program: def declare_type(self, decl: str) -> str | None: """Create or update a C type. Returns None on success, else the parse error. (Re-declaring a name updates it in place.)""" - payload = self.client.invoke("declare_type", decls=decl) + payload = self.client.call(remote_ops.declare_type, decls=decl) res = payload.get("result", []) if isinstance(payload, dict) else [] if res and isinstance(res[0], dict): return res[0].get("error") @@ -1646,20 +1702,32 @@ class Program: """Structured decompiler types for the function at ``ea`` (prototype + local variables). None if ``ea`` isn't a decompilable function.""" try: - r = self.client.invoke("func_types", addr=hex(ea)) + r = self.client.call(remote_ops.func_types, addr=hex(ea)) except IDAToolError: return None if not isinstance(r, dict) or r.get("error"): return None - lvars = [LVar(name=lv.get("name", ""), type=lv.get("type", ""), - is_arg=bool(lv.get("is_arg"))) - for lv in r.get("lvars", []) if isinstance(lv, dict)] - return FuncTypes(addr=_as_int(r.get("addr", hex(ea))), name=r.get("name", ""), - prototype=r.get("prototype", ""), lvars=lvars) + lvars = [ + LVar( + name=lv.get("name", ""), + type=lv.get("type", ""), + is_arg=bool(lv.get("is_arg")), + ) + for lv in r.get("lvars", []) + if isinstance(lv, dict) + ] + return FuncTypes( + addr=_as_int(r.get("addr", hex(ea))), + name=r.get("name", ""), + prototype=r.get("prototype", ""), + lvars=lvars, + ) def set_function_type(self, ea: int, signature: str) -> str | None: """Set a function's prototype. None on success, else an error string.""" - r = self.client.invoke("set_type", edits=[{"addr": hex(ea), "signature": signature}]) + r = self.client.call( + remote_ops.set_type, edits=[{"addr": hex(ea), "signature": signature}] + ) res = r.get("result", []) if isinstance(r, dict) else [] row = res[0] if res and isinstance(res[0], dict) else {} if row.get("ok"): @@ -1670,7 +1738,7 @@ class Program: """Current type info for a data item/global: {addr,name,type,size,is_func}. None if the operation fails or the address isn't mapped.""" try: - r = self.client.invoke("data_type", addr=hex(ea)) + r = self.client.call(remote_ops.data_type, addr=hex(ea)) except IDAToolError: return None if not isinstance(r, dict) or r.get("error"): @@ -1679,8 +1747,10 @@ class Program: def set_data_type(self, ea: int, decl: str) -> str | None: """Set a global/data item's type. None on success, else an error string.""" - r = self.client.invoke( - "set_type", edits=[{"kind": "global", "addr": hex(ea), "type": decl}]) + r = self.client.call( + remote_ops.set_type, + edits=[{"kind": "global", "addr": hex(ea), "type": decl}], + ) res = r.get("result", []) if isinstance(r, dict) else [] row = res[0] if res and isinstance(res[0], dict) else {} if row.get("ok"): @@ -1690,7 +1760,9 @@ class Program: def set_lvar_type(self, fn_ea: int, var: str, ty: str) -> str | None: """Set a decompiler local variable's type through ida-domain pseudocode. None on success, else an error string.""" - r = self.client.invoke("set_lvar_type", addr=hex(fn_ea), variable=var, type=ty) + r = self.client.call( + remote_ops.set_lvar_type, addr=hex(fn_ea), variable=var, type=ty + ) if isinstance(r, dict) and r.get("error"): return r["error"] if isinstance(r, dict) and not r.get("ok"): @@ -1701,7 +1773,7 @@ class Program: """Delete a named type. Returns None on success, else an error string. Returns a clear error instead of raising when the runtime cannot do it.""" try: - self.client.invoke("del_type", name=name) + self.client.call(remote_ops.del_type, name=name) return None except IDAToolError as e: msg = e.message @@ -1732,7 +1804,7 @@ class Program: self._pc_nums.pop(ea, None) self._decomp_maps.pop(ea, None) try: - self.client.invoke("force_recompile", items=[{"addr": hex(ea)}]) + self.client.call(remote_ops.force_recompile, items=[{"addr": hex(ea)}]) except Exception: # noqa: BLE001 -- refresh still refetches best-effort pass @@ -1749,22 +1821,15 @@ class Program: # Cached before a rename: names may be stale. Drop Hex-Rays' # cache so the refetch reflects the new names. try: - self.client.invoke("force_recompile", items=[{"addr": hex(ea)}]) + self.client.call( + remote_ops.force_recompile, items=[{"addr": hex(ea)}] + ) except Exception: # noqa: BLE001 pass - # Bound the decompile: a function Hex-Rays can't handle tends to stall - # near the client's default 30s timeout, and the transport retries a - # dropped connection up to max_retries+1 times, re-running the failing - # decompile each time. Cap it so the worst case stays well under the - # rpcclient socket timeout, and cache the failure below so a re-request - # returns instantly instead of re-grinding. + # The typed remote declaration carries a 15-second transport timeout, + # so a function Hex-Rays cannot handle does not stall the UI. try: - # Code Mode returns the complete JSON result directly; unlike the - # old MCP tool transport there is no structured-content envelope or - # out-of-band download URL to unwrap. - payload = self.client.invoke( - "decompile", addr=hex(ea), timeout=DECOMPILE_TIMEOUT - ) + payload = self.client.call(remote_ops.decompile, addr=hex(ea)) except Exception as e: # noqa: BLE001 -- surface as a failed decompile dec = Decompilation(ea, None, True, f"decompile error: {e}", False, None) with self._lock: @@ -1790,7 +1855,7 @@ class Program: self._name_gen += 1 models = list(self._disasm.values()) listings = list(self._listings.values()) - self._pc_nums.clear() # a reformat moves every literal on its line + self._pc_nums.clear() # a reformat moves every literal on its line for m in models: m.invalidate() for lm in listings: @@ -1832,6 +1897,34 @@ class Program: if self._listings.get(start) is lm: del self._listings[start] + def invalidate_external(self) -> None: + """Drop every cached view of an IDB changed by another client. + + An event may describe a rename, a byte patch, a new function, or a + segment move. Treating an unknown event as text-only risks displaying a + structurally impossible mix of old rows and new metadata, so the + external boundary deliberately invalidates all derived state. The app + debounces event bursts before reaching this method. + """ + with self._lock: + self._name_gen += 1 + models = list(self._disasm.values()) + self._indices.clear() + self._disasm.clear() + self._listings.clear() + self._decomp.clear() + self._pc_nums.clear() + self._decomp_maps.clear() + self._flowcharts.clear() + self._strings = None + self._linkage = None + self._segments_cache = None + self._sections = None + self._fileregions = None + self._hexmodel = None + for model in models: + model.invalidate() + # -- item / function structure edits (IDA c/d/u/p) --------------------- # @staticmethod def _first_result(payload) -> dict: @@ -1847,18 +1940,19 @@ class Program: Undefine first so it works even when the bytes are currently part of a data/align item — ``create_insn`` refuses to carve into a live item.""" try: - self.client.invoke("undefine", items=[{"addr": hex(ea)}]) + self.client.call(remote_ops.undefine, items=[{"addr": hex(ea)}]) except IDAToolError: pass # nothing defined here yet -> just try to create the insn res = self._first_result( - self.client.invoke("define_code", items=[{"addr": hex(ea)}])) + self.client.call(remote_ops.define_code, items=[{"addr": hex(ea)}]) + ) if res.get("error"): raise IDAToolError("define_code", f"@ {ea:#x}: {res['error']}") def decomp_error(self, ea: int) -> str: """Hex-Rays' own reason for refusing ``ea``, or "" if it won't say.""" try: - r = self.client.invoke("decomp_error", addr=hex(ea)) + r = self.client.call(remote_ops.decomp_error, addr=hex(ea)) except IDAToolError: return "" if not isinstance(r, dict): @@ -1877,19 +1971,22 @@ class Program: def thumb_scan(self, start: int, end: int, apply: bool = True) -> dict: """Find Thumb entry points from odd pointers in ``[start, end)``.""" - r = self.client.invoke("thumb_scan", start=hex(start), end=hex(end), - apply=bool(apply)) + r = self.client.call( + remote_ops.thumb_scan, start=hex(start), end=hex(end), apply=bool(apply) + ) if not isinstance(r, dict) or r.get("error"): - raise IDAToolError("thumb_scan", - f"@ {start:#x}: {(r or {}).get('error', 'failed')}") + raise IDAToolError( + "thumb_scan", f"@ {start:#x}: {(r or {}).get('error', 'failed')}" + ) return r def set_thumb(self, ea: int, mode: str = "toggle") -> dict: """Switch ARM/Thumb decoding at ``ea``. Returns the resulting state.""" - r = self.client.invoke("set_thumb", addr=hex(ea), mode=mode) + r = self.client.call(remote_ops.set_thumb, addr=hex(ea), mode=mode) if not isinstance(r, dict) or r.get("error"): - raise IDAToolError("set_thumb", - f"@ {ea:#x}: {(r or {}).get('error', 'failed')}") + raise IDAToolError( + "set_thumb", f"@ {ea:#x}: {(r or {}).get('error', 'failed')}" + ) return r def define_code_run(self, ea: int, limit: int = 20000) -> dict: @@ -1899,13 +1996,16 @@ class Program: provide the run operation. """ try: - r = self.client.invoke("define_code_run", addr=hex(ea), limit=int(limit)) + r = self.client.call( + remote_ops.define_code_run, addr=hex(ea), limit=int(limit) + ) except IDAToolError: self.define_code(ea) return {"count": 1, "stopped": "single", "end": hex(ea)} if not isinstance(r, dict) or r.get("error"): - raise IDAToolError("define_code_run", - f"@ {ea:#x}: {(r or {}).get('error', 'failed')}") + raise IDAToolError( + "define_code_run", f"@ {ea:#x}: {(r or {}).get('error', 'failed')}" + ) return r def define_func(self, ea: int) -> dict: @@ -1915,16 +2015,18 @@ class Program: falls back to a plain create for alternate clients. """ try: - r = self.client.invoke("define_func_run", addr=hex(ea)) + r = self.client.call(remote_ops.define_func_run, addr=hex(ea)) except IDAToolError: res = self._first_result( - self.client.invoke("define_func", items=[{"addr": hex(ea)}])) + self.client.call(remote_ops.define_func, items=[{"addr": hex(ea)}]) + ) if res.get("error"): raise IDAToolError("define_func", f"@ {ea:#x}: {res['error']}") return {"ok": True, "how": "legacy"} if not isinstance(r, dict) or not r.get("ok"): - raise IDAToolError("define_func", - f"@ {ea:#x}: {(r or {}).get('error', 'failed')}") + raise IDAToolError( + "define_func", f"@ {ea:#x}: {(r or {}).get('error', 'failed')}" + ) return r def undefine(self, ea: int, size: int | None = None) -> None: @@ -1932,7 +2034,7 @@ class Program: item: dict = {"addr": hex(ea)} if size: item["size"] = int(size) - res = self._first_result(self.client.invoke("undefine", items=[item])) + res = self._first_result(self.client.call(remote_ops.undefine, items=[item])) if res.get("error"): raise IDAToolError("undefine", f"@ {ea:#x}: {res['error']}") @@ -1942,24 +2044,29 @@ class Program: item: dict = {"addr": hex(ea), "type": type_decl} if name: item["name"] = name - res = self._first_result(self.client.invoke("make_data", items=[item])) + res = self._first_result(self.client.call(remote_ops.make_data, items=[item])) if res.get("ok") is False or res.get("error"): raise IDAToolError( - "make_data", f"@ {ea:#x}: {res.get('error') or 'rejected'}") + "make_data", f"@ {ea:#x}: {res.get('error') or 'rejected'}" + ) def make_string(self, ea: int, length: int = 0, kind: str = "c") -> str: """Create a string literal at ``ea`` (IDA's 'A'); auto-length when 0. Returns the decoded contents.""" - r = self.client.invoke("make_string", addr=hex(ea), length=int(length), kind=kind) + r = self.client.call( + remote_ops.make_string, addr=hex(ea), length=int(length), kind=kind + ) res = r if isinstance(r, dict) else {} if not res.get("ok"): raise IDAToolError( - "make_string", f"@ {ea:#x}: {res.get('error') or 'rejected'}") + "make_string", f"@ {ea:#x}: {res.get('error') or 'rejected'}" + ) return res.get("text", "") # -- literal display formats (IDA's 'o': hex / dec / char / offset) ---- # - def op_format(self, ea: int, mode: str = "cycle", col: int = -1, - n: int = -1) -> dict: + def op_format( + self, ea: int, mode: str = "cycle", col: int = -1, n: int = -1 + ) -> dict: """Change how the literal at ``ea`` is DISPLAYED in the listing. ``col`` is a column inside the rendered line, which is how the cursor @@ -1967,8 +2074,9 @@ class Program: ``cycle``/``back`` (step the stops that make sense for this value) or a format by name. ``show`` reports without changing anything. """ - r = self.client.invoke("op_format", addr=hex(ea), mode=str(mode), - col=int(col), n=int(n)) + r = self.client.call( + remote_ops.op_format, addr=hex(ea), mode=str(mode), col=int(col), n=int(n) + ) res = r if isinstance(r, dict) else {} if res.get("error"): raise IDAToolError("op_format", f"@ {ea:#x}: {res['error']}") @@ -1991,31 +2099,43 @@ class Program: if hit is not None and hit[1] == gen: return hit[0] try: - r = self.client.invoke("pc_nums", addr=hex(fn_ea)) + r = self.client.call(remote_ops.pc_nums, addr=hex(fn_ea)) except Exception: # noqa: BLE001 -- an older worker hasn't got the tool r = {} out: dict[int, list[tuple[int, int, str, int, int]]] = {} for rec in (r or {}).get("nums", []): try: out.setdefault(int(rec["line"]), []).append( - (int(rec["x0"]), int(rec["x1"]), str(rec.get("value", "")), - _as_int(rec["ea"]), int(rec.get("opnum", 0)))) + ( + int(rec["x0"]), + int(rec["x1"]), + str(rec.get("value", "")), + _as_int(rec["ea"]), + int(rec.get("opnum", 0)), + ) + ) except Exception: # noqa: BLE001 -- skip a malformed row continue with self._lock: self._pc_nums[fn_ea] = (out, gen) return out - def pc_num_format(self, fn_ea: int, mode: str = "cycle", line: int = -1, - col: int = -1) -> dict: + def pc_num_format( + self, fn_ea: int, mode: str = "cycle", line: int = -1, col: int = -1 + ) -> dict: """The same, for a number in the DECOMPILATION of ``fn_ea``. Hex-Rays keeps number formats of its own, per (address, operand) — the listing's format doesn't reach the pseudocode and vice versa, so this is a separate call rather than a flag on ``op_format``. """ - r = self.client.invoke("pc_num_format", addr=hex(fn_ea), mode=str(mode), - line=int(line), col=int(col)) + r = self.client.call( + remote_ops.pc_num_format, + addr=hex(fn_ea), + mode=str(mode), + line=int(line), + col=int(col), + ) res = r if isinstance(r, dict) else {} if res.get("error"): raise IDAToolError("pc_num_format", f"@ {fn_ea:#x}: {res['error']}") @@ -2043,22 +2163,30 @@ class Program: offset, page = 0, 2000 while True: try: - payload = self.client.invoke( - "list_strings", offset=offset, count=page, min_len=min_len, - refresh=(refresh and offset == 0)) + payload = self.client.call( + remote_ops.list_strings, + offset=offset, + count=page, + min_len=min_len, + refresh=(refresh and offset == 0), + ) except IDAToolError: return [] rows = payload.get("strings", []) if isinstance(payload, dict) else [] for r in rows: if not isinstance(r, dict): continue - out.append(StrLit( - addr=_as_int(r.get("addr", 0)), - text=r.get("text", ""), - length=int(r.get("len", 0) or 0), - type=r.get("type", "") or "", - )) - total = int(payload.get("total", 0) or 0) if isinstance(payload, dict) else 0 + out.append( + StrLit( + addr=_as_int(r.get("addr", 0)), + text=r.get("text", ""), + length=int(r.get("len", 0) or 0), + type=r.get("type", "") or "", + ) + ) + total = ( + int(payload.get("total", 0) or 0) if isinstance(payload, dict) else 0 + ) if len(rows) < page or len(out) >= total: break offset += len(rows) @@ -2074,27 +2202,39 @@ class Program: if hit is not None: return hit try: - payload = self.client.invoke("list_linkage", kind="both") + payload = self.client.call(remote_ops.list_linkage, kind="both") except IDAToolError: return ([], []) if not isinstance(payload, dict): return ([], []) - imps = [Linkage(addr=_as_int(r.get("addr", 0)), - name=link_name(r.get("name", "")), - module=r.get("module", "") or "", - raw=r.get("name", "") or "") - for r in payload.get("imports", []) if isinstance(r, dict)] - exps = [Linkage(addr=_as_int(r.get("addr", 0)), - name=link_name(r.get("name", "")), - ordinal=int(r.get("ordinal", 0) or 0), - raw=r.get("name", "") or "") - for r in payload.get("exports", []) if isinstance(r, dict)] + imps = [ + Linkage( + addr=_as_int(r.get("addr", 0)), + name=link_name(r.get("name", "")), + module=r.get("module", "") or "", + raw=r.get("name", "") or "", + ) + for r in payload.get("imports", []) + if isinstance(r, dict) + ] + exps = [ + Linkage( + addr=_as_int(r.get("addr", 0)), + name=link_name(r.get("name", "")), + ordinal=int(r.get("ordinal", 0) or 0), + raw=r.get("name", "") or "", + ) + for r in payload.get("exports", []) + if isinstance(r, dict) + ] out = ([i for i in imps if i.name], [e for e in exps if e.name]) with self._lock: self._linkage = out return out - def annotations(self, limit: int = 4000) -> tuple[list["Comment"], list["NamedItem"]]: + def annotations( + self, limit: int = 4000 + ) -> tuple[list["Comment"], list["NamedItem"]]: """``(comments, names)`` -- everything a person added to this database. Not cached: it is the *current* state of your work, and the one caller @@ -2102,44 +2242,64 @@ class Program: no such operation, so an alternate client degrades instead of breaking. """ try: - payload = self.client.invoke("list_annotations", limit=int(limit)) + payload = self.client.call(remote_ops.list_annotations, limit=int(limit)) except IDAToolError: return ([], []) if not isinstance(payload, dict): return ([], []) comments = [ - Comment(addr=_as_int(r.get("addr", 0)), text=str(r.get("text", "")), - repeatable=bool(r.get("repeatable")), - whole_func=bool(r.get("whole_func")), - line=str(r.get("line", "") or ""), - seg=str(r.get("seg", "") or ""), - func=(r.get("func") or None), - func_addr=(_as_int(r["func_addr"]) if r.get("func_addr") else None)) - for r in payload.get("comments", []) if isinstance(r, dict) and r.get("text")] + Comment( + addr=_as_int(r.get("addr", 0)), + text=str(r.get("text", "")), + repeatable=bool(r.get("repeatable")), + whole_func=bool(r.get("whole_func")), + line=str(r.get("line", "") or ""), + seg=str(r.get("seg", "") or ""), + func=(r.get("func") or None), + func_addr=(_as_int(r["func_addr"]) if r.get("func_addr") else None), + ) + for r in payload.get("comments", []) + if isinstance(r, dict) and r.get("text") + ] names = [ - NamedItem(addr=_as_int(r.get("addr", 0)), name=str(r.get("name", "")), - is_func=bool(r.get("func")), size=int(r.get("size", 0) or 0), - proto=(r.get("proto") or None), seg=str(r.get("seg", "") or "")) - for r in payload.get("names", []) if isinstance(r, dict) and r.get("name")] + NamedItem( + addr=_as_int(r.get("addr", 0)), + name=str(r.get("name", "")), + is_func=bool(r.get("func")), + size=int(r.get("size", 0) or 0), + proto=(r.get("proto") or None), + seg=str(r.get("seg", "") or ""), + ) + for r in payload.get("names", []) + if isinstance(r, dict) and r.get("name") + ] return (comments, names) - - def search(self, query: str, mode: str = "text", *, limit: int = 500, - regex: bool = False, case: bool = False, - ) -> tuple[list["SearchHit"], str | None, bool]: + def search( + self, + query: str, + mode: str = "text", + *, + limit: int = 500, + regex: bool = False, + case: bool = False, + ) -> tuple[list["SearchHit"], str | None, bool]: """Search the whole database. Returns ``(hits, error, truncated)``. A failed search is DATA (a message to show), not an exception: a bad regex or an unparsable byte pattern is something the user typed, and the palette wants to say so without unwinding. """ - op = "search_bytes" if mode == "bytes" else "search_text" + operation = ( + remote_ops.search_bytes if mode == "bytes" else remote_ops.search_text + ) args: dict = {"limit": int(limit), "case": bool(case)} if mode == "bytes": # Validate HERE, not just in the UI: IDA's find_bytes answers a # malformed pattern with zero hits and no error, which reads as # "not present" -- the most misleading answer a search can give. from .search import normalise_pattern, pattern_problem + problem = pattern_problem(query) if problem: return ([], problem, False) @@ -2148,29 +2308,32 @@ class Program: args["query"] = query args["regex"] = bool(regex) try: - payload = self.client.invoke(op, **args) + payload = self.client.call(operation, **args) except IDAToolError as e: return ([], str(e), False) if not isinstance(payload, dict): return ([], "the backend returned nothing searchable", False) hits = [ - SearchHit(addr=_as_int(r.get("addr", 0)), - head=_as_int(r.get("head", r.get("addr", 0))), - line=str(r.get("line", "") or ""), - func=(r.get("func") or None), - func_addr=(_as_int(r["func_addr"]) if r.get("func_addr") - else None), - seg=str(r.get("seg", "") or "")) - for r in payload.get("hits", []) if isinstance(r, dict)] + SearchHit( + addr=_as_int(r.get("addr", 0)), + head=_as_int(r.get("head", r.get("addr", 0))), + line=str(r.get("line", "") or ""), + func=(r.get("func") or None), + func_addr=(_as_int(r["func_addr"]) if r.get("func_addr") else None), + seg=str(r.get("seg", "") or ""), + ) + for r in payload.get("hits", []) + if isinstance(r, dict) + ] return (hits, payload.get("error") or None, bool(payload.get("truncated"))) def journal_get(self) -> str: """The findings journal blob stored in this database ('' if none).""" - payload = self.client.invoke("journal_get") + payload = self.client.call(remote_ops.journal_get) return str(payload.get("data", "")) if isinstance(payload, dict) else "" def journal_put(self, data: str) -> None: - self.client.invoke("journal_put", data=str(data)) + self.client.call(remote_ops.journal_put, data=str(data)) def decomp_map(self, ea: int) -> list[list[int]]: """Per-pseudocode-line instruction coverage for the split-view region @@ -2183,12 +2346,15 @@ class Program: if hit is not None and hit[1] == gen: return hit[0] try: - payload = self.client.invoke("decomp_map", addr=hex(ea)) + payload = self.client.call(remote_ops.decomp_map, addr=hex(ea)) except IDAToolError: return [] lines = payload.get("lines", []) if isinstance(payload, dict) else [] - out = [[_as_int(e) for e in (ln.get("eas") or [])] - for ln in lines if isinstance(ln, dict)] + out = [ + [_as_int(e) for e in (ln.get("eas") or [])] + for ln in lines + if isinstance(ln, dict) + ] with self._lock: self._decomp_maps[ea] = (out, gen) return out @@ -2213,7 +2379,7 @@ class Program: if hit is not None and hit[1] == gen: return hit[0] try: - payload = self.client.invoke("flowchart", addr=hex(ea)) + payload = self.client.call(remote_ops.flowchart, addr=hex(ea)) except IDAToolError: return None if not isinstance(payload, dict) or payload.get("error"): @@ -2224,10 +2390,14 @@ class Program: blocks = [] for b in raw: try: - blocks.append(BasicBlock( - id=int(b["id"]), start=_as_int(b["start"]), - end=_as_int(b["end"]), - succs=[(int(d), str(k)) for d, k in (b.get("succs") or [])])) + blocks.append( + BasicBlock( + id=int(b["id"]), + start=_as_int(b["start"]), + end=_as_int(b["end"]), + succs=[(int(d), str(k)) for d, k in (b.get("succs") or [])], + ) + ) except (KeyError, ValueError, TypeError): continue if not blocks: @@ -2239,8 +2409,9 @@ class Program: for b in blocks: # bisect, not a scan per block: a 400-block function against a few # thousand rows is a million comparisons done for nothing. - b.rows = rows[bisect.bisect_left(eas, b.start): - bisect.bisect_left(eas, b.end)] + b.rows = rows[ + bisect.bisect_left(eas, b.start) : bisect.bisect_left(eas, b.end) + ] fcv = Flowchart( func_ea=_as_int(f.get("addr", lo)), name=str(f.get("name") or f"sub_{lo:X}"), @@ -2283,11 +2454,12 @@ class Program: operand marks for free.""" out: list[Head] = [] addr = lo - for _ in range(64): # bounded: ~128k heads + for _ in range(64): # bounded: ~128k heads if addr >= hi: break - payload = self.client.invoke("heads", addr=hex(addr), end=hex(hi), - count=2000) + payload = self.client.call( + remote_ops.heads, addr=hex(addr), end=hex(hi), count=2000 + ) rows = payload.get("heads", []) if isinstance(payload, dict) else [] if not rows: break @@ -2315,27 +2487,34 @@ class Program: # -- cross-references & containing function --------------------------- # def function_of(self, ea: int) -> Func | None: """Return the function containing ``ea`` (resolves mid-function addrs).""" - payload = self.client.invoke("lookup_funcs", queries=[hex(ea)]) + payload = self.client.call(remote_ops.lookup_funcs, queries=[hex(ea)]) res = payload.get("result", []) if isinstance(payload, dict) else [] fn = res[0].get("fn") if res and isinstance(res[0], dict) else None return Func.from_raw(fn) if fn else None def xrefs_from(self, ea: int) -> list[Xref]: - payload = self.client.invoke( - "xref_query", + payload = self.client.call( + remote_ops.xref_query, queries=[{"addr": hex(ea), "direction": "from", "include_fn": True}], ) return _parse_xrefs(payload) def xrefs_to(self, ea: int, limit: int = 2000) -> list[Xref]: - q = [{"addr": hex(ea), "direction": "to", "include_fn": True, - "dedup": True, "count": limit}] + q = [ + { + "addr": hex(ea), + "direction": "to", + "include_fn": True, + "dedup": True, + "count": limit, + } + ] try: # xref_types adds a fine-grained `kind` (call/read/write/...) for the # xref dialog; fall back to xref_query (code/data only) if absent. - payload = self.client.invoke("xref_types", queries=q) + payload = self.client.call(remote_ops.xref_types, queries=q) except IDAToolError: - payload = self.client.invoke("xref_query", queries=q) + payload = self.client.call(remote_ops.xref_query, queries=q) return _parse_xrefs(payload) # -- address resolution ------------------------------------------------ # @@ -2353,7 +2532,7 @@ class Program: # (loc_/locret_): lookup_funcs would map a label to its *containing* # function's entry, so double-clicking a label jumped to the wrong place. try: - payload = self.client.invoke("resolve_names", queries=[s]) + payload = self.client.call(remote_ops.resolve_names, queries=[s]) res = payload.get("result", []) if isinstance(payload, dict) else [] ea = res[0].get("ea") if res and isinstance(res[0], dict) else None if ea: @@ -2363,7 +2542,7 @@ class Program: # Fall back to function-name resolution (also drives the 'did you mean' # suggestion when the name is unknown). try: - payload = self.client.invoke("lookup_funcs", queries=[s]) + payload = self.client.call(remote_ops.lookup_funcs, queries=[s]) except IDAToolError as e: raise KeyError(f"cannot resolve {target!r}: {e}") from e res = payload.get("result", []) if isinstance(payload, dict) else [] @@ -2389,8 +2568,10 @@ class Program: except Exception: # noqa: BLE001 -- suggestions are strictly optional return "" if not cands: - return (" (no function name contains it; it may be a data symbol or " - "not a function — pass an address like 0x1234)") + return ( + " (no function name contains it; it may be a data symbol or " + "not a function — pass an address like 0x1234)" + ) shown = cands[:5] names = ", ".join(f"{c.name} @ {c.addr:#x}" for c in shown) more = " …" if len(cands) > len(shown) else "" @@ -2401,7 +2582,9 @@ class Program: """Set (empty text clears) the comment at ``ea``; affects both the disasm and decompiler views. Returns the raw payload so the caller can surface a soft per-item error. The caller must invalidate/recompile to see it.""" - return self.client.invoke("set_comments", items=[{"addr": hex(ea), "comment": text}]) + return self.client.call( + remote_ops.set_comments, items=[{"addr": hex(ea), "comment": text}] + ) # -- invalidation (after edits) --------------------------------------- # def invalidate(self, ea: int) -> None: @@ -2430,14 +2613,16 @@ def _parse_xrefs(payload) -> list[Xref]: fn = d.get("fn") or {} frm = d.get("from", d.get("addr")) to = d.get("to") - out.append(Xref( - frm=_as_int(frm) if frm is not None else 0, - to=_as_int(to) if to is not None else None, - type=d.get("type", "?"), - fn_name=fn.get("name"), - fn_addr=_as_int(fn["addr"]) if fn.get("addr") else None, - kind=d.get("kind"), - )) + out.append( + Xref( + frm=_as_int(frm) if frm is not None else 0, + to=_as_int(to) if to is not None else None, + type=d.get("type", "?"), + fn_name=fn.get("name"), + fn_addr=_as_int(fn["addr"]) if fn.get("addr") else None, + kind=d.get("kind"), + ) + ) return out @@ -2447,13 +2632,15 @@ def _parse_decompilation(ea: int, payload) -> Decompilation: code = payload.get("code") error = payload.get("error") if not code: - return Decompilation(ea, None, True, error or "decompilation failed", - False, None) + return Decompilation( + ea, None, True, error or "decompilation failed", False, None + ) m = _TRUNC_RE.search(code) truncated = m is not None total_chars = int(m.group(1)) if m else len(code) refs = [ Ref(addr=_as_int(r["addr"]), name=r.get("name", ""), string=r.get("string")) - for r in payload.get("refs", []) if isinstance(r, dict) and "addr" in r + for r in payload.get("refs", []) + if isinstance(r, dict) and "addr" in r ] return Decompilation(ea, code, False, error, truncated, total_chars, refs) diff --git a/idatui/edit_ctl.py b/idatui/edit_ctl.py index 52566ca..89dfc69 100644 --- a/idatui/edit_ctl.py +++ b/idatui/edit_ctl.py @@ -19,6 +19,7 @@ The message handlers and the ``@work`` entry points stay on ``IdaTui``: Textual dispatches ``on_<message>`` by name on the DOMNode, and its worker machinery wants a DOMNode host. They are one-line delegates into here. """ + from __future__ import annotations import re @@ -26,10 +27,10 @@ from typing import TYPE_CHECKING from textual.widgets import DataTable -from . import diag +from . import diag, remote_ops from .errors import IDAToolError -if TYPE_CHECKING: # pragma: no cover +if TYPE_CHECKING: # pragma: no cover from .app import IdaTui _app_mod = None @@ -40,13 +41,18 @@ def _M(): global _app_mod if _app_mod is None: from . import app as _m + _app_mod = _m return _app_mod #: A C type wide enough for N bytes, for prefilling a retype/define prompt. -_BY_SIZE = {1: "unsigned __int8", 2: "unsigned __int16", - 4: "unsigned __int32", 8: "unsigned __int64"} +_BY_SIZE = { + 1: "unsigned __int8", + 2: "unsigned __int16", + 4: "unsigned __int32", + 8: "unsigned __int64", +} class EditController: @@ -169,16 +175,24 @@ class EditController: # If the cursor is on a symbol token (a call/branch target, a data # reference, or this head's own label) rename THAT symbol; otherwise # create/rename a label at the head's address (bare/undefined bytes). - if (word and app._looks_like_symbol(word) and word != mnem - and word.lower() not in M._ASM_KEYWORDS): + if ( + word + and app._looks_like_symbol(word) + and word != mnem + and word.lower() not in M._ASM_KEYWORDS + ): app.prompts.rename.show( f"rename '{word}' — Enter=apply Esc=cancel", - word, ctx=(msg.view, word, None)) + word, + ctx=(msg.view, word, None), + ) else: cur = head.name if (head is not None and head.name) else "" app.prompts.rename.show( f"name @ {ea:#x} — Enter=apply Esc=cancel", - cur, ctx=(msg.view, cur, ea)) + cur, + ctx=(msg.view, cur, ea), + ) return if not msg.name: app._status("nothing to rename under the cursor") @@ -186,11 +200,14 @@ class EditController: if self.is_pseudocode_label(msg.view, msg.name): app._status( f"can't rename pseudocode label '{msg.name}' " - "(Hex-Rays goto labels aren't renamable via the API)") + "(Hex-Rays goto labels aren't renamable via the API)" + ) return app.prompts.rename.show( f"rename '{msg.name}' — Enter=apply Esc=cancel", - msg.name, ctx=(msg.view, msg.name, None)) + msg.name, + ctx=(msg.view, msg.name, None), + ) def submit_rename(self, ctx, value: str) -> None: # type: ignore[no-untyped-def] view, old, addr = ctx @@ -242,7 +259,7 @@ class EditController: kind = "stack" batch = {"stack": {"func_addr": hex(cur.ea), "old": old, "new": new}} try: - res = prog.client.invoke("rename", batch=batch) + res = prog.client.call(remote_ops.rename, batch=batch) except IDAToolError as e: app.call_from_thread(app._status, f"rename failed: {e.message}") return @@ -275,8 +292,9 @@ class EditController: app = self.app assert app.program is not None try: - res = app.program.client.invoke( - "rename", batch={"data": {"addr": hex(addr), "new": name}}) + res = app.program.client.call( + remote_ops.rename, batch={"data": {"addr": hex(addr), "new": name}} + ) except IDAToolError as e: app.call_from_thread(app._status, f"name failed: {e.message}") return @@ -307,11 +325,11 @@ class EditController: lm = app.program.listing(addr) label = name if is_func_start else app.program.region_label(addr) idx = max(lm.ensure_ea(addr), 0) if lm is not None else 0 - app.call_from_thread(self.open_at_named, label, addr, idx, name, - is_func_start) + app.call_from_thread(self.open_at_named, label, addr, idx, name, is_func_start) - def open_at_named(self, label: str, addr: int, idx: int, name: str, - is_func_start: bool = False) -> None: + def open_at_named( + self, label: str, addr: int, idx: int, name: str, is_func_start: bool = False + ) -> None: app = self.app if is_func_start: app.program.bump_names() @@ -328,10 +346,11 @@ class EditController: pseudocode a comment is `// text` before the trailing /*0xEA*/ markers; C has no `//` operator, so the last `//` is unambiguously the comment.""" if isinstance(view, _M().DecompView) and 0 <= view.cursor < len(view._texts): - s = re.sub(r"(?:/\*\s*0x[0-9A-Fa-f]+\s*\*/\s*)+$", "", - view._texts[view.cursor]) + s = re.sub( + r"(?:/\*\s*0x[0-9A-Fa-f]+\s*\*/\s*)+$", "", view._texts[view.cursor] + ) i = s.rfind("//") - return s[i + 2:].strip() if i >= 0 else "" + return s[i + 2 :].strip() if i >= 0 else "" return "" def request_comment(self, msg) -> None: # type: ignore[no-untyped-def] @@ -349,7 +368,9 @@ class EditController: what = "function comment" if func_level else "comment" app.prompts.comment.show( f"{what} @ {ea:#x} — Enter=apply (empty=clear) Esc=cancel", - existing, ctx=(msg.view, ea, existing)) + existing, + ctx=(msg.view, ea, existing), + ) def submit_comment(self, ctx, value: str) -> None: # type: ignore[no-untyped-def] view, ea, existing = ctx @@ -369,10 +390,13 @@ class EditController: app.call_from_thread(app._status, f"comment failed: {e.message}") return data = res.get("result") if isinstance(res, dict) else None - if (isinstance(data, list) and data and isinstance(data[0], dict) - and data[0].get("error")): - app.call_from_thread(app._status, - f"comment failed: {data[0]['error']}") + if ( + isinstance(data, list) + and data + and isinstance(data[0], dict) + and data[0].get("error") + ): + app.call_from_thread(app._status, f"comment failed: {data[0]['error']}") return app.call_from_thread(self.after_comment, ea, text) @@ -429,30 +453,34 @@ class EditController: if dt is not None and not dt.get("is_func"): kind, subject = "data", tgt prefill = dt.get("type") or self.guess_data_type( - dt.get("size") or 0) + dt.get("size") or 0 + ) # 3) fall back to the current function itself if kind is None and ft is not None: kind, subject, prefill = "func", app._cur.ea, ft.prototype if kind is None: - app.call_from_thread(app._status, - "nothing to retype under the cursor") + app.call_from_thread(app._status, "nothing to retype under the cursor") return - app.call_from_thread(self.open_retype, view, kind, subject, - word or "", prefill) + app.call_from_thread(self.open_retype, view, kind, subject, word or "", prefill) - def open_retype(self, view, kind: str, subject: int, word: str, - prefill: str) -> None: # type: ignore[no-untyped-def] + def open_retype( + self, view, kind: str, subject: int, word: str, prefill: str + ) -> None: # type: ignore[no-untyped-def] label = "prototype" if kind == "func" else f"type for '{word}'" - self.app.prompts.retype.show(f"{label} — Enter=apply Esc=cancel", - prefill, ctx=(view, kind, subject, word)) + self.app.prompts.retype.show( + f"{label} — Enter=apply Esc=cancel", + prefill, + ctx=(view, kind, subject, word), + ) def submit_retype(self, ctx, value: str) -> None: # type: ignore[no-untyped-def] view, kind, subject, word = ctx if view is not None and value: self.app._do_retype(kind, subject, word, value) - def do_retype(self, kind: str, subject: int, word: str, - new: str) -> None: # worker context + def do_retype( + self, kind: str, subject: int, word: str, new: str + ) -> None: # worker context app = self.app assert app.program is not None if kind == "func": @@ -473,8 +501,9 @@ class EditController: app.program.bump_names() self.reload_active_code() app._dirty = True - app.journal.record("retype", getattr(app._cur, "ea", None), word, - {"kind": kind}) + app.journal.record( + "retype", getattr(app._cur, "ea", None), word, {"kind": kind} + ) what = "prototype" if kind == "func" else f"'{word}'" app._status(f"retyped {what} (Ctrl+S to save)") @@ -498,15 +527,17 @@ class EditController: f"data type @ {ea:#x} (e.g. int, char[16], my_struct)" " — Enter=apply Esc=cancel", self.default_data_type(head) if head is not None else "int", - ctx=(view, ea)) + ctx=(view, ea), + ) def submit_make_data(self, ctx, value: str) -> None: # type: ignore[no-untyped-def] view, ea = ctx if view is not None and value: self.app._do_make_data(ea, value, self.app._anchor()) - def do_make_data(self, ea: int, type_decl: str, - anchor=None) -> None: # worker context + def do_make_data( + self, ea: int, type_decl: str, anchor=None + ) -> None: # worker context app = self.app assert app.program is not None try: @@ -522,8 +553,7 @@ class EditController: lm = app.program.listing(ea) idx = max(lm.ensure_ea(ea), 0) if lm is not None else 0 _cur, top = app._anchor_rows(anchor, lm, ea) - app.call_from_thread( - app._open_at, ea, name, idx, False, -1, 0, True, None, top) + app.call_from_thread(app._open_at, ea, name, idx, False, -1, 0, True, None, top) app.call_from_thread(self.edit_done, anchor) # -- literal display formats (IDA 'o') --------------------------------- # @@ -554,8 +584,9 @@ class EditController: fn = view.loaded_ea if view.loaded_ea is not None else app._cur.ea app._do_op_format(msg.mode, "decomp", fn, view.cursor_x, view.cursor) - def do_op_format(self, mode: str, where: str, ea: int, col: int, - line: int = -1) -> None: # worker context + def do_op_format( + self, mode: str, where: str, ea: int, col: int, line: int = -1 + ) -> None: # worker context app = self.app assert app.program is not None try: @@ -584,7 +615,9 @@ class EditController: app.call_from_thread( app._status, f"{what}{fmt} {r.get('value') or ''}" - f" [{', '.join(r.get('choices', []))}]", True) + f" [{', '.join(r.get('choices', []))}]", + True, + ) return step = f"{prev} \u2192 {fmt}" if prev and prev != fmt else fmt desc = f"{what}{step}: {text[:96]}" @@ -629,13 +662,17 @@ class EditController: return app._do_edit_item(msg.kind, ea, app._anchor()) - def do_edit_item(self, kind: str, ea: int, - anchor=None) -> None: # worker context + def do_edit_item(self, kind: str, ea: int, anchor=None) -> None: # worker context app = self.app assert app.program is not None - verb = {"code": "defined code", "func": "created function", - "undef": "undefined", "string": "made string", - "thumb": "switched decoding", "thumbscan": "scanned"}[kind] + verb = { + "code": "defined code", + "func": "created function", + "undef": "undefined", + "string": "made string", + "thumb": "switched decoding", + "thumbscan": "scanned", + }[kind] try: if kind == "code": # Keep going until something stops it: one instruction is rarely @@ -646,20 +683,24 @@ class EditController: if n == 0 and why == "defined": # Already code/data here — a no-op, not a failure. Saying # "failed to create instruction" for it would be a lie. - app.call_from_thread( - app._status, f"already defined @ {ea:#x}") + app.call_from_thread(app._status, f"already defined @ {ea:#x}") return if n == 0: - raise IDAToolError("define_code", - f"@ {ea:#x}: Failed to create instruction") + raise IDAToolError( + "define_code", f"@ {ea:#x}: Failed to create instruction" + ) end = int(str(r.get("end", hex(ea))), 0) - reason = {"undecodable": "hit bytes that don't decode", - "flow": "control flow ends here", - "defined": "ran into existing code/data", - "segment": "end of segment", - "limit": "instruction limit"}.get(why, why) - verb = (f"defined {n} instruction{'s' if n != 1 else ''} " - f"({ea:#x}\u2013{end:#x}) \u2014 {reason}") + reason = { + "undecodable": "hit bytes that don't decode", + "flow": "control flow ends here", + "defined": "ran into existing code/data", + "segment": "end of segment", + "limit": "instruction limit", + }.get(why, why) + verb = ( + f"defined {n} instruction{'s' if n != 1 else ''} " + f"({ea:#x}\u2013{end:#x}) \u2014 {reason}" + ) elif kind == "thumbscan": # A vector table is a list of Thumb entry points that IDA won't # follow on a headerless image, because nothing tells it those @@ -668,11 +709,15 @@ class EditController: r = app.program.thumb_scan(ea, ea + 0x400) n, applied = int(r.get("n", 0)), int(r.get("applied", 0)) if not n: - verb = (f"no Thumb entry pointers in {ea:#x}\u2013{ea+0x400:#x}" - " (odd words pointing into the image)") + verb = ( + f"no Thumb entry pointers in {ea:#x}\u2013{ea + 0x400:#x}" + " (odd words pointing into the image)" + ) else: - verb = (f"{n} Thumb entr{'y' if n == 1 else 'ies'} found, " - f"{applied} disassembled") + verb = ( + f"{n} Thumb entr{'y' if n == 1 else 'ies'} found, " + f"{applied} disassembled" + ) elif kind == "thumb": # Switch the mode, then disassemble in it: flipping T and # leaving the bytes undefined shows nothing, and the reason you @@ -686,11 +731,16 @@ class EditController: verb += " (segment set to 32-bit; Thumb needs ARM32)" if r.get("db_64bit"): # Disassembly will look right and F5 will never work. - verb += (" \u26a0 this database is 64-bit, so Hex-Rays " - "won't decompile it \u2014 Ctrl+L and pick " - "arm:ARMv7-A") - verb += (f" \u2014 {n} instruction{'s' if n != 1 else ''}" - if n else " \u2014 still doesn't decode") + verb += ( + " \u26a0 this database is 64-bit, so Hex-Rays " + "won't decompile it \u2014 Ctrl+L and pick " + "arm:ARMv7-A" + ) + verb += ( + f" \u2014 {n} instruction{'s' if n != 1 else ''}" + if n + else " \u2014 still doesn't decode" + ) # falls through to the shared reload: same cache bump, same # anchor restore, same flash. That is the whole point of having # one path. @@ -698,9 +748,11 @@ class EditController: anchor.refresh_functions = True r = app.program.define_func(ea) if r.get("start") and r.get("end"): - verb = (f"created function {r['start']}\u2013{r['end']}" - + (" (end worked out from the code)" - if r.get("how") == "explicit-end" else "")) + verb = f"created function {r['start']}\u2013{r['end']}" + ( + " (end worked out from the code)" + if r.get("how") == "explicit-end" + else "" + ) elif kind == "string": s = app.program.make_string(ea) verb = f"made string ({s[:24]!r})" if s else verb @@ -726,13 +778,14 @@ class EditController: idx = 0 if ea == fn.addr else model.index_of_ea(ea) _cur, top = app._anchor_rows(anchor, model, ea) app.call_from_thread( - app._open_at, fn.addr, fn.name, idx, False, -1, 0, False, - None, top) + app._open_at, fn.addr, fn.name, idx, False, -1, 0, False, None, top + ) else: name = app.program.region_label(ea) lm = app.program.listing(ea) idx = max(lm.ensure_ea(ea), 0) if lm is not None else 0 _cur, top = app._anchor_rows(anchor, lm, ea) app.call_from_thread( - app._open_at, ea, name, idx, False, -1, 0, True, None, top) + app._open_at, ea, name, idx, False, -1, 0, True, None, top + ) app.call_from_thread(self.edit_done, anchor) diff --git a/idatui/pool.py b/idatui/pool.py index 465dff2..573aa75 100644 --- a/idatui/pool.py +++ b/idatui/pool.py @@ -225,6 +225,28 @@ class DatabasePool: self.evict(label, save=save, save_gui=save) self.active = None + def discard_changes(self, labels: list[str]) -> list[str]: + """Discard final managed sessions; return labels whose owner remains. + + A returned label is not an error: its client is attached to a GUI or a + still-shared worker, so releasing our lease transfers finalization to + that session's owner or remaining clients. + """ + transferred: list[str] = [] + for label in labels: + client = self._clients.get(label) + if client is not None and not client.discard_database(): + transferred.append(label) + return transferred + + def replace_client(self, label: str, old, new) -> bool: + """Replace one disconnected lease without changing residency policy.""" + if self._clients.get(label) is not old: + return False + self._clients[label] = new + self._touch(label) + return True + # -- introspection ------------------------------------------------------ # def status(self) -> list[dict]: """Per-binary residency for the switcher UI.""" diff --git a/idatui/remote_ops.py b/idatui/remote_ops.py new file mode 100644 index 0000000..91ffeb7 --- /dev/null +++ b/idatui/remote_ops.py @@ -0,0 +1,1683 @@ +"""Typed remote operations executed through ida-codemode.""" + +from __future__ import annotations +# ruff: noqa + +import threading +from collections.abc import Callable +from pathlib import Path +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from ida_domain import Database + + +def operation_label() -> str: + """Display attribution for the current call; ready for per-user context.""" + return "IDA TUI" + + +def data_type(db: Database, **a: Any) -> Any: + 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)} + return result + + +def declare_type(db: Database, **a: Any) -> Any: + 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} + return result + + +def decomp_error(db: Database, **a: Any) -> Any: + 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}" + return result + + +def define_code(db: Database, **a: Any) -> Any: + 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} + return result + + +def define_code_run(db: Database, **a: Any) -> Any: + 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, + } + return result + + +def define_func(db: Database, **a: Any) -> Any: + 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} + return result + + +def define_func_run(db: Database, **a: Any) -> Any: + 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}", + } + ) + return result + + +def del_type(db: Database, **a: Any) -> Any: + 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"}), + } + return result + + +def disasm(db: Database, **a: Any) -> Any: + 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), + } + return result + + +def file_regions(db: Database, **a: Any) -> Any: + 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} + return result + + +def flowchart(db: Database, **a: Any) -> Any: + 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, + } + return result + + +def force_recompile(db: Database, **a: Any) -> Any: + 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} + return result + + +def func_types(db: Database, **a: Any) -> Any: + 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, + } + return result + + +def get_bytes(db: Database, **a: Any) -> Any: + 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} + return result + + +def journal_get(db: Database, **a: Any) -> Any: + import ida_netnode + + n = ida_netnode.netnode(a.get("node", "$ idatui.journal")) + blob = n.getblob(0, "I") if ida_netnode.exist(n) else None + result = {"data": blob.decode("utf-8", "replace") if blob else ""} + return result + + +def journal_put(db: Database, **a: Any) -> Any: + import ida_netnode + + n = ida_netnode.netnode(a.get("node", "$ idatui.journal"), 0, True) + payload = (a.get("data") or "").encode("utf-8") + n.setblob(payload, 0, "I") + result = {"ok": True, "bytes": len(payload)} + return result + + +def list_annotations(db: Database, **a: Any) -> Any: + import ida_bytes, ida_funcs, ida_lines, ida_nalt, ida_name + import ida_segment, ida_typeinf, idautils + + limit = max(1, int(a.get("limit", 4000))) + max_scan = max(1000, int(a.get("max_scan", 2000000))) + comments, names = [], [] + scanned = 0 + + def _line(ea): + try: + txt = ida_lines.generate_disasm_line(ea, ida_lines.GENDSM_REMOVE_TAGS) + except Exception: + txt = "" + return " ".join((txt or "").split()) + + for ea, nm in idautils.Names(): + if len(names) >= limit: + break + if not nm or not ida_bytes.has_user_name(ida_bytes.get_flags(ea)): + continue + fn = ida_funcs.get_func(ea) + is_fn = fn is not None and int(fn.start_ea) == int(ea) + proto = None + if is_fn: + try: + ti = ida_typeinf.tinfo_t() + if ida_nalt.get_tinfo(ti, ea): + proto = str(ti) + except Exception: + proto = None + seg = ida_segment.getseg(ea) + names.append( + { + "addr": hex(int(ea)), + "name": nm, + "func": is_fn, + "size": (int(fn.end_ea - fn.start_ea) if is_fn else 0), + "proto": proto, + "seg": (ida_segment.get_segm_name(seg) if seg else ""), + } + ) + + for i in range(ida_segment.get_segm_qty()): + seg = ida_segment.getnseg(i) + if seg is None or len(comments) >= limit or scanned >= max_scan: + continue + for ea in idautils.Heads(seg.start_ea, seg.end_ea): + scanned += 1 + if len(comments) >= limit or scanned >= max_scan: + break + if not ida_bytes.has_cmt(ida_bytes.get_flags(ea)): + continue + for rep in (False, True): + text = ida_bytes.get_cmt(ea, rep) + if text: + fn = ida_funcs.get_func(ea) + comments.append( + { + "addr": hex(int(ea)), + "text": text, + "repeatable": rep, + "line": _line(ea), + "seg": ida_segment.get_segm_name(seg), + "func": ( + ida_funcs.get_func_name(fn.start_ea) if fn else None + ), + "func_addr": (hex(int(fn.start_ea)) if fn else None), + } + ) + + # Whole-function comments are not on the byte flags, so the scan cannot see them. + for fn_ea in idautils.Functions(): + fn = ida_funcs.get_func(fn_ea) + if fn is None or len(comments) >= limit: + continue + for rep in (False, True): + text = ida_funcs.get_func_cmt(fn, rep) + if text: + seg = ida_segment.getseg(fn_ea) + comments.append( + { + "addr": hex(int(fn_ea)), + "text": text, + "repeatable": rep, + "line": "", + "whole_func": True, + "seg": (ida_segment.get_segm_name(seg) if seg else ""), + "func": ida_funcs.get_func_name(fn_ea), + "func_addr": hex(int(fn_ea)), + } + ) + result = { + "comments": comments, + "names": names, + "scanned": scanned, + "truncated": ( + len(comments) >= limit or len(names) >= limit or scanned >= max_scan + ), + } + return result + + +def list_funcs(db: Database, **a: Any) -> Any: + 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)} + ] + } + return result + + +def list_linkage(db: Database, **a: Any) -> Any: + 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), + } + return result + + +def list_strings(db: Database, **a: Any) -> Any: + 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)} + return result + + +def lookup_funcs(db: Database, **a: Any) -> Any: + 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} + return result + + +def make_data(db: Database, **a: Any) -> Any: + 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} + return result + + +def make_string(db: Database, **a: Any) -> Any: + 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)} + return result + + +def read_raw(db: Database, **a: Any) -> Any: + 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)} + return result + + +def rename(db: Database, **a: Any) -> Any: + import idaapi, ida_hexrays, ida_name + + batch = a.get("batch") or {} + 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: + 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 + return result + + +def resolve_names(db: Database, **a: Any) -> Any: + 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} + return result + + +def search_bytes(db: Database, **a: Any) -> Any: + import ida_bytes, ida_funcs, ida_idaapi, ida_lines, ida_segment + + pat = str(a.get("pattern", "")).strip() + limit = max(1, int(a.get("limit", 500))) + lo = int(a.get("start", 0)) + hi = int(a.get("end", 0)) or ida_idaapi.BADADDR + flags = ida_bytes.BIN_SEARCH_FORWARD | ida_bytes.BIN_SEARCH_NOSHOW + if a.get("case"): + flags |= ida_bytes.BIN_SEARCH_CASE + rows, err, ea = [], None, lo + while len(rows) < limit: + try: + hit = ida_bytes.find_bytes(pat, range_start=ea, range_end=hi, flags=flags) + except Exception as exc: + err = str(exc) or exc.__class__.__name__ + break + if hit is None or hit == ida_idaapi.BADADDR: + break + head = ida_bytes.get_item_head(hit) + fn = ida_funcs.get_func(hit) + seg = ida_segment.getseg(hit) + try: + line = ( + ida_lines.generate_disasm_line(head, ida_lines.GENDSM_REMOVE_TAGS) or "" + ) + except Exception: + line = "" + rows.append( + { + "addr": hex(int(hit)), + "head": hex(int(head)), + "line": " ".join(line.split()), + "func": (ida_funcs.get_func_name(fn.start_ea) if fn else None), + "func_addr": (hex(int(fn.start_ea)) if fn else None), + "seg": (ida_segment.get_segm_name(seg) if seg else ""), + } + ) + ea = int(hit) + 1 + result = {"hits": rows, "error": err, "truncated": len(rows) >= limit} + return result + + +def search_structs(db: Database, **a: Any) -> Any: + 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} + return result + + +def search_text(db: Database, **a: Any) -> Any: + import ida_lines, ida_funcs, ida_segment, idautils + import re as _re + + q = str(a.get("query", "")) + limit = max(1, int(a.get("limit", 500))) + max_scan = max(1000, int(a.get("max_scan", 3000000))) + ci = (not a.get("case")) and q.islower() # smartcase, like the in-view search + rx, err = None, None + if a.get("regex"): + try: + rx = _re.compile(q, _re.I if ci else 0) + except Exception as exc: + err = "bad regex: " + str(exc) + needle = q.lower() if ci else q + rows, scanned = [], 0 + if err is None and q: + for i in range(ida_segment.get_segm_qty()): + seg = ida_segment.getnseg(i) + if seg is None or len(rows) >= limit or scanned >= max_scan: + continue + for ea in idautils.Heads(seg.start_ea, seg.end_ea): + scanned += 1 + if len(rows) >= limit or scanned >= max_scan: + break + try: + line = ( + ida_lines.generate_disasm_line(ea, ida_lines.GENDSM_REMOVE_TAGS) + or "" + ) + except Exception: + continue + # Match what the user SEES, not IDA's column padding: nobody types + # "call" + four spaces + "cs:getenv_ptr". + line = " ".join(line.split()) + hay = line.lower() if ci else line + if rx.search(line) if rx is not None else (needle in hay): + fn = ida_funcs.get_func(ea) + rows.append( + { + "addr": hex(int(ea)), + "head": hex(int(ea)), + "line": line, + "func": ( + ida_funcs.get_func_name(fn.start_ea) if fn else None + ), + "func_addr": (hex(int(fn.start_ea)) if fn else None), + "seg": ida_segment.get_segm_name(seg), + } + ) + result = { + "hits": rows, + "error": err, + "scanned": scanned, + "truncated": len(rows) >= limit or scanned >= max_scan, + } + return result + + +def set_comments(db: Database, **a: Any) -> Any: + import idaapi, idc, ida_hexrays + + rows = [] + for item in a.get("items", []): + addr_s = str(item.get("addr", "")) + text = str(item.get("comment") or "") + try: + 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": addr_s, "error": str(exc)}) + result = {"result": rows} + return result + + +def set_lvar_type(db: Database, **a: Any) -> Any: + 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}"} + return result + + +def set_thumb(db: Database, **a: Any) -> Any: + 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), + } + return result + + +def set_type(db: Database, **a: Any) -> Any: + 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} + return result + + +def survey_binary(db: Database, **a: Any) -> Any: + 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} + return result + + +def thumb_scan(db: Database, **a: Any) -> Any: + 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), + } + return result + + +def type_inspect(db: Database, **a: Any) -> Any: + 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} + return result + + +def undefine(db: Database, **a: Any) -> Any: + 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} + return result + + +def xref_query(db: Database, **a: Any) -> Any: + 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: + 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} + return result + + +def xref_types(db: Database, **a: Any) -> Any: + 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: + 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} + return result + + +def op_format(addr: str, mode: str = "cycle", col: int = -1, n: int = -1) -> dict: ... + + +def pc_nums(addr: str) -> dict: ... + + +def decompile(addr, include_addresses=True) -> dict: ... + + +def decomp_map(addr: str) -> dict: ... + + +def pc_num_format( + addr: str, + mode: str = "cycle", + line: int = -1, + col: int = -1, + ea: str = "", + opnum: int = -1, +) -> dict: ... + + +def segment_index( + addr: str, + end: str = "", + page_rows: int = 500, + detail: bool = False, +) -> dict: ... + + +def heads( + addr: str, + count: int = 200, + offset: int = 0, + end: str = "", + back: bool = False, + annotate: bool = False, + expect: str = "", + text: bool = True, +) -> dict: ... + + +def profile_remote(operation: str, args: dict[str, Any], reps: int = 5) -> dict: ... + + +OPERATIONS: dict[str, Callable[..., Any]] = { + "data_type": data_type, + "declare_type": declare_type, + "decomp_error": decomp_error, + "define_code": define_code, + "define_code_run": define_code_run, + "define_func": define_func, + "define_func_run": define_func_run, + "del_type": del_type, + "disasm": disasm, + "file_regions": file_regions, + "flowchart": flowchart, + "force_recompile": force_recompile, + "func_types": func_types, + "get_bytes": get_bytes, + "heads": heads, + "journal_get": journal_get, + "journal_put": journal_put, + "list_annotations": list_annotations, + "list_funcs": list_funcs, + "list_linkage": list_linkage, + "list_strings": list_strings, + "lookup_funcs": lookup_funcs, + "make_data": make_data, + "make_string": make_string, + "op_format": op_format, + "pc_nums": pc_nums, + "decompile": decompile, + "decomp_map": decomp_map, + "pc_num_format": pc_num_format, + "profile_remote": profile_remote, + "read_raw": read_raw, + "rename": rename, + "resolve_names": resolve_names, + "search_bytes": search_bytes, + "search_structs": search_structs, + "search_text": search_text, + "segment_index": segment_index, + "set_comments": set_comments, + "set_lvar_type": set_lvar_type, + "set_thumb": set_thumb, + "set_type": set_type, + "survey_binary": survey_binary, + "thumb_scan": thumb_scan, + "type_inspect": type_inspect, + "undefine": undefine, + "xref_query": xref_query, + "xref_types": xref_types, +} + +_MODULE_DECLARATIONS = frozenset( + ( + heads, + segment_index, + op_format, + pc_nums, + decompile, + decomp_map, + pc_num_format, + profile_remote, + ) +) +_BOUND: dict[Callable[..., Any], Any] | None = None +_BIND_LOCK = threading.Lock() + + +def _bindings() -> dict[Callable[..., Any], Any]: + global _BOUND + with _BIND_LOCK: + if _BOUND is not None: + return _BOUND + from ida_codemode import RemoteModule + + operations_module = RemoteModule( + Path(__file__), operation_label=operation_label, codec="json" + ) + tools_module = RemoteModule( + Path(__file__).with_name("remote_tools.py"), + operation_label=operation_label, + codec="json", + ) + bound: dict[Callable[..., Any], Any] = {} + for declaration in OPERATIONS.values(): + if declaration in _MODULE_DECLARATIONS: + bound[declaration] = tools_module.function( + declaration, + timeout=15.0 if declaration is decompile else None, + ) + else: + bound[declaration] = operations_module.function(declaration) + _BOUND = bound + return bound + + +def bind(function: Callable[..., Any]) -> Any: + """Return the lazily constructed remote callable for one declaration.""" + return _bindings()[function] diff --git a/idatui/remote_tools.py b/idatui/remote_tools.py index 0e26965..db04244 100644 --- a/idatui/remote_tools.py +++ b/idatui/remote_tools.py @@ -23,6 +23,7 @@ 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 @@ -134,7 +135,7 @@ def _idatui_head_row(ea, flags=None, text=True): """ f = ida_bytes.get_flags(ea) if flags is None else flags - cls = f & _MS_CLS # == is_code(f) / is_data(f), without the calls + cls = f & _MS_CLS # == is_code(f) / is_data(f), without the calls if cls == _FF_CODE: kind = "code" elif cls == _FF_DATA: @@ -211,10 +212,20 @@ _IDATUI_SPAN_KINDS = { # 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"), + "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"), @@ -241,7 +252,7 @@ _IDATUI_TAGS = None _IDATUI_OPND_TAGS = None -_IDATUI_CTL = None # re: a tag = one of three control chars plus its argument +_IDATUI_CTL = None # re: a tag = one of three control chars plus its argument _IDATUI_TAGINFO = None @@ -280,6 +291,7 @@ def _idatui_spans(line): _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 @@ -289,17 +301,18 @@ def _idatui_spans(line): 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)} + 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 + 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 + skip = 0 # characters of an address payload still due i, n = 0, len(parts) while i < n: txt = parts[i] @@ -317,11 +330,11 @@ def _idatui_spans(line): break pair = parts[i] i += 1 - if skip: # a tag INSIDE an address payload: 2 chars + 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 + if ch == esc: # escaped literal: keep the char it guards pend += pair[1] continue tag = pair[1] @@ -337,7 +350,7 @@ def _idatui_spans(line): 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 + opnd = o # operands nest: an inner colour keeps the operand elif stack: kind, opnd = stack.pop() else: @@ -359,7 +372,7 @@ def _idatui_spans(line): prev_space = False out.append([kind, txt, opnd]) continue - if not core: # the span is nothing but padding + if not core: # the span is nothing but padding if not prev_space: prev_space = True out.append([kind, " ", opnd]) @@ -396,7 +409,7 @@ def _idatui_spans(line): 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 + 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(): @@ -435,8 +448,17 @@ def _idatui_rows_digest(rows): 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)) + acc = hash( + ( + acc, + r.get("ea"), + r.get("kind"), + r.get("size"), + r.get("text"), + r.get("name"), + sh, + ) + ) return acc @@ -448,8 +470,12 @@ def _idatui_unknown_row(ea, size): if size <= 1: return _idatui_head_row(ea) - row = {"ea": hex(ea), "kind": "unknown", "size": int(size), - "text": f"db {size} dup(?)"} + 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 @@ -481,8 +507,7 @@ def _idatui_struct_member_rows(ea): 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}) + rows.append({"ea": hex(ea + off), "kind": "member", "size": sz, "text": text}) return rows @@ -494,8 +519,13 @@ def _idatui_func_header_rows(ea): 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}, + { + "ea": hex(ea), + "kind": "funchdr", + "size": 0, + "text": name + " proc", + "name": name, + }, ] @@ -504,8 +534,13 @@ def _idatui_func_footer_rows(ea, func): 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": "funchdr", + "size": 0, + "text": name + " endp", + "name": name, + }, {"ea": hex(ea), "kind": "sep", "size": 0, "text": "; " + "-" * 60}, ] @@ -543,9 +578,12 @@ def _idatui_segment_detail(addr, end, page_rows): except Exception: pass - K_CODE = _IDATUI_KIND_ID["code"]; K_DATA = _IDATUI_KIND_ID["data"] - K_UNK = _IDATUI_KIND_ID["unknown"]; K_SEP = _IDATUI_KIND_ID["sep"] - K_FUNC = _IDATUI_KIND_ID["funchdr"]; K_LABEL = _IDATUI_KIND_ID["label"] + K_CODE = _IDATUI_KIND_ID["code"] + K_DATA = _IDATUI_KIND_ID["data"] + K_UNK = _IDATUI_KIND_ID["unknown"] + K_SEP = _IDATUI_KIND_ID["sep"] + K_FUNC = _IDATUI_KIND_ID["funchdr"] + K_LABEL = _IDATUI_KIND_ID["label"] K_MEMBER = _IDATUI_KIND_ID["member"] eas = array.array("Q") @@ -569,8 +607,8 @@ def _idatui_segment_detail(addr, end, page_rows): # are the same number until a segment contains an undefined run) and then # silently yields pages that do not line up with a refetch. anchors = [] - page_phys = 0 # physical rows emitted into the page being filled - rows = 0 # logical rows so far (what the scrollbar counts) + page_phys = 0 # physical rows emitted into the page being filled + rows = 0 # logical rows so far (what the scrollbar counts) fn = None ea = ida_bytes.get_item_head(lo) while ea != BAD and ea < hi: @@ -584,7 +622,9 @@ def _idatui_segment_detail(addr, end, page_rows): nh = next_head(ea, hi) stop = nh if (nh != BAD and ea < nh <= hi) else hi run = stop - ea - ea_ap(ea); kind_ap(K_UNK); size_ap(run) + ea_ap(ea) + kind_ap(K_UNK) + size_ap(run) rows += run if run > 1 else 1 page_phys += len(eas) - before ea = stop @@ -593,23 +633,32 @@ def _idatui_segment_detail(addr, end, page_rows): fn = get_func(ea) at_start = fn is not None and fn.start_ea == ea if at_start: - for k in (K_SEP, K_SEP, K_FUNC): # blank, banner, `name proc` - ea_ap(ea); kind_ap(k); size_ap(0) + for k in (K_SEP, K_SEP, K_FUNC): # blank, banner, `name proc` + ea_ap(ea) + kind_ap(k) + size_ap(0) rows += 3 elif cls == _FF_CODE and get_ea_name(ea): - ea_ap(ea); kind_ap(K_LABEL); size_ap(0) + ea_ap(ea) + kind_ap(K_LABEL) + size_ap(0) rows += 1 - ea_ap(ea); kind_ap(K_CODE if cls == _FF_CODE else K_DATA) - size_ap(int(get_item_size(ea))); rows += 1 + ea_ap(ea) + kind_ap(K_CODE if cls == _FF_CODE else K_DATA) + size_ap(int(get_item_size(ea))) + rows += 1 if cls == _FF_DATA: for m in _idatui_struct_member_rows(ea): ea_ap(int(m["ea"], 16) if isinstance(m["ea"], str) else m["ea"]) kind_ap(_IDATUI_KIND_ID.get(m.get("kind", "member"), K_MEMBER)) - size_ap(int(m.get("size", 0) or 0)); rows += 1 + size_ap(int(m.get("size", 0) or 0)) + rows += 1 item_end = get_item_end(ea) if fn is not None and item_end >= fn.end_ea: - for k in (K_FUNC, K_SEP): # `name endp`, separator - ea_ap(ea); kind_ap(k); size_ap(0) + for k in (K_FUNC, K_SEP): # `name endp`, separator + ea_ap(ea) + kind_ap(k) + size_ap(0) rows += 2 page_phys += len(eas) - before ea = item_end if item_end > ea else ea + 1 @@ -618,19 +667,28 @@ def _idatui_segment_detail(addr, end, page_rows): # which turns a bytes object into its repr -- 4 characters per byte and # unparseable at the other end. Learned by watching 2.97MB arrive as 11.26MB. import base64 + b64 = base64.b64encode - return {"addr": hex(lo), "end": hex(hi), "rows": rows, "heads": len(eas), - "anchors": anchors, "kind_names": list(_IDATUI_KINDS), - "eas": b64(eas.tobytes()).decode(), - "kinds": b64(kinds.tobytes()).decode(), - "sizes": b64(sizes.tobytes()).decode()} + return { + "addr": hex(lo), + "end": hex(hi), + "rows": rows, + "heads": len(eas), + "anchors": anchors, + "kind_names": list(_IDATUI_KINDS), + "eas": b64(eas.tobytes()).decode(), + "kinds": b64(kinds.tobytes()).decode(), + "sizes": b64(sizes.tobytes()).decode(), + } def segment_index( addr: Annotated[str, "Any address in the segment to index"], end: Annotated[str, "Optional exclusive end address; default = segment end"] = "", page_rows: Annotated[int, "Rows between anchors (default 500)"] = 500, - detail: Annotated[bool, "Also return every row's ea/kind/size as packed arrays"] = False, + detail: Annotated[ + bool, "Also return every row's ea/kind/size as packed arrays" + ] = False, ) -> dict: """How many listing rows a segment has, and where to seek into it. @@ -662,6 +720,7 @@ def segment_index( if detail: return _idatui_segment_detail(addr, end, count) import ida_segment + seg = ida_segment.getseg(start) if not seg: return {"addr": str(addr), "error": "no segment", "rows": 0, "anchors": []} @@ -703,19 +762,24 @@ def segment_index( fn = get_func(ea) n = 1 if fn is not None and fn.start_ea == ea: - n += 3 # blank, banner, `proc` + n += 3 # blank, banner, `proc` elif cls == _FF_CODE and get_ea_name(ea): - n += 1 # loc_XXX label on its own row + n += 1 # loc_XXX label on its own row if cls == _FF_DATA: n += len(_idatui_struct_member_rows(ea)) item_end = get_item_end(ea) if fn is not None and item_end >= fn.end_ea: - n += 2 # `endp` + separator + n += 2 # `endp` + separator rows += n n_heads += 1 ea = item_end if item_end > ea else ea + 1 - return {"addr": hex(lo), "end": hex(hi), "rows": rows, - "heads": n_heads, "anchors": anchors} + return { + "addr": hex(lo), + "end": hex(hi), + "rows": rows, + "heads": n_heads, + "anchors": anchors, + } def heads( @@ -723,10 +787,21 @@ def heads( 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"] = "", - text: Annotated[bool, "Render each row's disassembly text (default true). False = a skeleton page: same rows, same addresses, no text"] = True, + 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", + ] = "", + text: Annotated[ + bool, + "Render each row's disassembly text (default true). False = a skeleton page: same rows, same addresses, no text", + ] = True, ) -> 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 @@ -740,10 +815,20 @@ def heads( try: start = parse_address(addr) except Exception as e: - return {"addr": str(addr), "error": str(e), "heads": [], "cursor": {"done": True}} + 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}} + return { + "addr": str(addr), + "error": "no segment", + "heads": [], + "cursor": {"done": True}, + } lo, hi = seg.start_ea, seg.end_ea if end: try: @@ -766,7 +851,9 @@ def heads( 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)} + 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 @@ -823,8 +910,9 @@ def heads( # 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}) + out.append( + {"ea": hex(e), "kind": "label", "size": 0, "text": nm + ":", "name": nm} + ) row = dict(row) row["name"] = None out.append(row) @@ -845,7 +933,7 @@ def heads( if len(rows) >= count: more = True break - f = get_flags(ea) # once per head, not once per consumer + 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} @@ -871,8 +959,18 @@ def heads( _IDATUI_FMT_CYCLE = ("hex", "dec", "bin", "char", "offset", "default") -_IDATUI_FMT_SETTABLE = ("hex", "dec", "oct", "bin", "char", "offset", "seg", - "float", "stack", "default") +_IDATUI_FMT_SETTABLE = ( + "hex", + "dec", + "oct", + "bin", + "char", + "offset", + "seg", + "float", + "stack", + "default", +) def _idatui_fmt_nibbles(): @@ -880,13 +978,20 @@ def _idatui_fmt_nibbles(): this module is injected into a file that is imported before a database is open.""" 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, + "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, } @@ -932,8 +1037,12 @@ def _idatui_op_value(ea, n): 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) + 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: @@ -991,9 +1100,9 @@ def _idatui_op_candidates(ea): F = ida_bytes.get_flags(ea) if ida_bytes.is_data(F): - return [0] # a data item's value is operand 0 + 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 + return [] # undefined bytes: IDA refuses a format outright insn = ida_ua.insn_t() if ida_ua.decode_insn(insn, ea) <= 0: return [] @@ -1037,7 +1146,7 @@ def _idatui_op_spans(ea, text): if not op: continue i = text.find(op, pos) - if i < 0: # duplicated operand text (mov eax, eax) + if i < 0: # duplicated operand text (mov eax, eax) i = text.find(op) if i < 0: continue @@ -1070,21 +1179,34 @@ def _idatui_apply_fmt(ea, n, fmt): 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) + 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 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, + 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, @@ -1126,19 +1248,27 @@ def op_format( # 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}"} + 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 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)") + 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] @@ -1148,9 +1278,12 @@ def op_format( # 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))] + 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 @@ -1160,10 +1293,18 @@ def op_format( 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} + 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: @@ -1176,31 +1317,51 @@ def op_format( 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)} + 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")} + 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} + 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} + 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")) + 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 @@ -1267,13 +1428,18 @@ def _idatui_lit_extent(plain, x): around the column, which cannot reach a ``)`` or a space.""" if x >= len(plain): return None - if plain[x] == "'": # a character constant: '-' + 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 ( + 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 @@ -1282,7 +1448,7 @@ def _idatui_lit_extent(plain, x): 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 + if lo > 0 and plain[lo - 1] == "-": # a unary minus is part of it lo -= 1 return (lo, hi) @@ -1321,26 +1487,36 @@ def _idatui_pc_nums(cf, sl): 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)) + 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 + 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 + 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 @@ -1367,31 +1543,36 @@ def pc_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": []} + 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": []} + 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"]), - }) + 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", + 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"] = "", @@ -1431,8 +1612,9 @@ def pc_num_format( 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)): + if rec["ea"] == want_ea and ( + int(opnum) < 0 or rec["opnum"] == int(opnum) + ): target, line = rec, i break if target: @@ -1446,26 +1628,42 @@ def pc_num_format( 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"} + 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"} + 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)] + 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()} + 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": @@ -1481,13 +1679,16 @@ def pc_num_format( 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["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["error"] = ( + f"unknown format {mode!r}; one of hex, dec, oct, char, default" + ) out["text"] = out["before"] return out @@ -1499,8 +1700,9 @@ def pc_num_format( ida_hexrays.user_numforms_erase(cf.numforms, it) if want != "default": nf = ida_hexrays.number_format_t(target["opnum"]) - nf.flags = ida_bytes.get_operand_flag(_idatui_fmt_nibbles()[want], - target["opnum"]) + nf.flags = ida_bytes.get_operand_flag( + _idatui_fmt_nibbles()[want], target["opnum"] + ) try: nf.org_nbytes = target["nbytes"] except Exception: @@ -1515,14 +1717,18 @@ def pc_num_format( 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") + 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) + 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"]) + 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}" @@ -1557,11 +1763,17 @@ def decompile(addr, include_addresses=True): 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}"} + 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}"} + 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(): @@ -1569,7 +1781,9 @@ def decompile(addr, include_addresses=True): 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): + 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: @@ -1595,9 +1809,13 @@ def decompile(addr, include_addresses=True): 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}) + refs.append( + { + "addr": hex(target), + "name": ida_name.get_name(target) or "", + "string": text, + } + ) return 0 try: @@ -1615,6 +1833,7 @@ def decomp_map( swept across the line's columns via get_line_item. Shape: {addr, lines:[{ea: primary|None, eas:[hex,...]}, ...]}.""" import ida_hexrays + try: ea = int(str(addr), 16) except ValueError: @@ -1691,3 +1910,32 @@ def decomp_map( eas.append(hex(e)) lines.append({"ea": eas[0] if eas else None, "eas": eas}) return {"addr": hex(func.start_ea), "lines": lines} + + +def profile_remote(operation, args, reps=5): + """Profile one persistent remote tool entirely inside IDA.""" + import cProfile + import io + import pstats + + functions = { + "heads": heads, + "segment_index": segment_index, + "op_format": op_format, + "pc_nums": pc_nums, + "decompile": decompile, + "decomp_map": decomp_map, + "pc_num_format": pc_num_format, + } + function = functions.get(str(operation)) + if function is None: + raise ValueError(f"unknown profile operation: {operation!r}") + profiler = cProfile.Profile() + for _ in range(max(1, int(reps))): + profiler.runcall(function, **dict(args)) + stats = pstats.Stats(profiler) + output = io.StringIO() + stats.stream = output + stats.sort_stats("tottime") + stats.print_stats(80) + return {"stats": output.getvalue(), "total": stats.total_tt, "reps": reps} |