From c9208de05d8583b677117fe43c9d3567e89eb2ce Mon Sep 17 00:00:00 2001 From: blasty Date: Fri, 7 Aug 2026 12:39:54 +0200 Subject: Rebase MISTER EXO's ida-codemode port onto the current tree Mechanical part of the port: the 27-file patch was cut against a base ~148 commits behind us, so it did not apply. Resolved 11 conflicts (all of them diff drift, not semantic clashes) and the three file deletions: - app.py: the patch re-inserted _do_rename/_do_name_addr/_seek_split etc. as "theirs" because our tree moved them to edit_ctl.py/trace_ctl.py. Kept ours and applied the real intent (WorkerClient->CodeModeClient, .call->.invoke, _open_worker_client->_open_database_client) at their current homes. - domain.py: kept Head as a NamedTuple -- the patch reverted it to a frozen dataclass, which the perf work measured at 2.9us vs 1.9us per row on a quarter-million-row walk. Dropped _fetch_output (no download_url under Code Mode) and its now-dead urllib/json imports. - pane.py: the patch's deletion swallowed our zellij support along with the worker-reaping block it meant to remove. Kept zellij, removed the reaping. - test_scenarios.py: the idb_save->save_database teardown hunk belongs to tests/_fixtures.py in our tree; applied it there and kept our pc_num_format scenario that the drift landed on. Three defects in the patch itself, fixed here: - It made "import idatui" hard-require ida_codemode, so every offline suite died at import -- including the pure ones (graph/index/trace) that are the house rule for "tests/run.py --fast". The import is now deferred and gated on the binding, which is also what lets the port's own contract tests inject a fake DatabaseHandle. - project.stage() inlined an ida_codemode.registry import and treated "library not installed" as "someone owns this database", which broke IDA-free project staging. Ownership lookup moved to codemode_client.database_owner(). - tests/test_codemode_client.py had no NEEDS_IDA marker, which tests/run.py rejects outright. Offline suite: 301 passed, 0 failed. Against master's 344 the whole delta is accounted for: -40 worker_client (module deleted), -18 launch sweep checks (behaviour deliberately removed) +3 guarding that it stays removed, +2 pool (GUI-save semantics), +13 new codemode_client contract tests. NOT yet done, and the port is not functional without it: the adapter is missing five operations our tree grew since the patch's base (flowchart, op_format, pc_nums, pc_num_format, survey_binary) and its "heads" predates back-walking and digest/expect. --- idatui/__init__.py | 5 +- idatui/app.py | 200 ++++---- idatui/codemode_client.py | 1164 +++++++++++++++++++++++++++++++++++++++++++++ idatui/domain.py | 209 ++++---- idatui/drive.py | 3 +- idatui/edit_ctl.py | 4 +- idatui/errors.py | 10 +- idatui/launch.py | 103 ++-- idatui/pane.py | 93 +--- idatui/pool.py | 100 ++-- idatui/project.py | 28 +- idatui/rpc.py | 4 +- idatui/worker.py | 301 ------------ idatui/worker_client.py | 266 ----------- 14 files changed, 1506 insertions(+), 984 deletions(-) create mode 100644 idatui/codemode_client.py delete mode 100644 idatui/worker.py delete mode 100644 idatui/worker_client.py (limited to 'idatui') diff --git a/idatui/__init__.py b/idatui/__init__.py index 2cbdde8..7da28b3 100644 --- a/idatui/__init__.py +++ b/idatui/__init__.py @@ -1,5 +1,4 @@ -"""idatui — a minimal keyboard-first TUI for IDA Pro, driving idalib via a -private unix-socket worker (idatui.worker / WorkerClient).""" +"""idatui — a keyboard-first TUI using shared IDA Code Mode databases.""" from .errors import ( IDAError, @@ -11,6 +10,7 @@ from .errors import ( IDASessionError, Session, ) +from .codemode_client import CodeModeClient from .domain import ( Program, FunctionIndex, @@ -25,6 +25,7 @@ from .domain import ( ) __all__ = [ + "CodeModeClient", "Program", "FunctionIndex", "DisasmModel", diff --git a/idatui/app.py b/idatui/app.py index 4a58b16..a4b03a1 100644 --- a/idatui/app.py +++ b/idatui/app.py @@ -10,8 +10,8 @@ Design notes: without ever materializing 52k lines in a widget. * All network/domain work runs in Textual worker threads; the UI never blocks. * An address-history stack backs Enter (follow) / Esc (back), IDA-style. -* On startup we bump the worker idle-TTL and run a keepalive heartbeat so the - session never gets reaped while we chill. +* Database lifecycle is lease-based through ida_codemode: matching GUI sessions + are reused, otherwise a shared managed idalib worker is opened on demand. """ from __future__ import annotations @@ -32,7 +32,7 @@ from textual import work from textual.app import App, ComposeResult from textual.binding import Binding from textual.command import DiscoveryHit, Hit, Provider -from textual.containers import Grid, Horizontal, Vertical, VerticalScroll +from textual.containers import Horizontal, Vertical, VerticalScroll from textual.geometry import Region, Size from textual.message import Message from textual.reactive import reactive @@ -53,8 +53,8 @@ from .trace_ctl import TraceController from .highlight import highlight_c from .errors import IDAToolError, IDAConnectionError -from .worker_client import WorkerClient -from .domain import DisasmModel, Func, Head, ListingModel, Program, Struct +from .codemode_client import CodeModeClient, registered_database +from .domain import Func, Head, ListingModel, Program, Struct # Styles for the disassembly listing. _S_ADDR = Style(color="#6b7684") @@ -173,8 +173,8 @@ _ADDR_MARK_STRIP_RE = re.compile(r"\s*/\*\s*0x[0-9A-Fa-f]+\s*\*/") @dataclass class BinaryState: """Everything that makes one project binary's session resumable across a - switch. Addresses outlive the worker, so nav history survives eviction; the - Program/index only survive while that worker is still resident.""" + switch. Addresses outlive a database lease, so nav history survives eviction; + the Program/index only survive while that lease remains resident.""" label: str program: object | None = None @@ -1085,9 +1085,9 @@ class ListingView(SearchMixin, NavMixin, ColumnCursor, ScrollView, can_focus=Tru def _span_segments(h: Head, fallback: Style): """Segments for a row's disassembly text. - Uses IDA's own token classification when the worker supplied it; falls - back to the old mnemonic/rest split so an older worker (or a row whose - spans didn't match the text) still renders. + Uses IDA's own token classification when Code Mode supplies it; falls + back to the mnemonic/rest split when spans are absent or disagree with + the plain text. """ if h.spans: return [Segment(t, _S_SPAN.get(k, fallback)) for k, t in h.spans] @@ -3493,12 +3493,16 @@ _HELP = ( class QuitScreen(ModalScreen): - """Asked before exiting with unsaved database changes. Dismisses with - "save", "discard" or None (stay).""" + """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. + """ BINDINGS = [ Binding("s", "save", "Save & quit"), - Binding("d", "discard", "Discard & quit"), + Binding("d", "discard", "Leave & quit"), Binding("escape,c", "cancel", "Cancel"), ] @@ -3515,7 +3519,7 @@ class QuitScreen(ModalScreen): for label in self._labels: body.append(f" \u2022 {label}\n", _S_LABEL) yield Static(body, id="quit-list") - yield Static("s save & quit d discard & quit Esc cancel", + yield Static("s save & quit d leave as-is & quit Esc cancel", id="quit-help") def action_save(self) -> None: @@ -4780,15 +4784,16 @@ class IdaTui(App): self._index = None # project-wide symbol/string index if project is not None: from .index import ProjectIndex - from .pool import WorkerPool - self._pool = WorkerPool(project, ttl=ttl) + from .pool import DatabasePool + self._pool = DatabasePool(project, ttl=ttl) self._index = ProjectIndex( os.path.join(project.index_dir, "project.db")) self._binary = project.refs[0].label open_path = project.refs[0].staged self._open_path = open_path self._ttl = ttl - self._load_args = load_args or "" # IDA switches for a headerless blob + self._load_args = load_args or "" # first-open options for a headerless blob + self._new_database = False # Ctrl+L asks Code Mode for a fresh IDB self._title = (os.path.basename(open_path) if open_path else "") #: Where we are in the execution trace, and everything that moves us. #: Owns the trace state; the _trace/_t/_trail_* properties below @@ -4797,7 +4802,7 @@ class IdaTui(App): self._do_keepalive = keepalive self._rpc_path = rpc_path self._rpc = None - self.client: WorkerClient | None = None + self.client: CodeModeClient | None = None self.program: Program | None = None self._loading_screen: LoadingScreen | None = None self._ka = None @@ -4900,8 +4905,8 @@ class IdaTui(App): if self._rpc_path: self._start_rpc() # A file no loader recognises has to be described before it can be - # opened, so ask BEFORE the worker starts — once IDA has made a database - # the answer is baked in and changing it means deleting the .i64. + # opened, so ask BEFORE Code Mode creates it — once IDA has made a database + # the answer is baked in and changing it requires a fresh-IDB reopen. if self._project is not None: ref = self._pending_load_ref() if ref is not None: @@ -4932,6 +4937,11 @@ class IdaTui(App): if os.path.exists(self._open_path + ".i64") or os.path.exists( os.path.splitext(self._open_path)[0] + ".i64"): return False + try: + if registered_database(self._open_path): + return False + except Exception: + pass # connect() will surface registry failures with full diagnostics return needs_load_options(self._open_path) def action_load_options(self) -> None: @@ -4945,7 +4955,11 @@ class IdaTui(App): forward. """ if not self._can_reload(): - self._status("nothing to reload") + if self.client is not None and self.client.backend == "gui": + self._status( + "reload unavailable for a GUI-owned database — reopen it in IDA") + else: + self._status("nothing to reload") return n = len(self._func_index) if self._func_index else 0 note = ("this image has no functions, so nothing is lost" @@ -4964,10 +4978,13 @@ class IdaTui(App): ref = self._project.by_label(self._binary) if ref is not None: path, label = ref.source, ref.label - # Drop the worker first: it holds the database open, and the .i64 can't - # be removed (or rebuilt) underneath a live one. - self._release_worker() - self._drop_database() + # Release our lease first. Code Mode waits for a managed worker's final + # lease grace, then creates the replacement IDB atomically. A GUI-backed + # database is rejected by _can_reload(): the TUI must never close it. + self._release_database() + self._new_database = True + if label is not None and self._pool is not None: + self._pool.recreate_on_next_open(label) self._reset_for_reload() self._load_args = "" if label is not None and self._project is not None: @@ -4979,7 +4996,9 @@ class IdaTui(App): self._pending_switch = None self._ask_load_options(path, label=label) - def _release_worker(self) -> None: + def _release_database(self) -> None: + if self.program is not None: + self.program.close() if self._pool is not None and self._binary is not None: try: self._pool.evict(self._binary, save=False) @@ -4993,23 +5012,6 @@ class IdaTui(App): self.client = None self.program = None - def _drop_database(self) -> None: - """Remove the .i64 (and any unpacked scratch) so the next open re-reads - the raw image with new options.""" - base = self._open_path - if self._project is not None and self._binary is not None: - ref = self._project.by_label(self._binary) - if ref is not None: - base = ref.staged - if not base: - return - for suffix in (".i64", ".id0", ".id1", ".id2", ".nam", ".til"): - for cand in (base + suffix, os.path.splitext(base)[0] + suffix): - try: - os.remove(cand) - except OSError: - pass - def _reset_for_reload(self) -> None: self._no_functions = False self._func_index = None @@ -5053,6 +5055,11 @@ class IdaTui(App): if os.path.exists(ref.db) or os.path.exists( os.path.splitext(ref.staged)[0] + ".i64"): return None # already analysed: the .i64 records how + try: + if registered_database(ref.staged, output_database=ref.db): + return None + except Exception: + pass from .formats import needs_load_options return ref if needs_load_options(ref.source) else None @@ -5106,10 +5113,6 @@ class IdaTui(App): asyncio.get_running_loop().create_task(_serve()) - async def on_unmount(self) -> None: - if self._rpc is not None: - await self._rpc.stop() - # -- status helper ----------------------------------------------------- # def _status(self, text: str, priority: bool = False) -> None: """Write the status bar. ``priority`` marks the RESULT of something the @@ -5174,9 +5177,10 @@ class IdaTui(App): # -- connection loss / recovery --------------------------------------- # def _handle_exception(self, error: BaseException) -> None: - """Intercept a lost-connection error from any worker so the whole app - doesn't die when the analysis server goes away (it can idle out, be - killed, or the box can sleep). Everything else crashes as usual.""" + """Intercept a lost Code Mode lease so the app can rediscover the DB. + + Everything unrelated to database connectivity crashes as usual. + """ from textual.worker import WorkerFailed orig = error.error if isinstance(error, WorkerFailed) else error if isinstance(orig, IDAConnectionError): @@ -5208,15 +5212,15 @@ class IdaTui(App): @work(thread=True, exclusive=True, group="reconnect") def _reconnect(self) -> None: - # The worker died (segfault -> dropped socket). Respawn it: it re-opens - # and re-analyzes the binary in a fresh process, then we rebuild. + # The registered instance disappeared. Rediscover it; Code Mode may find + # a GUI/replacement worker, then we rebuild caches against the new handle. try: if self._open_path is None: self.app.call_from_thread(self._reconnect_failed, "no binary to reopen") return - client = WorkerClient(self._open_path, ttl=self._ttl, - load_args=self._load_args) + client = CodeModeClient(self._open_path, ttl=self._ttl, + load_args=self._load_args) client.connect(progress=lambda m: self.app.call_from_thread( self._conn_note, m)) except Exception as e: # noqa: BLE001 @@ -5224,7 +5228,7 @@ class IdaTui(App): return self.app.call_from_thread(self._after_reconnect, client, Program(client)) - def _after_reconnect(self, client: "WorkerClient", program: "Program") -> None: + def _after_reconnect(self, client: "CodeModeClient", program: "Program") -> None: self.client = client self.program = program self._reconnecting = False @@ -5244,13 +5248,13 @@ class IdaTui(App): @work(thread=True, exclusive=True, group="connect") def _connect(self) -> None: try: - client = self._open_worker_client() + client = self._open_database_client() if client is None: return # the opener already reported + dismissed the overlay module = client.health().get("module", "?") if self._do_keepalive: - # Keep the session warm while we run; don't make it immortal, so - # it's reclaimed after the TUI closes. (No-op for the worker.) + # Compatibility shim: DatabaseHandle's SSE lease already owns + # liveness and heartbeat behavior. self._ka = client.keepalive(interval=120.0).start() program = Program(client) except Exception as e: # noqa: BLE001 @@ -5265,14 +5269,14 @@ class IdaTui(App): return self.client = client self.program = program - self.app.call_from_thread(self._status, f"{module} — loading functions…") + self._new_database = False + self.app.call_from_thread( + self._status, f"{module} [{client.backend}] — loading functions…") self._load_functions() - def _open_worker_client(self): # type: ignore[no-untyped-def] - """Our idalib-worker path: spawn the worker (it opens + analyzes the - binary in its own process) and connect. Returns the client, or None.""" - from .worker_client import WorkerClient - if self._pool is not None: # project mode: the pool owns the workers + def _open_database_client(self): # type: ignore[no-untyped-def] + """Attach through Code Mode, reusing a GUI or managed idalib database.""" + if self._pool is not None: # project mode: the pool owns the leases label = self._binary or self._project.refs[0].label client = self._pool.get(label, progress=lambda m: self.app.call_from_thread(self._status, m)) @@ -5283,14 +5287,15 @@ class IdaTui(App): return client if not self._open_path: self.app.call_from_thread( - self._status, "the worker backend needs a binary path") + self._status, "Code Mode needs a database or executable path") self.app.call_from_thread(self._dismiss_loading) return None base = os.path.basename(self._open_path) self.app.call_from_thread( - self._status, f"starting worker — initial auto-analysis of {base}…") - client = WorkerClient(self._open_path, ttl=self._ttl, - load_args=self._load_args) + self._status, f"discovering Code Mode database for {base}…") + client = CodeModeClient(self._open_path, ttl=self._ttl, + load_args=self._load_args, + new_database=self._new_database) client.connect(progress=lambda m: self.app.call_from_thread( self._status, m)) return client @@ -5362,7 +5367,7 @@ class IdaTui(App): @work(thread=True, exclusive=True, group="index") def _index_binary(self) -> None: """Fold this binary's symbols + strings into the project index, so it can - be searched later even when its worker is gone.""" + be searched later even when its Code Mode lease is gone.""" if self._index is None or self._project is None or self._binary is None: return ref = self._project.by_label(self._binary) @@ -5380,7 +5385,7 @@ class IdaTui(App): imps, exps = self.program.linkage() entries += [(KIND_IMPORT, i.addr, i.name) for i in imps] entries += [(KIND_EXPORT, e.addr, e.name) for e in exps] - except Exception: # noqa: BLE001 -- an old worker has no list_linkage + except Exception: # noqa: BLE001 -- indexing is best-effort pass try: n = self._index.reindex(self._binary, entries, source=ref.source) @@ -5465,7 +5470,13 @@ class IdaTui(App): cursor=0, push=True, is_region=True) def _can_reload(self) -> bool: - """Whether we're able to re-open this binary with different options.""" + """Whether Code Mode can replace this IDB with different options. + + A GUI database is owned by the user and has no remote close/rollback + route. Managed idalib databases can be released and reopened fresh. + """ + if self.client is not None and self.client.backend == "gui": + return False if self._project is not None and self._binary is not None: return True return bool(self._open_path) @@ -5654,6 +5665,9 @@ 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() elif choice == "save": @@ -5670,7 +5684,7 @@ class IdaTui(App): if self._pool is not None: self._pool.close_all(save=True) # saves each resident worker elif self.program is not None: - self.program.client.call("idb_save", timeout=600.0) + self.program.client.save_database() except Exception as e: # noqa: BLE001 -- still exit, but say so self.app.call_from_thread(self._status, f"save failed: {e}") self.app.call_from_thread(self._finish_exit) @@ -5710,7 +5724,7 @@ class IdaTui(App): self._ask_load_options(ref.source, label=label) return # Snapshot what we're leaving so coming back restores the view, then let - # the pool hand us a worker (spawning + evicting as the budget dictates). + # the pool hand us a lease (attaching + evicting as the budget dictates). if self._binary is not None: self._states[self._binary] = BinaryState( label=self._binary, program=self.program, @@ -5731,8 +5745,8 @@ class IdaTui(App): self.app.call_from_thread(self._switch_failed, label, str(e)) return st = self._states.get(label) - # The Program (and its caches) only survive while that worker does; a - # binary that was evicted comes back with a fresh one. Either way the nav + # The Program (and its caches) only survive while that lease does; an + # evicted binary reattaches. Either way the nav # history is just addresses, so it always survives. reuse = (st is not None and st.program is not None and getattr(st.program, "client", None) is client) @@ -5769,7 +5783,7 @@ class IdaTui(App): self._did_auto_land = False self._auto_land() return - # Cold (first visit, or the worker was evicted): rebuild the index, then + # Cold (first visit, or the lease was evicted): rebuild the index, then # land back where we were via _pending_restore. self._cur = None self._func_index = None @@ -5826,28 +5840,6 @@ class IdaTui(App): return self._goto_ea(addr, push=True) # land on the literal in the listing - def on_descendant_focus(self, event) -> None: # type: ignore[no-untyped-def] - """Keep ``_active`` in step with focus while split. - - Tab moves both together, but focus also moves on its own — a click, or a - pane focusing itself after a load — and then ``_active`` still names the - pane you're NOT in. Everything downstream trusts ``_active``: follow - resolves the word under that pane's cursor and pushes history for it, so - Enter in the pseudocode would follow something from the listing and the - next Esc got spent undoing it. - """ - if not self._split: - return - w = self.focused - mode = ("decomp" if isinstance(w, DecompView) - else "listing" if isinstance(w, ListingView) else None) - if mode is None or mode == self._active: - return - self._active = mode - self._sync_split(mode) # re-link the band from the new driver - if not self.query_one(DecompView).loading: - self._status_for_cur("split") # never clobber "decompiling…" - def action_toggle_view(self) -> None: """Tab: switch the code pane between disassembly and pseudocode (or leave the hex view back to the preferred code view).""" @@ -6356,7 +6348,7 @@ class IdaTui(App): def _cross_binary_impl(self, name: str) -> tuple[str, int] | None: """``(binary, addr)`` of a project binary that EXPORTS ``name``. - Reads the on-disk index, so a provider resolves even when its worker was + Reads the on-disk index, so a provider resolves even when its lease was evicted — the whole reason the index exists. """ if self._index is None or self._project is None or not name: @@ -6529,7 +6521,7 @@ class IdaTui(App): def _foreign_importers(self, subj: int, subj_name, fn): # type: ignore[no-untyped-def] """Project binaries that IMPORT the symbol at ``subj`` — the other half of the phase-3 join, read from the on-disk index so a caller shows up - whether or not its worker is resident. + whether or not its database lease is resident. Only for a symbol this binary actually exports: a local name that happens to collide with another binary's import isn't a caller of ours. @@ -6742,7 +6734,7 @@ class IdaTui(App): def _save(self) -> None: assert self.program is not None try: - self.program.client.call("idb_save", timeout=300.0) + self.program.client.save_database() except Exception as e: # noqa: BLE001 self.app.call_from_thread(self._status, f"save failed: {e}") return @@ -7727,7 +7719,9 @@ class IdaTui(App): "(c code · p func · u undefine · Enter follow)") # -- teardown ---------------------------------------------------------- # - def on_unmount(self) -> None: + async def on_unmount(self) -> None: + if self._rpc is not None: + await self._rpc.stop() if self._ka is not None: self._ka.stop() if self.program is not None: @@ -7739,7 +7733,7 @@ class IdaTui(App): elif self.client is not None: if self._save_on_exit is None and self._dirty: try: # unexpected teardown with edits: don't drop them - self.client.call("idb_save", timeout=600.0) + self.client.save_database() except Exception: # noqa: BLE001 pass self.client.close() diff --git a/idatui/codemode_client.py b/idatui/codemode_client.py new file mode 100644 index 0000000..77a0227 --- /dev/null +++ b/idatui/codemode_client.py @@ -0,0 +1,1164 @@ +"""Client adapter from ida-tui's domain operations to IDA Code Mode. + +``DatabaseHandle`` is the lifecycle boundary: it discovers an already-registered +GUI database, reuses a shared managed idalib worker, or starts one when needed. +The TUI never owns or terminates an IDA process. Closing this client releases +only its lease. + +The Code Mode transport intentionally exposes one broad operation, +``execute_python``. ``CodeModeClient.invoke`` turns the small, address-centric +operations needed by the paging layer into self-contained snippets. The +snippets prefer the public ``ida-domain`` ``db`` object. A handful of features +that ida-domain does not currently expose (IDA-coloured listing rows, creating +instructions, ARM T-state, and detailed Hex-Rays line maps/failures) use the +IDAPython modules that Code Mode deliberately makes importable. +""" +from __future__ import annotations + +import json +import os +import shlex +import threading +import time +from textwrap import dedent +from typing import Any + +from .errors import IDAConnectionError, IDATimeoutError, IDAToolError, Session + +# ida_codemode is imported EAGERLY-IF-PRESENT but never at hard import cost. +# +# The paging/graph/trace layers and their offline test suites must keep importing +# `idatui` on a machine with no IDA and no Code Mode installed -- that is the +# house rule the stdlib-only worker client used to satisfy for free, and +# `tests/run.py --fast` (257 checks, any python3) depends on it. A hard top-level +# import here makes the whole package unimportable, so the failure is deferred to +# the first operation that genuinely needs the library. +_CODEMODE_ERROR: Exception | None = None +try: + from ida_codemode.client import ( + ClientError, + DatabaseHandle, + InstanceDisconnectedError, + RemoteError, + ) + from ida_codemode.registry import ( + REGISTRY_DIR, + FileLock, + RegistryEntry, + canonical_path, + idb_key, + scan_instances, + ) + from ida_codemode.resolver import IdbBusy, expected_idb_path +except ImportError as _exc: # library absent: usable only for offline layers + _CODEMODE_ERROR = _exc + # Bound to None rather than left undefined so the names stay patchable: the + # offline contract tests inject a fake DatabaseHandle here. + ClientError = InstanceDisconnectedError = RemoteError = None # type: ignore[assignment,misc] + DatabaseHandle = RegistryEntry = FileLock = None # type: ignore[assignment,misc] + REGISTRY_DIR = canonical_path = idb_key = scan_instances = None # type: ignore[assignment] + IdbBusy = expected_idb_path = None # type: ignore[assignment] + + +def _require_codemode() -> None: + """Raise an actionable error when the Code Mode library is missing. + + Gated on the binding, not on the original import result, so a test that + injects a fake ``DatabaseHandle`` exercises the real adapter logic. + """ + if DatabaseHandle is None: + raise IDAConnectionError( + "ida-codemode-mcp is not installed in this environment " + f"({_CODEMODE_ERROR}). Install it (e.g. `uv sync`, or " + "`pip install -e ../ida-codemode-mcp`) so ida-tui can lease a " + "database.") from _CODEMODE_ERROR + + +def database_owner(idb_path: str, staged_path: str | None = None): + """The registry entry that owns ``idb_path``/``staged_path``, else None. + + Returns None when the Code Mode library is absent: with no library there is + no client in this environment that could be holding the database, and the + IDA-free layers (project staging) must keep working. Registry errors that + happen WITH the library installed still propagate -- those mean "we could + not determine ownership", which is not the same as "nobody owns it". + """ + if DatabaseHandle is None: + return None + expected_key = idb_key(idb_path) + staged = canonical_path(staged_path) if staged_path else None + for item in scan_instances(timeout=0.5): + entry = item.entry + if entry.idb_key == expected_key: + return entry + if staged and entry.exe_path and canonical_path(entry.exe_path) == staged: + return entry + return None + + +def registered_database(path: str, output_database: str | None = None) -> bool: + """Whether a live/lock-held Code Mode instance owns this target.""" + _require_codemode() + source = canonical_path(path) + expected = canonical_path(output_database) if output_database else expected_idb_path(source) + expected_key = idb_key(expected) + for instance in scan_instances(timeout=0.5): + entry = instance.entry + if entry.idb_key == expected_key: + return True + if not output_database and entry.backend == "gui" and entry.exe_path: + if canonical_path(entry.exe_path) == source: + return True + return False + + +class _NoopKeepAlive: + """Compatibility shim: the DatabaseHandle's SSE lease is the heartbeat.""" + + def __init__(self) -> None: + self.beats = self.failures = 0 + + def start(self) -> "_NoopKeepAlive": + return self + + def stop(self) -> None: + pass + + +def _parse_load_args(value: str) -> tuple[str | None, int | None, str | None]: + """Translate ida-tui's legacy first-open switches to Code Mode options. + + Code Mode has typed options for processor, natural loading address and file + type. It deliberately has no arbitrary command-line escape hatch; reject + switches we cannot represent instead of silently loading a blob wrongly. + """ + processor: str | None = None + loading_address: int | None = None + file_type: str | None = None + unsupported: list[str] = [] + try: + words = shlex.split(value or "", posix=os.name != "nt") + except ValueError as exc: + raise ValueError(f"invalid IDA load options: {exc}") from exc + for word in words: + if word.startswith("-p") and len(word) > 2: + processor = word[2:] + elif word.startswith("-b") and len(word) > 2: + try: + # IDA's -b is in 16-byte paragraphs. DatabaseHandle expects the + # natural address, which is the safer public API. + loading_address = int(word[2:], 16) << 4 + except ValueError as exc: + raise ValueError(f"invalid IDA loading address: {word!r}") from exc + elif word.startswith("-T") and len(word) > 2: + file_type = word[2:] + else: + unsupported.append(word) + if unsupported: + joined = " ".join(unsupported) + raise ValueError( + "ida-codemode cannot represent arbitrary IDA load options: " + f"{joined!r}; use processor/base/file type options instead" + ) + return processor, loading_address, file_type + + +def _script(args: dict[str, Any], body: str) -> str: + """Bind JSON arguments without interpolating user text into Python code.""" + encoded = json.dumps(args, ensure_ascii=False, separators=(",", ":")) + return f"import json\na = json.loads({encoded!r})\n{dedent(body).strip()}\n" + + +# Rich flat-listing generation is the largest ida-domain gap in this port. +# ida-domain can enumerate heads and render plain disassembly, but it does not +# expose undefined runs, IDA colour spans, function banners, or expanded UDT +# members. Keep that IDAPython-only logic isolated in this one operation. +_HEADS = r''' +import ida_bytes, ida_funcs, ida_idaapi, ida_lines, ida_name, ida_nalt, ida_segment, ida_typeinf +start = int(str(a["addr"]), 16) +count = max(1, min(int(a.get("count", 200)), 2000)) +offset = max(0, int(a.get("offset", 0))) +annotate = bool(a.get("annotate", False)) +seg = db.segments.get_at(start) +if seg is None: + result = {"addr": a["addr"], "error": "no segment", "heads": [], "cursor": {"done": True}} +else: + lo, hi = int(seg.start_ea), int(seg.end_ea) + if a.get("end"): + hi = min(hi, int(str(a["end"]), 16)) + + span_names = { + "insn": ("SCOLOR_INSN", "SCOLOR_KEYWORD", "SCOLOR_ASMDIR", "SCOLOR_MACRO"), + "reg": ("SCOLOR_REG",), + "num": ("SCOLOR_NUMBER", "SCOLOR_CHAR", "SCOLOR_BINPREF"), + "str": ("SCOLOR_STRING",), + "name": ("SCOLOR_DATNAME", "SCOLOR_CODNAME", "SCOLOR_LOCNAME", "SCOLOR_IMPNAME", + "SCOLOR_DEMNAME", "SCOLOR_LIBNAME", "SCOLOR_CNAME", "SCOLOR_DNAME", + "SCOLOR_CREF", "SCOLOR_DREF", "SCOLOR_CREFTAIL", "SCOLOR_DREFTAIL"), + "seg": ("SCOLOR_SEGNAME",), + "cmt": ("SCOLOR_AUTOCMT", "SCOLOR_REGCMT", "SCOLOR_RPTCMT", "SCOLOR_VOIDOP"), + "punct": ("SCOLOR_SYMBOL", "SCOLOR_ALTOP", "SCOLOR_HIDNAME"), + "err": ("SCOLOR_ERROR",), + } + tag_kinds = {} + for kind, names in span_names.items(): + for name in names: + value = getattr(ida_lines, name, None) + if isinstance(value, str) and value: + tag_kinds[value[0]] = kind + elif isinstance(value, int): + tag_kinds[chr(value)] = kind + + def spans(tagged): + on, off, esc = "\x01", "\x02", "\x03" + addr_tag = chr(getattr(ida_lines, "COLOR_ADDR", 0x28)) + addr_len = int(getattr(ida_lines, "COLOR_ADDR_SIZE", 16)) + out, stack, buf = [], [], [] + def flush(): + if buf: + out.append([stack[-1] if stack else "text", "".join(buf)]) + buf.clear() + i = 0 + while i < len(tagged): + ch = tagged[i] + if ch == on and i + 1 < len(tagged): + tag = tagged[i + 1] + if tag == addr_tag: + i += 2 + addr_len + continue + flush(); stack.append(tag_kinds.get(tag, "text")); i += 2; continue + if ch == off and i + 1 < len(tagged): + flush() + if stack: stack.pop() + i += 2; continue + if ch == esc and i + 1 < len(tagged): + buf.append(tagged[i + 1]); i += 2; continue + buf.append(ch); i += 1 + flush() + collapsed, previous_space = [], False + for kind, text in out: + acc = [] + for ch in text: + if ch.isspace(): + if previous_space: continue + acc.append(" "); previous_space = True + else: + acc.append(ch); previous_space = False + if acc: collapsed.append([kind, "".join(acc)]) + if collapsed: + collapsed[0][1] = collapsed[0][1].lstrip() + collapsed[-1][1] = collapsed[-1][1].rstrip() + return [[kind, text] for kind, text in collapsed if text] + + def row(ea): + flags = ida_bytes.get_flags(ea) + kind = "code" if ida_bytes.is_code(flags) else ("data" if ida_bytes.is_data(flags) else "unknown") + tagged = ida_lines.generate_disasm_line(ea, 0) or "" + text = " ".join(ida_lines.tag_remove(tagged).split()) if tagged else "" + item = {"ea": hex(ea), "kind": kind, "size": int(ida_bytes.get_item_size(ea)), "text": text} + if tagged: + rich = spans(tagged) + if " ".join("".join(x[1] for x in rich).split()) == text: + item["spans"] = rich + name = ida_name.get_ea_name(ea) + if name: item["name"] = name + return item + + def unknown_row(ea, size): + if size <= 1: return row(ea) + item = {"ea": hex(ea), "kind": "unknown", "size": int(size), "text": f"db {size} dup(?)"} + name = ida_name.get_ea_name(ea) + if name: item["name"] = name + return item + + def members(ea): + tif = db.types.get_at(ea) + if tif is None or not tif.is_udt(): return [] + answer = [] + for member in db.types.get_udt_members(tif): + type_text = member.type.dstr() or "" + text = f"+{member.offset:X} {member.name}" + (f" {type_text}" if type_text else "") + answer.append({"ea": hex(ea + member.offset), "kind": "member", + "size": int(member.size), "text": text}) + return answer + + def is_unknown(ea): + flags = ida_bytes.get_flags(ea) + return not (ida_bytes.is_code(flags) or ida_bytes.is_data(flags)) + def run_end(ea): + nxt = ida_bytes.next_head(ea, hi) + return nxt if nxt != ida_idaapi.BADADDR and ea < nxt <= hi else hi + def advance(ea): + if is_unknown(ea): return run_end(ea) + nxt = ida_bytes.get_item_end(ea) + return nxt if nxt > ea else ea + 1 + def rows_for(ea): + if is_unknown(ea): return [unknown_row(ea, run_end(ea) - ea)] + fn = db.functions.get_at(ea) if annotate else None + at_start = fn is not None and int(fn.start_ea) == ea + answer = [] + if at_start: + name = db.functions.get_name(fn) or f"sub_{ea:X}" + answer += [ + {"ea": hex(ea), "kind": "sep", "size": 0, "text": ""}, + {"ea": hex(ea), "kind": "sep", "size": 0, + "text": "; " + "=" * 15 + " S U B R O U T I N E " + "=" * 15}, + {"ea": hex(ea), "kind": "funchdr", "size": 0, + "text": name + " proc", "name": name}, + ] + item = row(ea) + if at_start: + item["name"] = None + elif annotate and item["kind"] == "code" and item.get("name"): + name = item["name"] + answer.append({"ea": hex(ea), "kind": "label", "size": 0, + "text": name + ":", "name": name}) + item["name"] = None + answer.append(item) + if item["kind"] == "data": answer += members(ea) + if fn is not None and ida_bytes.get_item_end(ea) >= int(fn.end_ea): + name = db.functions.get_name(fn) or f"sub_{int(fn.start_ea):X}" + answer += [ + {"ea": hex(ea), "kind": "funchdr", "size": 0, + "text": name + " endp", "name": name}, + {"ea": hex(ea), "kind": "sep", "size": 0, "text": "; " + "-" * 60}, + ] + return answer + + ea = ida_bytes.get_item_head(start) + if ea == ida_idaapi.BADADDR: ea = start + for _ in range(offset): + if ea >= hi: break + ea = advance(ea) + rows = [] + more = False + while ea != ida_idaapi.BADADDR and ea < hi: + if len(rows) >= count: + more = True; break + rows += rows_for(ea) + ea = advance(ea) + result = {"addr": a["addr"], "heads": rows, + "cursor": {"next": hex(ea)} if more else {"done": True}} +result +''' + + +_DECOMP_MAP_HELPER = r''' +def line_map(cfunc): + import ida_hexrays + answer = [] + for sl in cfunc.get_pseudocode(): + tagged, eas, seen = sl.line, [], set() + for x in range(len(tagged) + 1): + head = ida_hexrays.ctree_item_t(); item = ida_hexrays.ctree_item_t(); tail = ida_hexrays.ctree_item_t() + if not cfunc.get_line_item(tagged, x, False, head, item, tail): continue + text = item.dstr() or "" + try: ea = int(text.split(": ", 1)[0], 16) + except (ValueError, IndexError): continue + if ea not in seen: seen.add(ea); eas.append(ea) + answer.append(eas) + return answer +''' + + +_OPERATIONS: dict[str, str] = { + "list_funcs": r''' +import fnmatch +queries = a.get("queries") or [{}] +q = queries[0] +offset, count = max(0, int(q.get("offset", 0))), max(1, int(q.get("count", 500))) +pattern = str(q.get("filter") or "").lower() +if pattern and not any(ch in pattern for ch in "*?["): pattern = "*" + pattern + "*" +rows = [] +for fn in db.functions.get_all(): + name = db.functions.get_name(fn) or f"sub_{int(fn.start_ea):X}" + if pattern and not fnmatch.fnmatchcase(name.lower(), pattern): continue + rows.append({"addr": hex(int(fn.start_ea)), "name": name, + "size": int(fn.end_ea) - int(fn.start_ea)}) +page = rows[offset:offset + count] +result = {"result": [{"data": page, "next_offset": offset + len(page), "total": len(rows)}]} +result +''', + "disasm": r''' +ea = int(str(a["addr"]), 16) +fn = db.functions.get_at(ea) +if fn is None: + result = {"instructions": [], "total_instructions": 0, "instruction_count": 0} +else: + instructions = list(db.functions.get_instructions(fn)) + limit = max(1, int(a.get("max_instructions", len(instructions) or 1))) + rows = [{"addr": hex(int(insn.ea)), "instruction": db.instructions.get_disassembly(insn)} + for insn in instructions[:limit]] + result = {"instructions": rows, "total_instructions": len(instructions), + "instruction_count": len(instructions)} +result +''', + "file_regions": r''' +import idaapi +rows = [] +for seg in db.segments.get_all(): + try: file_off = int(idaapi.get_fileregion_offset(seg.start_ea)) + except Exception: file_off = -1 + if file_off < 0 or file_off >= (1 << 48): file_off = -1 + rows.append({"start": hex(int(seg.start_ea)), "end": hex(int(seg.end_ea)), + "file_off": file_off, "name": db.segments.get_name(seg) or ""}) +result = {"regions": rows} +result +''', + "read_raw": r''' +import ida_bytes +ea, size = int(str(a["addr"]), 16), max(0, int(a["size"])) +raw = ida_bytes.get_bytes(ea, size) or b"" +raw = raw[:size] + b"\xff" * max(0, size - len(raw)) +data = bytearray(raw) +for index, value in enumerate(data): + if value == 0xFF and not ida_bytes.is_loaded(ea + index): data[index] = 0 +result = {"addr": a["addr"], "hex": bytes(data).hex(), "n": len(data)} +result +''', + "get_bytes": r''' +rows = [] +for region in a.get("regions", []): + ea, size = int(str(region["addr"]), 16), int(region["size"]) + raw = db.bytes.get_bytes_at(ea, size) or b"" + rows.append({"addr": region["addr"], "data": " ".join(f"{b:02x}" for b in raw)}) +result = {"result": rows} +result +''', + "search_structs": r''' +needle = str(a.get("filter") or "").lower() +rows = [] +for tif in db.types.get_all(): + name = tif.get_type_name() or "" + if not name or needle not in name.lower() or not tif.is_udt(): continue + members = list(db.types.get_udt_members(tif)) + rows.append({"name": name, "size": int(tif.get_size()), "is_union": bool(tif.is_union()), + "cardinality": len(members), "ordinal": int(tif.get_ordinal())}) +result = {"result": rows} +result +''', + "type_inspect": r''' +rows = [] +for query in a.get("queries", []): + name = str(query.get("name") or "") + tif = db.types.get_by_name(name) + if tif is None: + rows.append({"name": name, "error": "type not found"}); continue + members = [{"name": m.name, "type": m.type.dstr() or str(m.type), + "offset": int(m.offset), "size": int(m.size)} + for m in db.types.get_udt_members(tif)] if tif.is_udt() else [] + rows.append({"name": name, "size": int(tif.get_size()), "is_union": bool(tif.is_union()), + "members": members}) +result = {"result": rows} +result +''', + "declare_type": r''' +import ida_typeinf +decls = a.get("decls", "") +if isinstance(decls, str): decls = [decls] +rows = [] +for declaration in decls: + try: + errors = int(db.types.parse_declarations(ida_typeinf.get_idati(), declaration)) + rows.append({"ok": errors == 0, **({} if errors == 0 else {"error": f"{errors} parse error(s)"})}) + except Exception as exc: + rows.append({"ok": False, "error": str(exc)}) +result = {"result": rows} +result +''', + "del_type": r''' +import ida_typeinf +name = str(a["name"]) +ok = bool(ida_typeinf.del_named_type(ida_typeinf.get_idati(), name, ida_typeinf.NTF_TYPE)) +result = {"name": name, "deleted": ok, **({} if ok else {"error": f"Type {name!r} not found or could not be deleted"})} +result +''', + "func_types": r''' +import ida_typeinf +ea = int(str(a["addr"]), 16) +fn = db.functions.get_at(ea) +if fn is None: + result = {"addr": a["addr"], "error": "no function at address"} +else: + pseudo = db.pseudocode.decompile(fn) + name = db.functions.get_name(fn) or "" + tif = pseudo.get_func_type() + try: prototype = ida_typeinf.print_tinfo("", 0, 0, ida_typeinf.PRTYPE_1LINE, tif, name, "") if tif else "" + except Exception: prototype = tif.dstr() if tif else "" + lvars = [{"name": var.name, "type": var.type_info.dstr() if var.type_info else "", + "is_arg": bool(var.is_arg)} for var in pseudo.local_variables] + result = {"addr": hex(int(fn.start_ea)), "name": name, + "prototype": (prototype or "").strip(), "lvars": lvars} +result +''', + "set_lvar_type": r''' +import ida_typeinf +ea, variable, declaration = int(str(a["addr"]), 16), str(a["variable"]), str(a["type"]) +fn = db.functions.get_at(ea) +if fn is None: + result = {"error": "no function at address"} +else: + pseudo = db.pseudocode.decompile(fn) + var = pseudo.find_local_variable(variable) + if var is None: + result = {"error": f"local variable {variable!r} not found"} + else: + try: + tif = db.types.parse_one_declaration(ida_typeinf.get_idati(), declaration) + accepted = bool(var.set_type(tif)) + saved = bool(pseudo.save_local_variable_info(var, save_type=True)) if accepted else False + result = {"addr": hex(int(fn.start_ea)), "variable": variable, + "type": declaration, "ok": accepted and saved} + except Exception as exc: + result = {"error": f"bad type {declaration!r}: {exc}"} +result +''', + "set_type": r''' +from ida_domain.types import TypeApplyFlags +rows = [] +for edit in a.get("edits", []): + ea = int(str(edit["addr"]), 16) + declaration = str(edit.get("signature") or edit.get("type") or "") + try: + ok = bool(db.types.apply_declaration_at(ea, declaration, TypeApplyFlags.DEFINITE)) + rows.append({"addr": hex(ea), "ok": ok, **({} if ok else {"error": "IDA rejected the type"})}) + except Exception as exc: + rows.append({"addr": hex(ea), "ok": False, "error": str(exc)}) +result = {"result": rows} +result +''', + "data_type": r''' +ea = int(str(a["addr"]), 16) +try: + tif = db.types.get_at(ea) + fn = db.functions.get_at(ea) + result = {"addr": hex(ea), "name": db.names.get_at(ea) or "", + "type": tif.dstr() if tif else "", "size": int(db.heads.size(ea)) if db.heads.is_head(ea) else 0, + "is_func": bool(fn)} +except Exception as exc: + result = {"addr": hex(ea), "error": str(exc)} +result +''', + "force_recompile": r''' +import ida_hexrays +rows = [] +for item in a.get("items", []): + ea = int(str(item["addr"]), 16) + ida_hexrays.mark_cfunc_dirty(ea, False) + rows.append({"addr": hex(ea), "ok": True}) +result = {"result": rows} +result +''', + "undefine": r''' +import ida_bytes +rows = [] +for item in a.get("items", []): + ea = int(str(item["addr"]), 16) + size = max(1, int(item.get("size") or ida_bytes.get_item_size(ea) or 1)) + ok = bool(ida_bytes.del_items(ea, ida_bytes.DELIT_SIMPLE, size)) + rows.append({"addr": hex(ea), "ok": ok, **({} if ok else {"error": "delete items failed"})}) +result = {"result": rows} +result +''', + "define_code": r''' +import ida_ua +rows = [] +for item in a.get("items", []): + ea = int(str(item["addr"]), 16); size = int(ida_ua.create_insn(ea)) + rows.append({"addr": hex(ea), "ok": size > 0, "size": size, + **({} if size > 0 else {"error": "instruction did not decode"})}) +result = {"result": rows} +result +''', + "define_func": r''' +rows = [] +for item in a.get("items", []): + ea = int(str(item["addr"]), 16); ok = bool(db.functions.create(ea)) + rows.append({"addr": hex(ea), "ok": ok, **({} if ok else {"error": "IDA refused the function"})}) +result = {"result": rows} +result +''', + "make_data": r''' +import ida_bytes, ida_idaapi, ida_typeinf +from ida_domain.types import TypeApplyFlags +rows = [] +for item in a.get("items", []): + ea, declaration = int(str(item["addr"]), 16), str(item["type"]) + try: + tif = db.types.parse_one_declaration(ida_typeinf.get_idati(), declaration) + size = max(1, int(tif.get_size())) + saved_names = [(addr, name) for addr, name in db.names.get_all() + if ea <= int(addr) < ea + size] + ida_bytes.del_items(ea, ida_bytes.DELIT_EXPAND | ida_bytes.DELIT_DELNAMES, + max(size, int(ida_bytes.get_item_size(ea) or 1))) + created = bool(ida_bytes.create_data(ea, ida_bytes.FF_BYTE, size, ida_idaapi.BADADDR)) + ok = created and bool(db.types.apply_at(tif, ea, TypeApplyFlags.DEFINITE)) + for address, name in saved_names: + db.names.set_name(int(address), name) + if ok and item.get("name"): ok = bool(db.names.set_name(ea, str(item["name"]))) + rows.append({"addr": hex(ea), "ok": ok, "size": size, + **({} if ok else {"error": "IDA rejected the data type"})}) + except Exception as exc: + rows.append({"addr": hex(ea), "ok": False, "error": str(exc)}) +result = {"result": rows} +result +''', + "make_string": r''' +from ida_domain.strings import StringType +ea, length = int(str(a["addr"]), 16), max(0, int(a.get("length", 0))) +kind = {"c": StringType.C, "c16": StringType.C_16, "c32": StringType.C_32, + "pascal": StringType.PASCAL}.get(str(a.get("kind", "c")).lower(), StringType.C) +import ida_bytes +try: + ida_bytes.del_items(ea, ida_bytes.DELIT_SIMPLE, length if length > 0 else 1) +except Exception: + pass +try: + ok = bool(db.bytes.create_string_at(ea, length or None, kind)) + text = db.bytes.get_string_at(ea) or "" if ok else "" + result = {"addr": hex(ea), "ok": ok, "size": int(db.heads.size(ea)) if ok else 0, "text": text} +except Exception as exc: + result = {"addr": hex(ea), "ok": False, "error": str(exc)} +result +''', + "list_strings": r''' +from ida_domain.strings import StringListConfig +offset, count, min_len = max(0, int(a.get("offset", 0))), max(1, int(a.get("count", 2000))), max(1, int(a.get("min_len", 4))) +if offset == 0 or a.get("refresh"): + from ida_domain.strings import StringType + db.strings.rebuild(StringListConfig(string_types=list(StringType), min_len=min_len, + only_ascii_7bit=False)) +items = list(db.strings.get_all()) +page = items[offset:offset + count] +rows = [] +for item in page: + try: text = str(item) + except Exception: text = item.contents.decode("utf-8", "replace") if item.contents else "" + rows.append({"addr": hex(int(item.address)), "text": text, "len": int(item.length), "type": item.type.name}) +result = {"strings": rows, "total": len(items), "next_offset": offset + len(rows)} +result +''', + "list_linkage": r''' +imports = [{"addr": hex(int(item.address)), "name": item.name, "module": item.module_name} + for item in db.imports.get_all_imports() if item.name] +exports = [{"addr": hex(int(item.address)), "name": item.name, "ordinal": int(item.ordinal)} + for item in db.entries.get_all() if item.name] +result = {"imports": imports, "exports": exports, + "n_imports": len(imports), "n_exports": len(exports)} +result +''', + "lookup_funcs": r''' +rows = [] +for query in a.get("queries", []): + raw = str(query) + try: ea = int(raw, 16) + except ValueError: + fn = db.functions.get_by_name(raw); ea = int(fn.start_ea) if fn else None + else: fn = db.functions.get_at(ea) + if fn is None: + rows.append({"query": raw, "fn": None}) + else: + rows.append({"query": raw, "fn": {"addr": hex(int(fn.start_ea)), + "name": db.functions.get_name(fn) or f"sub_{int(fn.start_ea):X}", + "size": int(fn.end_ea) - int(fn.start_ea)}}) +result = {"result": rows} +result +''', + "resolve_names": r''' +import ida_idaapi, ida_name +rows = [] +for query in a.get("queries", []): + name = str(query).strip(); ea = ida_name.get_name_ea(ida_idaapi.BADADDR, name) + rows.append({"query": name, "ea": hex(int(ea)) if ea != ida_idaapi.BADADDR else None}) +result = {"result": rows} +result +''', + "xref_types": r''' +queries = a.get("queries") or [] +all_results = [] +for query in queries: + ea, direction = int(str(query["addr"]), 16), str(query.get("direction", "both")) + refs = [] + if direction in ("to", "both"): refs += list(db.xrefs.to_ea(ea)) + if direction in ("from", "both"): refs += list(db.xrefs.from_ea(ea)) + rows, seen = [], set() + for ref in refs: + key = (int(ref.from_ea), int(ref.to_ea), int(ref.type)) + if query.get("dedup") and key in seen: continue + seen.add(key) + fn = db.functions.get_at(int(ref.from_ea)) + kind = ("call" if ref.is_call else "jump" if ref.is_jump else "flow" if ref.is_flow + else "read" if ref.is_read else "write" if ref.is_write else ref.type.name.lower()) + row = {"from": hex(int(ref.from_ea)), "to": hex(int(ref.to_ea)), + "type": "code" if ref.is_code else "data", "kind": kind} + if query.get("include_fn") and fn is not None: + row["fn"] = {"addr": hex(int(fn.start_ea)), "name": db.functions.get_name(fn) or ""} + rows.append(row) + if len(rows) >= int(query.get("count", 2000)): break + all_results.append({"data": rows}) +result = {"result": all_results} +result +''', + "xref_query": r''' +queries = a.get("queries") or [] +all_results = [] +for query in queries: + ea, direction = int(str(query["addr"]), 16), str(query.get("direction", "both")) + refs = [] + if direction in ("to", "both"): refs += list(db.xrefs.to_ea(ea)) + if direction in ("from", "both"): refs += list(db.xrefs.from_ea(ea)) + rows = [] + for ref in refs[:int(query.get("count", 2000))]: + fn = db.functions.get_at(int(ref.from_ea)) + row = {"from": hex(int(ref.from_ea)), "to": hex(int(ref.to_ea)), + "type": "code" if ref.is_code else "data"} + if query.get("include_fn") and fn is not None: + row["fn"] = {"addr": hex(int(fn.start_ea)), "name": db.functions.get_name(fn) or ""} + rows.append(row) + all_results.append({"data": rows}) +result = {"result": all_results} +result +''', + "set_comments": r''' +rows = [] +for item in a.get("items", []): + ea, text = int(str(item["addr"]), 16), str(item.get("comment") or "") + try: + if text: ok = bool(db.comments.set_at(ea, text)) + else: db.comments.delete_at(ea); ok = True + rows.append({"addr": hex(ea), "ok": ok}) + except Exception as exc: + rows.append({"addr": hex(ea), "ok": False, "error": str(exc)}) +result = {"result": rows} +result +''', + "rename": r''' +import ida_idaapi, ida_name, ida_typeinf +batch = a.get("batch") or {} +out = {}; ok_count = failed = 0 +for category, edit in batch.items(): + try: + if category == "func": + ea, new = int(str(edit["addr"]), 16), str(edit["name"]) + fn = db.functions.get_at(ea); ok = bool(fn and db.functions.set_name(fn, new)) + elif category == "data": + new = str(edit.get("new") or "") + if edit.get("addr") is not None: ea = int(str(edit["addr"]), 16) + else: ea = int(ida_name.get_name_ea(ida_idaapi.BADADDR, str(edit.get("old") or ""))) + ok = bool(db.names.set_name(ea, new)) + elif category in ("local", "stack"): + ea, old, new = int(str(edit["func_addr"]), 16), str(edit["old"]), str(edit["new"]) + pseudo = db.pseudocode.decompile(ea); var = pseudo.find_local_variable(old) + if var is None: ok = False + else: + var.set_user_name(new) + ok = bool(pseudo.save_local_variable_info(var, save_name=True)) + else: + raise ValueError(f"unsupported rename category: {category}") + row = {"ok": ok, **({} if ok else {"error": "IDA rejected the name"})} + except Exception as exc: + row = {"ok": False, "error": str(exc)} + out[category] = [row] + if row["ok"]: ok_count += 1 + else: failed += 1 +out["summary"] = {"ok": ok_count, "failed": failed} +result = out +result +''', +} + + +_OPERATIONS["decompile"] = _DECOMP_MAP_HELPER + r''' +ea = int(str(a["addr"]), 16) +fn = db.functions.get_at(ea) +if fn is None: + result = {"error": f"no function at {ea:#x}"} +else: + pseudo = db.pseudocode.decompile(fn) + mapping = line_map(pseudo.raw_cfunc) + plain = pseudo.to_text() + marked = [line + (f" /*0x{eas[0]:X}*/" if eas else "") + for line, eas in zip(plain, mapping)] + import ida_name + refs, seen = [], set() + for expr in pseudo.find_objects(): + target = int(expr.obj_ea) + if target in seen or not (db.is_valid_ea(target) or db.is_private_ea(target)): continue + seen.add(target) + name = expr.obj_name or ida_name.get_name(target) or "" + try: string = db.bytes.get_string_at(target) if db.is_valid_ea(target) else None + except Exception: string = None + refs.append({"addr": hex(target), "name": name, "string": string}) + result = {"addr": hex(int(fn.start_ea)), "code": "\n".join(marked), "refs": refs} +result +''' + +_OPERATIONS["decomp_map"] = _DECOMP_MAP_HELPER + r''' +ea = int(str(a["addr"]), 16) +fn = db.functions.get_at(ea) +if fn is None: + result = {"error": f"no function at {ea:#x}"} +else: + pseudo = db.pseudocode.decompile(fn) + mapping = line_map(pseudo.raw_cfunc) + result = {"addr": hex(int(fn.start_ea)), + "lines": [{"ea": hex(eas[0]) if eas else None, + "eas": [hex(item) for item in eas]} for eas in mapping]} +result +''' + +_OPERATIONS["define_code_run"] = r''' +import ida_bytes, ida_idp, ida_segment, ida_ua, idaapi +ea, limit = int(str(a["addr"]), 16), max(1, min(int(a.get("limit", 20000)), 200000)) +seg = ida_segment.getseg(ea) +if seg is None: + result = {"addr": a["addr"], "error": "no segment", "count": 0} +else: + start, count, stopped, hi = ea, 0, "limit", int(seg.end_ea) + while count < limit: + if ea >= hi: stopped = "segment"; break + flags = ida_bytes.get_flags(ea) + if ida_bytes.is_code(flags) or ida_bytes.is_data(flags): stopped = "defined"; break + size = int(ida_ua.create_insn(ea)) + if size <= 0: stopped = "undecodable"; break + count += 1 + insn = ida_ua.insn_t() + if ida_ua.decode_insn(insn, ea) > 0: + try: is_ret = bool(ida_idp.is_ret_insn(insn)) + except Exception: is_ret = False + if is_ret or (insn.get_canon_feature() & idaapi.CF_STOP): + ea += size; stopped = "flow"; break + ea += size + result = {"start": hex(start), "end": hex(ea), "count": count, "stopped": stopped} +result +''' + +_OPERATIONS["define_func_run"] = r''' +import ida_bytes, ida_funcs, ida_segment +ea = int(str(a["addr"]), 16) +fn = db.functions.get_at(ea) +if fn is not None and int(fn.start_ea) == ea: + result = {"addr": hex(ea), "ok": True, "start": hex(ea), "end": hex(int(fn.end_ea)), "how": "existed"} +else: + automatic = bool(db.functions.create(ea)) + if not automatic: + seg = db.segments.get_at(ea); end = ea; hi = int(seg.end_ea) if seg else ea + while end < hi and ida_bytes.is_code(ida_bytes.get_flags(end)): + nxt = int(ida_bytes.get_item_end(end)) + if nxt <= end: break + end = nxt + ok = bool(end > ea and ida_funcs.add_func(ea, end)) + else: ok = True + fn = db.functions.get_at(ea) + result = ({"addr": hex(ea), "ok": True, "start": hex(int(fn.start_ea)), + "end": hex(int(fn.end_ea)), "how": "auto" if automatic else "explicit-end"} + if ok and fn is not None else + {"addr": hex(ea), "ok": False, "error": f"IDA refused a function at {ea:#x}"}) +result +''' + +_OPERATIONS["set_thumb"] = r''' +import ida_bytes, ida_ida, ida_idp, ida_segment, ida_segregs +ea = int(str(a["addr"]), 16); treg = ida_idp.str2reg("T") +seg = ida_segment.getseg(ea) +if treg is None or treg < 0: + result = {"addr": hex(ea), "error": "no T register (not an ARM database)"} +elif seg is None: + result = {"addr": hex(ea), "error": "no segment"} +else: + current = ida_segregs.get_sreg(ea, treg) + current = 0 if current in (None, 0xFFFFFFFF, -1) else int(current) + want = {"on": 1, "off": 0}.get(str(a.get("mode", "toggle")).lower(), 0 if current else 1) + changed = False + if want and seg.bitness != 1: + ida_segment.set_segm_addressing(seg, 1); changed = True + size = max(int(ida_bytes.get_item_size(ea)), 2) + ida_bytes.del_items(ea, 0, size) + ok = bool(ida_segregs.split_sreg_range(ea, treg, want, ida_segregs.SR_user)) + now = ida_segregs.get_sreg(ea, treg) + result = {"addr": hex(ea), "thumb": bool(now), "was": bool(current), "ok": ok, + "bitness": ida_segment.getseg(ea).bitness, "forced_32bit": changed, + "db_64bit": bool(ida_ida.inf_get_app_bitness() == 64 and want)} +result +''' + +_OPERATIONS["thumb_scan"] = r''' +import ida_bytes, ida_funcs, ida_idp, ida_segment, ida_segregs, ida_ua +lo, hi = int(str(a["start"]), 16), int(str(a["end"]), 16) +apply, limit = bool(a.get("apply", True)), int(a.get("limit", 512)) +treg = ida_idp.str2reg("T"); found = []; applied = 0; cursor = lo +while cursor + 4 <= hi and len(found) < limit: + at = cursor; value = int(ida_bytes.get_dword(cursor)); cursor += 4 + if not value & 1: continue + target = value & ~1; seg = ida_segment.getseg(target) + if seg is None or not (seg.perm & ida_segment.SEGPERM_EXEC or seg.perm == 0): continue + flags = ida_bytes.get_flags(target) + if ida_bytes.is_data(flags): continue + item = {"at": hex(at), "value": hex(value), "target": hex(target), + "was_code": bool(ida_bytes.is_code(flags))}; found.append(item) + if not apply: continue + if treg is not None and treg >= 0: ida_segregs.split_sreg_range(target, treg, 1, ida_segregs.SR_user) + if not ida_bytes.is_code(ida_bytes.get_flags(target)): + ida_bytes.del_items(target, 0, 2) + if ida_ua.create_insn(target) <= 0: item["decoded"] = False; continue + item["decoded"] = True; item["function"] = bool(db.functions.get_at(target) or db.functions.create(target)); applied += 1 +result = {"start": hex(lo), "end": hex(hi), "found": found, "applied": applied, "n": len(found)} +result +''' + +_OPERATIONS["decomp_error"] = r''' +import ida_hexrays, ida_ida +ea = int(str(a["addr"]), 16); fn = db.functions.get_at(ea) +result = {"addr": hex(ea), "bitness": ida_ida.inf_get_app_bitness()} +if fn is None: + result["reason"] = "no function here" +else: + try: + failure = ida_hexrays.hexrays_failure_t(); cfunc = ida_hexrays.decompile_func(fn, failure) + if cfunc is not None: result["reason"] = "" + else: + result.update({"reason": failure.desc() or f"error {failure.code}", + "code": int(failure.code), "errea": hex(int(failure.errea))}) + except Exception as exc: result["reason"] = f"{type(exc).__name__}: {exc}" +result +''' + + +class CodeModeClient: + """A leased GUI/idalib database accessed through ``ida_codemode``.""" + + def __init__( + self, + binary_path: str, + *, + ttl: int = 0, + load_args: str = "", + processor: str | None = None, + loading_address: int | None = None, + file_type: str | None = None, + output_database: str | None = None, + spawn: bool = True, + new_database: bool = False, + ) -> None: + del ttl # managed-worker lifetime is lease-based, not idle-TTL based + self._path = os.path.abspath(os.path.expanduser(binary_path)) + parsed_processor, parsed_address, parsed_file_type = _parse_load_args(load_args) + self._processor = processor or parsed_processor + self._loading_address = loading_address if loading_address is not None else parsed_address + self._file_type = file_type or parsed_file_type + self._output_database = output_database + self._spawn = spawn + self._new_database = new_database + self._handle: DatabaseHandle | None = None + self._last_entry: RegistryEntry | None = None + self._connect_lock = threading.Lock() + + def connect(self, timeout: float = 1800.0, progress=None) -> "CodeModeClient": + _require_codemode() + with self._connect_lock: + if self._handle is not None and self._handle.connected: + return self + if progress: + progress(f"discovering Code Mode database for {os.path.basename(self._path)}…") + try: + # A Ctrl+L reload releases its current managed-worker lease, but + # that worker remains registered during Code Mode's final-lease + # grace period. Retry only that known handoff window. A GUI or + # another long-lived client remains busy and yields a clear + # failure rather than being modified underneath its owner. + deadline = time.monotonic() + min(timeout, 60.0) + while True: + try: + handle = DatabaseHandle.open( + self._path, + spawn=self._spawn, + timeout=max(0.1, timeout), + output_database=self._output_database, + processor=self._processor, + loading_address=self._loading_address, + file_type=self._file_type, + new_database=self._new_database, + ) + break + except IdbBusy: + if not self._new_database or time.monotonic() >= deadline: + raise + if progress: + progress("waiting for the previous Code Mode lease to close…") + # Remember the record before managed shutdown withdraws + # its JSON. The lifetime lock remains held until IDA has + # actually closed the IDB; waiting on it avoids racing a + # replacement worker into the old process's file lock. + expected = canonical_path( + self._output_database or expected_idb_path(self._path) + ) + owners = [item.entry for item in scan_instances(timeout=0.5) + if item.entry.idb_key == idb_key(expected)] + if owners: + self._wait_for_entry_release( + owners[0], max(0.0, deadline - time.monotonic()) + ) + else: + time.sleep(0.2) + if progress: + backend = handle.entry.backend + progress(f"attached to {backend} database; waiting for auto-analysis…") + handle.wait_autoanalysis(timeout=timeout) + except Exception as exc: # normalize the dependency's transport errors + raise self._connection_error(exc) from exc + self._handle = handle + self._last_entry = handle.entry + return self + + @staticmethod + def _connection_error(exc: BaseException) -> IDAConnectionError: + return IDAConnectionError(str(exc) or type(exc).__name__) + + @property + def connected(self) -> bool: + return self._handle is not None and self._handle.connected + + @property + def pid(self) -> int | None: + return self._handle.entry.pid if self._handle is not None else None + + @property + def backend(self) -> str | None: + return self._handle.entry.backend if self._handle is not None else None + + def execute_python(self, code: str, *, timeout: float | None = None) -> Any: + if not self.connected: + self.connect() + handle = self._handle + if handle is None: + raise IDAConnectionError("Code Mode database is not connected") + try: + response = handle.execute_python(code, timeout=timeout) + except RemoteError as exc: + details = exc.details or {} + message = str(exc) + if details.get("traceback"): + message += f"\n{details['traceback']}" + if exc.code == "operation_timeout": + raise IDATimeoutError(message) from exc + raise IDAToolError("execute_python", message) from exc + except (InstanceDisconnectedError, ClientError) as exc: + raise self._connection_error(exc) from exc + if not isinstance(response, dict) or "result" not in response: + raise IDAToolError("execute_python", "Code Mode returned an invalid execution result") + return response["result"] + + def invoke(self, operation: str, *, timeout: float | None = None, **args) -> Any: + """Execute one TUI domain operation through Code Mode.""" + if operation in ("idb_save", "save"): + return self.save_database() + if operation in ("server_health", "ping", "health", "state"): + return self.health() + body = _HEADS if operation == "heads" else _OPERATIONS.get(operation) + if body is None: + raise IDAToolError(operation, f"unknown ida-tui Code Mode operation: {operation}") + try: + return self.execute_python(_script(args, body), timeout=timeout) + except IDAToolError as exc: + if exc.tool == "execute_python": + raise IDAToolError(operation, exc.message) from exc + raise + + # Temporary source compatibility for external drivers/tests that used the + # old WorkerClient. Application code uses the accurately named invoke(). + call = invoke + + def save_database(self) -> dict[str, Any]: + if not self.connected: + self.connect() + handle = self._handle + if handle is None: + raise IDAConnectionError("Code Mode database is not connected") + try: + return handle.save_database() + except RemoteError as exc: + raise IDAToolError("save_database", str(exc)) from exc + except (InstanceDisconnectedError, ClientError) as exc: + raise self._connection_error(exc) from exc + + def health(self) -> dict[str, Any]: + if not self.connected: + self.connect() + assert self._handle is not None + entry = self._handle.entry + module = os.path.basename(entry.exe_path or entry.idb_path or self._path) + return { + "ok": self._handle.connected, + "module": module, + "backend": entry.backend, + "record_id": entry.record_id, + "input_path": entry.exe_path, + "idb_path": entry.idb_path, + } + + def keepalive(self, interval: float = 120.0) -> _NoopKeepAlive: + del interval + return _NoopKeepAlive() + + def resolve_db(self) -> str: + if not self.connected: + self.connect() + assert self._handle is not None + return self._handle.entry.record_id + + def set_db(self, db: str | None) -> None: + del db # one handle is permanently bound to one registered database + + def list_sessions(self) -> list[Session]: + if not self.connected: + self.connect() + assert self._handle is not None + entry = self._handle.entry + path = entry.exe_path or entry.idb_path or self._path + return [Session(session_id=entry.record_id, filename=os.path.basename(path), + input_path=path, is_active=True)] + + def close(self, grace: float = 0.0) -> None: + del grace + with self._connect_lock: + handle, self._handle = self._handle, None + if handle is not None: + self._last_entry = handle.entry + handle.close() # release our lease; never close a GUI/other client's DB + + @staticmethod + def _wait_for_entry_release(entry: "RegistryEntry", timeout: float) -> bool: + _require_codemode() + path = REGISTRY_DIR / f"{entry.record_id}.lock" + deadline = time.monotonic() + max(0.0, timeout) + while True: + lock = FileLock(path) + try: + if lock.try_acquire(): + return True + except OSError: + pass + finally: + lock.close() + if time.monotonic() >= deadline: + return False + time.sleep(min(0.1, deadline - time.monotonic())) + + def wait_released(self, timeout: float = 45.0) -> bool: + """Wait until a managed instance releases its lifetime lock. + + Normal application shutdown must not wait: another client may retain the + worker. This is an explicit test/maintenance helper for deleting a + temporary IDB safely after this client closes. GUI instances return + ``False`` immediately because clients never own their lifetime. + """ + entry = self._last_entry + if entry is None or entry.backend != "idalib": + return False + return self._wait_for_entry_release(entry, timeout) + + def __enter__(self) -> "CodeModeClient": + return self.connect() + + def __exit__(self, *exc) -> None: + self.close() diff --git a/idatui/domain.py b/idatui/domain.py index aedade9..6fe73c2 100644 --- a/idatui/domain.py +++ b/idatui/domain.py @@ -1,19 +1,15 @@ -"""Domain / paging layer: address-centric models over the raw MCP client. +"""Domain / paging layer: address-centric models over IDA Code Mode. This is where the "millions of lines" problem is solved, so the TUI widgets only ever see a viewport-sized slice. Every hard-won constraint from ``docs/PAGING_FINDINGS.md`` is encoded here: -* Per-call caps are silent (over the cap the server returns 10, not a clamp), so - we clamp page sizes ourselves: ``LIST_PAGE`` / ``DISASM_BLOCK`` <= the caps. -* ``next_offset`` is unreliable; we paginate by advancing ``len(data)``. -* ``disasm offset=N`` is O(N) with no resumable cursor, so windowed disassembly - is **block-cached** (revisits are free) and **prefetches** the next block on a - background thread (the client is concurrency-safe). -* ``include_total`` scans the whole function (~200ms on monsters); totals are - fetched once and cached. -* ``decompile`` can hard-fail on huge functions as a *soft* error (``code`` is - null); that is surfaced as data, not an exception. +* Page sizes remain bounded so remote execution returns viewport-scale JSON. +* Pagination advances by the number of rows actually returned. +* Deep head walks are block-cached (revisits are free) and neighboring blocks + prefetch through the thread-safe Code Mode client. +* Expensive function totals are fetched once and cached. +* Decompilation failures are surfaced as data, not application crashes. Everything here is synchronous and thread-safe. The TUI runs these calls from Textual worker threads; the internal prefetch pool is separate and small. @@ -22,10 +18,8 @@ Textual worker threads; the internal prefetch pool is separate and small. from __future__ import annotations import bisect -import json import re import threading -import urllib.request from concurrent.futures import ThreadPoolExecutor from collections.abc import Sequence from dataclasses import dataclass, field, replace @@ -36,7 +30,7 @@ from . import diag from .errors import IDAToolError if TYPE_CHECKING: # type hint only - from .worker_client import WorkerClient # noqa: F401 + from .codemode_client import CodeModeClient # Clamps derived from measured caps (list ~700, disasm ~500). Margin included. LIST_PAGE = 500 @@ -66,7 +60,7 @@ class Func: def from_raw(cls, d: dict) -> "Func": addr = _as_int(d["addr"]) name = d.get("name") - # An unnamed function (server returns null/empty) must still have a + # An unnamed function must still have a # usable string name — synthesize IDA's sub_ADDR so every consumer # (palette, sort, rename prefill) can treat name as a str. if not name: @@ -93,7 +87,7 @@ class Line: class Head(NamedTuple): - """One flat-listing item (from the ``heads`` server tool): a code + """One flat-listing item (from the Code Mode ``heads`` operation): a code instruction, a data item, or an undefined byte run. A ``NamedTuple`` rather than a dataclass because this is by far the @@ -114,7 +108,7 @@ class Head(NamedTuple): name: str | None = None raw: bytes | None = None # opcode/item bytes (filled in for code by the model) #: [(kind, text)] from IDA's own colour tags — mnem/reg/num/name/str/punct/… - #: None when the worker didn't provide them (older worker, or the spans + #: None when Code Mode didn't provide them (or the spans #: disagreed with the plain text, in which case the text wins). #: #: Held exactly as it came off the wire, and **read-only**. The worker @@ -301,7 +295,7 @@ class FunctionIndex: """A lazily-paginated, cached view of the function list. Loads pages of ``LIST_PAGE`` on demand, advancing by ``len(data)`` (never by - ``next_offset``). A single index instance corresponds to one server-side + ``next_offset``). A single index instance corresponds to one remote ``filter`` glob (``None`` = all functions). """ @@ -321,7 +315,7 @@ class FunctionIndex: query: dict = {"offset": offset, "count": LIST_PAGE} if self.filter: query["filter"] = self.filter - data = _query_data(self._prog.client.call("list_funcs", queries=[query])) + data = _query_data(self._prog.client.invoke("list_funcs", queries=[query])) added = 0 with self._lock: for d in data: @@ -431,7 +425,7 @@ 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.call( + payload = self._prog.client.invoke( "disasm", addr=hex(self.ea), max_instructions=1, include_total=True ) total = payload.get("total_instructions") @@ -492,7 +486,7 @@ class DisasmModel: # 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.call( + payload = self._prog.client.invoke( "heads", addr=hex(self.ea), offset=b * self.BLOCK, count=self.BLOCK + 1, **self._end_kw(), ) @@ -633,15 +627,15 @@ class ListingModel: """A flat, IDA-style disassembly *listing* over one segment: code, data and undefined heads interleaved, unlike ``DisasmModel`` (one function, code only). - Backed by the injected ``heads`` server tool, which walks item heads and - renders each via ``generate_disasm_line``. The segment is walked lazily in + Backed by the Code Mode adapter's ``heads`` operation, which walks item heads + and renders each via ``generate_disasm_line``. The segment is walked lazily in forward pages (``FunctionIndex`` style); line index == position in the walked head list. Random access to an address is O(distance-from-seg-start) the first time (then cached) — the same tradeoff as ``disasm offset=N``. Grows on demand as the viewport scrolls. Synchronous + thread-safe. """ - PAGE = 500 # heads per server call (well under the tool's 2000 cap) + PAGE = 500 # viewport-scale heads per Code Mode execution def __init__(self, program: "Program", seg_start: int, seg_end: int, name: str | None = None): @@ -754,7 +748,7 @@ class ListingModel: if self._done or self._next is None: return 0 frm = self._next - payload = self._prog.client.call( + payload = self._prog.client.invoke( "heads", addr=hex(frm), count=self.PAGE, annotate=True) rows = payload.get("heads", []) if isinstance(payload, dict) else [] cur = payload.get("cursor", {}) if isinstance(payload, dict) else {} @@ -1019,7 +1013,7 @@ class ListingModel: # the expectation rather than asking first means a page that HAS changed # still costs one round trip. try: - payload = self._prog.client.call( + payload = self._prog.client.invoke( "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 @@ -1233,7 +1227,7 @@ class HexModel: class Program: """The bound analysis session: models, caches, and a small prefetch pool.""" - def __init__(self, client: "WorkerClient", prefetch_workers: int = 2): + def __init__(self, client: "CodeModeClient", prefetch_workers: int = 2): self.client = client self._pool = ThreadPoolExecutor( max_workers=prefetch_workers, thread_name_prefix="idatui-prefetch" @@ -1257,7 +1251,7 @@ class Program: self._sections: list[tuple[int, int, str]] | None = None self._fileregions: list[tuple[int, int, int]] | None = None self._hexmodel: "HexModel | None" = None - self._no_read_raw = False # set if the server lacks the read_raw tool + self._no_read_raw = False # compatibility fallback for alternate clients self._lock = threading.Lock() # -- prefetch plumbing ------------------------------------------------- # @@ -1284,17 +1278,14 @@ class Program: """Sorted raw segment map [(start, end, file_off, name)] — the single source for sections()/file_regions()/image_range. Cached. - Uses the injected ``file_regions`` tool (a plain segment walk, ~ms). - This deliberately AVOIDS ``survey_binary``, which also computes function - counts / strings / stats and takes *seconds* on a large IDB (it was the - cause of the multi-second hex-pane open). Falls back to survey_binary - only if the injected tool is missing. + Uses the Code Mode adapter's ``file_regions`` operation (a plain segment + walk, ~ms), avoiding broad binary surveys on the hex-pane open path. """ if self._segments_cache is not None: return self._segments_cache segs: list[tuple[int, int, int, str]] = [] try: - r = self.client.call("file_regions") + r = self.client.invoke("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"]), @@ -1303,7 +1294,7 @@ class Program: segs = [] if not segs: # older server without file_regions -> survey_binary (slow) try: - sb = self.client.call("survey_binary") + sb = self.client.invoke("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, @@ -1345,8 +1336,7 @@ class Program: def file_regions(self) -> list[tuple[int, int, int]]: """Sorted [(start, end, file_off)] mapping loaded segments to raw file - offsets (file_off == -1 for non-file-backed, e.g. .bss). Cached; needs - the injected ``file_regions`` server tool.""" + offsets (file_off == -1 for non-file-backed, e.g. .bss). Cached.""" if self._fileregions is not None: return self._fileregions regions = [(s, e, fo) for s, e, fo, _nm in self._segments()] @@ -1364,15 +1354,14 @@ class Program: def read_bytes(self, ea: int, n: int) -> bytes: """Raw bytes [ea, ea+n) from IDA (gaps read as zero). - Fast path: the injected ``read_raw`` tool returns one contiguous hex - string (C-speed both ends). Falls back to the stock ``get_bytes`` (a - per-byte '0x..'-with-spaces string) on an older server without it. + The Code Mode adapter returns one contiguous hex string (C-speed in IDA). + A legacy ``get_bytes`` decoding fallback remains for alternate clients. """ if n <= 0: return b"" if not self._no_read_raw: try: - r = self.client.call("read_raw", addr=hex(ea), size=int(n)) + r = self.client.invoke("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) @@ -1386,7 +1375,7 @@ class Program: except (ValueError, KeyError): pass # malformed hex -> fall through to the legacy decoder try: - r = self.client.call("get_bytes", regions=[{"addr": hex(ea), "size": int(n)}]) + r = self.client.invoke("get_bytes", regions=[{"addr": hex(ea), "size": int(n)}]) except IDAToolError: return b"\x00" * n res = r.get("result", []) if isinstance(r, dict) else [] @@ -1437,7 +1426,7 @@ class Program: def list_structs(self, filter: str = "") -> list[Struct]: """All local structs/unions (optionally name-substring filtered), sorted by name.""" - payload = self.client.call("search_structs", filter=filter) + payload = self.client.invoke("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") @@ -1447,9 +1436,9 @@ class Program: def struct_source(self, name: str) -> str: """A C definition for ``name`` reconstructed from its member layout - (the server exposes members, not printable source). Faithful to IDA's + (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.call( + payload = self.client.invoke( "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 {} @@ -1472,7 +1461,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.call("declare_type", decls=decl) + payload = self.client.invoke("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") @@ -1481,10 +1470,9 @@ class Program: # -- function / variable types ---------------------------------------- # def func_types(self, ea: int) -> FuncTypes | None: """Structured decompiler types for the function at ``ea`` (prototype + - local variables). None if ``ea`` isn't a decompilable function. Requires - the injected ``func_types`` server tool.""" + local variables). None if ``ea`` isn't a decompilable function.""" try: - r = self.client.call("func_types", addr=hex(ea)) + r = self.client.invoke("func_types", addr=hex(ea)) except IDAToolError: return None if not isinstance(r, dict) or r.get("error"): @@ -1497,7 +1485,7 @@ class Program: 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.call("set_type", edits=[{"addr": hex(ea), "signature": signature}]) + r = self.client.invoke("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"): @@ -1506,9 +1494,9 @@ class Program: def data_type(self, ea: int) -> dict | None: """Current type info for a data item/global: {addr,name,type,size,is_func}. - None if the tool is unavailable or the address isn't mapped.""" + None if the operation fails or the address isn't mapped.""" try: - r = self.client.call("data_type", addr=hex(ea)) + r = self.client.invoke("data_type", addr=hex(ea)) except IDAToolError: return None if not isinstance(r, dict) or r.get("error"): @@ -1517,7 +1505,7 @@ 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.call( + r = self.client.invoke( "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 {} @@ -1526,9 +1514,9 @@ class Program: return row.get("error") or "failed to set the type" def set_lvar_type(self, fn_ea: int, var: str, ty: str) -> str | None: - """Set a decompiler local variable's type (via the injected server tool). + """Set a decompiler local variable's type through ida-domain pseudocode. None on success, else an error string.""" - r = self.client.call("set_lvar_type", addr=hex(fn_ea), variable=var, type=ty) + r = self.client.invoke("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"): @@ -1537,15 +1525,14 @@ class Program: def delete_type(self, name: str) -> str | None: """Delete a named type. Returns None on success, else an error string. - Requires a server-side ``del_type`` tool; if absent, a clear message is - returned instead of raising.""" + Returns a clear error instead of raising when the runtime cannot do it.""" try: - self.client.call("del_type", name=name) + self.client.invoke("del_type", name=name) return None except IDAToolError as e: msg = e.message if "not found" in msg.lower() and "del_type" in msg: - return "delete needs a 'del_type' tool on the ida-pro-mcp server" + return "the connected Code Mode runtime cannot delete local types" return msg # -- disassembly ------------------------------------------------------- # @@ -1559,13 +1546,7 @@ class Program: # -- decompilation ----------------------------------------------------- # def decompile(self, ea: int, refresh: bool = False) -> Decompilation: - """Full pseudocode for a function. - - The server truncates responses over 50KB (strings clipped to 1000 - chars) but caches the full output and exposes it at - ``_meta.ida_mcp.download_url``. We transparently fetch that so the view - always gets the complete body, not a 1KB stub. - """ + """Full pseudocode for a function, returned directly by Code Mode.""" if not refresh: with self._lock: hit = self._decomp.get(ea) @@ -1574,10 +1555,10 @@ class Program: dec, hit_gen = hit if hit_gen == gen: return dec - # Cached before a rename: names may be stale. Drop the server's - # Hex-Rays cache so the refetch reflects the new names. + # Cached before a rename: names may be stale. Drop Hex-Rays' + # cache so the refetch reflects the new names. try: - self.client.call("force_recompile", items=[{"addr": hex(ea)}]) + self.client.invoke("force_recompile", items=[{"addr": hex(ea)}]) except Exception: # noqa: BLE001 pass # Bound the decompile: a function Hex-Rays can't handle tends to stall @@ -1587,7 +1568,10 @@ class Program: # rpcclient socket timeout, and cache the failure below so a re-request # returns instantly instead of re-grinding. try: - envelope = self.client.call_envelope( + # 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 ) except Exception as e: # noqa: BLE001 -- surface as a failed decompile @@ -1595,15 +1579,6 @@ class Program: with self._lock: self._decomp[ea] = (dec, self._name_gen) return dec - result = envelope.get("result", {}) - payload = result.get("structuredContent") - if payload is None: # fall back to text content - payload = self.client._extract_payload("decompile", result) - meta = (result.get("_meta") or {}).get("ida_mcp") - if isinstance(meta, dict) and meta.get("download_url"): - full = self._fetch_output(meta["download_url"]) - if isinstance(full, dict) and full.get("code"): - payload = full dec = _parse_decompilation(ea, payload) with self._lock: self._decomp[ea] = (dec, self._name_gen) @@ -1681,18 +1656,18 @@ 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.call("undefine", items=[{"addr": hex(ea)}]) + self.client.invoke("undefine", items=[{"addr": hex(ea)}]) except IDAToolError: pass # nothing defined here yet -> just try to create the insn res = self._first_result( - self.client.call("define_code", items=[{"addr": hex(ea)}])) + self.client.invoke("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.call("decomp_error", addr=hex(ea)) + r = self.client.invoke("decomp_error", addr=hex(ea)) except IDAToolError: return "" if not isinstance(r, dict): @@ -1711,7 +1686,7 @@ 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.call("thumb_scan", start=hex(start), end=hex(end), + r = self.client.invoke("thumb_scan", start=hex(start), end=hex(end), apply=bool(apply)) if not isinstance(r, dict) or r.get("error"): raise IDAToolError("thumb_scan", @@ -1720,7 +1695,7 @@ class Program: def set_thumb(self, ea: int, mode: str = "toggle") -> dict: """Switch ARM/Thumb decoding at ``ea``. Returns the resulting state.""" - r = self.client.call("set_thumb", addr=hex(ea), mode=mode) + r = self.client.invoke("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')}") @@ -1729,11 +1704,11 @@ class Program: def define_code_run(self, ea: int, limit: int = 20000) -> dict: """Disassemble consecutively from ``ea`` until something stops it. - Falls back to a single instruction when the worker predates the tool, so - an old worker degrades to the previous behaviour instead of failing. + Falls back to a single instruction for alternate clients that do not + provide the run operation. """ try: - r = self.client.call("define_code_run", addr=hex(ea), limit=int(limit)) + r = self.client.invoke("define_code_run", addr=hex(ea), limit=int(limit)) except IDAToolError: self.define_code(ea) return {"count": 1, "stopped": "single", "end": hex(ea)} @@ -1745,14 +1720,14 @@ class Program: def define_func(self, ea: int) -> dict: """Create a function starting at ``ea`` (IDA's 'p'). - Prefers the injected tool, which works out the end when IDA can't; - falls back to the plain one for an older worker. + Prefers the Code Mode operation, which works out the end when IDA can't; + falls back to a plain create for alternate clients. """ try: - r = self.client.call("define_func_run", addr=hex(ea)) + r = self.client.invoke("define_func_run", addr=hex(ea)) except IDAToolError: res = self._first_result( - self.client.call("define_func", items=[{"addr": hex(ea)}])) + self.client.invoke("define_func", items=[{"addr": hex(ea)}])) if res.get("error"): raise IDAToolError("define_func", f"@ {ea:#x}: {res['error']}") return {"ok": True, "how": "legacy"} @@ -1766,7 +1741,7 @@ class Program: item: dict = {"addr": hex(ea)} if size: item["size"] = int(size) - res = self._first_result(self.client.call("undefine", items=[item])) + res = self._first_result(self.client.invoke("undefine", items=[item])) if res.get("error"): raise IDAToolError("undefine", f"@ {ea:#x}: {res['error']}") @@ -1776,7 +1751,7 @@ class Program: item: dict = {"addr": hex(ea), "type": type_decl} if name: item["name"] = name - res = self._first_result(self.client.call("make_data", items=[item])) + res = self._first_result(self.client.invoke("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'}") @@ -1784,7 +1759,7 @@ class Program: 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.call("make_string", addr=hex(ea), length=int(length), kind=kind) + r = self.client.invoke("make_string", addr=hex(ea), length=int(length), kind=kind) res = r if isinstance(r, dict) else {} if not res.get("ok"): raise IDAToolError( @@ -1801,7 +1776,7 @@ 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.call("op_format", addr=hex(ea), mode=str(mode), + r = self.client.invoke("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"): @@ -1825,7 +1800,7 @@ class Program: if hit is not None and hit[1] == gen: return hit[0] try: - r = self.client.call("pc_nums", addr=hex(fn_ea)) + r = self.client.invoke("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]]] = {} @@ -1848,7 +1823,7 @@ class Program: 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.call("pc_num_format", addr=hex(fn_ea), mode=str(mode), + r = self.client.invoke("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"): @@ -1865,18 +1840,6 @@ class Program: sec = None return f"{sec} @ {ea:#x}" if sec else f" @ {ea:#x}" - @staticmethod - def _fetch_output(url: str, timeout: float = 15.0): - """GET the server's cached full-output blob (plain HTTP, not MCP).""" - try: - with urllib.request.urlopen(url, timeout=timeout) as r: - return json.loads(r.read().decode("utf-8", "replace")) - except Exception as e: # noqa: BLE001 -- fall back to the truncated preview - # The user gets CLIPPED pseudocode with no indication that a fetch - # failed rather than the function genuinely being that short. - diag.note(f"decompile: full-body fetch {url}", e) - return None - def strings(self, min_len: int = 4, refresh: bool = False) -> list[StrLit]: """Every string literal in the binary (IDA's Shift+F12 list), paged in full and cached. ``[]`` if the tool is unavailable.""" @@ -1889,7 +1852,7 @@ class Program: offset, page = 0, 2000 while True: try: - payload = self.client.call( + payload = self.client.invoke( "list_strings", offset=offset, count=page, min_len=min_len, refresh=(refresh and offset == 0)) except IDAToolError: @@ -1914,13 +1877,13 @@ class Program: def linkage(self) -> tuple[list[Linkage], list[Linkage]]: """``(imports, exports)`` for this binary, cached. ``([], [])`` if the - tool is unavailable — an old worker must not break the caller.""" + operation is unavailable — an alternate client must not break the caller.""" with self._lock: hit = self._linkage if hit is not None: return hit try: - payload = self.client.call("list_linkage", kind="both") + payload = self.client.invoke("list_linkage", kind="both") except IDAToolError: return ([], []) if not isinstance(payload, dict): @@ -1951,7 +1914,7 @@ class Program: if hit is not None and hit[1] == gen: return hit[0] try: - payload = self.client.call("decomp_map", addr=hex(ea)) + payload = self.client.invoke("decomp_map", addr=hex(ea)) except IDAToolError: return [] lines = payload.get("lines", []) if isinstance(payload, dict) else [] @@ -1981,7 +1944,7 @@ class Program: if hit is not None and hit[1] == gen: return hit[0] try: - payload = self.client.call("flowchart", addr=hex(ea)) + payload = self.client.invoke("flowchart", addr=hex(ea)) except IDAToolError: return None if not isinstance(payload, dict) or payload.get("error"): @@ -2054,7 +2017,7 @@ class Program: for _ in range(64): # bounded: ~128k heads if addr >= hi: break - payload = self.client.call("heads", addr=hex(addr), end=hex(hi), + payload = self.client.invoke("heads", addr=hex(addr), end=hex(hi), count=2000) rows = payload.get("heads", []) if isinstance(payload, dict) else [] if not rows: @@ -2083,13 +2046,13 @@ 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.call("lookup_funcs", queries=[hex(ea)]) + payload = self.client.invoke("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.call( + payload = self.client.invoke( "xref_query", queries=[{"addr": hex(ea), "direction": "from", "include_fn": True}], ) @@ -2101,9 +2064,9 @@ class Program: 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.call("xref_types", queries=q) + payload = self.client.invoke("xref_types", queries=q) except IDAToolError: - payload = self.client.call("xref_query", queries=q) + payload = self.client.invoke("xref_query", queries=q) return _parse_xrefs(payload) # -- address resolution ------------------------------------------------ # @@ -2121,17 +2084,17 @@ 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.call("resolve_names", queries=[s]) + payload = self.client.invoke("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: return _as_int(ea) except IDAToolError: - pass # older server without resolve_names -> fall back below + pass # alternate client without resolve_names -> fall back below # Fall back to function-name resolution (also drives the 'did you mean' # suggestion when the name is unknown). try: - payload = self.client.call("lookup_funcs", queries=[s]) + payload = self.client.invoke("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 [] @@ -2169,7 +2132,7 @@ 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.call("set_comments", items=[{"addr": hex(ea), "comment": text}]) + return self.client.invoke("set_comments", items=[{"addr": hex(ea), "comment": text}]) # -- invalidation (after edits) --------------------------------------- # def invalidate(self, ea: int) -> None: diff --git a/idatui/drive.py b/idatui/drive.py index 1c7d8c2..825111c 100644 --- a/idatui/drive.py +++ b/idatui/drive.py @@ -121,7 +121,8 @@ def cmd_pc(c, args): lines = d["code"].splitlines() if needle: nlow = needle.lower() - lines = [f"{i:4} {l}" for i, l in enumerate(lines) if nlow in l.lower()] + lines = [f"{i:4} {line}" for i, line in enumerate(lines) + if nlow in line.lower()] return "\n".join(lines) or f"(no line matches {needle!r})" return d["code"] diff --git a/idatui/edit_ctl.py b/idatui/edit_ctl.py index 435fab9..80e1109 100644 --- a/idatui/edit_ctl.py +++ b/idatui/edit_ctl.py @@ -242,7 +242,7 @@ class EditController: kind = "stack" batch = {"stack": {"func_addr": hex(cur.ea), "old": old, "new": new}} try: - res = prog.client.call("rename", batch=batch) + res = prog.client.invoke("rename", batch=batch) except IDAToolError as e: app.call_from_thread(app._status, f"rename failed: {e.message}") return @@ -274,7 +274,7 @@ class EditController: app = self.app assert app.program is not None try: - res = app.program.client.call( + res = app.program.client.invoke( "rename", batch={"data": {"addr": hex(addr), "new": name}}) except IDAToolError as e: app.call_from_thread(app._status, f"name failed: {e.message}") diff --git a/idatui/errors.py b/idatui/errors.py index 29b09ae..aaf2dc5 100644 --- a/idatui/errors.py +++ b/idatui/errors.py @@ -1,10 +1,8 @@ -"""Transport-agnostic error hierarchy and the Session model. +"""TUI-facing error hierarchy and lightweight database session model. -These were originally defined in client.py (the ida-pro-mcp HTTP client), but the -idalib worker path (worker_client / domain / app) needs the same exception types -and Session dataclass without dragging in the HTTP transport. They live here so -both backends share one definition; client.py re-exports them for backwards -compatibility with the (deprecated) mcp tooling and the stress tests. +The Code Mode adapter normalizes ``ida_codemode.client`` transport and execution +errors into these types so the domain and Textual layers do not depend on HTTP or +registry implementation details. """ from __future__ import annotations diff --git a/idatui/launch.py b/idatui/launch.py index 35221a1..0c53987 100644 --- a/idatui/launch.py +++ b/idatui/launch.py @@ -1,15 +1,13 @@ -"""One-shot launcher: ``ida-tui foo.elf`` and you're in the TUI. +"""One-shot launcher for the IDA Code Mode-backed TUI. -Spawns a private idalib worker (``idatui.worker``) that opens + auto-analyzes -THIS binary in its own process, talking to the TUI over a unix socket. No shared -supervisor, no HTTP: everything slow (open + analysis) happens behind the TUI's -loading overlay. +A path first resolves to a registered GUI database; when none matches, Code Mode +reuses or starts a managed idalib worker. With no path, a single registered +database is selected automatically. -Usage: +Usage:: - ida-tui /path/to/binary # open a binary and drive it - -Extras: --ttl, --no-keepalive, --rpc (all forwarded to the TUI). + ida-tui /path/to/binary + ida-tui # attach when exactly one database is registered """ from __future__ import annotations @@ -17,12 +15,6 @@ import argparse import os import sys -# The unpacked working-copy files IDA writes next to a `.i64` while a database is -# open. A hard-killed worker leaves them behind and the `.i64` then refuses to -# reopen ("Failed to open database"). Safe to delete when nothing holds the DB. -_LOCK_SUFFIXES = (".id0", ".id1", ".id2", ".nam", ".til") - - def _load_args(load: dict) -> str: """``load`` as IDA switches, for the single-binary path (no project ref). @@ -41,35 +33,19 @@ def _log(msg: str) -> None: print(f"ida-tui: {msg}", file=sys.stderr) -def _sweep_locks(binary: str) -> int: - """Remove stale unpacked DB files next to ``binary``. Returns how many. - - Never touches the ``.i64`` -- that is the real database, and nothing is - saved unless ``idb_save`` was called -- and never the input file itself. - The second guard is not theoretical: ``.til`` is both an unpacked-DB suffix - and the extension of an IDA type library, so ``ida-tui mylib.til`` swept its - own argument out of existence. Same for anything named ``*.id0``/``*.nam``. - """ - keep = os.path.abspath(binary) - stem = os.path.splitext(binary)[0] - n = 0 - for base in (binary, stem): # IDA may key on the full name or the stem - for suf in _LOCK_SUFFIXES: - victim = base + suf - if os.path.abspath(victim) == keep: - continue # that's what the user asked us to open - try: - os.remove(victim) - n += 1 - except OSError: - pass - return n +def _registered_databases() -> tuple[list[dict], list[dict]]: + """Ready and blocked Code Mode registrations, with normalized errors.""" + try: + from ida_codemode.registry import discover_instances + return discover_instances() + except Exception as exc: # discovery diagnostics belong at the CLI boundary + return [], [{"error": str(exc)}] def main(argv: list[str] | None = None) -> int: p = argparse.ArgumentParser( prog="ida-tui", - description="Open a binary in the IDA TUI (private idalib worker).") + description="Open a registered GUI or managed idalib database in the IDA TUI.") p.add_argument("binary", nargs="*", help="binary to open and analyze (several with --project " "creates/extends that project)") @@ -77,9 +53,9 @@ def main(argv: list[str] | None = None) -> int: help="open a multi-binary project (created from the given " "binaries if FILE doesn't exist)") p.add_argument("--ttl", type=int, default=1800, - help="worker idle-TTL seconds (default 1800)") + help="deprecated compatibility option (Code Mode uses leases)") p.add_argument("--no-keepalive", action="store_true", - help="do not run the keepalive heartbeat") + help="deprecated compatibility option (the lease is the heartbeat)") p.add_argument("--rpc", metavar="PATH", help="listen for RPC on this unix socket (puppeteer the TUI)") p.add_argument("--trace", metavar="FILE", @@ -94,7 +70,7 @@ def main(argv: list[str] | None = None) -> int: g.add_argument("--base", metavar="ADDR", help="load address, e.g. 0x8000000 (any base; NOT paragraphs)") g.add_argument("--ida-args", metavar="STR", dest="ida_args", - help="extra IDA command-line switches, passed through as-is") + help="legacy switches; only Code Mode-representable -p/-b/-T are accepted") args = p.parse_args(argv) load: dict = {} @@ -152,27 +128,42 @@ def main(argv: list[str] | None = None) -> int: _log(str(e)) return 2 else: - if len(args.binary) != 1: - _log("give exactly one binary, or use --project for several") + ready, blocked = _registered_databases() + if len(args.binary) > 1: + _log("give at most one binary, or use --project for several") return 2 - binary = os.path.abspath(os.path.expanduser(args.binary[0])) - if not os.path.isfile(binary): - _log(f"no such file: {binary}") + if args.binary: + binary = os.path.abspath(os.path.expanduser(args.binary[0])) + key = os.path.normcase(os.path.realpath(binary)) + registered = any( + key == os.path.normcase(os.path.realpath(str(item.get(field) or ""))) + for item in ready for field in ("exe_path", "idb_path") + if item.get(field) + ) + if not os.path.isfile(binary) and not registered: + _log(f"no such file or registered database: {binary}") + return 2 + elif len(ready) == 1: + item = ready[0] + binary = str(item.get("exe_path") or item.get("idb_path") or "") + _log(f"attaching to registered {item.get('backend')} database: {binary}") + elif not ready: + detail = f" ({blocked[0].get('error')})" if blocked else "" + _log(f"no registered Code Mode database; pass a binary path{detail}") return 2 - if not os.access(os.path.dirname(binary), os.W_OK): - _log(f"directory not writable (IDA writes a .i64 there): " - f"{os.path.dirname(binary)}") + else: + _log("several Code Mode databases are registered; pass one of these paths:") + for item in ready: + _log(f" {item.get('exe_path') or item.get('idb_path')} " + f"[{item.get('backend')}, {item.get('record_id')}]") return 2 - swept = _sweep_locks(binary) # a crashed worker can leave the DB wedged - if swept: - _log(f"cleared {swept} stale lock file(s) from a crashed worker") - # Hand off to the TUI (imported late so --help works without textual). It - # spawns the worker behind its loading overlay while auto-analysis runs. + # Hand off to the TUI (imported late so --help works without Textual). Code + # Mode discovery/opening happens behind its loading overlay. try: from .app import IdaTui except ImportError as e: - _log(f"the TUI needs textual; run with ~/ida-venv/bin/python ({e})") + _log(f"TUI dependencies are missing; run `uv sync` ({e})") return 1 # Ask the terminal about graphics support NOW: the query needs a reply from # stdin, and once Textual starts it reads stdin on its own thread and would diff --git a/idatui/pane.py b/idatui/pane.py index 2592581..c31a93c 100644 --- a/idatui/pane.py +++ b/idatui/pane.py @@ -23,9 +23,9 @@ per pane in the registry, so stop/list/capture/keys keep working across both python -m idatui.pane capture --pane python -m idatui.pane keys --pane Escape -Requires: running inside tmux or zellij. Each pane spawns its own private idalib -worker (no shared supervisor). Uses ~/ida-venv/bin/python for the TUI (needs -textual) unless --python / IDATUI_PYTHON says otherwise. +Requires: running inside tmux or zellij. Each pane leases a registered GUI or +shared managed idalib database through Code Mode. Uses ~/ida-venv/bin/python for +the TUI (needs textual) unless --python / IDATUI_PYTHON says otherwise. """ from __future__ import annotations @@ -33,7 +33,6 @@ import argparse import json import os import secrets -import signal import subprocess import sys import time @@ -251,35 +250,9 @@ def _pane_keys(pane: str, keys: list[str], mux: str | None = None) -> None: subprocess.run(["tmux", "send-keys", "-t", pane, *keys], check=True) -# --------------------------------------------------------------------------- # -# idalib worker reaping -# -# ``pane stop`` kills the TUI pane, but a hard-killed pane can leave its private -# idalib worker (idatui/worker.py) running. A worker is only *safe* to reap when -# no idatui pane is live (then every worker is orphaned), which avoids killing an -# in-use analyser. -# --------------------------------------------------------------------------- # -_WORKER_PATTERN = r"idatui/worker\.py" - - -def _worker_pids() -> list[int]: - """PIDs of our private per-pane idalib worker processes (idatui/worker.py), - never our own PID.""" - try: - out = subprocess.run(["pgrep", "-f", _WORKER_PATTERN], - capture_output=True, text=True) - except OSError: - return [] - me = os.getpid() - pids: list[int] = [] - for tok in out.stdout.split(): - try: - pid = int(tok) - except ValueError: - continue - if pid != me: - pids.append(pid) - return pids +# Code Mode owns database process lifetime: a closed pane drops its lease at the +# socket/kernel boundary and Code Mode decides whether a managed worker still +# has clients. There is nothing for the pane layer to reap. def _count_live_panes() -> int: @@ -288,20 +261,9 @@ def _count_live_panes() -> int: def _reap_orphan_workers(force: bool = False) -> int: - """Kill leaked idalib workers when it is safe (no live pane) or ``force``. - - Returns the number of workers signalled. Best-effort; never raises. - """ - if not force and _count_live_panes() > 0: - return 0 - reaped = 0 - for pid in _worker_pids(): - try: - os.kill(pid, signal.SIGKILL) - reaped += 1 - except OSError: - pass - return reaped + """Compatibility no-op: Code Mode workers are shared and lease-managed.""" + del force + return 0 # --------------------------------------------------------------------------- # @@ -332,16 +294,8 @@ def spawn(args) -> int: print(f"error: no such project: {project}", file=sys.stderr) return 2 - # Reap workers leaked by previously-stopped/crashed panes so we don't spawn - # into a full IDA_MCP_MAX_WORKERS (which makes the new TUI hang forever, - # never reaching ready). No-op while any pane is live. - reaped = _reap_orphan_workers() - if reaped: - print(f"reaped {reaped} orphaned idalib worker(s) before spawn", - file=sys.stderr) - - # the command the pane runs: the launcher spawns a private idalib worker for - # this binary and becomes the TUI, so kill-pane tears the whole thing down. + # The pane owns only the TUI. Code Mode's lease cleanup handles crashes; + # kill-pane must never reap a shared GUI/idalib database. if project is not None: # launch takes: --project FILE [binaries...]; extra binaries are added to # the project (and a missing project file is created from them). @@ -393,9 +347,8 @@ def _wait_ready(sock: str, timeout: float, pane: str, stuck_after: float = 45.0, mux: str | None = None) -> dict[str, Any]: """Poll the socket + ping until the TUI reports ready (or timeout). - Emits a one-time hint to stderr if it's still not ready after ``stuck_after`` - seconds, so a wedged idalib worker / full worker pool surfaces a diagnostic - instead of an unexplained silent hang. + Emits a one-time hint if Code Mode discovery/opening is still not ready after + ``stuck_after`` seconds. """ start = time.time() deadline = start + timeout @@ -416,9 +369,8 @@ def _wait_ready(sock: str, timeout: float, pane: str, warned = True why = ("RPC socket not created yet" if not os.path.exists(sock) else "TUI up but analysis not ready") - print(f"still waiting ({int(time.time() - start)}s): {why}. If this " - f"hangs, the idalib worker may be stuck — try " - f"`python -m idatui.pane reap`.", file=sys.stderr) + print(f"still waiting ({int(time.time() - start)}s): {why}. " + f"Check Code Mode registrations and worker logs.", file=sys.stderr) time.sleep(0.4) last = dict(last) last["ready"] = False @@ -515,13 +467,9 @@ def list_panes(args) -> int: def reap(args) -> int: - """Kill leaked idalib workers (safe when no pane is live; --force overrides).""" - live = _count_live_panes() - n = _reap_orphan_workers(force=args.force) - print(json.dumps({"reaped_workers": n, "live_panes": live, "forced": args.force})) - if n == 0 and not args.force and live > 0: - print(f"note: {live} live pane(s) — not reaping in-use workers; pass " - f"--force to reap anyway", file=sys.stderr) + """Deprecated no-op; shared Code Mode workers are managed by leases.""" + print(json.dumps({"reaped_workers": 0, "live_panes": _count_live_panes(), + "forced": args.force, "deprecated": True})) return 0 @@ -624,9 +572,8 @@ def main(argv: list[str]) -> int: ls.add_argument("--prune", action="store_true", help="drop dead panes (and their sockets)") ls.set_defaults(fn=list_panes) - rp = sub.add_parser("reap", help="kill leaked idalib workers (frees worker slots)") - rp.add_argument("--force", action="store_true", - help="reap even while panes are live (may kill an in-use analyser)") + rp = sub.add_parser("reap", help="deprecated no-op (Code Mode uses shared leases)") + rp.add_argument("--force", action="store_true", help=argparse.SUPPRESS) rp.set_defaults(fn=reap) cp = sub.add_parser("capture", help="print a pane's visible screen") diff --git a/idatui/pool.py b/idatui/pool.py index ae37c25..465dff2 100644 --- a/idatui/pool.py +++ b/idatui/pool.py @@ -1,23 +1,16 @@ -"""WorkerPool — keeps a live idalib worker per project binary, within a budget. +"""DatabasePool — LRU leases on Code Mode databases for a project. -One worker process holds exactly one database (idalib is single-DB and -main-thread-only), so a project with N binaries means up to N processes. They are -not cheap and they do not share: a worker on ``bash`` measures ~126 MB RSS / -117 MB PSS, and the database working set dominates for anything larger -(``libcrypto.so.3``'s ``.i64`` alone is 72 MB). +Code Mode may bind a lease to an existing IDA GUI or to a shared managed idalib +worker. The pool therefore owns *client interest*, never an IDA process. Releasing +an LRU entry persists managed IDBs but does not implicitly save a GUI, then closes +only this TUI's lease; other clients and GUI sessions remain alive. Managed workers exit themselves after their final lease. -Residency is therefore bounded by a **memory budget**, not a worker count — a -count is the wrong knob when one project holds both a 50 KB helper and a 6 MB -crypto library. Workers are spawned lazily on first use, kept resident while they -fit, and least-recently-used ones evicted when they don't. Eviction **saves the -database first**, so coming back is a load rather than a re-analysis. - -The pool never evicts the active binary, nor anything pinned. +The historical memory budget remains useful for managed idalib instances, while +GUI process memory is only advisory. The active and pinned databases are never +released to satisfy it. """ from __future__ import annotations -import os - from .project import BinaryRef, Project #: Fallback budget if /proc/meminfo can't be read (MB). @@ -36,11 +29,11 @@ def _total_ram_mb() -> int: def _pss_mb(pid: int | None) -> int: - """Proportional set size of a worker, in MB. + """Proportional set size of the leased instance process, in MB. - PSS (not RSS) is the honest per-worker cost: it splits shared pages between - the processes mapping them. In practice workers share very little, so the two - are close, but PSS is what makes summing across workers meaningful. + PSS is useful for managed idalib workers. For GUI/shared processes it is only + advisory because the TUI neither owns all that memory nor controls process + exit. """ if not pid: return 0 @@ -54,13 +47,19 @@ def _pss_mb(pid: int | None) -> int: return 0 -def _default_spawn(ref: BinaryRef, ttl: int): # pragma: no cover - needs idalib - from .worker_client import WorkerClient - return WorkerClient(ref.staged, ttl=ttl, load_args=ref.load_args) +def _default_spawn(ref: BinaryRef, ttl: int, *, new_database: bool = False): # pragma: no cover - needs IDA + from .codemode_client import CodeModeClient + return CodeModeClient( + ref.staged, + ttl=ttl, + load_args=ref.load_args, + output_database=ref.db, + new_database=new_database, + ) -class WorkerPool: - """Live workers for a project's binaries, keyed by label.""" +class DatabasePool: + """Live Code Mode database leases, keyed by project label.""" def __init__(self, project: Project, *, budget_mb: int | None = None, ttl: int = 1800, spawn=None, mem_fn=None) -> None: @@ -71,6 +70,7 @@ class WorkerPool: self._clients: dict[str, object] = {} self._lru: list[str] = [] # least-recently-used first self._pinned: set[str] = set() + self._recreate: set[str] = set() # Ctrl+L: next attachment creates a fresh IDB self.active: str | None = None # never evicted if budget_mb is None: ram = _total_ram_mb() @@ -98,11 +98,11 @@ class WorkerPool: # -- acquire ----------------------------------------------------------- # def get(self, label: str, progress=None): - """A live client for ``label``, spawning it (and making room) if needed. + """A live client for ``label``, attaching or spawning as needed. - Staging and the scratch sweep happen here: a worker killed hard last time - leaves unpacked ``.id0/.id1/...`` behind, and the database then refuses to - reopen. Nothing else holds this DB (one worker per label), so it is safe. + Do not sweep IDA scratch files here: a registered GUI or another Code + Mode client may own the database. Code Mode's registry locks and health + probes are the authority for safe discovery and stale-record cleanup. """ client = self._clients.get(label) if client is not None: @@ -118,28 +118,29 @@ class WorkerPool: note(f"staging {ref.label}\u2026") self.project.stage(ref) - self.project.sweep_scratch(ref) note(f"opening {ref.label}\u2026") - client = self._spawn(ref, self._ttl) + fresh = label in self._recreate + client = (_default_spawn(ref, self._ttl, new_database=fresh) + if self._spawn is _default_spawn else self._spawn(ref, self._ttl)) connect = getattr(client, "connect", None) if connect is not None: connect(progress=progress) if progress is not None else connect() self._clients[label] = client + self._recreate.discard(label) self._lru.append(label) self._enforce_budget(protect=label) return client def prewarm(self, label: str, progress=None) -> bool: - """Spawn a worker for ``label`` only if it fits the budget AS IT STANDS. + """Attach a database for ``label`` only if it fits the current budget. Pre-warming must never cost residency: evicting a binary the user actually visited to speculatively load one they haven't is a straight downgrade, and the eviction would also throw away that binary's caches. So this refuses rather than making room, and returns False. - The cost of a worker that doesn't exist yet can only be estimated; the - largest resident one is the best evidence available (they are all the - same program with a different database). With nothing resident we have + The cost of a database not attached yet can only be estimated; the + largest resident instance is the best evidence available. With nothing resident we have no evidence at all, so we allow one — that is the case where the budget is certainly free. """ @@ -153,13 +154,19 @@ class WorkerPool: return False self.get(label, progress=progress) # get() enforces the budget protecting the NEW label; if that had to - # evict, our estimate was wrong and the speculative worker is the one + # evict, our estimate was wrong and the speculative lease is the one # that should go — never a binary the user chose. if self.memory_mb() > self.budget_mb and label != self.active: self.evict(label) return False return True + def recreate_on_next_open(self, label: str) -> None: + """Request a fresh IDB after the current lease has been released.""" + if self.project.by_label(label) is None: + raise KeyError(f"no such binary in the project: {label}") + self._recreate.add(label) + def _touch(self, label: str) -> None: if label in self._lru: self._lru.remove(label) @@ -171,16 +178,21 @@ class WorkerPool: self._touch(label) # -- release ----------------------------------------------------------- # - def evict(self, label: str, save: bool = True) -> bool: - """Drop a resident worker, persisting its database first.""" + def evict(self, label: str, save: bool = True, + save_gui: bool = False) -> bool: + """Release a resident lease, persisting a managed database first. + + A budget-driven eviction must not save somebody's GUI implicitly. GUI + saves are reserved for an explicit/defensive ``close_all(save=True)``. + """ client = self._clients.pop(label, None) if client is None: return False if label in self._lru: self._lru.remove(label) - if save: + if save and (save_gui or getattr(client, "backend", None) != "gui"): try: # persist analysis + edits so the next open is a load - client.call("idb_save") + client.save_database() except Exception: # noqa: BLE001 -- evict regardless pass try: @@ -198,7 +210,7 @@ class WorkerPool: return None def _enforce_budget(self, protect: str | None = None) -> int: - """Evict LRU workers until the pool fits its budget. Returns how many.""" + """Release LRU leases until the pool fits its budget. Returns how many.""" n = 0 while self.memory_mb() > self.budget_mb: victim = self._evictable(protect) @@ -210,7 +222,7 @@ class WorkerPool: def close_all(self, save: bool = True) -> None: for label in list(self._clients): - self.evict(label, save=save) + self.evict(label, save=save, save_gui=save) self.active = None # -- introspection ------------------------------------------------------ # @@ -231,5 +243,9 @@ class WorkerPool: return out def __repr__(self) -> str: # pragma: no cover - debug aid - return (f"") + + +# Source compatibility for callers that imported the pre-Code-Mode name. +WorkerPool = DatabasePool diff --git a/idatui/project.py b/idatui/project.py index e2fc542..53f3b0a 100644 --- a/idatui/project.py +++ b/idatui/project.py @@ -23,7 +23,8 @@ firmware image, a cleaned build tree). A source whose size/mtime no longer matches the staged copy is re-staged, and its now-stale database is dropped (the DB describes the old bytes). -stdlib-only, like the domain/worker layers — the TUI is the only Textual consumer. +The model has no IDA imports. Staging consults ida_codemode's registry before +replacing files so it never mutates a database owned by a GUI/shared worker. """ from __future__ import annotations @@ -56,7 +57,7 @@ class BinaryRef: #: it is, a raw firmware image doesn't, and IDA defaults to metapc at 0. processor: str = "" # IDA processor name: arm, armb, mipsb, metapc, … base: int = 0 # load address (natural, e.g. 0x8000000) - ida_args: str = "" # escape hatch: extra IDA command-line switches + ida_args: str = "" # legacy -p/-b/-T switches accepted by Code Mode adapter @property def db(self) -> str: @@ -310,13 +311,25 @@ class Project: """Ensure ``ref`` is staged in the sidecar; returns the staged path. Re-staging a changed source drops its database: the DB describes the old - bytes, so keeping it would silently mismatch the disassembly (any renames - in it are lost, which is why callers should say so out loud). + bytes. Refuse while Code Mode reports a GUI/idalib owner; replacing a + staged executable or IDB underneath a shared live instance is corruption. """ if not os.path.isfile(ref.source): raise ProjectError(f"no such binary: {ref.source}") if not self.is_stale(ref): return ref.staged + try: + from .codemode_client import database_owner + owner = database_owner(ref.db, ref.staged) + except Exception as exc: + raise ProjectError( + f"cannot verify Code Mode ownership before staging {ref.label}: {exc}" + ) from exc + if owner is not None: + raise ProjectError( + f"cannot restage {ref.label}: Code Mode instance {owner.record_id} " + f"still owns {owner.idb_path}; close/release it first" + ) os.makedirs(self.bin_dir, exist_ok=True) tmp = ref.staged + ".staging" _unlink(tmp) @@ -338,10 +351,11 @@ class Project: return out def sweep_scratch(self, ref: BinaryRef) -> int: - """Delete IDA's unpacked working files (never the ``.i64``) for ``ref``. + """Delete unpacked working files (never the ``.i64``) for maintenance. - A hard-killed worker leaves them behind and the database then refuses to - reopen. Only safe when no worker holds it. + Runtime paths no longer call this: Code Mode instances are shared, so a + registry owner may still be using these files. Callers must independently + prove that no GUI/idalib instance owns the database. """ return sum(1 for suf in SCRATCH_SUFFIXES if _unlink(ref.staged + suf)) diff --git a/idatui/rpc.py b/idatui/rpc.py index 698e4cf..225aeb2 100644 --- a/idatui/rpc.py +++ b/idatui/rpc.py @@ -757,7 +757,7 @@ class RpcServer: batch = {"func": ops, "allow_overwrite": bool(overwrite)} # The worker is blocking and single-threaded; off the event loop it goes, # or the TUI freezes for the length of the batch. - res = await asyncio.to_thread(app.program.client.call, "rename", batch=batch) + res = await asyncio.to_thread(app.program.client.invoke, "rename", batch=batch) summary = res.get("summary", {}) if isinstance(res, dict) else {} failed = [r for r in (res.get("func") or []) if isinstance(r, dict) and r.get("error")] if isinstance(res, dict) else [] @@ -768,7 +768,7 @@ class RpcServer: # this a batch import leaves pseudocode calling sub_98C0 forever while # the listing (and every readback) says memset. try: - await asyncio.to_thread(app.program.client.call, "force_recompile") + await asyncio.to_thread(app.program.client.invoke, "force_recompile") except Exception: # noqa: BLE001 -- older worker without the tool pass app.program.bump_names() diff --git a/idatui/worker.py b/idatui/worker.py deleted file mode 100644 index 556e69a..0000000 --- a/idatui/worker.py +++ /dev/null @@ -1,301 +0,0 @@ -"""idatui's own idalib worker — the replacement for the ida-pro-mcp supervisor. - -Opens ONE database in-process (on the main thread, as idalib requires) and -serves ida-pro-mcp's *tool functions* over a unix socket with length-prefixed -pickle. Same tool implementations as the MCP path (we call -``MCP_SERVER.tools.methods[name](**args)`` directly), so return shapes are -byte-identical — but with ~50us/call instead of the HTTP path's ~5ms, and no -supervisor / HTTP / JSON / 50KB-truncation machinery. - - python -m idatui.worker - -The socket only appears once the database is open + analyzed, so a client can -poll ``connect()`` to know when the worker is ready. Requests are served -serially on the main thread (idalib is single-threaded; every tool runs inline -through its own execute_sync, which is a no-op on the main thread). - -Protocol (both directions length-prefixed: 4-byte big-endian len + pickle): - request = (tool_name: str, kwargs: dict) - response = (ok: bool, result_or_error) - tool_name == "__shutdown__" ends the worker. -""" -from __future__ import annotations - -import os -import pickle -import socket -import struct -import sys -import threading -import time -import uuid - -#: Seconds a single tool call may run before it is cancelled. 0 disables the -#: deadline entirely. -TOOL_TIMEOUT_SEC = float(os.environ.get("IDATUI_TOOL_TIMEOUT_SEC") or 60) - -# ida-pro-mcp enforces its own tool deadline by installing a `sys.setprofile` -# hook for the duration of every call, so that a pure-python loop inside a tool -# body can be interrupted. That hook runs a python function on EVERY python call -# and return -- and our tools are exactly the call-heavy kind: `heads` renders -# hundreds of items per request and measured 92us/row with the hook against -# 28us/row without it. A 3.3x tax on the whole backend to bound loops that are -# already bounded by their `count` argument. -# -# So: turn the upstream mechanism off and re-arm the half that does the real -# work ourselves (see _Deadline). ida_kernwin.set_cancelled() is what actually -# frees the IDA main thread -- decompile, auto_wait, find_bytes and friends poll -# user_cancelled() and bail within a poll cycle -- and it costs nothing until it -# fires. -os.environ["IDA_MCP_TOOL_TIMEOUT_SEC"] = "0" - - -class _Deadline: - """A single watchdog thread that cancels a tool call which overruns. - - Arming is two attribute writes, because it is on the path of every call the - TUI makes (a scroll is dozens of them). The watchdog polls instead of being - signalled for the same reason: waking a thread per call costs more than the - 0.25s of granularity it buys on a 60s deadline. - """ - - TICK = 0.25 - - def __init__(self, seconds: float) -> None: - import ida_kernwin - self._kernwin = ida_kernwin - self.seconds = seconds - self._until: float | None = None - t = threading.Thread(target=self._run, name="idatui-deadline", - daemon=True) - t.start() - - def _run(self) -> None: - while True: - time.sleep(self.TICK) - until = self._until - if until is not None and time.monotonic() >= until: - self._until = None - # THREAD_SAFE in the IDA SDK; upstream fires it off a Timer too. - self._kernwin.set_cancelled() - - def arm(self) -> None: - # Clear unconditionally: the flag is sticky, and one left set would make - # every later user_cancelled() true forever. - self._kernwin.clr_cancelled() - self._until = time.monotonic() + self.seconds - - def disarm(self) -> None: - self._until = None - - -# --------------------------------------------------------------------------- # -# framing -# --------------------------------------------------------------------------- # -def _recvn(sock: socket.socket, n: int) -> bytes | None: - buf = bytearray() - while len(buf) < n: - chunk = sock.recv(n - len(buf)) - if not chunk: - return None - buf += chunk - return bytes(buf) - - -def send(sock: socket.socket, obj) -> None: - data = pickle.dumps(obj, protocol=pickle.HIGHEST_PROTOCOL) - sock.sendall(struct.pack(">I", len(data)) + data) - - -def recv(sock: socket.socket): - hdr = _recvn(sock, 4) - if hdr is None: - return None - (n,) = struct.unpack(">I", hdr) - body = _recvn(sock, n) - return None if body is None else pickle.loads(body) - - -# --------------------------------------------------------------------------- # -# worker -# --------------------------------------------------------------------------- # -def _ensure_tools_injected() -> None: - """Inject idatui's custom tools (heads/read_raw/resolve_names/func_types/...) - into the installed ida_pro_mcp, idempotently, so the worker is self-sufficient - (nothing else has to inject these tools first). Must run BEFORE - ida_pro_mcp.ida_mcp is imported (the injected code lives in api_types.py).""" - import importlib.util - repo = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) - patch = os.path.join(repo, "server", "patch_server.py") - if not os.path.exists(patch): - return - try: - spec = importlib.util.spec_from_file_location("_idatui_patch", patch) - mod = importlib.util.module_from_spec(spec) - spec.loader.exec_module(mod) # IDA-free; just defines + patches api_types - mod.main() - except Exception as e: # noqa: BLE001 -- tools may already be present - sys.stderr.write(f"idatui: tool injection skipped: {e}\n") - - -def _has_database(binpath: str) -> bool: - """Whether IDA already has a database for ``binpath``. - - IDA names it ``.i64`` (keeping the extension), but a database made - from ``foo.bin`` can also appear as ``foo.i64`` depending on how it was - created — check both, because guessing wrong here means re-passing load - switches to an existing database, which fails the open. - """ - return (os.path.exists(binpath + ".i64") - or os.path.exists(os.path.splitext(binpath)[0] + ".i64")) - - -def _open_and_register(binpath: str, load_args: str = ""): - """Open the DB (main thread) then import ida-pro-mcp so every @tool registers - against this live database. Returns (tools_dict, module_name, save_fn). - - ``load_args`` is passed to IDA as command-line switches, which is the only - way to tell it how to read a headerless blob: a raw firmware image has no - format to detect, so without ``-p`` it loads as metapc at 0 and - finds nothing. Ignored once a database exists — the .i64 already records how - it was loaded, and re-passing conflicting switches is how you corrupt one. - """ - _ensure_tools_injected() # before any ida_pro_mcp import - import idapro - idapro.enable_console_messages(False) - args = load_args or None - if args and _has_database(binpath): - # The .i64 already records how this image was loaded. Passing the - # switches again on reopen makes IDA fail outright (rc != 0) — the load - # options belong to the FIRST open only. - args = None - if idapro.open_database(binpath, run_auto_analysis=True, - args=args): # nonzero == failure - if args: - # With load switches in play they are the likeliest culprit by far: - # IDA refuses an unknown -p name with no diagnostic of its own, so - # saying "the database is locked" here sends people hunting a - # problem they don't have. - raise RuntimeError( - f"failed to open {binpath} with load options {args!r}: IDA " - f"rejected them \u2014 an unknown processor name is the usual " - f"cause (see tools/verify_procs.py for the valid ones)") - raise RuntimeError( - f"failed to open {binpath}: the .i64 is likely held by a running " - f"ida-mcp worker (try: pkill -f idalib) or wedged from a crash " - f"(delete its .id0/.id1/.id2/.nam/.til next to the binary)") - import ida_auto - ida_auto.auto_wait() # block until auto-analysis settles (match ida-mcp) - - # importing the package registers all api_*/patched tools against MCP_SERVER - from ida_pro_mcp.ida_mcp import MCP_SERVER # noqa: WPS433 - - import ida_nalt - module = os.path.basename(ida_nalt.get_root_filename() or binpath) - - def save(): - import idc - try: - idc.save_database(idc.get_idb_path(), 0) - except Exception: # noqa: BLE001 - import ida_loader, ida_pro # noqa: WPS433 - ida_loader.save_database(idc.get_idb_path(), 0) - - return MCP_SERVER.tools.methods, module, save - - -def serve(sockpath: str, binpath: str, load_args: str = "") -> None: - tools, module, save = _open_and_register(binpath, load_args) - sid = uuid.uuid4().hex[:8] - deadline = _Deadline(TOOL_TIMEOUT_SEC) if TOOL_TIMEOUT_SEC > 0 else None - - def dispatch(name: str, args: dict): - args = dict(args) - args.pop("database", None) # single-DB worker: no session routing - # session-management shims (were the supervisor's job): - if name in ("idb_open",): - return {"success": True, - "session": {"session_id": sid, "module": module, - "input_path": binpath}} - if name in ("idb_save", "save"): - save() - return {"success": True} - if name in ("server_health", "ping", "health", "state"): - return {"module": module, "ok": True, "session_id": sid} - if name in ("idb_list",): - return {"sessions": [{"session_id": sid, "module": module, - "input_path": binpath}]} - fn = tools.get(name) - if fn is None: - raise KeyError(f"unknown tool: {name!r}") - if deadline is None: - result = fn(**args) - else: - deadline.arm() - try: - result = fn(**args) - finally: - deadline.disarm() - # Match the MCP server's structuredContent: a dict passes through, any - # other return (list/scalar) is wrapped as {"result": ...}. domain.py - # parses that exact shape (e.g. lookup_funcs -> payload["result"]). - return result if isinstance(result, dict) else {"result": result} - - try: - os.unlink(sockpath) - except OSError: - pass - srv = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) - srv.bind(sockpath) - srv.listen(8) - try: - while True: - conn, _ = srv.accept() - try: - while True: - req = recv(conn) - if req is None: - break - name, args = req - if name == "__shutdown__": - return - try: - send(conn, (True, dispatch(name, args))) - except Exception as e: # noqa: BLE001 -- report, keep serving - send(conn, (False, f"{type(e).__name__}: {e}")) - except (ConnectionError, OSError): - pass - finally: - conn.close() - finally: - try: - import idapro - idapro.close_database(save=False) - except Exception: # noqa: BLE001 - pass - try: - os.unlink(sockpath) - except OSError: - pass - - -def main(argv=None) -> None: - argv = argv if argv is not None else sys.argv[1:] - if len(argv) < 2: - sys.stderr.write( - "usage: python -m idatui.worker [ida-load-args]\n") - raise SystemExit(2) - try: - serve(argv[0], argv[1], argv[2] if len(argv) > 2 else "") - except SystemExit: - raise - except BaseException as e: # noqa: BLE001 -- surface a clean cause + code 1 - import traceback - sys.stderr.write(f"\nWORKER-FATAL: {type(e).__name__}: {e}\n") - traceback.print_exc() - sys.stderr.flush() - raise SystemExit(1) - - -if __name__ == "__main__": - main() diff --git a/idatui/worker_client.py b/idatui/worker_client.py deleted file mode 100644 index 173cd9d..0000000 --- a/idatui/worker_client.py +++ /dev/null @@ -1,266 +0,0 @@ -"""WorkerClient — a drop-in replacement for ``IDAClient`` backed by our own -idalib worker (``idatui.worker``) over a unix socket instead of ida-pro-mcp's -HTTP/JSON transport. - -It exposes exactly the surface the app/domain use on the client -(``call``/``call_envelope``/``connect``/``set_db``/``resolve_db``/ -``list_sessions``/``health``/``keepalive``/``close``) and returns byte-identical -payloads (the worker calls the same tool functions), so ``domain.py`` and the -app are unchanged — you just construct a WorkerClient instead of an IDAClient. - -Concurrency: the app fires calls from several worker threads over one client; -the worker is single-threaded, so calls are serialized under a lock (the worker -processes one tool at a time anyway — and at ~50us/call that's free). -""" -from __future__ import annotations - -import os -import socket -import subprocess -import sys -import threading -import time -import uuid -from typing import Any - -from .errors import IDAToolError, IDAConnectionError, Session -from .worker import recv as _recv -from .worker import send as _send - -_WORKER_PY = os.path.join(os.path.dirname(os.path.abspath(__file__)), "worker.py") -_worker_python_cache: str | None = None - - -def _find_worker_python() -> str: - """A python that can import ``ida_pro_mcp`` (and thus idalib) — NOT necessarily - the TUI's python. On a typical box the TUI runs under a venv that has textual - + idalib but not ida_pro_mcp, while the system python has idalib + - ida_pro_mcp. Override with IDATUI_WORKER_PYTHON.""" - global _worker_python_cache - if _worker_python_cache: - return _worker_python_cache - override = os.environ.get("IDATUI_WORKER_PYTHON") - candidates = [override] if override else [] - candidates += ["/usr/bin/python", "/usr/bin/python3", sys.executable] - for py in candidates: - if not py or not os.path.exists(py): - continue - try: - r = subprocess.run([py, "-c", "import ida_pro_mcp"], - capture_output=True, timeout=30) - if r.returncode == 0: - _worker_python_cache = py - return py - except Exception: # noqa: BLE001 - continue - return sys.executable # last resort; the worker will report the real error - - -class _NoopKeepAlive: - """The worker is ours and never idles out, so keepalive is a no-op.""" - - def __init__(self) -> None: - self.beats = self.failures = 0 - - def start(self): - return self - - def stop(self) -> None: - pass - - -class WorkerClient: - def __init__(self, binary_path: str, *, ttl: int = 0, - python: str | None = None, load_args: str = "") -> None: - self._bin = os.path.abspath(os.path.expanduser(binary_path)) - self._load_args = load_args or "" # IDA switches for a headerless blob - self._python = python or _find_worker_python() - tag = f"{os.getpid()}-{uuid.uuid4().hex[:8]}" - self._sock_path = f"/tmp/idatui-worker-{tag}.sock" - self._log_path = f"/tmp/idatui-worker-{tag}.log" - self._proc: subprocess.Popen | None = None - self._sock: socket.socket | None = None - self._sid = uuid.uuid4().hex[:8] - self._lock = threading.Lock() # serialize socket use - self._spawn_lock = threading.Lock() - #: Set by close(). A dropped socket is respawned on the next call (the - #: worker segfaulted and we want it back); a CLOSED one must not be. - #: Teardown and binary-switch both close while @work threads are still - #: in flight, so without this, quitting during a decompile spawned a - #: fresh idalib worker that re-opened the database nobody was looking - #: at any more -- a stray process holding the .i64 we just released. - self._closed = False - - # -- lifecycle --------------------------------------------------------- # - def connect(self, timeout: float = 1800.0, progress=None) -> "WorkerClient": - """Spawn the worker (opens + analyzes the DB) and connect once ready.""" - with self._spawn_lock: - self._closed = False # an explicit reconnect revives this client - if self._sock is not None: - return self - if self._proc is None or self._proc.poll() is not None: - # run worker.py as a SCRIPT (not -m idatui.worker) so we don't - # import the textual-dependent idatui package __init__ under the - # IDA python, which usually has no textual. - argv = [self._python, _WORKER_PY, self._sock_path, self._bin] - if self._load_args: - argv.append(self._load_args) - self._proc = subprocess.Popen( - argv, - stdout=open(self._log_path, "wb"), - stderr=subprocess.STDOUT, stdin=subprocess.DEVNULL, - ) - deadline = time.time() + timeout - t0 = time.time() - # Poll fast at first, then back off. A flat 0.2s cost every caller a - # fifth of a second even when the worker was ready in milliseconds - # (a small binary, or a seeded .i64), which is most of the time in - # the tests and noticeable on a re-open. - # - # Backing off all the way to 0.2s was too eager: a seeded database - # is ready at ~250ms, by which point the delay has grown to 134ms, so - # every open waited ~350ms whatever the binary -- the same number for - # a 47KB `echo` and a 1.2MB `bash`, which is what gives a polling - # artefact away. Cap the backoff at 25ms instead: the overshoot on a - # fast open is bounded by that, and 40 probes a second is nothing - # next to an auto-analysis that runs for minutes. - # - # Do NOT be tempted to hold the 5ms rate instead. This poll runs on a - # background thread while the UI thread is drawing, and 200 wakeups a - # second through a cold analysis cost enough GIL time to delay the - # app's own startup -- it left the loading overlay up long enough for - # project mode's first keypress to land on it. - delay = 0.005 - while time.time() < deadline: - try: - s = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) - s.connect(self._sock_path) - self._sock = s - return self - except OSError: - if self._proc.poll() is not None: - raise IDAConnectionError( - f"worker exited (code {self._proc.returncode}): " - f"{self._log_tail()} [full log: {self._log_path}]") - if progress: - progress(f"auto-analyzing {os.path.basename(self._bin)}… " - f"({int(time.time() - t0)}s)") - time.sleep(delay) - delay = min(delay * 1.6, 0.025) - raise IDAConnectionError("worker did not become ready in time") - - @property - def pid(self) -> int | None: - """The worker process id (for memory accounting), or None if not spawned.""" - return self._proc.pid if self._proc is not None else None - - def close(self, grace: float = 20.0) -> None: - """Shut the worker down cleanly. - - After ``__shutdown__`` the worker still has to ``close_database()``, which - re-packs the ``.i64`` and removes the unpacked ``.id0/.id1/...`` scratch. - Signalling it before that finishes is what leaves databases wedged, so - wait out the grace period first and only escalate if it really is stuck. - """ - with self._lock: - s = self._sock - self._sock = None - self._closed = True - if s is not None: - try: - _send(s, ("__shutdown__", {})) - except Exception: # noqa: BLE001 - pass - try: - s.close() - except Exception: # noqa: BLE001 - pass - if self._proc is not None: - try: - self._proc.wait(timeout=grace) # let it close the DB properly - except Exception: # noqa: BLE001 -- TimeoutExpired: it's stuck - try: - self._proc.terminate() - self._proc.wait(timeout=5) - except Exception: # noqa: BLE001 - try: - self._proc.kill() - except Exception: # noqa: BLE001 - pass - - # -- the call surface -------------------------------------------------- # - def call(self, tool: str, *, timeout: float | None = None, **args) -> Any: - if self._closed: - raise IDAConnectionError( - f"{tool}: this worker was closed (call connect() to revive it)") - if self._sock is None: - self.connect() - with self._lock: - s = self._sock - if s is None: - raise IDAConnectionError("worker connection is closed") - try: - _send(s, (tool, args)) - reply = _recv(s) - except (OSError, ConnectionError) as e: - self._sock = None - raise IDAConnectionError(f"worker transport failed: {e}") from e - if reply is None: - self._sock = None - raise IDAConnectionError("worker closed the connection") - ok, payload = reply - if not ok: - raise IDAToolError(tool, str(payload)) - return payload - - def call_envelope(self, tool: str, *, timeout: float | None = None, - **args) -> dict: - # domain.decompile() reads result.structuredContent — mirror that shape. - return {"result": {"structuredContent": self.call(tool, timeout=timeout, - **args)}} - - # -- session shims (single-DB worker) --------------------------------- # - def set_db(self, db: str | None) -> None: - if db: - self._sid = db - - def resolve_db(self) -> str: - return self._sid - - def list_sessions(self) -> list[Session]: - return [Session(session_id=self._sid, - filename=os.path.basename(self._bin), - input_path=self._bin, is_active=True)] - - def health(self) -> dict: - try: - return self.call("server_health") - except IDAToolError: - return {"module": os.path.basename(self._bin), "ok": True} - - def keepalive(self, interval: float = 120.0) -> _NoopKeepAlive: - return _NoopKeepAlive() - - def _log_tail(self, n: int = 400) -> str: - """Last meaningful line(s) of the worker log (skip IDA's licence banner), - so a startup crash surfaces the real cause instead of just 'code 1'.""" - try: - with open(self._log_path, encoding="utf-8", errors="replace") as f: - lines = [ln.strip() for ln in f if ln.strip()] - except OSError: - return "(no worker log)" - # the worker prints a clean 'WORKER-FATAL: ...' line on a startup crash - for ln in reversed(lines): - if ln.startswith("WORKER-FATAL:"): - return ln[len("WORKER-FATAL:"):].strip()[-n:] - skip = ("thank you", "licensed to", "[mcp]", "ida ", "hex-rays") - meaningful = [ln for ln in lines - if not any(s in ln.lower() for s in skip)] - return " | ".join((meaningful or lines)[-3:])[-n:] - - # context manager parity with IDAClient - def __enter__(self) -> "WorkerClient": - return self.connect() - - def __exit__(self, *exc) -> None: - self.close() -- cgit v1.3.1-sl0p From 6e58e9e4ea72771619ded48b0ac2390a58d31cfb Mon Sep 17 00:00:00 2001 From: blasty Date: Fri, 7 Aug 2026 12:56:27 +0200 Subject: codemode: fix DatabaseHandle.open kwarg, and check kwargs against the real signature ida-codemode is now cloned at ../ida-codemode (0.3.1) and installed into ~/ida-venv, so the adapter can be checked against the library instead of against assumptions. First thing it found: connect() passed loading_address=, which DatabaseHandle.open() does not have. The real parameter is image_base, and it already wants the natural 16-byte-aligned address we compute, so this is a rename. Every connect would have died with TypeError on the first call. The port's own contract test could not catch it: its fake handle takes **kwargs, so any keyword at all looks accepted. The test now also validates the keywords we send against inspect.signature(DatabaseHandle.open) when the library is importable, and skips that one check when it is not. Offline suite: 302 passed with the library installed, 302 without it. --- idatui/codemode_client.py | 5 ++++- tests/test_codemode_client.py | 24 +++++++++++++++++++++++- 2 files changed, 27 insertions(+), 2 deletions(-) (limited to 'idatui') diff --git a/idatui/codemode_client.py b/idatui/codemode_client.py index 77a0227..b43bfca 100644 --- a/idatui/codemode_client.py +++ b/idatui/codemode_client.py @@ -975,7 +975,10 @@ class CodeModeClient: timeout=max(0.1, timeout), output_database=self._output_database, processor=self._processor, - loading_address=self._loading_address, + # DatabaseHandle calls this image_base and wants the + # natural (16-byte aligned) address; it does the + # conversion to IDA's paragraph-based -b itself. + image_base=self._loading_address, file_type=self._file_type, new_database=self._new_database, ) diff --git a/tests/test_codemode_client.py b/tests/test_codemode_client.py index 0954918..2f70ba3 100644 --- a/tests/test_codemode_client.py +++ b/tests/test_codemode_client.py @@ -76,6 +76,21 @@ class FakeDatabaseHandle: return FakeHandle(path) +def _open_kwargs_are_real(sent: dict): + """(ok, detail) for the kwargs the adapter passes to DatabaseHandle.open. + + Skips (passes) when ida_codemode is not installed, so the file stays pure. + """ + try: + import inspect + from ida_codemode.client import DatabaseHandle as Real + except ImportError: + return True, "ida_codemode not installed - signature not checked" + accepted = set(inspect.signature(Real.open).parameters) + unknown = sorted(set(sent) - accepted) + return not unknown, f"open() rejects {unknown}" + + def main() -> int: proc, base, file_type = _parse_load_args("-parm:ARMv7-M -b800000 -TRaw") check("legacy switches map to typed Code Mode options", @@ -103,8 +118,15 @@ def main() -> int: FakeDatabaseHandle.opened == path and handle is not None) check("typed loader options cross the dependency boundary", FakeDatabaseHandle.kwargs["processor"] == "arm:ARMv7-A" - and FakeDatabaseHandle.kwargs["loading_address"] == 0x1000, + and FakeDatabaseHandle.kwargs["image_base"] == 0x1000, FakeDatabaseHandle.kwargs) + # A fake that swallows **kwargs cannot catch a keyword the real + # library does not have -- which is exactly how this port shipped + # `loading_address` (the real name is `image_base`) and would have + # raised TypeError on the very first connect. Check the names we + # send against the real signature whenever it is importable. + check("every open() keyword exists in the real library", + *_open_kwargs_are_real(FakeDatabaseHandle.kwargs)) check("connect waits for Code Mode autoanalysis", handle.waited == 42, getattr(handle, "waited", None)) check("progress distinguishes discovery and backend attachment", -- cgit v1.3.1-sl0p From 5de318682c527976f3a298d38a0dc947c3a8c0dd Mon Sep 17 00:00:00 2001 From: blasty Date: Fri, 7 Aug 2026 13:02:27 +0200 Subject: codemode: restore the graph view and pseudocode comments Verified live now: ida-codemode 0.3.1 spawns a managed idalib worker on this box, so the pilot suite runs against the port. flowchart: the port simply does not have the operation, so domain.get_flowchart returned None and every graph key reported 'no control-flow graph for this function'. Ported ours onto ida_gdl (ida-domain exposes no basic-block or edge-kind surface). Blocks stay address RANGES, never text -- that is what lets graph boxes reuse the listing's own rows. Graph suite: 0 -> 50 passed. set_comments: the port set only the disassembly comment via db.comments.set_at(), so a comment never appeared in the pseudocode. A Hex-Rays comment is anchored to a ctree location and an anchor the ctree does not own is discarded as an orphan, so the itp slot must be searched until one sticks, and the entry ea is a function comment instead. Ported that logic back. survey_binary: added as the (caught) fallback domain.py expects behind file_regions, so the fallback path is real rather than always empty. --- idatui/codemode_client.py | 102 +++++++++++++++++++++++++++++++++++++++++++--- 1 file changed, 97 insertions(+), 5 deletions(-) (limited to 'idatui') diff --git a/idatui/codemode_client.py b/idatui/codemode_client.py index b43bfca..0eac29a 100644 --- a/idatui/codemode_client.py +++ b/idatui/codemode_client.py @@ -719,16 +719,60 @@ for query in queries: result = {"result": all_results} result ''', + # A comment must land in BOTH views, and the pseudocode half is not a + # simple set: db.comments.set_at() alone leaves the pseudocode unchanged. + # Hex-Rays comments are anchored to a ctree location (treeloc_t), and an + # anchor the ctree does not actually own is dropped as an "orphan" -- so the + # itp slot has to be searched until one sticks, exactly as IDA's own UI does. + # Without it a comment silently never appears in the decompilation. "set_comments": r''' +import idaapi, idc, ida_hexrays rows = [] for item in a.get("items", []): - ea, text = int(str(item["addr"]), 16), str(item.get("comment") or "") + addr_s = str(item.get("addr", "")) + text = str(item.get("comment") or "") try: - if text: ok = bool(db.comments.set_at(ea, text)) - else: db.comments.delete_at(ea); ok = True - rows.append({"addr": hex(ea), "ok": ok}) + ea = int(addr_s, 16) + if not idaapi.set_cmt(ea, text, False): + rows.append({"addr": addr_s, + "error": f"Failed to set disassembly comment at {hex(ea)}"}) + continue + if not ida_hexrays.init_hexrays_plugin(): + rows.append({"addr": addr_s}); continue + try: + cfunc = ida_hexrays.decompile(ea) + except Exception: + cfunc = None + if cfunc is None: + rows.append({"addr": addr_s}); continue + if ea == cfunc.entry_ea: + # The signature line carries no ctree item: it is a function comment. + idc.set_func_cmt(ea, text, True) + cfunc.refresh_func_ctext() + rows.append({"addr": addr_s}); continue + eamap = cfunc.get_eamap() + if ea not in eamap: + rows.append({"addr": addr_s, + "error": f"Failed to set decompiler comment at {hex(ea)}"}) + continue + nearest_ea = eamap[ea][0].ea + if cfunc.has_orphan_cmts(): + cfunc.del_orphan_cmts(); cfunc.save_user_cmts() + tl = idaapi.treeloc_t(); tl.ea = nearest_ea + placed = False + for itp in range(idaapi.ITP_SEMI, idaapi.ITP_COLON): + tl.itp = itp + cfunc.set_user_cmt(tl, text) + cfunc.save_user_cmts() + cfunc.refresh_func_ctext() + if not cfunc.has_orphan_cmts(): + placed = True; break + cfunc.del_orphan_cmts(); cfunc.save_user_cmts() + rows.append({"addr": addr_s} if placed else + {"addr": addr_s, + "error": f"Failed to set decompiler comment at {hex(ea)}"}) except Exception as exc: - rows.append({"addr": hex(ea), "ok": False, "error": str(exc)}) + rows.append({"addr": addr_s, "error": str(exc)}) result = {"result": rows} result ''', @@ -923,6 +967,54 @@ else: result ''' +# The graph view's only backend call. Blocks are address RANGES, never text: +# the client re-renders them with `heads`, so boxes reuse the exact listing rows +# (colours, operand marks, trail painting) instead of growing a second renderer. +# +# ida-domain exposes no basic-block/edge-kind surface, so this stays on ida_gdl. +_OPERATIONS["flowchart"] = r''' +import ida_funcs, ida_gdl +ea = int(str(a["addr"]), 16) +fn = ida_funcs.get_func(ea) +if fn is None: + result = {"addr": hex(ea), "error": "no function at that address", "blocks": []} +else: + fc = ida_gdl.FlowChart(fn, flags=ida_gdl.FC_PREDS) + index, order = {}, [] + for bb in fc: + index[bb.start_ea] = len(order) + order.append(bb) + blocks = [] + for bb in order: + sl = [s for s in bb.succs() if s.start_ea in index] + succs = [] + for s in sl: + # Edge kind is what the graph view colours by: an n-way dispatch is + # "switch", a successor that is literally the next address falls + # through, anything else is a taken branch. + if len(sl) > 2: kind = "switch" + elif s.start_ea == bb.end_ea: kind = "fall" + else: kind = "jump" + succs.append([index[s.start_ea], kind]) + blocks.append({"id": index[bb.start_ea], "start": hex(int(bb.start_ea)), + "end": hex(int(bb.end_ea)), "succs": succs}) + result = {"addr": hex(ea), + "func": {"addr": hex(int(fn.start_ea)), "end": hex(int(fn.end_ea)), + "name": ida_funcs.get_func_name(fn.start_ea) or ""}, + "entry": index.get(fn.start_ea, 0), "blocks": blocks} +result +''' + +# Only ever reached as domain.py's fallback when file_regions yields nothing. +_OPERATIONS["survey_binary"] = r''' +segments = [] +for seg in db.segments.get_all(): + segments.append({"start": hex(int(seg.start_ea)), "end": hex(int(seg.end_ea)), + "name": db.segments.get_name(seg) or ""}) +result = {"segments": segments} +result +''' + class CodeModeClient: """A leased GUI/idalib database accessed through ``ida_codemode``.""" -- cgit v1.3.1-sl0p From 3edb58cc0b93fdd90b68c6aacc8af91874142099 Mon Sep 17 00:00:00 2001 From: blasty Date: Fri, 7 Aug 2026 13:10:53 +0200 Subject: codemode: port the xref tools' real contract, order included The pseudocode follow's address fallback broke: following a call landed on the NEXT LINE instead of the callee (decomp_nav's stale-name check, cur=0x20dd want=0x2060). The port's xref_query returned rows in raw IDA order, and at a call site IDA yields the ordinary-flow xref (fl_F, the next instruction) before the call xref (fl_CN), so 'first code xref' picked the fall-through. The tool ida-tui was written against sorts rows by the far-end address and dedups by default; sorted, 0x2060 precedes 0x210e and the follow is correct. That ordering is load-bearing, so it is now part of the port rather than an accident of the old implementation. Also fixed: the port attached 'fn' to ref.from_ea for both directions, where a from-xref must describe its TARGET (the xref dialog shows the wrong function otherwise), and the envelope was missing direction/addr/total/next_offset/resolved_addr. xref_types (ours, the kind badges in the xref dialog) is ported verbatim and deliberately stays UNsorted -- that dialog lists xrefs in IDA's own order. decomp_nav, follow_xrefs, xref_labels, decomp_follow_self: 15 passed, 0 failed. --- idatui/codemode_client.py | 143 +++++++++++++++++++++++++++++++++++----------- 1 file changed, 111 insertions(+), 32 deletions(-) (limited to 'idatui') diff --git a/idatui/codemode_client.py b/idatui/codemode_client.py index 0eac29a..18504a5 100644 --- a/idatui/codemode_client.py +++ b/idatui/codemode_client.py @@ -673,49 +673,128 @@ for query in a.get("queries", []): result = {"result": rows} result ''', + # Ours: the coarse code/data type plus a fine `kind` (call/jump/flow, + # read/write/offset/text/info) that the xref dialog draws its badges from. + # Deliberately NOT sorted -- the dialog lists xrefs in IDA's own order. "xref_types": r''' +import idaapi, idautils, ida_bytes, ida_funcs, ida_xref +code_kind = {ida_xref.fl_CF: "call", ida_xref.fl_CN: "call", ida_xref.fl_JF: "jump", + ida_xref.fl_JN: "jump", ida_xref.fl_F: "flow"} +data_kind = {ida_xref.dr_O: "offset", ida_xref.dr_W: "write", ida_xref.dr_R: "read", + ida_xref.dr_T: "text", ida_xref.dr_I: "info"} +def _kind(xr): + return (code_kind if xr.iscode else data_kind).get(xr.type, "code" if xr.iscode else "data") +def _fn(ea): + f = ida_funcs.get_func(ea) + return {"addr": hex(int(f.start_ea)), "name": ida_funcs.get_func_name(f.start_ea) or ""} if f else None queries = a.get("queries") or [] all_results = [] for query in queries: - ea, direction = int(str(query["addr"]), 16), str(query.get("direction", "both")) - refs = [] - if direction in ("to", "both"): refs += list(db.xrefs.to_ea(ea)) - if direction in ("from", "both"): refs += list(db.xrefs.from_ea(ea)) - rows, seen = [], set() - for ref in refs: - key = (int(ref.from_ea), int(ref.to_ea), int(ref.type)) - if query.get("dedup") and key in seen: continue - seen.add(key) - fn = db.functions.get_at(int(ref.from_ea)) - kind = ("call" if ref.is_call else "jump" if ref.is_jump else "flow" if ref.is_flow - else "read" if ref.is_read else "write" if ref.is_write else ref.type.name.lower()) - row = {"from": hex(int(ref.from_ea)), "to": hex(int(ref.to_ea)), - "type": "code" if ref.is_code else "data", "kind": kind} - if query.get("include_fn") and fn is not None: - row["fn"] = {"addr": hex(int(fn.start_ea)), "name": db.functions.get_name(fn) or ""} - rows.append(row) - if len(rows) >= int(query.get("count", 2000)): break - all_results.append({"data": rows}) + query = query if isinstance(query, dict) else {"addr": query} + raw = str(query.get("addr", "")).strip() + direction = str(query.get("direction", "to") or "to").lower() + include_fn = bool(query.get("include_fn", True)) + dedup = bool(query.get("dedup", True)) + try: count = int(query.get("count", 2000) or 2000) + except (TypeError, ValueError): count = 2000 + try: target = int(raw, 16) + except ValueError: target = idaapi.get_name_ea(idaapi.BADADDR, raw) + rows = [] + if target is not None and target != idaapi.BADADDR and ida_bytes.is_mapped(target): + if direction in ("to", "both"): + for xr in idautils.XrefsTo(target, 0): + row = {"direction": "to", "addr": hex(int(xr.frm)), "from": hex(int(xr.frm)), + "to": hex(int(target)), "type": "code" if xr.iscode else "data", "kind": _kind(xr)} + if include_fn: row["fn"] = _fn(xr.frm) + rows.append(row) + if direction in ("from", "both"): + for xr in idautils.XrefsFrom(target, 0): + row = {"direction": "from", "addr": hex(int(xr.to)), "from": hex(int(target)), + "to": hex(int(xr.to)), "type": "code" if xr.iscode else "data", "kind": _kind(xr)} + if include_fn: row["fn"] = _fn(xr.to) + rows.append(row) + if dedup: + seen, deduped = set(), [] + for r in rows: + k = (r["direction"], r["from"], r["to"], r["kind"]) + if k in seen: continue + seen.add(k); deduped.append(r) + rows = deduped + rows = rows[:count] + all_results.append({"query": raw, "data": rows, "next_offset": None}) result = {"result": all_results} result ''', + # Mirrors the tool ida-tui was written against, ORDER INCLUDED. The rows are + # sorted by the far-end address and deduped by default, and the pseudocode + # follow's address fallback silently depends on it: at a call site the raw + # IDA order yields the ordinary-flow xref (the next instruction) first, so an + # unsorted result makes "follow the call" land on the following line instead. "xref_query": r''' +import idaapi, idautils, ida_bytes, ida_funcs +def _fn(ea): + f = ida_funcs.get_func(ea) + return {"addr": hex(int(f.start_ea)), "name": ida_funcs.get_func_name(f.start_ea) or ""} if f else None queries = a.get("queries") or [] all_results = [] for query in queries: - ea, direction = int(str(query["addr"]), 16), str(query.get("direction", "both")) - refs = [] - if direction in ("to", "both"): refs += list(db.xrefs.to_ea(ea)) - if direction in ("from", "both"): refs += list(db.xrefs.from_ea(ea)) - rows = [] - for ref in refs[:int(query.get("count", 2000))]: - fn = db.functions.get_at(int(ref.from_ea)) - row = {"from": hex(int(ref.from_ea)), "to": hex(int(ref.to_ea)), - "type": "code" if ref.is_code else "data"} - if query.get("include_fn") and fn is not None: - row["fn"] = {"addr": hex(int(fn.start_ea)), "name": db.functions.get_name(fn) or ""} - rows.append(row) - all_results.append({"data": rows}) + raw = str(query.get("addr", "")).strip() + direction = str(query.get("direction", "both") or "both").lower() + if direction not in ("to", "from", "both"): direction = "both" + xref_type = str(query.get("xref_type", "any") or "any").lower() + if xref_type not in ("any", "code", "data"): xref_type = "any" + include_fn = bool(query.get("include_fn", True)) + dedup = bool(query.get("dedup", True)) + sort_by = str(query.get("sort_by", "addr") or "addr") + descending = bool(query.get("descending", False)) + try: offset = max(0, int(query.get("offset", 0) or 0)) + except (TypeError, ValueError): offset = 0 + try: count = max(0, min(int(query.get("count", 200) or 200), 5000)) + except (TypeError, ValueError): count = 200 + try: + try: target = int(raw, 16) + except ValueError: + target = idaapi.get_name_ea(idaapi.BADADDR, raw) + if target == idaapi.BADADDR: raise ValueError(f"Failed to resolve address/name: {raw}") + if not ida_bytes.is_mapped(target): raise ValueError(f"Address not mapped: {raw}") + rows = [] + if direction in ("to", "both"): + for xr in idautils.XrefsTo(target, 0): + kind = "code" if xr.iscode else "data" + if xref_type != "any" and kind != xref_type: continue + row = {"direction": "to", "addr": hex(int(xr.frm)), "from": hex(int(xr.frm)), + "to": hex(int(target)), "type": kind} + if include_fn: row["fn"] = _fn(xr.frm) + rows.append(row) + if direction in ("from", "both"): + for xr in idautils.XrefsFrom(target, 0): + kind = "code" if xr.iscode else "data" + if xref_type != "any" and kind != xref_type: continue + row = {"direction": "from", "addr": hex(int(xr.to)), "from": hex(int(target)), + "to": hex(int(xr.to)), "type": kind} + if include_fn: row["fn"] = _fn(xr.to) + rows.append(row) + if dedup: + seen, deduped = set(), [] + for row in rows: + key = (row["direction"], row["from"], row["to"], row["type"]) + if key in seen: continue + seen.add(key); deduped.append(row) + rows = deduped + if sort_by == "type": + rows.sort(key=lambda r: (str(r.get("type", "")), int(str(r["addr"]), 16)), reverse=descending) + else: + rows.sort(key=lambda r: int(str(r["addr"]), 16), reverse=descending) + page = rows[offset:offset + count] if count else rows[offset:] + nxt = offset + len(page) + all_results.append({"target": raw, "resolved_addr": hex(int(target)), "direction": direction, + "xref_type": xref_type, "data": page, + "next_offset": nxt if nxt < len(rows) else None, + "total": len(rows), "error": None}) + except Exception as exc: + all_results.append({"target": raw, "resolved_addr": None, "direction": direction, + "xref_type": xref_type, "data": [], "next_offset": None, + "total": 0, "error": str(exc)}) result = {"result": all_results} result ''', -- cgit v1.3.1-sl0p From ddeb7afea4553e73a7c1decd78fde968e40d186e Mon Sep 17 00:00:00 2001 From: blasty Date: Fri, 7 Aug 2026 13:34:03 +0200 Subject: codemode: carry over the listing + operand-format tools, and stop reshipping them This closes the five operations the port was missing and restores the listing's own tooling instead of a re-implementation of it. idatui/remote_tools.py is the port's IDAPython island: `heads` (the continuous listing) and `op_format`/`pc_nums`/`pc_num_format` (`o`/`O`), extracted verbatim from the BODY that server/patch_server.py used to inject. They are real, diffable source shipped to the database process as text, not string literals, because this is the most performance-tuned and behaviour-sensitive code in the project. Why carry `heads` over rather than keep the port's version: the port's rewrite emitted no per-operand extents ("ops"), so no keypress could show which literal it would reformat (opfmt_highlight had no two-operand row to find); it had no digest/`expect` support, so every page was re-sent after any edit; and its span walk was the per-character loop ours had already been rewritten out of. It also dropped struct-member expansion sizing and the func banner/label rows' exact shapes. The library is installed ONCE per database process (sys.modules, keyed by a hash of the source) and then called by name. Code Mode's execute_python builds a fresh namespace per call, so a library exec'd inline is rebuilt every time and its module-level caches thrown away -- the per-line render lru_cache in particular, which the perf work sized to 65536 entries. Installing it once took `heads` count=200 from 181ms to 92ms; the cache reports 211 hits on a second call where it previously reported none. (Extraction footgun recorded: ast FunctionDef.lineno points at `def`, not at the decorators, so a naive slice silently drops @lru_cache.) Also ported: flowchart, survey_binary, and the xref contract. Live pilot suite on targets/echo: 301 passed, 0 failed -- identical to master. Known, quantified, and NOT fixed here: Code Mode's transport is much slower than the unix-socket worker for the listing's paging. heads count=200 is 2.6ms on master vs 92ms here, count=500 is 6.3ms vs 214ms. Roughly half of that is to_jsonable + HTTP framing per call and is inherent to the architecture; the empty round trip alone is 2ms. The digest/`expect` path (unchanged pages) is the main mitigation and is restored. --- idatui/codemode_client.py | 250 +++------ idatui/remote_tools.py | 1367 +++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 1442 insertions(+), 175 deletions(-) create mode 100644 idatui/remote_tools.py (limited to 'idatui') diff --git a/idatui/codemode_client.py b/idatui/codemode_client.py index 18504a5..069cf98 100644 --- a/idatui/codemode_client.py +++ b/idatui/codemode_client.py @@ -15,11 +15,13 @@ IDAPython modules that Code Mode deliberately makes importable. """ from __future__ import annotations +import hashlib import json import os import shlex import threading import time +from pathlib import Path from textwrap import dedent from typing import Any @@ -169,180 +171,6 @@ def _script(args: dict[str, Any], body: str) -> str: return f"import json\na = json.loads({encoded!r})\n{dedent(body).strip()}\n" -# Rich flat-listing generation is the largest ida-domain gap in this port. -# ida-domain can enumerate heads and render plain disassembly, but it does not -# expose undefined runs, IDA colour spans, function banners, or expanded UDT -# members. Keep that IDAPython-only logic isolated in this one operation. -_HEADS = r''' -import ida_bytes, ida_funcs, ida_idaapi, ida_lines, ida_name, ida_nalt, ida_segment, ida_typeinf -start = int(str(a["addr"]), 16) -count = max(1, min(int(a.get("count", 200)), 2000)) -offset = max(0, int(a.get("offset", 0))) -annotate = bool(a.get("annotate", False)) -seg = db.segments.get_at(start) -if seg is None: - result = {"addr": a["addr"], "error": "no segment", "heads": [], "cursor": {"done": True}} -else: - lo, hi = int(seg.start_ea), int(seg.end_ea) - if a.get("end"): - hi = min(hi, int(str(a["end"]), 16)) - - span_names = { - "insn": ("SCOLOR_INSN", "SCOLOR_KEYWORD", "SCOLOR_ASMDIR", "SCOLOR_MACRO"), - "reg": ("SCOLOR_REG",), - "num": ("SCOLOR_NUMBER", "SCOLOR_CHAR", "SCOLOR_BINPREF"), - "str": ("SCOLOR_STRING",), - "name": ("SCOLOR_DATNAME", "SCOLOR_CODNAME", "SCOLOR_LOCNAME", "SCOLOR_IMPNAME", - "SCOLOR_DEMNAME", "SCOLOR_LIBNAME", "SCOLOR_CNAME", "SCOLOR_DNAME", - "SCOLOR_CREF", "SCOLOR_DREF", "SCOLOR_CREFTAIL", "SCOLOR_DREFTAIL"), - "seg": ("SCOLOR_SEGNAME",), - "cmt": ("SCOLOR_AUTOCMT", "SCOLOR_REGCMT", "SCOLOR_RPTCMT", "SCOLOR_VOIDOP"), - "punct": ("SCOLOR_SYMBOL", "SCOLOR_ALTOP", "SCOLOR_HIDNAME"), - "err": ("SCOLOR_ERROR",), - } - tag_kinds = {} - for kind, names in span_names.items(): - for name in names: - value = getattr(ida_lines, name, None) - if isinstance(value, str) and value: - tag_kinds[value[0]] = kind - elif isinstance(value, int): - tag_kinds[chr(value)] = kind - - def spans(tagged): - on, off, esc = "\x01", "\x02", "\x03" - addr_tag = chr(getattr(ida_lines, "COLOR_ADDR", 0x28)) - addr_len = int(getattr(ida_lines, "COLOR_ADDR_SIZE", 16)) - out, stack, buf = [], [], [] - def flush(): - if buf: - out.append([stack[-1] if stack else "text", "".join(buf)]) - buf.clear() - i = 0 - while i < len(tagged): - ch = tagged[i] - if ch == on and i + 1 < len(tagged): - tag = tagged[i + 1] - if tag == addr_tag: - i += 2 + addr_len - continue - flush(); stack.append(tag_kinds.get(tag, "text")); i += 2; continue - if ch == off and i + 1 < len(tagged): - flush() - if stack: stack.pop() - i += 2; continue - if ch == esc and i + 1 < len(tagged): - buf.append(tagged[i + 1]); i += 2; continue - buf.append(ch); i += 1 - flush() - collapsed, previous_space = [], False - for kind, text in out: - acc = [] - for ch in text: - if ch.isspace(): - if previous_space: continue - acc.append(" "); previous_space = True - else: - acc.append(ch); previous_space = False - if acc: collapsed.append([kind, "".join(acc)]) - if collapsed: - collapsed[0][1] = collapsed[0][1].lstrip() - collapsed[-1][1] = collapsed[-1][1].rstrip() - return [[kind, text] for kind, text in collapsed if text] - - def row(ea): - flags = ida_bytes.get_flags(ea) - kind = "code" if ida_bytes.is_code(flags) else ("data" if ida_bytes.is_data(flags) else "unknown") - tagged = ida_lines.generate_disasm_line(ea, 0) or "" - text = " ".join(ida_lines.tag_remove(tagged).split()) if tagged else "" - item = {"ea": hex(ea), "kind": kind, "size": int(ida_bytes.get_item_size(ea)), "text": text} - if tagged: - rich = spans(tagged) - if " ".join("".join(x[1] for x in rich).split()) == text: - item["spans"] = rich - name = ida_name.get_ea_name(ea) - if name: item["name"] = name - return item - - def unknown_row(ea, size): - if size <= 1: return row(ea) - item = {"ea": hex(ea), "kind": "unknown", "size": int(size), "text": f"db {size} dup(?)"} - name = ida_name.get_ea_name(ea) - if name: item["name"] = name - return item - - def members(ea): - tif = db.types.get_at(ea) - if tif is None or not tif.is_udt(): return [] - answer = [] - for member in db.types.get_udt_members(tif): - type_text = member.type.dstr() or "" - text = f"+{member.offset:X} {member.name}" + (f" {type_text}" if type_text else "") - answer.append({"ea": hex(ea + member.offset), "kind": "member", - "size": int(member.size), "text": text}) - return answer - - def is_unknown(ea): - flags = ida_bytes.get_flags(ea) - return not (ida_bytes.is_code(flags) or ida_bytes.is_data(flags)) - def run_end(ea): - nxt = ida_bytes.next_head(ea, hi) - return nxt if nxt != ida_idaapi.BADADDR and ea < nxt <= hi else hi - def advance(ea): - if is_unknown(ea): return run_end(ea) - nxt = ida_bytes.get_item_end(ea) - return nxt if nxt > ea else ea + 1 - def rows_for(ea): - if is_unknown(ea): return [unknown_row(ea, run_end(ea) - ea)] - fn = db.functions.get_at(ea) if annotate else None - at_start = fn is not None and int(fn.start_ea) == ea - answer = [] - if at_start: - name = db.functions.get_name(fn) or f"sub_{ea:X}" - answer += [ - {"ea": hex(ea), "kind": "sep", "size": 0, "text": ""}, - {"ea": hex(ea), "kind": "sep", "size": 0, - "text": "; " + "=" * 15 + " S U B R O U T I N E " + "=" * 15}, - {"ea": hex(ea), "kind": "funchdr", "size": 0, - "text": name + " proc", "name": name}, - ] - item = row(ea) - if at_start: - item["name"] = None - elif annotate and item["kind"] == "code" and item.get("name"): - name = item["name"] - answer.append({"ea": hex(ea), "kind": "label", "size": 0, - "text": name + ":", "name": name}) - item["name"] = None - answer.append(item) - if item["kind"] == "data": answer += members(ea) - if fn is not None and ida_bytes.get_item_end(ea) >= int(fn.end_ea): - name = db.functions.get_name(fn) or f"sub_{int(fn.start_ea):X}" - answer += [ - {"ea": hex(ea), "kind": "funchdr", "size": 0, - "text": name + " endp", "name": name}, - {"ea": hex(ea), "kind": "sep", "size": 0, "text": "; " + "-" * 60}, - ] - return answer - - ea = ida_bytes.get_item_head(start) - if ea == ida_idaapi.BADADDR: ea = start - for _ in range(offset): - if ea >= hi: break - ea = advance(ea) - rows = [] - more = False - while ea != ida_idaapi.BADADDR and ea < hi: - if len(rows) >= count: - more = True; break - rows += rows_for(ea) - ea = advance(ea) - result = {"addr": a["addr"], "heads": rows, - "cursor": {"next": hex(ea)} if more else {"done": True}} -result -''' - - _DECOMP_MAP_HELPER = r''' def line_map(cfunc): import ida_hexrays @@ -1046,6 +874,72 @@ else: result ''' +# `heads` and the operand-format tools are the port's IDAPython island: the +# continuous listing's presentation model (undefined runs, colour spans, operand +# extents, banners, struct members, the digest protocol) and IDA/Hex-Rays number +# formats have no ida-domain surface. Rather than paraphrase ~1100 lines of +# performance-tuned, behaviour-sensitive code into string literals, they stay +# real, diffable source in idatui/remote_tools.py and are shipped to the database +# process as text. Read once at import; the file ships beside this module. +_REMOTE_LIB = (Path(__file__).with_name("remote_tools.py")).read_text(encoding="utf-8") + +#: Versioned by content, so editing remote_tools.py re-installs it instead of +#: silently running the copy a long-lived worker already has. +_REMOTE_MODULE = "_idatui_remote_" + hashlib.sha1( + _REMOTE_LIB.encode("utf-8")).hexdigest()[:12] + +#: Sent back when the database process has not got the library yet; the client +#: installs it and retries once. Amortised, a worker receives it exactly once. +_NEED_LIB = "__idatui_needs_remote_lib__" + +#: Installs the library as a real module in the database process. Persisting it +#: in sys.modules is what makes the module-level caches (the tag maps, and the +#: line-render lru_cache the listing's throughput depends on) survive between +#: calls -- execute_python builds a fresh namespace every time, so a library +#: exec'd inline is rebuilt, and its caches thrown away, on every single call. +_INSTALL_LIB = f''' +import sys, types +_m = types.ModuleType({_REMOTE_MODULE!r}) +exec(compile(a["source"], {_REMOTE_MODULE!r}, "exec"), _m.__dict__) +sys.modules[{_REMOTE_MODULE!r}] = _m +result = True +result +''' + + +def _remote_op(call: str) -> str: + """A snippet that calls one of the carried-over tools by its real signature. + + Costs one short request: the library is imported from the database process's + own sys.modules, not shipped again. + """ + return (f"import sys\n" + f"_m = sys.modules.get({_REMOTE_MODULE!r})\n" + f"result = {{{_NEED_LIB!r}: True}} if _m is None else _m.{call}\n" + f"result\n") + + +_OPERATIONS["op_format"] = _remote_op( + 'op_format(addr=a["addr"], mode=a.get("mode", "cycle"),' + ' col=int(a.get("col", -1)), n=int(a.get("n", -1)))') +_OPERATIONS["pc_nums"] = _remote_op('pc_nums(addr=a["addr"])') +_OPERATIONS["pc_num_format"] = _remote_op( + 'pc_num_format(addr=a["addr"], mode=a.get("mode", "cycle"),' + ' line=int(a.get("line", -1)), col=int(a.get("col", -1)),' + ' ea=a.get("ea", ""), opnum=int(a.get("opnum", -1)))') + +# The listing walker itself. Replaces the port's re-implementation, which +# rendered no per-operand extents (so no keypress could say which literal it +# would reformat) and had no digest/expect support (so every page was re-sent +# after any edit), and whose span walk was the per-character loop our own +# version had already been rewritten to avoid. +_HEADS = _remote_op( + 'heads(addr=a["addr"], count=int(a.get("count", 200)),' + ' offset=int(a.get("offset", 0)), end=a.get("end", ""),' + ' back=bool(a.get("back", False)), annotate=bool(a.get("annotate", False)),' + ' expect=a.get("expect", ""))') + + # The graph view's only backend call. Blocks are address RANGES, never text: # the client re-renders them with `heads`, so boxes reuse the exact listing rows # (colours, operand marks, trail painting) instead of growing a second renderer. @@ -1232,7 +1126,13 @@ class CodeModeClient: if body is None: raise IDAToolError(operation, f"unknown ida-tui Code Mode operation: {operation}") try: - return self.execute_python(_script(args, body), timeout=timeout) + answer = self.execute_python(_script(args, body), timeout=timeout) + if isinstance(answer, dict) and answer.get(_NEED_LIB): + # First call against this database process (or a restarted one). + self.execute_python(_script({"source": _REMOTE_LIB}, _INSTALL_LIB), + timeout=timeout) + answer = self.execute_python(_script(args, body), timeout=timeout) + return answer except IDAToolError as exc: if exc.tool == "execute_python": raise IDAToolError(operation, exc.message) from exc diff --git a/idatui/remote_tools.py b/idatui/remote_tools.py new file mode 100644 index 0000000..41de0de --- /dev/null +++ b/idatui/remote_tools.py @@ -0,0 +1,1367 @@ +"""The IDAPython ida-tui runs inside the Code Mode sandbox. + +Two features have no ida-domain surface at all and are carried over VERBATIM +from the tools ida-tui was developed against (`server/patch_server.py`'s +injected BODY, which the Code Mode port deletes): + +* `heads` -- the continuous listing. ida-domain enumerates defined heads and + renders plain disassembly; the listing also needs coalesced undefined runs, + IDA colour-tag spans, PER-OPERAND EXTENTS, function banners, code labels and + expanded struct members, plus the digest/`expect` protocol the paging layer + uses to skip re-sending a page that has not changed. +* `op_format` / `pc_nums` / `pc_num_format` -- `o`/`O`. IDA's operand types and + Hex-Rays' per-(ea, opnum) numforms are separate sets, and neither is exposed. + +Keeping the originals rather than paraphrasing them is deliberate: this is the +most performance-tuned and most behaviour-sensitive code in the project (the +span walker is a single regex pass because a per-character loop was the most +expensive thing the listing did, and the cycle only offers stops that change +what you see). A re-implementation drifts from it silently. + +This file is SOURCE SHIPPED AS TEXT to the database process; it is never +imported here, because the ida_* modules do not exist in the TUI's interpreter. +`codemode_client` reads it and prepends it to the relevant snippets. Keep it +self-contained: no relative imports, nothing beyond what Code Mode provides. +""" +# ruff: noqa +import re as _re + +from typing import Annotated # the extracted tool signatures still carry these + + +class IDAError(Exception): + """The MCP host's error type; the tools raise/catch it by name.""" + + +def parse_address(addr): + """ida_pro_mcp.utils.parse_address: hex/decimal string, int, or a symbol.""" + if isinstance(addr, int): + return addr + try: + return int(addr, 0) + except ValueError: + import idaapi + ea = idaapi.get_name_ea(idaapi.BADADDR, str(addr).strip()) + if ea != idaapi.BADADDR: + return ea + raise IDAError(f"Not found: {addr!r}") + + +#: Byte-identical to ida_pro_mcp.utils._STRING_OR_SPACES_RE: the pseudocode +#: column coordinates the client holds depend on collapsing exactly the same way. +_IDATUI_STRING_OR_SPACES_RE = _re.compile( + r'"(?:[^"\\]|\\.)*"' # double-quoted string + r"|'(?:[^'\\]|\\.)*'" # single-quoted string / char + r"|[ \t]{2,}" # run of 2+ whitespace (outside strings) +) + + +def compact_whitespace(line: str) -> str: + """ida_pro_mcp.utils.compact_whitespace: collapse runs of 2+ spaces/tabs to + one, preserving string literals.""" + stripped = line.lstrip(" \t") + if not stripped: + return line + lead = line[: len(line) - len(stripped)] + + def _repl(m): + s = m.group() + if s[0] in ('"', "'"): + return s # preserve string content + return " " + + return lead + _IDATUI_STRING_OR_SPACES_RE.sub(_repl, stripped) + + +def _idatui_head_row(ea, flags=None): + """One flat-listing row for the head at ``ea``: kind (code/data/unknown), + byte size, rendered text, and any symbol name. + + ``flags`` lets a caller that already asked for them say so -- the walk in + ``heads`` used to fetch them three times per head (here, in _is_unknown from + _advance, and again from _rows_for). + """ + import ida_bytes + import ida_lines + import ida_name + + f = ida_bytes.get_flags(ea) if flags is None else flags + if ida_bytes.is_code(f): + kind = "code" + elif ida_bytes.is_data(f): + kind = "data" + else: + kind = "unknown" + line = ida_lines.generate_disasm_line(ea, 0) + text, spans, ops = _idatui_line_parts(line) if line else ("", None, None) + row = { + "ea": hex(ea), + "kind": kind, + "size": int(ida_bytes.get_item_size(ea)), + "text": text, + } + if spans is not None: + row["spans"] = spans + # Where each operand sits in `text`. Comes out of the same tag walk + # (free), and is what lets the client show WHICH literal a keypress + # would reformat before you press it. + if ops: + row["ops"] = ops + nm = ida_name.get_ea_name(ea) + if nm: + row["name"] = nm + return row + + +import functools as _idatui_functools + + +import os as _idatui_os + + +_IDATUI_LINE_CACHE = int(_idatui_os.environ.get("IDATUI_LINE_CACHE") or 65536) + + +def _idatui_line_parts(line): + """``(text, spans, ops)`` for one tagged disassembly line -- memoised. + + A function of the tagged line and nothing else, so the same line always + gives the same answer: a rename changes the line, which changes the key. + And listings repeat themselves hard -- 196k lines of bash are 53k distinct + ones, so a 16k-entry cache serves ~70% of them and takes the per-line cost + from 10.4us to 3.9us. This is the most expensive thing the backend does per + listing row, and a jump to an address near the end of a big binary walks + hundreds of thousands of them. + + ``spans`` is None when the tag walk and the plain text disagree about what + the line says (then the text wins and the row renders unhighlighted). + + The returned lists are SHARED between every row that has the same line; + treat them as read-only. Pickle notices the sharing too, so a page of + repetitive disassembly also serialises smaller. + """ + import ida_lines + text = " ".join(ida_lines.tag_remove(line).split()) # collapse the padding + spans, ops = _idatui_spans(line) + # Built from the SAME line as `text`, then whitespace-collapsed identically, + # so the two can never disagree about what the row says. + joined = "".join([t for _k, t in spans]) + if " ".join(joined.split()) != text: + return (text, None, None) + return (text, spans, ops) + + +def _idatui_head_row(ea, flags=None): + """One flat-listing row for the head at ``ea``: kind (code/data/unknown), + byte size, rendered text, and any symbol name. + + ``flags`` lets a caller that already asked for them say so -- the walk in + ``heads`` used to fetch them three times per head (here, in _is_unknown from + _advance, and again from _rows_for). + """ + import ida_bytes + import ida_lines + import ida_name + + f = ida_bytes.get_flags(ea) if flags is None else flags + if ida_bytes.is_code(f): + kind = "code" + elif ida_bytes.is_data(f): + kind = "data" + else: + kind = "unknown" + line = ida_lines.generate_disasm_line(ea, 0) + text, spans, ops = _idatui_line_parts(line) if line else ("", None, None) + row = { + "ea": hex(ea), + "kind": kind, + "size": int(ida_bytes.get_item_size(ea)), + "text": text, + } + if spans is not None: + row["spans"] = spans + # Where each operand sits in `text`. Comes out of the same tag walk + # (free), and is what lets the client show WHICH literal a keypress + # would reformat before you press it. + if ops: + row["ops"] = ops + nm = ida_name.get_ea_name(ea) + if nm: + row["name"] = nm + return row + + +import functools as _idatui_functools + + +import os as _idatui_os + + +_IDATUI_LINE_CACHE = int(_idatui_os.environ.get("IDATUI_LINE_CACHE") or 65536) + + +@_idatui_functools.lru_cache(maxsize=_IDATUI_LINE_CACHE) +def _idatui_line_parts(line): + """``(text, spans, ops)`` for one tagged disassembly line -- memoised. + + A function of the tagged line and nothing else, so the same line always + gives the same answer: a rename changes the line, which changes the key. + And listings repeat themselves hard -- 196k lines of bash are 53k distinct + ones, so a 16k-entry cache serves ~70% of them and takes the per-line cost + from 10.4us to 3.9us. This is the most expensive thing the backend does per + listing row, and a jump to an address near the end of a big binary walks + hundreds of thousands of them. + + ``spans`` is None when the tag walk and the plain text disagree about what + the line says (then the text wins and the row renders unhighlighted). + + The returned lists are SHARED between every row that has the same line; + treat them as read-only. Pickle notices the sharing too, so a page of + repetitive disassembly also serialises smaller. + """ + import ida_lines + text = " ".join(ida_lines.tag_remove(line).split()) # collapse the padding + spans, ops = _idatui_spans(line) + # Built from the SAME line as `text`, then whitespace-collapsed identically, + # so the two can never disagree about what the row says. + joined = "".join([t for _k, t in spans]) + if " ".join(joined.split()) != text: + return (text, None, None) + return (text, spans, ops) + + +_IDATUI_SPAN_KINDS = { + "insn": ("SCOLOR_INSN", "SCOLOR_KEYWORD", "SCOLOR_ASMDIR", "SCOLOR_MACRO"), + "reg": ("SCOLOR_REG",), + "num": ("SCOLOR_NUMBER", "SCOLOR_CHAR", "SCOLOR_BINPREF"), + "str": ("SCOLOR_STRING",), + # NB the real constant names: DATNAME/CODNAME, not "DNAME". Guessing here + # fails silently — an unmapped tag renders as plain body text, so symbols + # just quietly aren't blue and nothing tells you why. + "name": ("SCOLOR_DATNAME", "SCOLOR_CODNAME", "SCOLOR_LOCNAME", + "SCOLOR_IMPNAME", "SCOLOR_DEMNAME", "SCOLOR_LIBNAME", + "SCOLOR_CNAME", "SCOLOR_DNAME", + "SCOLOR_CREF", "SCOLOR_DREF", "SCOLOR_CREFTAIL", "SCOLOR_DREFTAIL"), + "seg": ("SCOLOR_SEGNAME",), + "cmt": ("SCOLOR_AUTOCMT", "SCOLOR_REGCMT", "SCOLOR_RPTCMT", "SCOLOR_VOIDOP"), + "punct": ("SCOLOR_SYMBOL", "SCOLOR_ALTOP", "SCOLOR_HIDNAME"), + "err": ("SCOLOR_ERROR",), +} + + +def _idatui_tag_map(): + """{tag character: kind}, built once from whatever this IDA actually has.""" + import ida_lines + out = {} + for kind, names in _IDATUI_SPAN_KINDS.items(): + for n in names: + v = getattr(ida_lines, n, None) + if isinstance(v, str) and v: + out[v[0]] = kind + elif isinstance(v, int): + out[chr(v)] = kind + return out + + +_IDATUI_TAGS = None + + +_IDATUI_OPND_TAGS = None + + +_IDATUI_CTL = None # re: a tag = one of three control chars plus its argument + + +_IDATUI_TAGINFO = None + + +def _idatui_opnd_tag_map(): + """{tag character: operand index}. IDA wraps each operand of a disassembly + line in COLOR_OPND1..8, so the line already says where operand N starts and + ends -- no need to re-render operands with print_operand to find out (and + the two agree exactly; checked over thousands of instructions).""" + import ida_lines + out = {} + for i in range(1, 9): + v = getattr(ida_lines, "COLOR_OPND%d" % i, None) + if isinstance(v, int): + out[chr(v)] = i - 1 + elif isinstance(v, str) and v: + out[v[0]] = i - 1 + return out + + +def _idatui_spans(line): + """(spans, ops) for a tagged disasm line. + + ``spans`` is [[kind, text], ...] with colour tags resolved; ``ops`` is + [[start, end, n], ...], the extent of each operand in the SAME (collapsed) + coordinates the row's ``text`` uses -- which is what lets a cursor column + name the operand it is standing on. + + Unknown tags become 'text' rather than being dropped: a processor module can + emit a colour we don't classify, and losing the characters would corrupt the + line.""" + global _IDATUI_TAGS, _IDATUI_OPND_TAGS, _IDATUI_CTL, _IDATUI_TAGINFO + import ida_lines + if _IDATUI_TAGS is None: + _IDATUI_TAGS = _idatui_tag_map() + if _IDATUI_OPND_TAGS is None: + _IDATUI_OPND_TAGS = _idatui_opnd_tag_map() + if _IDATUI_CTL is None: + import re as _re + # One capturing split gives [text, tag, text, tag, ..., text] in a + # single C pass. A per-character python loop over the line used to be + # the most expensive thing the `heads` tool did, and a line is ~54 + # characters but only ~13 tags -- everything between two tags is already + # exactly one span's worth of text. + _IDATUI_CTL = _re.compile("([\\x01\\x02\\x03](?s:.))") + if _IDATUI_TAGINFO is None: + _IDATUI_TAGINFO = { + tag: (_IDATUI_TAGS.get(tag, "text"), _IDATUI_OPND_TAGS.get(tag)) + for tag in set(_IDATUI_TAGS) | set(_IDATUI_OPND_TAGS)} + taginfo = _IDATUI_TAGINFO + plain_tag = ("text", None) + on, off, esc = "\x01", "\x02", "\x03" + addr_tag = chr(getattr(ida_lines, "COLOR_ADDR", 0x28)) + addr_len = int(getattr(ida_lines, "COLOR_ADDR_SIZE", 16)) + parts = _IDATUI_CTL.split(line) + spans, stack = [], [] # stack entries: (kind, operand index|None) + kind, opnd = "text", None # state the current run of text belongs to + pend = "" + skip = 0 # characters of an address payload still due + i, n = 0, len(parts) + while i < n: + txt = parts[i] + i += 1 + if skip: + if len(txt) <= skip: + skip -= len(txt) + txt = "" + else: + txt = txt[skip:] + skip = 0 + if txt: + pend += txt + if i >= n: + break + pair = parts[i] + i += 1 + if skip: # a tag INSIDE an address payload: 2 chars + skip = skip - 2 if skip > 2 else 0 + continue + ch = pair[0] + if ch == esc: # escaped literal: keep the char it guards + pend += pair[1] + continue + tag = pair[1] + if ch == on and tag == addr_tag: + # An embedded target address, not display text: 16 hex digits that + # must not reach the screen. Deliberately NOT a span boundary. + skip = addr_len + continue + if pend: + spans.append([kind, pend, opnd]) + pend = "" + if ch == on: + stack.append((kind, opnd)) + kind, o = taginfo.get(tag, plain_tag) + if o is not None: + opnd = o # operands nest: an inner colour keeps the operand + elif stack: + kind, opnd = stack.pop() + else: + kind, opnd = "text", None + if pend: + spans.append([kind, pend, opnd]) + # Collapse IDA's column padding EXACTLY as the plain text does. A run of + # spaces can straddle two spans, so the leading space of a span is dropped + # when the previous one ended in space — otherwise the spans and `text` + # disagree about the line and the row silently loses its highlighting. + # ``" ".join(txt.split())`` splits on exactly what str.isspace() calls + # whitespace, which is what the character walk this replaces tested. + out, prev_space = [], False + for kind, txt, opnd in spans: + core = " ".join(txt.split()) + if core == txt: + # Nothing to collapse and no edge whitespace -- which is the common + # case ("mov", "rax", ", ") and skips both isspace() probes below. + prev_space = False + out.append([kind, txt, opnd]) + continue + if not core: # the span is nothing but padding + if not prev_space: + prev_space = True + out.append([kind, " ", opnd]) + continue + acc = core + if txt[0].isspace() and not prev_space: + acc = " " + acc + if txt[-1].isspace(): + acc += " " + prev_space = acc[-1] == " " + out.append([kind, acc, opnd]) + while out and out[0][1] == " ": + out.pop(0) + while out and out[-1][1] == " ": + out.pop() + if out and out[0][1].startswith(" "): + out[0][1] = out[0][1].lstrip() + if out and out[-1][1].endswith(" "): + out[-1][1] = out[-1][1].rstrip() + out = [s for s in out if s[1]] + # Operand extents, in the coordinates of the collapsed text these spans + # spell out. Adjacent spans of the same operand merge, so an operand like + # ``[rbp+var_40]`` (five differently-coloured tokens) comes back as ONE + # range -- which is the thing a cursor is inside of, and the thing a format + # change applies to. + ops, pos, cur, start = [], 0, None, 0 + for _kind, txt, opnd in out: + if opnd != cur: + if cur is not None and pos > start: + ops.append([start, pos, cur]) + cur, start = opnd, pos + pos += len(txt) + if cur is not None and pos > start: + ops.append([start, pos, cur]) + text = "".join(t for _k, t, _o in out) + trimmed = [] + for lo, hi, k in ops: # don't let a range own trailing space + while hi > lo and text[hi - 1].isspace(): + hi -= 1 + while lo < hi and text[lo].isspace(): + lo += 1 + if hi > lo: + trimmed.append([lo, hi, k]) + return [[k, t] for k, t, _o in out], trimmed + + +def _idatui_rows_digest(rows): + """A value that changes whenever any of ``rows`` would render differently. + + Covers everything a client keeps off a row: address, kind, size, the plain + text, the symbol name and the colour spans (which is what makes it exact + rather than a heuristic -- two lines can collapse to the same text and still + be coloured differently). + + Uses the interpreter's own ``hash``, deliberately. It never has to mean + anything outside this process: the client stores what a page hashed to when + it loaded it and hands the same number back to ask whether the page still + hashes to that. One worker, one process, one hash seed. + """ + acc = 0 + # The per-line render is memoised, so one spans list is shared by every row + # that says the same thing -- about 45% of them within a page. Hash each + # distinct list once and key that by identity, rather than rebuilding a + # tuple of tuples per row (which is the exact cost that was measured and + # removed from the client side for the same reason). + seen = {} + for r in rows: + sp = r.get("spans") + if sp is None: + sh = None + else: + key = id(sp) + sh = seen.get(key) + if sh is None: + sh = seen[key] = hash(tuple(map(tuple, sp))) + acc = hash((acc, r.get("ea"), r.get("kind"), r.get("size"), + r.get("text"), r.get("name"), sh)) + return acc + + +def _idatui_unknown_row(ea, size): + """One collapsed row for a run of ``size`` undefined bytes starting at + ``ea``. A single byte is rendered normally (shows its value); a longer run + collapses to ``db N dup(?)`` so a big .bss/gap doesn't explode into millions + of one-byte rows.""" + import ida_name + + if size <= 1: + return _idatui_head_row(ea) + row = {"ea": hex(ea), "kind": "unknown", "size": int(size), + "text": f"db {size} dup(?)"} + nm = ida_name.get_ea_name(ea) + if nm: + row["name"] = nm + return row + + +def _idatui_struct_member_rows(ea): + """Indented member rows for a struct-typed data item at ``ea`` (expansion), + or [] if it isn't a struct. Top-level fields only.""" + import ida_nalt + import ida_typeinf + import idaapi + + tif = ida_typeinf.tinfo_t() + if not (ida_nalt.get_tinfo(tif, ea) and tif.is_udt()): + return [] + udt = ida_typeinf.udt_type_data_t() + if not tif.get_udt_details(udt): + return [] + rows = [] + for m in udt: + off = m.begin() // 8 + try: + mtype = m.type._print() or "" + except Exception: + mtype = "" + try: + sz = int(m.type.get_size()) + if sz == idaapi.BADSIZE: + sz = 0 + except Exception: + sz = 0 + name = m.name or "" + text = f"+{off:X} {name}" + (f" {mtype}" if mtype else "") + rows.append({"ea": hex(ea + off), "kind": "member", "size": sz, + "text": text}) + return rows + + +def _idatui_func_header_rows(ea): + """IDA-style subroutine banner rows shown just before a function's entry.""" + import ida_funcs + + name = ida_funcs.get_func_name(ea) or "sub_%X" % ea + bar = "=" * 15 + " S U B R O U T I N E " + "=" * 15 + return [ + {"ea": hex(ea), "kind": "sep", "size": 0, "text": ""}, + {"ea": hex(ea), "kind": "sep", "size": 0, "text": "; " + bar}, + {"ea": hex(ea), "kind": "funchdr", "size": 0, + "text": name + " proc", "name": name}, + ] + + +def _idatui_func_footer_rows(ea, func): + """End-of-function marker shown just after a function's last item.""" + import ida_funcs + + name = ida_funcs.get_func_name(func.start_ea) or "sub_%X" % func.start_ea + return [ + {"ea": hex(ea), "kind": "funchdr", "size": 0, + "text": name + " endp", "name": name}, + {"ea": hex(ea), "kind": "sep", "size": 0, "text": "; " + "-" * 60}, + ] + + +def heads( + addr: Annotated[str, "Start address or name to walk from"], + count: Annotated[int, "Max heads to return (default 200, max 2000)"] = 200, + offset: Annotated[int, "Skip first N heads from addr (default 0)"] = 0, + end: Annotated[str, "Optional exclusive end address; default = segment end"] = "", + back: Annotated[bool, "Walk backwards: return the count heads ENDING just before addr, in forward order"] = False, + annotate: Annotated[bool, "Emit IDA-style function boundary banner rows (kind sep/funchdr)"] = False, + expect: Annotated[str, "Digest a caller already holds: the rows are omitted when they still hash to it"] = "", +) -> dict: + """Walk item heads from ``addr`` as a flat listing: every head is rendered + (code OR data OR undefined) via generate_disasm_line and stepped with + next_head/prev_head. Unlike ``disasm`` (code-only, bails at the first data + byte) this shows db/dw/dd/... lines for data and undefined regions — IDA's + real disassembly view. Address-paged: page forward by re-calling with + ``addr`` = the returned cursor.next; page up with ``back=true``.""" + import ida_bytes + import ida_segment + import idaapi + + count = 2000 if count > 2000 else (1 if count < 1 else count) + offset = max(int(offset), 0) + try: + start = parse_address(addr) + except Exception as e: + return {"addr": str(addr), "error": str(e), "heads": [], "cursor": {"done": True}} + seg = ida_segment.getseg(start) + if not seg: + return {"addr": str(addr), "error": "no segment", "heads": [], "cursor": {"done": True}} + lo, hi = seg.start_ea, seg.end_ea + if end: + try: + hi = min(hi, parse_address(end)) + except Exception: + pass + + rows = [] + if back: + # Collect up to (count+offset) heads strictly before `start`, then take + # the window closest to `start`, returned in forward order. + walk = [] + cur = ida_bytes.prev_head(start, lo) + while cur != idaapi.BADADDR and cur >= lo and len(walk) < count + offset: + walk.append(cur) + cur = ida_bytes.prev_head(cur, lo) + walk.reverse() + chosen = walk[: len(walk) - offset] if offset else walk + chosen = chosen[-count:] + rows = [_idatui_head_row(e) for e in chosen] + first = chosen[0] if chosen else start + pea = ida_bytes.prev_head(first, lo) + cursor = {"done": True} if pea == idaapi.BADADDR or pea < lo else {"prev": hex(pea)} + return {"addr": str(addr), "heads": rows, "cursor": cursor} + + # Walk by item END (not next_head): next_head SKIPS undefined bytes, but a + # flat listing must show them (IDA renders undefined as `db ?` lines, and + # navigating to an unmarked address must land ON it). Defined items advance + # by get_item_end; a run of undefined bytes is COLLAPSED into one row (its + # end found in O(1) via next_head, which skips undefined) so a large .bss or + # gap doesn't explode into millions of one-byte rows. + def _is_unknown_f(f): + return not (ida_bytes.is_code(f) or ida_bytes.is_data(f)) + + def _run_end(e): + """End (exclusive) of the undefined run starting at ``e``.""" + nh = ida_bytes.next_head(e, hi) + return nh if (nh != idaapi.BADADDR and e < nh <= hi) else hi + + def _advance(e, f): + if _is_unknown_f(f): + return _run_end(e) + nxt = ida_bytes.get_item_end(e) + return nxt if nxt > e else e + 1 + + # The function the walk is currently inside, reused while it stays inside. + # get_func is ~0.5us and the walk asks per head; a head is nearly always in + # the same function as the one before it. Only ever consulted when ``e`` + # falls in [start_ea, end_ea), so a tail chunk elsewhere cannot be + # misattributed -- checked against get_func over 437k heads of + # bash/ls_ttl/echo with zero disagreements. + fn_cache = [None] + + def _func_at(e): + cur = fn_cache[0] + if cur is not None and cur.start_ea <= e < cur.end_ea: + return cur + cur = idaapi.get_func(e) + fn_cache[0] = cur + return cur + + def _rows_for(e, f): + if _is_unknown_f(f): + return [_idatui_unknown_row(e, _run_end(e) - e)] + func = _func_at(e) if annotate else None + at_start = func is not None and func.start_ea == e + out = [] + if at_start: + out.extend(_idatui_func_header_rows(e)) + row = _idatui_head_row(e, f) + if at_start: + row = dict(row) + row["name"] = None # the name is shown on the proc header line + elif annotate and row.get("kind") == "code" and row.get("name"): + # A code label (loc_XXX/jump target) gets its OWN line at depth 0, + # like IDA; strip it from the instruction row below. + nm = row["name"] + out.append({"ea": hex(e), "kind": "label", "size": 0, + "text": nm + ":", "name": nm}) + row = dict(row) + row["name"] = None + out.append(row) + if row.get("kind") == "data": + out.extend(_idatui_struct_member_rows(e)) # expand struct fields + if func is not None and ida_bytes.get_item_end(e) >= func.end_ea: + out.extend(_idatui_func_footer_rows(e, func)) + return out + + ea = ida_bytes.get_item_head(start) + get_flags = ida_bytes.get_flags + for _ in range(offset): + if ea >= hi or ea == idaapi.BADADDR: + break + ea = _advance(ea, get_flags(ea)) + more = False + while ea != idaapi.BADADDR and ea < hi: + if len(rows) >= count: + more = True + break + f = get_flags(ea) # once per head, not once per consumer + rows.extend(_rows_for(ea, f)) # a struct head expands into member rows + ea = _advance(ea, f) + cursor = {"next": hex(ea)} if more else {"done": True} + dig = _idatui_rows_digest(rows) + out = {"addr": str(addr), "cursor": cursor, "digest": dig, "count": len(rows)} + # ``expect`` says "I already hold a page that hashed to this". The rows are + # built either way -- generate_disasm_line is the floor and there is no way + # to know a line is unchanged without rendering it -- but pickling several + # hundred rows with their colour spans, unpickling them and rebuilding Heads + # is about 40% of what a page costs, and after a rename almost every page + # comes back identical. + # + # It carries the expected value rather than being a yes/no "digest mode" so + # that a page which HAS changed still costs one round trip: asking first and + # fetching afterwards made every changed page two. + if not (expect and str(dig) == expect): + out["heads"] = rows + return out + + +_IDATUI_FMT_CYCLE = ("hex", "dec", "bin", "char", "offset", "default") + + +_IDATUI_FMT_SETTABLE = ("hex", "dec", "oct", "bin", "char", "offset", "seg", + "float", "stack", "default") + + +def _idatui_fmt_nibbles(): + """{format name: IDA operand-type nibble}. Built on call, not at import: + this module is injected into a file that is imported before a database is + open.""" + import ida_bytes + return { + "default": ida_bytes.FF_N_VOID, "hex": ida_bytes.FF_N_NUMH, + "dec": ida_bytes.FF_N_NUMD, "char": ida_bytes.FF_N_CHAR, + "seg": ida_bytes.FF_N_SEG, "offset": ida_bytes.FF_N_OFF, + "bin": ida_bytes.FF_N_NUMB, "oct": ida_bytes.FF_N_NUMO, + "enum": ida_bytes.FF_N_ENUM, "forced": ida_bytes.FF_N_FOP, + "stroff": ida_bytes.FF_N_STRO, "stack": ida_bytes.FF_N_STK, + "float": ida_bytes.FF_N_FLT, "custom": ida_bytes.FF_N_CUST, + } + + +def _idatui_fmt_name(nib): + for name, v in _idatui_fmt_nibbles().items(): + if v == nib: + return name + return "default" + + +def _idatui_op_fmt(ea, n): + """The format operand ``n`` of the item at ``ea`` is currently displayed in. + + Reads the nibble IDA keeps per operand rather than guessing from the text -- + ``1`` renders identically in hex and decimal, so the rendered line cannot + answer this.""" + import ida_bytes + F = ida_bytes.get_flags(ea) + nib = (F >> ida_bytes.get_operand_type_shift(int(n))) & 0xF + return _idatui_fmt_name(nib) + + +def _idatui_op_value(ea, n): + """(value, byte width) of operand ``n``, or (None, 0) if it hasn't got one. + + The value is what decides which formats are OFFERED: a character constant + for 0x38A9 or an offset to an unmapped address are stops worth skipping.""" + import ida_bytes + import ida_ua + + F = ida_bytes.get_flags(ea) + if ida_bytes.is_code(F): + insn = ida_ua.insn_t() + if ida_ua.decode_insn(insn, ea) <= 0: + return None, 0 + try: + op = insn.ops[int(n)] + except Exception: + return None, 0 + if op.type == ida_ua.o_void: + return None, 0 + v = op.value if op.type == ida_ua.o_imm else op.addr + try: + size = int(ida_ua.get_dtype_size(op.dtype)) + except Exception: + size = 0 + return int(v), size + size = int(ida_bytes.get_item_size(ea)) + read = {1: ida_bytes.get_byte, 2: ida_bytes.get_word, + 4: ida_bytes.get_dword, 8: ida_bytes.get_qword}.get(size) + if read is None: + return None, size + try: + return int(read(ea)), size + except Exception: + return None, size + + +def _idatui_printable(v): + """Whether ``v`` would actually render as a character constant. IDA accepts + op_chr on anything and then prints the number anyway, so a cycle that offers + 'char' for 0x18 has a stop where nothing visibly happens.""" + if v is None or v < 0 or v > 0xFFFFFFFF: + return False + bs, x = [], int(v) + while True: + bs.append(x & 0xFF) + x >>= 8 + if not x: + break + return all(0x20 <= b <= 0x7E or b in (9, 10, 13) for b in bs) + + +def _idatui_offset_worth(v): + """Whether 'offset' is worth OFFERING as a cycle stop for value ``v``. + + Making an offset is not free: IDA invents a dummy name at the target + (``off_18``) and that name STAYS once you cycle past it. So the ring only + stops there when the target is already something you could name -- a symbol, + a function, or an item something else references. In a PIE at base 0 half + the small constants in a function are 'mapped' (they land in the ELF + header); ``sub rsp, 18h`` is not a reference and must not offer to become + one on the way past. + + An explicit request still converts anything mapped: that's a decision, not a + keypress that happened to land here. After it, the target HAS a name, so the + ring includes the stop from then on.""" + import ida_bytes + import ida_name + + return bool(v and ida_bytes.is_mapped(v) and ida_name.get_ea_name(v)) + + +def _idatui_op_candidates(ea): + """Operand indices at ``ea`` whose display format is worth changing. + + Immediates and displacements -- the literals. Deliberately NOT: + + * branch targets (o_near/o_far), or every jump on the listing would offer to + become a bare number, on a view you navigate by label; + * memory references (o_mem), e.g. x86-64's RIP-relative ``lea rdi, name``. + IDA prints those from the reference, not from the operand's number format, + so setting one is accepted and changes nothing on screen -- a keypress + that appears to do nothing is worse than one that says it can't. + + An explicit ``n`` still reaches them; this is what a bare cursor picks.""" + import ida_bytes + import ida_ua + + F = ida_bytes.get_flags(ea) + if ida_bytes.is_data(F): + return [0] # a data item's value is operand 0 + if not ida_bytes.is_code(F): + return [] # undefined bytes: IDA refuses a format outright + insn = ida_ua.insn_t() + if ida_ua.decode_insn(insn, ea) <= 0: + return [] + want = (ida_ua.o_imm, ida_ua.o_displ) + out = [] + for i in range(len(insn.ops)): + op = insn.ops[i] + if op.type == ida_ua.o_void: + break + if op.type in want: + out.append(i) + return out + + +def _idatui_op_spans(ea, text): + """[(start, end, n)] -- where each operand sits inside ``text`` (the + whitespace-collapsed line the TUI shows), so a cursor column can name the + operand it is standing on. + + Read out of IDA's own COLOR_OPND markers on the line, which is both free + (the line is generated anyway) and exact. print_operand is kept as a + fallback for a processor module that emits no operand markers -- it agrees + with the tags where both exist, but it re-renders every operand to say so. + """ + import ida_lines + import ida_ua + + line = ida_lines.generate_disasm_line(ea, 0) + if line: + _spans, ops = _idatui_spans(line) + if ops: + return [tuple(o) for o in ops] + + out, pos = [], 0 + for n in range(8): + try: + raw = ida_ua.print_operand(ea, n) + except Exception: + raw = None + if not raw: + continue + op = " ".join(ida_lines.tag_remove(raw).split()) + if not op: + continue + i = text.find(op, pos) + if i < 0: # duplicated operand text (mov eax, eax) + i = text.find(op) + if i < 0: + continue + out.append((i, i + len(op), n)) + pos = i + len(op) + return out + + +def _idatui_line_text(ea): + import ida_lines + line = ida_lines.generate_disasm_line(ea, 0) + return " ".join(ida_lines.tag_remove(line).split()) if line else "" + + +def _idatui_op_text(ea, text, n): + """How operand ``n`` reads on the line, for a message that names it.""" + for lo, hi, i in _idatui_op_spans(ea, text): + if i == int(n): + return text[lo:hi].strip() + return "" + + +def _idatui_apply_fmt(ea, n, fmt): + """Set operand ``n``'s display format. Returns (ok, error).""" + import ida_bytes + import ida_offset + import idaapi + + n = int(n) + if fmt == "default": + return bool(ida_bytes.clr_op_type(ea, n)), "" + if fmt == "offset": + base = ida_offset.calc_offset_base(ea, n) + if base in (idaapi.BADADDR, None) or base < 0: + base = 0 + return bool(ida_offset.op_plain_offset(ea, n, base)), "" + fn = {"hex": ida_bytes.op_hex, "dec": ida_bytes.op_dec, + "oct": ida_bytes.op_oct, "bin": ida_bytes.op_bin, + "char": ida_bytes.op_chr, "seg": ida_bytes.op_seg, + "float": ida_bytes.op_flt, "stack": ida_bytes.op_stkvar}.get(fmt) + if fn is None: + return False, (f"can't set {fmt!r} from a name alone" + if fmt in _idatui_fmt_nibbles() else + f"unknown format {fmt!r}") + return bool(fn(ea, n)), "" + + +def op_format( + addr: Annotated[str, "Address of the instruction or data item"], + mode: Annotated[str, "cycle | back | show | hex | dec | oct | bin | char | offset | stack | default"] = "cycle", + col: Annotated[int, "Cursor column inside the rendered line (-1: first literal)"] = -1, + n: Annotated[int, "Operand index; -1 derives it from ``col``"] = -1, +) -> dict: + """Change how a literal is DISPLAYED (IDA's 'o' family): hex, decimal, + binary, character, or an offset to the address it names. + + The value in the bytes never changes -- only the representation IDA renders + and remembers. ``cycle``/``back`` step the stops that make sense for THIS + operand: 'char' is skipped unless the value prints as one, 'offset' unless + the target is already named, so no press is ever a no-op you have to press + again. ``show`` reports without changing anything. + + A format the ring can't hold (a stack variable, an enum) is reported in + ``warn`` on the way out, with what to do about it -- ``mode`` takes any of + the names above outright, which is also how you put one back. + + Which operand: ``n`` if given, else the one under ``col`` (a column in the + whitespace-collapsed line, as ``heads`` renders it), else the first literal + on the line.""" + import ida_bytes + + try: + ea = ida_bytes.get_item_head(parse_address(addr)) + except Exception as e: + return {"addr": str(addr), "error": str(e)} + + before = _idatui_line_text(ea) + cands = _idatui_op_candidates(ea) + n = int(n) + if n < 0: + n = -1 + if int(col) >= 0: + for lo, hi, i in _idatui_op_spans(ea, before): + if not (lo <= int(col) < hi): + continue + if i in cands: + n = i + break + # The cursor IS on an operand, just not one with a format. The + # client highlights what the cursor is on, so quietly moving to + # a different operand would make that highlight a lie -- say + # which one can be changed instead. + where = before[lo:hi].strip() + alt = (f"; the literal on this line is operand {cands[0]} " + f"({_idatui_op_text(ea, before, cands[0])})" + if cands else "") + return {"addr": hex(ea), "n": i, "text": before, + "error": f"operand {i} ({where}) has no format to " + f"change{alt}"} + if n < 0: + if not cands: + F = ida_bytes.get_flags(ea) + why = ("no literal on this line to reformat" + if ida_bytes.is_code(F) or ida_bytes.is_data(F) else + "undefined bytes have no format to change -- define " + "them first ('d' makes data, 'c' makes code)") + return {"addr": hex(ea), "text": before, "error": why} + n = cands[0] + + cur = _idatui_op_fmt(ea, n) + value, width = _idatui_op_value(ea, n) + mapped = value is not None and value != 0 and ida_bytes.is_mapped(value) + # The ring is a property of the OPERAND, not of what you last pressed: every + # stop is one that changes what you see for this value, and it is the same + # ring at every step, so a lap always comes home. + choices = [f for f in _IDATUI_FMT_CYCLE + if (f != "char" or _idatui_printable(value)) + and (f != "offset" or _idatui_offset_worth(value))] + # A stack variable is deliberately NOT a stop: ``[rbp+var_40]`` is a frame + # member, not a way of writing a number, and IDA's own "is this a stack + # variable" test isn't exposed to Python here (calc_stkvar_struc_offset + # happily answers for ``[r14+8]`` too, which would put a bogus stop in the + # ring). Leaving one is reported instead, with the command that undoes it. + lossy = cur not in choices and cur != "default" + + mode = str(mode or "cycle").lower() + if mode == "show": + return {"addr": hex(ea), "n": n, "format": cur, "prev": cur, + "choices": choices, "text": before, "before": before, + "value": None if value is None else hex(value), + "width": width, "applied": False} + if mode in ("cycle", "back"): + step = 1 if mode == "cycle" else -1 + if cur in choices: + want = choices[(choices.index(cur) + step) % len(choices)] + else: + # Standing on a format the ring can't hold (an enum names a type a + # nibble doesn't record): enter the ring at its end, don't skip a + # stop working out where we "would have" been. + want = choices[0] if step > 0 else choices[-1] + else: + want = mode + if want not in _idatui_fmt_nibbles(): + return {"addr": hex(ea), "n": n, "text": before, + "error": f"unknown format {mode!r}; one of " + + ", ".join(_IDATUI_FMT_SETTABLE)} + if want == "offset" and not mapped: + return {"addr": hex(ea), "n": n, "text": before, "format": cur, + "error": (f"{'0x%x' % value if value is not None else 'this operand'}" + " isn't a mapped address -- an offset to it would" + " invent a name for nothing")} + + ok, err = _idatui_apply_fmt(ea, n, want) + if err: + return {"addr": hex(ea), "n": n, "text": before, "format": cur, + "error": err} + got = _idatui_op_fmt(ea, n) + out = {"addr": hex(ea), "n": n, "prev": cur, "format": got, + "requested": want, "applied": bool(ok), "choices": choices, + "before": before, "text": _idatui_line_text(ea), + "value": None if value is None else hex(value), "width": width} + if not ok: + out["error"] = f"IDA refused {want} on operand {n}" + elif lossy: + out["warn"] = ( + f"operand {n} was {cur} and the ring has no stop there -- " + + (f"'{cur}' sets it again" if cur in _IDATUI_FMT_SETTABLE else + f"{cur} names a type this can't put back, reassign it by hand")) + return out + + +_IDATUI_PC_FMT_CYCLE = ("hex", "dec", "oct", "char", "default") + + +def _idatui_compact(line): + """The ida-pro-mcp whitespace collapse the pseudocode is served through, so + a column in what the client SHOWS can be mapped back to Hex-Rays' line.""" + try: + from ida_pro_mcp.ida_mcp.utils import compact_whitespace + return compact_whitespace(line) + except Exception: + import re as _re + stripped = line.lstrip(" \t") + lead = line[: len(line) - len(stripped)] + return lead + _re.sub(r"[ \t]{2,}", " ", stripped) + + +def _idatui_compact_col(plain, compact, col): + """The inverse of ``_idatui_uncompact_col``: a column in Hex-Rays' own line, + expressed in the collapsed line the client shows.""" + j = 0 + for i in range(min(int(col), len(plain))): + if j < len(compact) and plain[i] == compact[j]: + j += 1 + return j + + +def _idatui_uncompact_col(plain, compact, col): + """Map a column in the collapsed line back to the same character in the + original. The transform only ever DELETES spaces, so walking both in step + and skipping what vanished is exact.""" + i = 0 + for j in range(min(int(col), len(compact))): + c = compact[j] + while i < len(plain) and plain[i] != c: + i += 1 + i += 1 + return min(i, max(len(plain) - 1, 0)) + + +_IDATUI_LIT_CHARS = frozenset("0123456789abcdefABCDEFxXuUlL") + + +def _idatui_lit_extent(plain, x): + """The [start, end) of the literal token containing column ``x``. + + Hex-Rays says WHICH item a column belongs to, but not how wide the printed + literal is -- and it attributes neighbouring punctuation to the same item, + so ``if ( a1 > 1 )`` reports the closing paren as part of the number. The + identity comes from the ctree; the extent is the run of literal characters + around the column, which cannot reach a ``)`` or a space.""" + if x >= len(plain): + return None + if plain[x] == "'": # a character constant: '-' + end = plain.find("'", x + 1) + return (x, end + 1) if end > x else None + lo = plain.rfind("'", 0, x) + if lo >= 0 and plain.find("'", x) > x and "'" in plain[lo:x] and \ + plain[lo:x].count("'") == 1 and " " not in plain[lo:x]: + return (lo, plain.find("'", x) + 1) # inside 'c' + if plain[x] not in _IDATUI_LIT_CHARS: + return None + lo = x + while lo > 0 and plain[lo - 1] in _IDATUI_LIT_CHARS: + lo -= 1 + hi = x + while hi < len(plain) and plain[hi] in _IDATUI_LIT_CHARS: + hi += 1 + if lo > 0 and plain[lo - 1] == "-": # a unary minus is part of it + lo -= 1 + return (lo, hi) + + +def _idatui_pc_nums(cf, sl): + """Every number literal on one pseudocode line, as + [{x0, x1, ea, opnum, value, nbytes, fmt}]. + + Asks Hex-Rays what each column belongs to rather than pattern-matching the + text: a regex over ``v6 = a1 - 1;`` has to guess which of those characters + are a literal, and ``v11`` looks like one.""" + import ida_bytes + import ida_hexrays + import ida_lines + import idaapi + + plain = ida_lines.tag_remove(sl.line) + out = [] + x = 0 + while x < len(plain): + ch = plain[x] + if ch not in _IDATUI_LIT_CHARS and ch != "'": + x += 1 + continue + head, item, tail = (ida_hexrays.ctree_item_t() for _ in range(3)) + if not cf.get_line_item(sl.line, x, True, head, item, tail): + x += 1 + continue + if item.citype != ida_hexrays.VDI_EXPR: + x += 1 + continue + e = item.e + if e.op != ida_hexrays.cot_num: + x += 1 + continue + extent = _idatui_lit_extent(plain, x) + if extent is None: + x += 1 + continue + nf = e.n.nf + opnum = ord(nf.opnum) if isinstance(nf.opnum, str) else int(nf.opnum) + nbytes = (ord(nf.org_nbytes) if isinstance(nf.org_nbytes, str) + else int(nf.org_nbytes)) + ea = int(e.ea) + if ea == idaapi.BADADDR: + x = extent[1] + continue # synthesised: nothing to key on + nib = (nf.flags >> ida_bytes.get_operand_type_shift(opnum)) & 0xF + # Whether this format is the USER's or Hex-Rays' own guess. The nibble + # can't say: an untouched number reads back as whatever it happens to + # be printed as, and cycling from there would skip that stop forever + # (default already looks like it) and never come back to it. + loc = ida_hexrays.operand_locator_t(ea, opnum) + user = (ida_hexrays.user_numforms_find(cf.numforms, loc) + != ida_hexrays.user_numforms_end(cf.numforms)) + out.append({"x0": extent[0], "x1": extent[1], "ea": ea, + "opnum": opnum, "value": int(e.n._value), + "nbytes": nbytes, "user": user, + "fmt": _idatui_fmt_name(nib) if user else "default", + "shown": _idatui_fmt_name(nib)}) + x = extent[1] # past this literal, not into it + return out + + +def pc_nums( + addr: Annotated[str, "Function address (or any address inside it)"], +) -> dict: + """Every number literal in a function's pseudocode, as + [{line, x0, x1, ea, opnum, value, fmt, user}]. + + One call per decompilation, so a client can show WHICH literal the cursor is + on (and reformat exactly that one) without a round trip per cursor move. + Columns are in the same collapsed coordinates the decompile tool serves its + text in, i.e. what the client actually displays.""" + import ida_hexrays + import ida_lines + import idaapi + + if not ida_hexrays.init_hexrays_plugin(): + return {"addr": str(addr), "error": "no decompiler", "nums": []} + try: + f = idaapi.get_func(parse_address(addr)) + except Exception as e: + return {"addr": str(addr), "error": str(e), "nums": []} + if f is None: + return {"addr": str(addr), "error": "no function here", "nums": []} + try: + cf = ida_hexrays.decompile(f.start_ea) + except Exception as e: + return {"addr": hex(f.start_ea), "error": f"decompile failed: {e}", + "nums": []} + if cf is None: + return {"addr": hex(f.start_ea), "error": "decompilation failed", + "nums": []} + sv = cf.get_pseudocode() + out = [] + for i in range(len(sv)): + plain = ida_lines.tag_remove(sv[i].line) + compact = _idatui_compact(plain) + for rec in _idatui_pc_nums(cf, sv[i]): + out.append({ + "line": i, + "x0": _idatui_compact_col(plain, compact, rec["x0"]), + "x1": _idatui_compact_col(plain, compact, rec["x1"]), + "ea": hex(rec["ea"]), "opnum": rec["opnum"], + "value": hex(rec["value"]), "fmt": rec["fmt"], + "shown": rec["shown"], "user": bool(rec["user"]), + }) + return {"addr": hex(f.start_ea), "nums": out, "lines": len(sv)} + + +def pc_num_format( + addr: Annotated[str, "Function address (or any address inside it)"], + mode: Annotated[str, "cycle | back | show | hex | dec | oct | char | default"] = "cycle", + line: Annotated[int, "0-based pseudocode line index"] = -1, + col: Annotated[int, "Cursor column in the DISPLAYED line (-1: first literal)"] = -1, + ea: Annotated[str, "Address of the number instead of line/col"] = "", + opnum: Annotated[int, "Operand number, with ``ea``"] = -1, +) -> dict: + """Change how a number is displayed in the DECOMPILATION (Hex-Rays keeps its + own number formats, per (address, operand), independent of the listing). + + Same stops as ``op_format`` minus the two C can't express: binary (no such + literal -- IDA takes the format and prints decimal anyway) and offset (it + makes the function stop decompiling). Returns the re-rendered line, and + marks the function dirty so the next decompile is the new text.""" + import ida_hexrays + import ida_lines + import idaapi + + if not ida_hexrays.init_hexrays_plugin(): + return {"addr": str(addr), "error": "no decompiler"} + try: + f = idaapi.get_func(parse_address(addr)) + except Exception as e: + return {"addr": str(addr), "error": str(e)} + if f is None: + return {"addr": str(addr), "error": "no function here"} + try: + cf = ida_hexrays.decompile(f.start_ea) + except Exception as e: + return {"addr": hex(f.start_ea), "error": f"decompile failed: {e}"} + if cf is None: + return {"addr": hex(f.start_ea), "error": "decompilation failed"} + + sv = cf.get_pseudocode() + line = int(line) + target = None + if ea: + try: + want_ea = parse_address(ea) + except Exception as e: + return {"addr": hex(f.start_ea), "error": str(e)} + for i in range(len(sv)): + for rec in _idatui_pc_nums(cf, sv[i]): + if rec["ea"] == want_ea and (int(opnum) < 0 + or rec["opnum"] == int(opnum)): + target, line = rec, i + break + if target: + break + elif 0 <= line < len(sv): + nums = _idatui_pc_nums(cf, sv[line]) + if nums: + if int(col) >= 0: + plain = ida_lines.tag_remove(sv[line].line) + x = _idatui_uncompact_col(plain, _idatui_compact(plain), int(col)) + target = next((r for r in nums if r["x0"] <= x < r["x1"]), None) + target = target or nums[0] + else: + return {"addr": hex(f.start_ea), + "error": f"line {line} is outside the {len(sv)}-line decompilation"} + if target is None: + return {"addr": hex(f.start_ea), "line": line, + "text": (ida_lines.tag_remove(sv[line].line).strip() + if 0 <= line < len(sv) else ""), + "error": "no number literal on this line"} + + cur, value = target["fmt"], target["value"] + choices = [c for c in _IDATUI_PC_FMT_CYCLE + if c != "char" or _idatui_printable(value)] + # Same rule as the listing: one ring per literal, every step. A format the + # ring can't hold (an enum set in the GUI) is reported on the way out + # instead of being kept for one lap and then lost. + lossy = cur not in choices and cur != "default" + out = {"addr": hex(f.start_ea), "ea": hex(target["ea"]), + "opnum": target["opnum"], "line": line, "prev": cur, + "format": cur, "shown": target["shown"], "choices": choices, + "value": hex(value), + "before": ida_lines.tag_remove(sv[line].line).strip()} + + mode = str(mode or "cycle").lower() + if mode == "show": + out["text"] = out["before"] + out["applied"] = False + return out + if mode in ("cycle", "back"): + step = 1 if mode == "cycle" else -1 + if cur in choices: + want = choices[(choices.index(cur) + step) % len(choices)] + else: + want = choices[0] if step > 0 else choices[-1] + else: + want = mode + if want in ("bin", "offset", "stack", "seg", "float"): + out["error"] = (f"Hex-Rays has no {want} format for a number " + f"-- set it on the listing instead") + out["text"] = out["before"] + return out + if want not in ("hex", "dec", "oct", "char", "default"): + out["error"] = (f"unknown format {mode!r}; one of hex, dec, oct, " + f"char, default") + out["text"] = out["before"] + return out + + loc = ida_hexrays.operand_locator_t(target["ea"], target["opnum"]) + it = ida_hexrays.user_numforms_find(cf.numforms, loc) + if it != ida_hexrays.user_numforms_end(cf.numforms): + # std::map::insert is a no-op on an existing key, so a format already + # set here would silently win over the new one. + ida_hexrays.user_numforms_erase(cf.numforms, it) + if want != "default": + import ida_bytes + nf = ida_hexrays.number_format_t(target["opnum"]) + nf.flags = ida_bytes.get_operand_flag(_idatui_fmt_nibbles()[want], + target["opnum"]) + try: + nf.org_nbytes = target["nbytes"] + except Exception: + pass + ida_hexrays.user_numforms_insert(cf.numforms, loc, nf) + cf.save_user_numforms() + try: + ida_hexrays.mark_cfunc_dirty(f.start_ea) + except Exception: + pass + + out["format"] = want + out["applied"] = True + if lossy: + out["warn"] = (f"this number was {cur}, which names a type a radix " + f"can't put back -- reassign it in IDA") + try: + cf2 = ida_hexrays.decompile(f.start_ea, + flags=ida_hexrays.DECOMP_NO_CACHE) + sv2 = cf2.get_pseudocode() if cf2 is not None else None + out["text"] = (ida_lines.tag_remove(sv2[line].line).strip() + if sv2 is not None and line < len(sv2) else out["before"]) + except Exception as e: + out["text"] = out["before"] + out["warn"] = f"re-render failed: {e}" + return out -- cgit v1.3.1-sl0p From df4a42e25881531469a0a8399f457ecf58cab099 Mon Sep 17 00:00:00 2001 From: blasty Date: Fri, 7 Aug 2026 13:50:42 +0200 Subject: codemode: rename takes a LIST of edits per category, not just one Found by tests/test_rawimage_rpc.py, which the earlier runs had not covered: every rename_many check failed with {"ok": 0, "failed": 2, "errors": [{"addr": null, "error": "list indices must be integers or slices, not str"}]} The port's rename read each category as a single edit (edit["addr"]), but the batch shape is {func: [{addr,name}, ...], data: [...], local/stack: [...]} -- a list per category, with a single dict accepted as shorthand. Indexing the list with "addr" raised, and because the whole category was one try block the error came back attached to addr=null, naming nothing. That is the entire point of the rename_many RPC verb: a firmware image arrives with hundreds of names from a loader map or an emulator's symbols.json, and applying them one at a time costs a navigation plus two prompt round trips each. Only the single-rename UI path worked. Now mirrors the real tool: one row per EDIT (addr/old/name plus a per-row error), a summary counting edits rather than categories, conflict detection before the write, and dry_run/allow_overwrite/stop_on_error. Renaming a function refreshes Hex-Rays' ctext, whose cache is per function and persisted in the .i64 -- without it the pseudocode keeps calling the old name forever while every other readback reports the new one. Clearing a label with an empty new name is kept as a real request (the scenarios revert with it) rather than being rejected as a missing argument. tests/test_rawimage_rpc.py: 14 passed/7 failed -> 21 passed, 0 failed. --- .gitignore | 1 + idatui/codemode_client.py | 145 +++++++++++++++++++++++++++++++++++++--------- 2 files changed, 120 insertions(+), 26 deletions(-) (limited to 'idatui') diff --git a/.gitignore b/.gitignore index 65cc5f3..ff80117 100644 --- a/.gitignore +++ b/.gitignore @@ -19,3 +19,4 @@ bin/ # core dumps (idalib/SWIG can segfault under differential probes) core core.* +.fastfeedback/ diff --git a/idatui/codemode_client.py b/idatui/codemode_client.py index 069cf98..5e8089b 100644 --- a/idatui/codemode_client.py +++ b/idatui/codemode_client.py @@ -683,36 +683,129 @@ for item in a.get("items", []): result = {"result": rows} result ''', + # Every category takes EITHER one edit or a LIST of them, and the answer is + # one row per edit. The port accepted only a single dict, so any batch path + # (rpc rename_many applying a whole symbol file, which is the entire point of + # that verb) died with "list indices must be integers or slices, not str" and + # reported the failure against addr=null. Mirrors the real tool: conflict + # detection before the write, dry_run/allow_overwrite/stop_on_error, per-row + # addr/old/name, and a summary counting EDITS rather than categories. "rename": r''' -import ida_idaapi, ida_name, ida_typeinf +import idaapi, ida_hexrays, ida_name batch = a.get("batch") or {} -out = {}; ok_count = failed = 0 -for category, edit in batch.items(): - try: - if category == "func": - ea, new = int(str(edit["addr"]), 16), str(edit["name"]) - fn = db.functions.get_at(ea); ok = bool(fn and db.functions.set_name(fn, new)) - elif category == "data": - new = str(edit.get("new") or "") - if edit.get("addr") is not None: ea = int(str(edit["addr"]), 16) - else: ea = int(ida_name.get_name_ea(ida_idaapi.BADADDR, str(edit.get("old") or ""))) - ok = bool(db.names.set_name(ea, new)) - elif category in ("local", "stack"): - ea, old, new = int(str(edit["func_addr"]), 16), str(edit["old"]), str(edit["new"]) - pseudo = db.pseudocode.decompile(ea); var = pseudo.find_local_variable(old) - if var is None: ok = False +dry_run = bool(batch.get("dry_run", False)) +allow_overwrite = bool(batch.get("allow_overwrite", False)) +stop_on_error = bool(batch.get("stop_on_error", False)) + +def _items(value): + if value is None: return [] + if isinstance(value, dict): return [value] + if isinstance(value, list): return [i for i in value if isinstance(i, dict)] + return [] + +def _set_name_checked(ea, new): + conflict = idaapi.get_name_ea(idaapi.BADADDR, new) + if conflict != idaapi.BADADDR and conflict != ea and not allow_overwrite: + return False, f"can't rename at {hex(ea)} as {new!r}: name already used at {hex(conflict)}" + if dry_run: + return True, None + flags = idaapi.SN_CHECK + if allow_overwrite: flags |= int(getattr(idaapi, "SN_FORCE", 0)) + if not idaapi.set_name(ea, new, flags): + return False, (f"Rename failed at {hex(ea)}: IDA rejected name {new!r} " + "(invalid identifier or internal conflict)") + return True, None + +def _refresh_ctext(fn_addr): + # A renamed function must invalidate Hex-Rays' cache, which is per function + # and persisted in the .i64: without this the pseudocode keeps calling the + # old name forever while every other readback reports the new one. + if not ida_hexrays.init_hexrays_plugin(): return + failure = ida_hexrays.hexrays_failure_t() + cfunc = ida_hexrays.decompile_func(fn_addr, failure, ida_hexrays.DECOMP_WARNINGS) + if cfunc: cfunc.refresh_func_ctext() + +out = {}; ok_count = failed = 0; halted = False +for category in ("func", "data", "local", "stack"): + if category not in batch: continue + rows = [] + for edit in _items(batch.get(category)): + try: + if category == "func": + addr_text = edit.get("addr") or edit.get("func_addr") or edit.get("func") + new = edit.get("name") or edit.get("new") or edit.get("new_name") + if not addr_text or not new: + row = {"addr": addr_text, "name": new, + "error": "Function rename requires addr + name"} + else: + ea = int(str(addr_text), 16) + fn = idaapi.get_func(ea) + if fn is None: + row = {"addr": addr_text, "name": new, "error": "Function not found"} + else: + old = idaapi.get_name(fn.start_ea) or None + ok, err = _set_name_checked(fn.start_ea, str(new)) + row = {"addr": addr_text, "old": old, "name": str(new)} + if err: row["error"] = err + if dry_run: row["dry_run"] = True + if ok and not dry_run: _refresh_ctext(fn.start_ea) + elif category == "data": + addr_text = edit.get("addr") + old = edit.get("old") or edit.get("old_name") + new = edit.get("new") or edit.get("new_name") or edit.get("name") + if not new and new != "": + row = {"old": old, "new": None, + "error": "Global rename requires target and new name"} + else: + if addr_text is not None: + ea = int(str(addr_text), 16) + old = old or (idaapi.get_name(ea) or None) + else: + ea = idaapi.get_name_ea(idaapi.BADADDR, str(old or "")) + if ea == idaapi.BADADDR: + row = {"old": old, "new": str(new), "error": f"Global {old!r} not found"} + else: + # An empty new name CLEARS the label; that is a real + # request (tests revert with it), not a missing argument. + if str(new) == "": + ok = bool(ida_name.set_name(ea, "", idaapi.SN_CHECK)) + err = None if ok else f"Failed to clear the name at {hex(ea)}" + else: + ok, err = _set_name_checked(ea, str(new)) + row = {"addr": hex(ea), "old": old, "new": str(new)} + if err: row["error"] = err + if dry_run: row["dry_run"] = True else: - var.set_user_name(new) - ok = bool(pseudo.save_local_variable_info(var, save_name=True)) - else: - raise ValueError(f"unsupported rename category: {category}") - row = {"ok": ok, **({} if ok else {"error": "IDA rejected the name"})} - except Exception as exc: - row = {"ok": False, "error": str(exc)} - out[category] = [row] - if row["ok"]: ok_count += 1 - else: failed += 1 + fa, old, new = edit.get("func_addr"), edit.get("old"), edit.get("new") + if not fa or not old or not new: + row = {"old": old, "new": new, + "error": f"{category} rename requires func_addr + old + new"} + else: + ea = int(str(fa), 16) + pseudo = db.pseudocode.decompile(ea) + var = pseudo.find_local_variable(str(old)) + if var is None: + row = {"func_addr": fa, "old": old, "new": new, + "error": f"no local {old!r} in that function"} + elif dry_run: + row = {"func_addr": fa, "old": old, "new": new, "dry_run": True} + else: + var.set_user_name(str(new)) + ok = bool(pseudo.save_local_variable_info(var, save_name=True)) + row = {"func_addr": fa, "old": old, "new": new} + if not ok: row["error"] = "IDA rejected the local variable name" + except Exception as exc: + row = {"addr": edit.get("addr"), "error": str(exc)} + rows.append(row) + if row.get("error"): failed += 1 + else: ok_count += 1 + if row.get("error") and stop_on_error: + halted = True; break + out[category] = rows + if halted: break out["summary"] = {"ok": ok_count, "failed": failed} +if dry_run: out["summary"]["dry_run"] = True +if halted: out["summary"]["halted"] = True result = out result ''', -- cgit v1.3.1-sl0p From ecc58d7725db6d4929ae3e299dea4f0202a81946 Mon Sep 17 00:00:00 2001 From: blasty Date: Fri, 7 Aug 2026 14:08:12 +0200 Subject: codemode: three defects the A/B benchmark found in the decompiler path Benchmarking the port against master op-by-op (rather than only asking whether tests pass) turned up three real bugs, all in the most user-visible path: opening pseudocode. 1. decompile was doing decomp_map's job. It called the full per-column line map purely to fill in each line's /*0xEA*/ anchor. The tool ida-tui was written against takes ONE get_line_item at column 0 per line; the port took one per COLUMN, i.e. thousands of get_line_item+dstr() calls per function instead of one per line. Every pseudocode open cost the same as opening the split view. Carried the real implementation over: 1888ms -> 53ms. 2. decomp_map used the pre-optimisation line map. Ours memoises obj_id -> ea for the whole function (commit 853d90c: dstr() was 79% of the tool, and consecutive columns report the same ctree item), the port's did not. 1925ms -> 287ms. 3. _idatui_compact imported ida_pro_mcp on every call. Under Code Mode that package is not installed in the database process, so the import failed every time -- and a FAILED import is never cached, so each one re-searched the whole of sys.path: 422 failed imports per pc_nums call, which was most of its runtime. 1428ms -> 257ms. The same bug was a correctness bug hiding behind the perf bug: the fallback path collapsed whitespace INSIDE string literals, where the real function preserves it. Pseudocode columns are served in those coordinates, so on any line containing a string with two spaces, every literal's mark and every reformat would have been placed on the wrong column. It never fired on master because ida_pro_mcp is installed there. Now calls the byte-identical module-level shim directly, with the deviation from the extracted original documented in place. Narrow verification: decomp/split_view/opfmt/follow/comment/structs scenarios, 72 passed, 0 failed. Full gate running separately. --- idatui/codemode_client.py | 65 ++------------- idatui/remote_tools.py | 199 +++++++++++++++++++++++++++++++++++++++++++--- 2 files changed, 198 insertions(+), 66 deletions(-) (limited to 'idatui') diff --git a/idatui/codemode_client.py b/idatui/codemode_client.py index 5e8089b..3bf01cc 100644 --- a/idatui/codemode_client.py +++ b/idatui/codemode_client.py @@ -171,24 +171,6 @@ def _script(args: dict[str, Any], body: str) -> str: return f"import json\na = json.loads({encoded!r})\n{dedent(body).strip()}\n" -_DECOMP_MAP_HELPER = r''' -def line_map(cfunc): - import ida_hexrays - answer = [] - for sl in cfunc.get_pseudocode(): - tagged, eas, seen = sl.line, [], set() - for x in range(len(tagged) + 1): - head = ida_hexrays.ctree_item_t(); item = ida_hexrays.ctree_item_t(); tail = ida_hexrays.ctree_item_t() - if not cfunc.get_line_item(tagged, x, False, head, item, tail): continue - text = item.dstr() or "" - try: ea = int(text.split(": ", 1)[0], 16) - except (ValueError, IndexError): continue - if ea not in seen: seen.add(ea); eas.append(ea) - answer.append(eas) - return answer -''' - - _OPERATIONS: dict[str, str] = { "list_funcs": r''' import fnmatch @@ -812,45 +794,6 @@ result } -_OPERATIONS["decompile"] = _DECOMP_MAP_HELPER + r''' -ea = int(str(a["addr"]), 16) -fn = db.functions.get_at(ea) -if fn is None: - result = {"error": f"no function at {ea:#x}"} -else: - pseudo = db.pseudocode.decompile(fn) - mapping = line_map(pseudo.raw_cfunc) - plain = pseudo.to_text() - marked = [line + (f" /*0x{eas[0]:X}*/" if eas else "") - for line, eas in zip(plain, mapping)] - import ida_name - refs, seen = [], set() - for expr in pseudo.find_objects(): - target = int(expr.obj_ea) - if target in seen or not (db.is_valid_ea(target) or db.is_private_ea(target)): continue - seen.add(target) - name = expr.obj_name or ida_name.get_name(target) or "" - try: string = db.bytes.get_string_at(target) if db.is_valid_ea(target) else None - except Exception: string = None - refs.append({"addr": hex(target), "name": name, "string": string}) - result = {"addr": hex(int(fn.start_ea)), "code": "\n".join(marked), "refs": refs} -result -''' - -_OPERATIONS["decomp_map"] = _DECOMP_MAP_HELPER + r''' -ea = int(str(a["addr"]), 16) -fn = db.functions.get_at(ea) -if fn is None: - result = {"error": f"no function at {ea:#x}"} -else: - pseudo = db.pseudocode.decompile(fn) - mapping = line_map(pseudo.raw_cfunc) - result = {"addr": hex(int(fn.start_ea)), - "lines": [{"ea": hex(eas[0]) if eas else None, - "eas": [hex(item) for item in eas]} for eas in mapping]} -result -''' - _OPERATIONS["define_code_run"] = r''' import ida_bytes, ida_idp, ida_segment, ida_ua, idaapi ea, limit = int(str(a["addr"]), 16), max(1, min(int(a.get("limit", 20000)), 200000)) @@ -877,6 +820,7 @@ else: result ''' + _OPERATIONS["define_func_run"] = r''' import ida_bytes, ida_funcs, ida_segment ea = int(str(a["addr"]), 16) @@ -901,6 +845,7 @@ else: result ''' + _OPERATIONS["set_thumb"] = r''' import ida_bytes, ida_ida, ida_idp, ida_segment, ida_segregs ea = int(str(a["addr"]), 16); treg = ida_idp.str2reg("T") @@ -926,6 +871,7 @@ else: result ''' + _OPERATIONS["thumb_scan"] = r''' import ida_bytes, ida_funcs, ida_idp, ida_segment, ida_segregs, ida_ua lo, hi = int(str(a["start"]), 16), int(str(a["end"]), 16) @@ -950,6 +896,7 @@ result = {"start": hex(lo), "end": hex(hi), "found": found, "applied": applied, result ''' + _OPERATIONS["decomp_error"] = r''' import ida_hexrays, ida_ida ea = int(str(a["addr"]), 16); fn = db.functions.get_at(ea) @@ -1016,6 +963,10 @@ _OPERATIONS["op_format"] = _remote_op( 'op_format(addr=a["addr"], mode=a.get("mode", "cycle"),' ' col=int(a.get("col", -1)), n=int(a.get("n", -1)))') _OPERATIONS["pc_nums"] = _remote_op('pc_nums(addr=a["addr"])') +_OPERATIONS["decompile"] = _remote_op( + 'decompile(addr=a["addr"],' + ' include_addresses=bool(a.get("include_addresses", True)))') +_OPERATIONS["decomp_map"] = _remote_op('decomp_map(addr=a["addr"])') _OPERATIONS["pc_num_format"] = _remote_op( 'pc_num_format(addr=a["addr"], mode=a.get("mode", "cycle"),' ' line=int(a.get("line", -1)), col=int(a.get("col", -1)),' diff --git a/idatui/remote_tools.py b/idatui/remote_tools.py index 41de0de..6fb6436 100644 --- a/idatui/remote_tools.py +++ b/idatui/remote_tools.py @@ -1050,15 +1050,26 @@ _IDATUI_PC_FMT_CYCLE = ("hex", "dec", "oct", "char", "default") def _idatui_compact(line): """The ida-pro-mcp whitespace collapse the pseudocode is served through, so - a column in what the client SHOWS can be mapped back to Hex-Rays' line.""" - try: - from ida_pro_mcp.ida_mcp.utils import compact_whitespace - return compact_whitespace(line) - except Exception: - import re as _re - stripped = line.lstrip(" \t") - lead = line[: len(line) - len(stripped)] - return lead + _re.sub(r"[ \t]{2,}", " ", stripped) + a column in what the client SHOWS can be mapped back to Hex-Rays' line. + + DEVIATION FROM THE EXTRACTED ORIGINAL, deliberately: this used to be + ``from ida_pro_mcp.ida_mcp.utils import compact_whitespace`` inside a + try/except, with a plain ``[ \\t]{2,}`` regex as the fallback. Under Code + Mode ida_pro_mcp is not installed in the database process, so BOTH halves + of that were wrong: + + * the import failed on every call, and a failed import is never cached, so + each one re-searched the whole of sys.path -- 422 failures per pc_nums + call, which was the majority of its runtime; + * the fallback collapses runs of spaces INSIDE STRING LITERALS, which the + real function preserves. Pseudocode columns are served in these + coordinates, so a line containing a string with two spaces would have put + every literal's mark, and every reformat, on the wrong column. + + The module-level shim above is byte-identical to the original regex, so + call it directly. + """ + return compact_whitespace(line) def _idatui_compact_col(plain, compact, col): @@ -1365,3 +1376,173 @@ def pc_num_format( out["text"] = out["before"] out["warn"] = f"re-render failed: {e}" return out + + +def decompile(addr, include_addresses=True): + """Pseudocode for the function at ``addr``, plus the objects it references. + + Faithful to the tool ida-tui was written against, and in particular to its + COST: the per-line address anchor comes from ONE ``get_line_item`` at column + 0 per line. The Code Mode port asked for the full per-column line map (what + ``decomp_map`` is for) purely to fill in that anchor, which is thousands of + ``get_line_item``+``dstr()`` calls per function instead of one per line, and + made every pseudocode open cost the same as opening the split view. + + Text is whitespace-collapsed exactly as the client displays it, because + ``pc_nums`` reports literal columns in those coordinates. + """ + import ida_bytes + import ida_hexrays + import ida_lines + import ida_name + import idaapi + + try: + ea = parse_address(addr) + except Exception as e: + return {"addr": str(addr), "code": None, "error": str(e)} + fn = idaapi.get_func(ea) + if fn is None: + return {"addr": str(addr), "code": None, "error": f"no function at {ea:#x}"} + if not ida_hexrays.init_hexrays_plugin(): + return {"addr": hex(int(fn.start_ea)), "code": None, "error": "no decompiler"} + failure = ida_hexrays.hexrays_failure_t() + try: + cfunc = ida_hexrays.decompile_func(fn, failure) + except Exception as e: + return {"addr": hex(int(fn.start_ea)), "code": None, + "error": f"Decompilation failed at {ea:#x}: {e}"} + if cfunc is None: + return {"addr": hex(int(fn.start_ea)), "code": None, + "error": failure.desc() or f"Decompilation failed at {ea:#x}"} + + lines = [] + for sl in cfunc.get_pseudocode(): + head = ida_hexrays.ctree_item_t() + item = ida_hexrays.ctree_item_t() + tail = ida_hexrays.ctree_item_t() + line_ea = None + if include_addresses and cfunc.get_line_item(sl.line, 0, False, head, item, tail): + parts = (item.dstr() or "").split(": ") + if len(parts) == 2: + try: + line_ea = int(parts[0], 16) + except ValueError: + line_ea = None + text = compact_whitespace(ida_lines.tag_remove(sl.line)) + lines.append(f"{text} /*{line_ea:#x}*/" if line_ea is not None else text) + + refs, seen = [], set() + + class _RefVisitor(ida_hexrays.ctree_visitor_t): + def __init__(self): + ida_hexrays.ctree_visitor_t.__init__(self, ida_hexrays.CV_FAST) + + def visit_expr(self, e): + if e.op == ida_hexrays.cot_obj: + target = int(e.obj_ea) + if target != idaapi.BADADDR and target not in seen: + seen.add(target) + try: + raw = ida_bytes.get_strlit_contents(target, -1, 0) + text = raw.decode("utf-8", "replace") if raw else None + except Exception: + text = None + refs.append({"addr": hex(target), + "name": ida_name.get_name(target) or "", + "string": text}) + return 0 + + try: + _RefVisitor().apply_to(cfunc.body, None) + except Exception: + pass + return {"addr": hex(int(fn.start_ea)), "code": "\n".join(lines), "refs": refs} + + +def decomp_map( + addr: Annotated[str, "Function address or name"], +) -> dict: + """Per-pseudocode-line instruction coverage for the split view's region + highlight: for each line, the set of EAs the decompiler attributes to it, + swept across the line's columns via get_line_item. Shape: + {addr, lines:[{ea: primary|None, eas:[hex,...]}, ...]}.""" + import ida_hexrays + import idaapi + try: + ea = int(str(addr), 16) + except ValueError: + ea = idaapi.get_name_ea(idaapi.BADADDR, str(addr).strip()) + func = idaapi.get_func(ea) + if not func: + return {"error": f"no function at {addr}"} + try: + cfunc = ida_hexrays.decompile(func.start_ea) + except Exception as e: # noqa: BLE001 + return {"error": f"decompile failed: {e}"} + if cfunc is None: + return {"error": "decompile failed"} + import ida_lines + # Three things this loop must not do, each measured on real functions (the 25 + # largest of bash went 68.3s -> 6.5s; echo's 60 largest 5.4s -> 0.6s, with + # byte-identical output): + # + # * allocate ctree_item_t's per COLUMN. They are SWIG objects and this is + # the innermost loop; one per call is enough, and head/tail are never + # read, so don't ask for them at all. + # * sweep the TAGGED length. ``x`` is a screen column but ``sl.line`` still + # carries IDA's colour tags, so a 23-column line was swept 124 times. + # * call dstr() per column. It formats a whole 'EA: description' string -- + # 24us a call, which is 79% of this tool. Comparing against the PREVIOUS + # column's item id is not enough: items interleave, so `foo(a, b)` flips + # call -> arg -> call -> arg and every flip re-formats an item already + # seen (106 594 calls for 15 417 lines of bash). Memoise id -> ea for the + # whole function instead: obj_id is unique within a cfunc, so the same id + # always yields the same string, and the result is deduped by ``seen`` + # anyway. Items with no ctree node (it is None) have no id to key on and + # still pay per occurrence. + item = ida_hexrays.ctree_item_t() + tag_remove = ida_lines.tag_remove + get_line_item = cfunc.get_line_item + ea_of_id = {} + lines = [] + for sl in cfunc.get_pseudocode(): + line = sl.line + eas, seen = [], set() + prev_id = None + for x in range(len(tag_remove(line)) + 1): + if not get_line_item(line, x, False, None, item, None): + continue + it = item.it + if it is not None: + oid = it.obj_id + if oid == prev_id: + continue + prev_id = oid + if oid in ea_of_id: + e = ea_of_id[oid] + if e is not None and e not in seen: + seen.add(e) + eas.append(hex(e)) + continue + else: + oid = None + prev_id = None + # Match the /*ea*/ marker's source (decompile_function_safe): the + # item's dstr() is 'EA: description'; get_ea() reports a different ea. + e = None + dstr = item.dstr() + if dstr: + parts = dstr.split(": ", 1) + if len(parts) == 2: + try: + e = int(parts[0], 16) + except ValueError: + e = None + if oid is not None: + ea_of_id[oid] = e + if e is not None and e not in seen: + seen.add(e) + eas.append(hex(e)) + lines.append({"ea": eas[0] if eas else None, "eas": eas}) + return {"addr": hex(func.start_ea), "lines": lines} -- cgit v1.3.1-sl0p From 21971d0a9ce9b8a2fb1db2ec67303696cea9cbbb Mon Sep 17 00:00:00 2001 From: blasty Date: Fri, 7 Aug 2026 14:32:56 +0200 Subject: tests: blob_ui 39.8s -> 4.1s, and a fatal it was hiding Three separate wastes, all of the same family: waiting on a guess instead of a signal, and paying for work that never had to be repeated. 1. The 64KB blob was built with os.urandom into a fresh TemporaryDirectory on every run. New bytes at a new path means the pristine-database cache can never apply, so full auto-analysis of 64KB of AArch64-decoded noise was paid every single run. It is now built from a seeded PRNG at a stable path (tests/.synthetic/, gitignored) and staged through the existing cache. Determinism is also a correctness fix: whether 64KB of chance bytes contains something IDA reads as a function is luck, and this suite asserts "and really has no functions". 2. `wait(lambda: lst.model is not old, ..., 30)` after commenting. The perf work made an item edit KEEP the listing's walk and re-render in place, so the model object is never replaced and this waited out its full 30s timeout on every run -- and then "commenting leaves the view where it was" passed vacuously, because nothing had happened at all. A test that burns 30s to check nothing is worse than no test. 3. Two `pause(2.0)`/`pause(2.5)` after a carve, replaced with settle() on a real condition. The second one deliberately has NO predicate: that spot is random data, so the carve may legitimately produce nothing, and "the row became code" would never hold -- gating on it cost another 30s timeout. What that check is about is the VIEW not moving, so the gate is "the app finished reacting". Fixing (1) exposed a real bug in the client, fixed here too: reopening a database that already exists while passing loader switches is FATAL in IDA -- FATAL ERROR: Switch '-b400' can be used only when loading a new file which kills the worker before it can report anything. Loader switches describe an IMPORT and are recorded in the database they produce, so they are now sent only when there is an import to describe. This was never reachable from the old suite (a fresh random blob never had a database to reopen), but it is reachable by any user who opens a raw blob with --ida-args twice. 30 passed, 0 failed. --- .gitignore | 1 + idatui/codemode_client.py | 25 +++++++++-- tests/_fixtures.py | 30 +++++++++++++ tests/test_blob_ui.py | 109 +++++++++++++++++++++++++++++++++++----------- 4 files changed, 136 insertions(+), 29 deletions(-) (limited to 'idatui') diff --git a/.gitignore b/.gitignore index ff80117..9a7b5b2 100644 --- a/.gitignore +++ b/.gitignore @@ -20,3 +20,4 @@ bin/ core core.* .fastfeedback/ +tests/.synthetic/ diff --git a/idatui/codemode_client.py b/idatui/codemode_client.py index 3bf01cc..9d91335 100644 --- a/idatui/codemode_client.py +++ b/idatui/codemode_client.py @@ -1062,6 +1062,17 @@ class CodeModeClient: self._last_entry: RegistryEntry | None = None self._connect_lock = threading.Lock() + def _database_exists(self) -> bool: + """Whether the IDB this open would target is already on disk. + + Its loader switches are baked in, so they must not be sent again. + """ + try: + target = self._output_database or expected_idb_path(self._path) + except Exception: # noqa: BLE001 -- resolver unavailable: assume fresh + return False + return bool(target) and os.path.exists(target) + def connect(self, timeout: float = 1800.0, progress=None) -> "CodeModeClient": _require_codemode() with self._connect_lock: @@ -1078,17 +1089,25 @@ class CodeModeClient: deadline = time.monotonic() + min(timeout, 60.0) while True: try: + # Loader switches describe how to IMPORT a raw file and + # are recorded in the database it produces. Sending them + # again for a database that already exists is a FATAL + # error in IDA itself ("Switch '-b400' can be used only + # when loading a new file"), which kills the worker + # before it can report anything useful. So: describe the + # import only when there is an import to describe. + fresh = self._new_database or not self._database_exists() handle = DatabaseHandle.open( self._path, spawn=self._spawn, timeout=max(0.1, timeout), output_database=self._output_database, - processor=self._processor, + processor=self._processor if fresh else None, # DatabaseHandle calls this image_base and wants the # natural (16-byte aligned) address; it does the # conversion to IDA's paragraph-based -b itself. - image_base=self._loading_address, - file_type=self._file_type, + image_base=self._loading_address if fresh else None, + file_type=self._file_type if fresh else None, new_database=self._new_database, ) break diff --git a/tests/_fixtures.py b/tests/_fixtures.py index c7a45d3..0b72856 100644 --- a/tests/_fixtures.py +++ b/tests/_fixtures.py @@ -37,6 +37,36 @@ def cache_is_fresh(binary: str) -> bool: return os.path.exists(c) and os.path.getmtime(c) >= os.path.getmtime(binary) +#: Generated targets live here so their pristine caches survive between runs. +#: Gitignored; safe to delete (the next run rebuilds both). +SYNTHETIC_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), ".synthetic") + + +def synthetic(name: str, build) -> str: + """A generated binary at a STABLE path, rebuilt only when its bytes change. + + Generated targets used to be written into a fresh TemporaryDirectory on + every run, which quietly defeated the whole pristine-cache scheme: a new + path with new bytes every time means auto-analysis is paid in full, every + run, forever. `test_blob_ui`'s 64KB blob cost ~40s a run that way. + + ``build()`` must be DETERMINISTIC and return bytes. That is also what makes + the suites reproducible: a blob built from os.urandom can, by luck, contain + something IDA reads as a function, and then a test asserting "no functions" + fails for reasons no one can reproduce. + """ + os.makedirs(SYNTHETIC_DIR, exist_ok=True) + path = os.path.join(SYNTHETIC_DIR, name) + data = build() + if not os.path.exists(path) or open(path, "rb").read() != data: + with open(path, "wb") as fh: # content changed -> cache is stale + fh.write(data) + for stale in (cache_path(path), path + ".i64"): + if os.path.exists(stale): + os.remove(stale) + return path + + async def build_pristine(binary: str, cache: str, app_factory) -> None: """Analyse ``binary`` once and keep the database as a golden copy. diff --git a/tests/test_blob_ui.py b/tests/test_blob_ui.py index f60e914..1055e5a 100644 --- a/tests/test_blob_ui.py +++ b/tests/test_blob_ui.py @@ -17,10 +17,13 @@ import sys import tempfile sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) from textual.widgets import Input, Static # noqa: E402 from idatui.app import ConfirmScreen, IdaTui, ListingView # noqa: E402 +from _fixtures import staged, synthetic # noqa: E402 +from idatui._sync import settle # noqa: E402 PASS = FAIL = 0 @@ -46,28 +49,63 @@ async def wait(pred, pilot, t=240.0): return False -async def run() -> int: - with tempfile.TemporaryDirectory() as tmp: - # Random bytes so IDA finds no functions... but with REAL AArch64 - # instructions planted at a known offset. Whether arbitrary random bytes - # happen to decode is chance, and a test that depends on chance tells you - # nothing on the run where it fails. - data = bytearray(os.urandom(64 * 1024)) +#: File offset of the planted instruction run -> ea 0x4000 + PLANTED. +PLANTED = 0x40 + + +def _blob_bytes() -> bytes: + """A DETERMINISTIC pseudo-random blob with real AArch64 instructions planted. + + Seeded, not os.urandom: the bytes must be identical every run or the + pristine-database cache can never apply (this suite used to pay ~40s of + auto-analysis per run because the content, and the path, changed each time). + Determinism also removes a genuine flake -- whether 64KB of chance bytes + contains something IDA reads as a function is luck, and "and really has no + functions" is asserted below. + """ + import random + data = bytearray(random.Random(0xB10BCAFE).randbytes(64 * 1024)) # -parm puts IDA in AArch64 mode, so these are A64 encodings; the ARM32 # spelling of a nop (0xE1A00000) is NOT decodable there and made this # test fail for a reason that had nothing to do with what it checks. - planted = 0x40 # file offset -> ea 0x4040 - for k, insn in enumerate((0xD503201F, # nop - 0xD503201F, # nop - 0xD65F03C0)): # ret <- the run must stop here - data[planted + k * 4:planted + k * 4 + 4] = insn.to_bytes(4, "little") - blob = os.path.join(tmp, "rnd.bin") - with open(blob, "wb") as f: - f.write(bytes(data)) + for k, insn in enumerate((0xD503201F, # nop + 0xD503201F, # nop + 0xD65F03C0)): # ret <- the run must stop here + data[PLANTED + k * 4:PLANTED + k * 4 + 4] = insn.to_bytes(4, "little") + return bytes(data) + + +#: -parm puts IDA in AArch64 mode (the ARM32 spelling of a nop is not decodable +#: there); -b400 sets the image base. The cached database must be built with the +#: SAME switches, so both go through one factory. +BLOB_ARGS = "-parm -b400" + +def _blob_app(path): + return IdaTui(open_path=path, keepalive=False, load_args=BLOB_ARGS) + + +def head_at(lst, ea): + """The listing row for ``ea`` off the LIVE model, or None. + + Always re-reads ``lst.model``: an edit may rebuild the model, and holding + the old object shows pre-edit rows -- which looks exactly like the edit + silently failing. + """ + m = lst.model + if m is None: + return None + i = m.index_of_ea(ea) + return m.get(i) if i is not None and i >= 0 else None + + +async def run() -> int: + blob_src = synthetic("rnd.bin", _blob_bytes) + # staged() analyses once ever and copies the result in on later runs. + async with staged(blob_src, _blob_app) as blob: # Skip the dialog by answering up front; this test is about what # happens AFTER a described blob turns out to contain nothing. - app = IdaTui(open_path=blob, keepalive=False, load_args="-parm -b400") + app = _blob_app(blob) async with app.run_test(size=(140, 44)) as pilot: ok = await wait(lambda: app._func_index is not None and app._func_index.complete, pilot) @@ -126,7 +164,7 @@ async def run() -> int: i >= 0 and m.get(i).ea == 0x4021, f"row={i} ea={m.get(i).ea if i >= 0 else None}") - target = 0x4000 + planted # a NOP we put there ourselves + target = 0x4000 + PLANTED # a NOP we put there ourselves lst.cursor = m.index_of_ea(target) lst._scroll_cursor_into_view() await pilot.pause(0.1) @@ -134,12 +172,18 @@ async def run() -> int: lst._cursor_ea() == target, f"{lst._cursor_ea():#x} want {target:#x}") await pilot.press("c") - await pilot.pause(2.0) - # Defining an item REBUILDS the listing model, so re-read it from the - # view: holding the old object shows the pre-edit rows and looks - # exactly like the edit silently failing. + # settle(), not a fixed sleep AND not a bare predicate: an edit can + # look done for a moment and then be replaced when a queued listing + # rebuild lands, so the gate has to be "the row is code AND the app + # has stopped working". settle() is the same helper the app's own + # RPC layer uses, so tests and driver agree on what "done" means. + await settle(app, lambda: (lambda h: h is not None and h.kind == "code")( + head_at(lst, target)), timeout=30) + # Re-read the model: defining an item rebuilds it, and holding the + # old object shows pre-edit rows -- which looks exactly like the + # edit silently failing. m = lst.model - h = m.get(m.index_of_ea(target)) + h = head_at(lst, target) check("`c` on a chosen byte carves an instruction there", h is not None and h.kind == "code", f"kind={h.kind if h else None} text={h.text if h else None!r}") @@ -184,9 +228,17 @@ async def run() -> int: for ch in "note": await pilot.press(ch) await pilot.press("enter") - await wait(lambda: lst.model is not old and lst.model is not None, - pilot, 30) - await pilot.pause(0.4) + # Wait for the COMMENT ITSELF to show up, not for the model object to + # be replaced: a comment now re-renders the listing in place (the + # walk is kept), so `model is not old` never becomes true and this + # burned its full 30s timeout on every run -- after which the check + # below passed vacuously, because nothing had happened at all. + # The prompt closing plus quiescence is the real end of the edit. + # (The listing re-renders its text lazily, so the comment is not + # necessarily visible in model rows the moment the worker returns -- + # which is why this waits for the app, not for the text.) + await settle(app, lambda: not app.query_one("#comment", Input).display, + timeout=30) check("commenting leaves the view where it was", lst.model.get(round(lst.scroll_offset.y)).ea == ctop and lst._cursor_ea() == ccur, @@ -208,7 +260,12 @@ async def run() -> int: check("scrolled somewhere with rows above us", round(lst.scroll_offset.y) > 0, f"top={lst.scroll_offset.y}") await pilot.press("c") - await pilot.pause(2.5) + # No predicate here on purpose: this spot is random data, so the + # carve may legitimately produce nothing and "the row became code" + # would never hold (it timed out for 30s and then passed anyway). + # What is being checked is that the VIEW did not move, so the gate + # is simply "the app has finished reacting". + await settle(app, timeout=30) m2 = lst.model top_after = m2.get(round(lst.scroll_offset.y)).ea check("carving leaves the scroll position where it was", -- cgit v1.3.1-sl0p From 55e9d9fdca4baafce728659d43614e203bfbea9a Mon Sep 17 00:00:00 2001 From: blasty Date: Fri, 7 Aug 2026 15:02:05 +0200 Subject: codemode: close the performance gap with the old worker (heads 35x -> 2.2x) Two changes, both about work that was never ours to do, found by profiling the A/B benchmark rather than guessing. 1. Serialise inside the database process. Code Mode runs to_jsonable() over whatever a snippet returns, walking the entire structure to make it JSON-safe. Our answers are already JSON-safe and they are large: a 200-row listing page is ~10k small objects, and walking them cost 66ms of the page's 92ms -- 114x what json.dumps of the very same data costs (0.58ms). Snippets now return one pre-serialised string, so that walk is O(1) and the client parses a payload it was going to parse anyway. heads(200): 92ms -> 24.7ms. 2. Detach the runtime's trace hook while our snippet runs. ida_codemode.runtime wraps every execute_python in sys.settrace(timeout_trace) to enforce deadlines, and timeout_trace RETURNS ITSELF -- which switches on LINE tracing in every frame it sees. Every line of every function we call pays a Python-level callback. Measured here: ida_bytes.get_flags 0.106us untraced 5.49us traced 52x (plain idalib, no Code Mode: 0.119us -- i.e. untraced == native) heads(200 rows) 2.0ms untraced 20.2ms traced 10x That one hook was the entire residual gap against the old unix-socket worker. The snippet now detaches it and restores it in a finally. What that gives up, stated plainly: the deadline is no longer enforced for a pure-Python loop inside our snippet. The runtime's other cancellation path -- a threading.Timer calling ida_kernwin.set_cancelled() -- does not go through the trace and still fires, so a long IDA operation remains interruptible, and every operation here is bounded by its own count/limit argument. Set IDATUI_CODEMODE_TRACE=1 to keep the stock behaviour. Against the worker backend, same box, targets/echo (worker -> codemode): heads_200 2.65ms -> 5.85ms 2.2x (was 35x) heads_500 6.21ms -> 10.67ms 1.7x heads expect-hit 2.16ms -> 4.61ms 2.1x disasm_200 9.03ms -> 5.25ms 0.6x faster decompile_cold 162.62ms -> 30.66ms 0.2x faster decompile_warm 30.02ms -> 25.39ms 0.8x faster decomp_map 45.82ms -> 47.85ms 1.0x parity pc_nums 19.90ms -> 22.56ms 1.1x parity rename_func 254.08ms -> 255.54ms 1.0x parity connect 550.0ms -> 410.0ms 0.7x faster What is left is the transport floor: an empty execute_python round trip is 2.0ms, so trivial calls (data_type 0.07ms -> 2.63ms, force_recompile, a single xref query) look like 40x while being 2.5ms of wall clock. Reducing those needs fewer calls, not faster ones -- the digest/expect path already does that for the listing, which is where call volume actually is. Full suite: 788 passed, 0 failed, 115.3s (was 146.3s; the pilot alone went 80.9s -> 62.2s). --- idatui/codemode_client.py | 71 +++++++++++++++++++++++++++++++++++++++++++---- 1 file changed, 66 insertions(+), 5 deletions(-) (limited to 'idatui') diff --git a/idatui/codemode_client.py b/idatui/codemode_client.py index 9d91335..88d7b31 100644 --- a/idatui/codemode_client.py +++ b/idatui/codemode_client.py @@ -22,7 +22,7 @@ import shlex import threading import time from pathlib import Path -from textwrap import dedent +from textwrap import dedent, indent from typing import Any from .errors import IDAConnectionError, IDATimeoutError, IDAToolError, Session @@ -165,10 +165,63 @@ def _parse_load_args(value: str) -> tuple[str | None, int | None, str | None]: return processor, loading_address, file_type +#: Key of the pre-serialised payload envelope. See _script(). +_PACKED = "__idatui_json__" + +#: Serialise the answer INSIDE the database process and hand back one string. +#: +#: Code Mode runs to_jsonable() over whatever a snippet returns, walking the +#: whole structure to make it JSON-safe. Our answers are already JSON-safe, and +#: they are big: a 200-row listing page is ~10k small objects, which costs 66ms +#: to walk -- 72% of the page's total cost, and 114x what json.dumps of the very +#: same data costs (0.58ms). Returning a STRING makes that walk O(1); the client +#: parses it, which it was going to do at the transport layer anyway. +_PACK_EPILOGUE = ( + '\n{"' + _PACKED + '": json.dumps(result, separators=(",", ":"), default=str)}\n' +) + + +#: Keep Code Mode's per-line trace hook installed while our snippet runs. +#: Set IDATUI_CODEMODE_TRACE=1 to restore the stock behaviour. +_KEEP_TRACE = os.environ.get("IDATUI_CODEMODE_TRACE", "") not in ("", "0") + + def _script(args: dict[str, Any], body: str) -> str: - """Bind JSON arguments without interpolating user text into Python code.""" + """Bind JSON arguments without interpolating user text into Python code. + + Also runs the body with Code Mode's trace hook detached, which is worth an + order of magnitude. The runtime wraps every execute_python in + sys.settrace(timeout_trace), and that trace function RETURNS ITSELF, which + turns on line tracing in every frame it sees -- so every line of every + function we call pays a Python-level callback. Measured on this box: + ida_bytes.get_flags is 0.106us untraced (0.119us in a plain idalib process) + and 5.49us traced, 52x; a 200-row listing page is 2.0ms untraced and 20.2ms + traced. That single hook was the whole residual gap against the old worker. + + What this gives up: the deadline is no longer enforced for a pure-Python + loop inside our snippet. The runtime's OTHER cancellation path -- a + threading.Timer that calls ida_kernwin.set_cancelled() -- is independent of + the trace and still fires, so a long IDA operation is still interruptible; + and every operation here is bounded by its own count/limit argument. The + trace is restored in a finally, so a raising snippet cannot leak the change. + """ encoded = json.dumps(args, ensure_ascii=False, separators=(",", ":")) - return f"import json\na = json.loads({encoded!r})\n{dedent(body).strip()}\n" + head = f"import json\na = json.loads({encoded!r})\n" + if _KEEP_TRACE: + return f"{head}{dedent(body).strip()}\n{_PACK_EPILOGUE}" + return ( + f"{head}" + "import sys\n" + "_idatui_trace = sys.gettrace()\n" + "sys.settrace(None)\n" + "try:\n" + f"{indent(dedent(body).strip(), ' ')}\n" + ' _idatui_packed = {"' + _PACKED + '": json.dumps(' + 'result, separators=(",", ":"), default=str)}\n' + "finally:\n" + " sys.settrace(_idatui_trace)\n" + "_idatui_packed\n" + ) _OPERATIONS: dict[str, str] = { @@ -1179,6 +1232,13 @@ class CodeModeClient: raise IDAToolError("execute_python", "Code Mode returned an invalid execution result") return response["result"] + @staticmethod + def _unpack(answer: Any) -> Any: + """Undo _PACK_EPILOGUE. Anything else passes through untouched.""" + if isinstance(answer, dict) and _PACKED in answer: + return json.loads(answer[_PACKED]) + return answer + def invoke(self, operation: str, *, timeout: float | None = None, **args) -> Any: """Execute one TUI domain operation through Code Mode.""" if operation in ("idb_save", "save"): @@ -1189,12 +1249,13 @@ class CodeModeClient: if body is None: raise IDAToolError(operation, f"unknown ida-tui Code Mode operation: {operation}") try: - answer = self.execute_python(_script(args, body), timeout=timeout) + answer = self._unpack(self.execute_python(_script(args, body), timeout=timeout)) if isinstance(answer, dict) and answer.get(_NEED_LIB): # First call against this database process (or a restarted one). self.execute_python(_script({"source": _REMOTE_LIB}, _INSTALL_LIB), timeout=timeout) - answer = self.execute_python(_script(args, body), timeout=timeout) + answer = self._unpack( + self.execute_python(_script(args, body), timeout=timeout)) return answer except IDAToolError as exc: if exc.tool == "execute_python": -- cgit v1.3.1-sl0p