diff options
Diffstat (limited to 'idatui')
| -rw-r--r-- | idatui/__init__.py | 5 | ||||
| -rw-r--r-- | idatui/app.py | 200 | ||||
| -rw-r--r-- | idatui/codemode_client.py | 1362 | ||||
| -rw-r--r-- | idatui/domain.py | 209 | ||||
| -rw-r--r-- | idatui/drive.py | 3 | ||||
| -rw-r--r-- | idatui/edit_ctl.py | 4 | ||||
| -rw-r--r-- | idatui/errors.py | 10 | ||||
| -rw-r--r-- | idatui/launch.py | 103 | ||||
| -rw-r--r-- | idatui/pane.py | 93 | ||||
| -rw-r--r-- | idatui/pool.py | 100 | ||||
| -rw-r--r-- | idatui/project.py | 28 | ||||
| -rw-r--r-- | idatui/remote_tools.py | 1548 | ||||
| -rw-r--r-- | idatui/rpc.py | 4 | ||||
| -rw-r--r-- | idatui/worker.py | 301 | ||||
| -rw-r--r-- | idatui/worker_client.py | 266 |
15 files changed, 3252 insertions, 984 deletions
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..88d7b31 --- /dev/null +++ b/idatui/codemode_client.py @@ -0,0 +1,1362 @@ +"""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 hashlib +import json +import os +import shlex +import threading +import time +from pathlib import Path +from textwrap import dedent, indent +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 + + +#: 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. + + 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=(",", ":")) + 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] = { + "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 +''', + # Ours: the coarse code/data type plus a fine `kind` (call/jump/flow, + # read/write/offset/text/info) that the xref dialog draws its badges from. + # Deliberately NOT sorted -- the dialog lists xrefs in IDA's own order. + "xref_types": r''' +import idaapi, idautils, ida_bytes, ida_funcs, ida_xref +code_kind = {ida_xref.fl_CF: "call", ida_xref.fl_CN: "call", ida_xref.fl_JF: "jump", + ida_xref.fl_JN: "jump", ida_xref.fl_F: "flow"} +data_kind = {ida_xref.dr_O: "offset", ida_xref.dr_W: "write", ida_xref.dr_R: "read", + ida_xref.dr_T: "text", ida_xref.dr_I: "info"} +def _kind(xr): + return (code_kind if xr.iscode else data_kind).get(xr.type, "code" if xr.iscode else "data") +def _fn(ea): + f = ida_funcs.get_func(ea) + return {"addr": hex(int(f.start_ea)), "name": ida_funcs.get_func_name(f.start_ea) or ""} if f else None +queries = a.get("queries") or [] +all_results = [] +for query in queries: + query = query if isinstance(query, dict) else {"addr": query} + raw = str(query.get("addr", "")).strip() + direction = str(query.get("direction", "to") or "to").lower() + include_fn = bool(query.get("include_fn", True)) + dedup = bool(query.get("dedup", True)) + try: count = int(query.get("count", 2000) or 2000) + except (TypeError, ValueError): count = 2000 + try: target = int(raw, 16) + except ValueError: target = idaapi.get_name_ea(idaapi.BADADDR, raw) + rows = [] + if target is not None and target != idaapi.BADADDR and ida_bytes.is_mapped(target): + if direction in ("to", "both"): + for xr in idautils.XrefsTo(target, 0): + row = {"direction": "to", "addr": hex(int(xr.frm)), "from": hex(int(xr.frm)), + "to": hex(int(target)), "type": "code" if xr.iscode else "data", "kind": _kind(xr)} + if include_fn: row["fn"] = _fn(xr.frm) + rows.append(row) + if direction in ("from", "both"): + for xr in idautils.XrefsFrom(target, 0): + row = {"direction": "from", "addr": hex(int(xr.to)), "from": hex(int(target)), + "to": hex(int(xr.to)), "type": "code" if xr.iscode else "data", "kind": _kind(xr)} + if include_fn: row["fn"] = _fn(xr.to) + rows.append(row) + if dedup: + seen, deduped = set(), [] + for r in rows: + k = (r["direction"], r["from"], r["to"], r["kind"]) + if k in seen: continue + seen.add(k); deduped.append(r) + rows = deduped + rows = rows[:count] + all_results.append({"query": raw, "data": rows, "next_offset": None}) +result = {"result": all_results} +result +''', + # Mirrors the tool ida-tui was written against, ORDER INCLUDED. The rows are + # sorted by the far-end address and deduped by default, and the pseudocode + # follow's address fallback silently depends on it: at a call site the raw + # IDA order yields the ordinary-flow xref (the next instruction) first, so an + # unsorted result makes "follow the call" land on the following line instead. + "xref_query": r''' +import idaapi, idautils, ida_bytes, ida_funcs +def _fn(ea): + f = ida_funcs.get_func(ea) + return {"addr": hex(int(f.start_ea)), "name": ida_funcs.get_func_name(f.start_ea) or ""} if f else None +queries = a.get("queries") or [] +all_results = [] +for query in queries: + raw = str(query.get("addr", "")).strip() + direction = str(query.get("direction", "both") or "both").lower() + if direction not in ("to", "from", "both"): direction = "both" + xref_type = str(query.get("xref_type", "any") or "any").lower() + if xref_type not in ("any", "code", "data"): xref_type = "any" + include_fn = bool(query.get("include_fn", True)) + dedup = bool(query.get("dedup", True)) + sort_by = str(query.get("sort_by", "addr") or "addr") + descending = bool(query.get("descending", False)) + try: offset = max(0, int(query.get("offset", 0) or 0)) + except (TypeError, ValueError): offset = 0 + try: count = max(0, min(int(query.get("count", 200) or 200), 5000)) + except (TypeError, ValueError): count = 200 + try: + try: target = int(raw, 16) + except ValueError: + target = idaapi.get_name_ea(idaapi.BADADDR, raw) + if target == idaapi.BADADDR: raise ValueError(f"Failed to resolve address/name: {raw}") + if not ida_bytes.is_mapped(target): raise ValueError(f"Address not mapped: {raw}") + rows = [] + if direction in ("to", "both"): + for xr in idautils.XrefsTo(target, 0): + kind = "code" if xr.iscode else "data" + if xref_type != "any" and kind != xref_type: continue + row = {"direction": "to", "addr": hex(int(xr.frm)), "from": hex(int(xr.frm)), + "to": hex(int(target)), "type": kind} + if include_fn: row["fn"] = _fn(xr.frm) + rows.append(row) + if direction in ("from", "both"): + for xr in idautils.XrefsFrom(target, 0): + kind = "code" if xr.iscode else "data" + if xref_type != "any" and kind != xref_type: continue + row = {"direction": "from", "addr": hex(int(xr.to)), "from": hex(int(target)), + "to": hex(int(xr.to)), "type": kind} + if include_fn: row["fn"] = _fn(xr.to) + rows.append(row) + if dedup: + seen, deduped = set(), [] + for row in rows: + key = (row["direction"], row["from"], row["to"], row["type"]) + if key in seen: continue + seen.add(key); deduped.append(row) + rows = deduped + if sort_by == "type": + rows.sort(key=lambda r: (str(r.get("type", "")), int(str(r["addr"]), 16)), reverse=descending) + else: + rows.sort(key=lambda r: int(str(r["addr"]), 16), reverse=descending) + page = rows[offset:offset + count] if count else rows[offset:] + nxt = offset + len(page) + all_results.append({"target": raw, "resolved_addr": hex(int(target)), "direction": direction, + "xref_type": xref_type, "data": page, + "next_offset": nxt if nxt < len(rows) else None, + "total": len(rows), "error": None}) + except Exception as exc: + all_results.append({"target": raw, "resolved_addr": None, "direction": direction, + "xref_type": xref_type, "data": [], "next_offset": None, + "total": 0, "error": str(exc)}) +result = {"result": all_results} +result +''', + # A comment must land in BOTH views, and the pseudocode half is not a + # simple set: db.comments.set_at() alone leaves the pseudocode unchanged. + # Hex-Rays comments are anchored to a ctree location (treeloc_t), and an + # anchor the ctree does not actually own is dropped as an "orphan" -- so the + # itp slot has to be searched until one sticks, exactly as IDA's own UI does. + # Without it a comment silently never appears in the decompilation. + "set_comments": r''' +import idaapi, idc, ida_hexrays +rows = [] +for item in a.get("items", []): + addr_s = str(item.get("addr", "")) + text = str(item.get("comment") or "") + try: + ea = int(addr_s, 16) + if not idaapi.set_cmt(ea, text, False): + rows.append({"addr": addr_s, + "error": f"Failed to set disassembly comment at {hex(ea)}"}) + continue + if not ida_hexrays.init_hexrays_plugin(): + rows.append({"addr": addr_s}); continue + try: + cfunc = ida_hexrays.decompile(ea) + except Exception: + cfunc = None + if cfunc is None: + rows.append({"addr": addr_s}); continue + if ea == cfunc.entry_ea: + # The signature line carries no ctree item: it is a function comment. + idc.set_func_cmt(ea, text, True) + cfunc.refresh_func_ctext() + rows.append({"addr": addr_s}); continue + eamap = cfunc.get_eamap() + if ea not in eamap: + rows.append({"addr": addr_s, + "error": f"Failed to set decompiler comment at {hex(ea)}"}) + continue + nearest_ea = eamap[ea][0].ea + if cfunc.has_orphan_cmts(): + cfunc.del_orphan_cmts(); cfunc.save_user_cmts() + tl = idaapi.treeloc_t(); tl.ea = nearest_ea + placed = False + for itp in range(idaapi.ITP_SEMI, idaapi.ITP_COLON): + tl.itp = itp + cfunc.set_user_cmt(tl, text) + cfunc.save_user_cmts() + cfunc.refresh_func_ctext() + if not cfunc.has_orphan_cmts(): + placed = True; break + cfunc.del_orphan_cmts(); cfunc.save_user_cmts() + rows.append({"addr": addr_s} if placed else + {"addr": addr_s, + "error": f"Failed to set decompiler comment at {hex(ea)}"}) + except Exception as exc: + rows.append({"addr": addr_s, "error": str(exc)}) +result = {"result": rows} +result +''', + # Every category takes EITHER one edit or a LIST of them, and the answer is + # one row per edit. The port accepted only a single dict, so any batch path + # (rpc rename_many applying a whole symbol file, which is the entire point of + # that verb) died with "list indices must be integers or slices, not str" and + # reported the failure against addr=null. Mirrors the real tool: conflict + # detection before the write, dry_run/allow_overwrite/stop_on_error, per-row + # addr/old/name, and a summary counting EDITS rather than categories. + "rename": r''' +import idaapi, ida_hexrays, ida_name +batch = a.get("batch") or {} +dry_run = bool(batch.get("dry_run", False)) +allow_overwrite = bool(batch.get("allow_overwrite", False)) +stop_on_error = bool(batch.get("stop_on_error", False)) + +def _items(value): + if value is None: return [] + if isinstance(value, dict): return [value] + if isinstance(value, list): return [i for i in value if isinstance(i, dict)] + return [] + +def _set_name_checked(ea, new): + conflict = idaapi.get_name_ea(idaapi.BADADDR, new) + if conflict != idaapi.BADADDR and conflict != ea and not allow_overwrite: + return False, f"can't rename at {hex(ea)} as {new!r}: name already used at {hex(conflict)}" + if dry_run: + return True, None + flags = idaapi.SN_CHECK + if allow_overwrite: flags |= int(getattr(idaapi, "SN_FORCE", 0)) + if not idaapi.set_name(ea, new, flags): + return False, (f"Rename failed at {hex(ea)}: IDA rejected name {new!r} " + "(invalid identifier or internal conflict)") + return True, None + +def _refresh_ctext(fn_addr): + # A renamed function must invalidate Hex-Rays' cache, which is per function + # and persisted in the .i64: without this the pseudocode keeps calling the + # old name forever while every other readback reports the new one. + if not ida_hexrays.init_hexrays_plugin(): return + failure = ida_hexrays.hexrays_failure_t() + cfunc = ida_hexrays.decompile_func(fn_addr, failure, ida_hexrays.DECOMP_WARNINGS) + if cfunc: cfunc.refresh_func_ctext() + +out = {}; ok_count = failed = 0; halted = False +for category in ("func", "data", "local", "stack"): + if category not in batch: continue + rows = [] + for edit in _items(batch.get(category)): + try: + if category == "func": + addr_text = edit.get("addr") or edit.get("func_addr") or edit.get("func") + new = edit.get("name") or edit.get("new") or edit.get("new_name") + if not addr_text or not new: + row = {"addr": addr_text, "name": new, + "error": "Function rename requires addr + name"} + else: + ea = int(str(addr_text), 16) + fn = idaapi.get_func(ea) + if fn is None: + row = {"addr": addr_text, "name": new, "error": "Function not found"} + else: + old = idaapi.get_name(fn.start_ea) or None + ok, err = _set_name_checked(fn.start_ea, str(new)) + row = {"addr": addr_text, "old": old, "name": str(new)} + if err: row["error"] = err + if dry_run: row["dry_run"] = True + if ok and not dry_run: _refresh_ctext(fn.start_ea) + elif category == "data": + addr_text = edit.get("addr") + old = edit.get("old") or edit.get("old_name") + new = edit.get("new") or edit.get("new_name") or edit.get("name") + if not new and new != "": + row = {"old": old, "new": None, + "error": "Global rename requires target and new name"} + else: + if addr_text is not None: + ea = int(str(addr_text), 16) + old = old or (idaapi.get_name(ea) or None) + else: + ea = idaapi.get_name_ea(idaapi.BADADDR, str(old or "")) + if ea == idaapi.BADADDR: + row = {"old": old, "new": str(new), "error": f"Global {old!r} not found"} + else: + # An empty new name CLEARS the label; that is a real + # request (tests revert with it), not a missing argument. + if str(new) == "": + ok = bool(ida_name.set_name(ea, "", idaapi.SN_CHECK)) + err = None if ok else f"Failed to clear the name at {hex(ea)}" + else: + ok, err = _set_name_checked(ea, str(new)) + row = {"addr": hex(ea), "old": old, "new": str(new)} + if err: row["error"] = err + if dry_run: row["dry_run"] = True + else: + fa, old, new = edit.get("func_addr"), edit.get("old"), edit.get("new") + if not fa or not old or not new: + row = {"old": old, "new": new, + "error": f"{category} rename requires func_addr + old + new"} + else: + ea = int(str(fa), 16) + pseudo = db.pseudocode.decompile(ea) + var = pseudo.find_local_variable(str(old)) + if var is None: + row = {"func_addr": fa, "old": old, "new": new, + "error": f"no local {old!r} in that function"} + elif dry_run: + row = {"func_addr": fa, "old": old, "new": new, "dry_run": True} + else: + var.set_user_name(str(new)) + ok = bool(pseudo.save_local_variable_info(var, save_name=True)) + row = {"func_addr": fa, "old": old, "new": new} + if not ok: row["error"] = "IDA rejected the local variable name" + except Exception as exc: + row = {"addr": edit.get("addr"), "error": str(exc)} + rows.append(row) + if row.get("error"): failed += 1 + else: ok_count += 1 + if row.get("error") and stop_on_error: + halted = True; break + out[category] = rows + if halted: break +out["summary"] = {"ok": ok_count, "failed": failed} +if dry_run: out["summary"]["dry_run"] = True +if halted: out["summary"]["halted"] = True +result = out +result +''', +} + + +_OPERATIONS["define_code_run"] = r''' +import ida_bytes, ida_idp, ida_segment, ida_ua, idaapi +ea, limit = int(str(a["addr"]), 16), max(1, min(int(a.get("limit", 20000)), 200000)) +seg = ida_segment.getseg(ea) +if seg is None: + result = {"addr": a["addr"], "error": "no segment", "count": 0} +else: + start, count, stopped, hi = ea, 0, "limit", int(seg.end_ea) + while count < limit: + if ea >= hi: stopped = "segment"; break + flags = ida_bytes.get_flags(ea) + if ida_bytes.is_code(flags) or ida_bytes.is_data(flags): stopped = "defined"; break + size = int(ida_ua.create_insn(ea)) + if size <= 0: stopped = "undecodable"; break + count += 1 + insn = ida_ua.insn_t() + if ida_ua.decode_insn(insn, ea) > 0: + try: is_ret = bool(ida_idp.is_ret_insn(insn)) + except Exception: is_ret = False + if is_ret or (insn.get_canon_feature() & idaapi.CF_STOP): + ea += size; stopped = "flow"; break + ea += size + result = {"start": hex(start), "end": hex(ea), "count": count, "stopped": stopped} +result +''' + + +_OPERATIONS["define_func_run"] = r''' +import ida_bytes, ida_funcs, ida_segment +ea = int(str(a["addr"]), 16) +fn = db.functions.get_at(ea) +if fn is not None and int(fn.start_ea) == ea: + result = {"addr": hex(ea), "ok": True, "start": hex(ea), "end": hex(int(fn.end_ea)), "how": "existed"} +else: + automatic = bool(db.functions.create(ea)) + if not automatic: + seg = db.segments.get_at(ea); end = ea; hi = int(seg.end_ea) if seg else ea + while end < hi and ida_bytes.is_code(ida_bytes.get_flags(end)): + nxt = int(ida_bytes.get_item_end(end)) + if nxt <= end: break + end = nxt + ok = bool(end > ea and ida_funcs.add_func(ea, end)) + else: ok = True + fn = db.functions.get_at(ea) + result = ({"addr": hex(ea), "ok": True, "start": hex(int(fn.start_ea)), + "end": hex(int(fn.end_ea)), "how": "auto" if automatic else "explicit-end"} + if ok and fn is not None else + {"addr": hex(ea), "ok": False, "error": f"IDA refused a function at {ea:#x}"}) +result +''' + + +_OPERATIONS["set_thumb"] = r''' +import ida_bytes, ida_ida, ida_idp, ida_segment, ida_segregs +ea = int(str(a["addr"]), 16); treg = ida_idp.str2reg("T") +seg = ida_segment.getseg(ea) +if treg is None or treg < 0: + result = {"addr": hex(ea), "error": "no T register (not an ARM database)"} +elif seg is None: + result = {"addr": hex(ea), "error": "no segment"} +else: + current = ida_segregs.get_sreg(ea, treg) + current = 0 if current in (None, 0xFFFFFFFF, -1) else int(current) + want = {"on": 1, "off": 0}.get(str(a.get("mode", "toggle")).lower(), 0 if current else 1) + changed = False + if want and seg.bitness != 1: + ida_segment.set_segm_addressing(seg, 1); changed = True + size = max(int(ida_bytes.get_item_size(ea)), 2) + ida_bytes.del_items(ea, 0, size) + ok = bool(ida_segregs.split_sreg_range(ea, treg, want, ida_segregs.SR_user)) + now = ida_segregs.get_sreg(ea, treg) + result = {"addr": hex(ea), "thumb": bool(now), "was": bool(current), "ok": ok, + "bitness": ida_segment.getseg(ea).bitness, "forced_32bit": changed, + "db_64bit": bool(ida_ida.inf_get_app_bitness() == 64 and want)} +result +''' + + +_OPERATIONS["thumb_scan"] = r''' +import ida_bytes, ida_funcs, ida_idp, ida_segment, ida_segregs, ida_ua +lo, hi = int(str(a["start"]), 16), int(str(a["end"]), 16) +apply, limit = bool(a.get("apply", True)), int(a.get("limit", 512)) +treg = ida_idp.str2reg("T"); found = []; applied = 0; cursor = lo +while cursor + 4 <= hi and len(found) < limit: + at = cursor; value = int(ida_bytes.get_dword(cursor)); cursor += 4 + if not value & 1: continue + target = value & ~1; seg = ida_segment.getseg(target) + if seg is None or not (seg.perm & ida_segment.SEGPERM_EXEC or seg.perm == 0): continue + flags = ida_bytes.get_flags(target) + if ida_bytes.is_data(flags): continue + item = {"at": hex(at), "value": hex(value), "target": hex(target), + "was_code": bool(ida_bytes.is_code(flags))}; found.append(item) + if not apply: continue + if treg is not None and treg >= 0: ida_segregs.split_sreg_range(target, treg, 1, ida_segregs.SR_user) + if not ida_bytes.is_code(ida_bytes.get_flags(target)): + ida_bytes.del_items(target, 0, 2) + if ida_ua.create_insn(target) <= 0: item["decoded"] = False; continue + item["decoded"] = True; item["function"] = bool(db.functions.get_at(target) or db.functions.create(target)); applied += 1 +result = {"start": hex(lo), "end": hex(hi), "found": found, "applied": applied, "n": len(found)} +result +''' + + +_OPERATIONS["decomp_error"] = r''' +import ida_hexrays, ida_ida +ea = int(str(a["addr"]), 16); fn = db.functions.get_at(ea) +result = {"addr": hex(ea), "bitness": ida_ida.inf_get_app_bitness()} +if fn is None: + result["reason"] = "no function here" +else: + try: + failure = ida_hexrays.hexrays_failure_t(); cfunc = ida_hexrays.decompile_func(fn, failure) + if cfunc is not None: result["reason"] = "" + else: + result.update({"reason": failure.desc() or f"error {failure.code}", + "code": int(failure.code), "errea": hex(int(failure.errea))}) + except Exception as exc: result["reason"] = f"{type(exc).__name__}: {exc}" +result +''' + +# `heads` and the operand-format tools are the port's IDAPython island: the +# continuous listing's presentation model (undefined runs, colour spans, operand +# extents, banners, struct members, the digest protocol) and IDA/Hex-Rays number +# formats have no ida-domain surface. Rather than paraphrase ~1100 lines of +# performance-tuned, behaviour-sensitive code into string literals, they stay +# real, diffable source in idatui/remote_tools.py and are shipped to the database +# process as text. Read once at import; the file ships beside this module. +_REMOTE_LIB = (Path(__file__).with_name("remote_tools.py")).read_text(encoding="utf-8") + +#: Versioned by content, so editing remote_tools.py re-installs it instead of +#: silently running the copy a long-lived worker already has. +_REMOTE_MODULE = "_idatui_remote_" + hashlib.sha1( + _REMOTE_LIB.encode("utf-8")).hexdigest()[:12] + +#: Sent back when the database process has not got the library yet; the client +#: installs it and retries once. Amortised, a worker receives it exactly once. +_NEED_LIB = "__idatui_needs_remote_lib__" + +#: Installs the library as a real module in the database process. Persisting it +#: in sys.modules is what makes the module-level caches (the tag maps, and the +#: line-render lru_cache the listing's throughput depends on) survive between +#: calls -- execute_python builds a fresh namespace every time, so a library +#: exec'd inline is rebuilt, and its caches thrown away, on every single call. +_INSTALL_LIB = f''' +import sys, types +_m = types.ModuleType({_REMOTE_MODULE!r}) +exec(compile(a["source"], {_REMOTE_MODULE!r}, "exec"), _m.__dict__) +sys.modules[{_REMOTE_MODULE!r}] = _m +result = True +result +''' + + +def _remote_op(call: str) -> str: + """A snippet that calls one of the carried-over tools by its real signature. + + Costs one short request: the library is imported from the database process's + own sys.modules, not shipped again. + """ + return (f"import sys\n" + f"_m = sys.modules.get({_REMOTE_MODULE!r})\n" + f"result = {{{_NEED_LIB!r}: True}} if _m is None else _m.{call}\n" + f"result\n") + + +_OPERATIONS["op_format"] = _remote_op( + 'op_format(addr=a["addr"], mode=a.get("mode", "cycle"),' + ' col=int(a.get("col", -1)), n=int(a.get("n", -1)))') +_OPERATIONS["pc_nums"] = _remote_op('pc_nums(addr=a["addr"])') +_OPERATIONS["decompile"] = _remote_op( + 'decompile(addr=a["addr"],' + ' include_addresses=bool(a.get("include_addresses", True)))') +_OPERATIONS["decomp_map"] = _remote_op('decomp_map(addr=a["addr"])') +_OPERATIONS["pc_num_format"] = _remote_op( + 'pc_num_format(addr=a["addr"], mode=a.get("mode", "cycle"),' + ' line=int(a.get("line", -1)), col=int(a.get("col", -1)),' + ' ea=a.get("ea", ""), opnum=int(a.get("opnum", -1)))') + +# The listing walker itself. Replaces the port's re-implementation, which +# rendered no per-operand extents (so no keypress could say which literal it +# would reformat) and had no digest/expect support (so every page was re-sent +# after any edit), and whose span walk was the per-character loop our own +# version had already been rewritten to avoid. +_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. +# +# 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``.""" + + 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 _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: + 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: + # 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 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 if fresh else None, + file_type=self._file_type if fresh else None, + 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"] + + @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"): + 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: + answer = self._unpack(self.execute_python(_script(args, body), timeout=timeout)) + if isinstance(answer, dict) and answer.get(_NEED_LIB): + # First call against this database process (or a restarted one). + self.execute_python(_script({"source": _REMOTE_LIB}, _INSTALL_LIB), + timeout=timeout) + answer = self._unpack( + self.execute_python(_script(args, body), timeout=timeout)) + return answer + except IDAToolError as exc: + if exc.tool == "execute_python": + raise IDAToolError(operation, exc.message) from exc + raise + + # 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"<no function> @ {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 <pane> python -m idatui.pane keys --pane <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"<WorkerPool {len(self._clients)}/{len(self.project.refs)} resident " + return (f"<DatabasePool {len(self._clients)}/{len(self.project.refs)} resident " f"{self.memory_mb()}/{self.budget_mb}MB active={self.active}>") + + +# 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/remote_tools.py b/idatui/remote_tools.py new file mode 100644 index 0000000..6fb6436 --- /dev/null +++ b/idatui/remote_tools.py @@ -0,0 +1,1548 @@ +"""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. + + 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): + """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 + + +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} 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 <sock_path> <binary_path> - -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 ``<file>.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<processor>`` 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 <sock> <binary> [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() |