aboutsummaryrefslogtreecommitdiffstats
path: root/idatui
diff options
context:
space:
mode:
Diffstat (limited to 'idatui')
-rw-r--r--idatui/__init__.py5
-rw-r--r--idatui/app.py200
-rw-r--r--idatui/codemode_client.py1164
-rw-r--r--idatui/domain.py209
-rw-r--r--idatui/drive.py3
-rw-r--r--idatui/edit_ctl.py4
-rw-r--r--idatui/errors.py10
-rw-r--r--idatui/launch.py103
-rw-r--r--idatui/pane.py93
-rw-r--r--idatui/pool.py100
-rw-r--r--idatui/project.py28
-rw-r--r--idatui/rpc.py4
-rw-r--r--idatui/worker.py301
-rw-r--r--idatui/worker_client.py266
14 files changed, 1506 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..77a0227
--- /dev/null
+++ b/idatui/codemode_client.py
@@ -0,0 +1,1164 @@
+"""Client adapter from ida-tui's domain operations to IDA Code Mode.
+
+``DatabaseHandle`` is the lifecycle boundary: it discovers an already-registered
+GUI database, reuses a shared managed idalib worker, or starts one when needed.
+The TUI never owns or terminates an IDA process. Closing this client releases
+only its lease.
+
+The Code Mode transport intentionally exposes one broad operation,
+``execute_python``. ``CodeModeClient.invoke`` turns the small, address-centric
+operations needed by the paging layer into self-contained snippets. The
+snippets prefer the public ``ida-domain`` ``db`` object. A handful of features
+that ida-domain does not currently expose (IDA-coloured listing rows, creating
+instructions, ARM T-state, and detailed Hex-Rays line maps/failures) use the
+IDAPython modules that Code Mode deliberately makes importable.
+"""
+from __future__ import annotations
+
+import json
+import os
+import shlex
+import threading
+import time
+from textwrap import dedent
+from typing import Any
+
+from .errors import IDAConnectionError, IDATimeoutError, IDAToolError, Session
+
+# ida_codemode is imported EAGERLY-IF-PRESENT but never at hard import cost.
+#
+# The paging/graph/trace layers and their offline test suites must keep importing
+# `idatui` on a machine with no IDA and no Code Mode installed -- that is the
+# house rule the stdlib-only worker client used to satisfy for free, and
+# `tests/run.py --fast` (257 checks, any python3) depends on it. A hard top-level
+# import here makes the whole package unimportable, so the failure is deferred to
+# the first operation that genuinely needs the library.
+_CODEMODE_ERROR: Exception | None = None
+try:
+ from ida_codemode.client import (
+ ClientError,
+ DatabaseHandle,
+ InstanceDisconnectedError,
+ RemoteError,
+ )
+ from ida_codemode.registry import (
+ REGISTRY_DIR,
+ FileLock,
+ RegistryEntry,
+ canonical_path,
+ idb_key,
+ scan_instances,
+ )
+ from ida_codemode.resolver import IdbBusy, expected_idb_path
+except ImportError as _exc: # library absent: usable only for offline layers
+ _CODEMODE_ERROR = _exc
+ # Bound to None rather than left undefined so the names stay patchable: the
+ # offline contract tests inject a fake DatabaseHandle here.
+ ClientError = InstanceDisconnectedError = RemoteError = None # type: ignore[assignment,misc]
+ DatabaseHandle = RegistryEntry = FileLock = None # type: ignore[assignment,misc]
+ REGISTRY_DIR = canonical_path = idb_key = scan_instances = None # type: ignore[assignment]
+ IdbBusy = expected_idb_path = None # type: ignore[assignment]
+
+
+def _require_codemode() -> None:
+ """Raise an actionable error when the Code Mode library is missing.
+
+ Gated on the binding, not on the original import result, so a test that
+ injects a fake ``DatabaseHandle`` exercises the real adapter logic.
+ """
+ if DatabaseHandle is None:
+ raise IDAConnectionError(
+ "ida-codemode-mcp is not installed in this environment "
+ f"({_CODEMODE_ERROR}). Install it (e.g. `uv sync`, or "
+ "`pip install -e ../ida-codemode-mcp`) so ida-tui can lease a "
+ "database.") from _CODEMODE_ERROR
+
+
+def database_owner(idb_path: str, staged_path: str | None = None):
+ """The registry entry that owns ``idb_path``/``staged_path``, else None.
+
+ Returns None when the Code Mode library is absent: with no library there is
+ no client in this environment that could be holding the database, and the
+ IDA-free layers (project staging) must keep working. Registry errors that
+ happen WITH the library installed still propagate -- those mean "we could
+ not determine ownership", which is not the same as "nobody owns it".
+ """
+ if DatabaseHandle is None:
+ return None
+ expected_key = idb_key(idb_path)
+ staged = canonical_path(staged_path) if staged_path else None
+ for item in scan_instances(timeout=0.5):
+ entry = item.entry
+ if entry.idb_key == expected_key:
+ return entry
+ if staged and entry.exe_path and canonical_path(entry.exe_path) == staged:
+ return entry
+ return None
+
+
+def registered_database(path: str, output_database: str | None = None) -> bool:
+ """Whether a live/lock-held Code Mode instance owns this target."""
+ _require_codemode()
+ source = canonical_path(path)
+ expected = canonical_path(output_database) if output_database else expected_idb_path(source)
+ expected_key = idb_key(expected)
+ for instance in scan_instances(timeout=0.5):
+ entry = instance.entry
+ if entry.idb_key == expected_key:
+ return True
+ if not output_database and entry.backend == "gui" and entry.exe_path:
+ if canonical_path(entry.exe_path) == source:
+ return True
+ return False
+
+
+class _NoopKeepAlive:
+ """Compatibility shim: the DatabaseHandle's SSE lease is the heartbeat."""
+
+ def __init__(self) -> None:
+ self.beats = self.failures = 0
+
+ def start(self) -> "_NoopKeepAlive":
+ return self
+
+ def stop(self) -> None:
+ pass
+
+
+def _parse_load_args(value: str) -> tuple[str | None, int | None, str | None]:
+ """Translate ida-tui's legacy first-open switches to Code Mode options.
+
+ Code Mode has typed options for processor, natural loading address and file
+ type. It deliberately has no arbitrary command-line escape hatch; reject
+ switches we cannot represent instead of silently loading a blob wrongly.
+ """
+ processor: str | None = None
+ loading_address: int | None = None
+ file_type: str | None = None
+ unsupported: list[str] = []
+ try:
+ words = shlex.split(value or "", posix=os.name != "nt")
+ except ValueError as exc:
+ raise ValueError(f"invalid IDA load options: {exc}") from exc
+ for word in words:
+ if word.startswith("-p") and len(word) > 2:
+ processor = word[2:]
+ elif word.startswith("-b") and len(word) > 2:
+ try:
+ # IDA's -b is in 16-byte paragraphs. DatabaseHandle expects the
+ # natural address, which is the safer public API.
+ loading_address = int(word[2:], 16) << 4
+ except ValueError as exc:
+ raise ValueError(f"invalid IDA loading address: {word!r}") from exc
+ elif word.startswith("-T") and len(word) > 2:
+ file_type = word[2:]
+ else:
+ unsupported.append(word)
+ if unsupported:
+ joined = " ".join(unsupported)
+ raise ValueError(
+ "ida-codemode cannot represent arbitrary IDA load options: "
+ f"{joined!r}; use processor/base/file type options instead"
+ )
+ return processor, loading_address, file_type
+
+
+def _script(args: dict[str, Any], body: str) -> str:
+ """Bind JSON arguments without interpolating user text into Python code."""
+ encoded = json.dumps(args, ensure_ascii=False, separators=(",", ":"))
+ return f"import json\na = json.loads({encoded!r})\n{dedent(body).strip()}\n"
+
+
+# Rich flat-listing generation is the largest ida-domain gap in this port.
+# ida-domain can enumerate heads and render plain disassembly, but it does not
+# expose undefined runs, IDA colour spans, function banners, or expanded UDT
+# members. Keep that IDAPython-only logic isolated in this one operation.
+_HEADS = r'''
+import ida_bytes, ida_funcs, ida_idaapi, ida_lines, ida_name, ida_nalt, ida_segment, ida_typeinf
+start = int(str(a["addr"]), 16)
+count = max(1, min(int(a.get("count", 200)), 2000))
+offset = max(0, int(a.get("offset", 0)))
+annotate = bool(a.get("annotate", False))
+seg = db.segments.get_at(start)
+if seg is None:
+ result = {"addr": a["addr"], "error": "no segment", "heads": [], "cursor": {"done": True}}
+else:
+ lo, hi = int(seg.start_ea), int(seg.end_ea)
+ if a.get("end"):
+ hi = min(hi, int(str(a["end"]), 16))
+
+ span_names = {
+ "insn": ("SCOLOR_INSN", "SCOLOR_KEYWORD", "SCOLOR_ASMDIR", "SCOLOR_MACRO"),
+ "reg": ("SCOLOR_REG",),
+ "num": ("SCOLOR_NUMBER", "SCOLOR_CHAR", "SCOLOR_BINPREF"),
+ "str": ("SCOLOR_STRING",),
+ "name": ("SCOLOR_DATNAME", "SCOLOR_CODNAME", "SCOLOR_LOCNAME", "SCOLOR_IMPNAME",
+ "SCOLOR_DEMNAME", "SCOLOR_LIBNAME", "SCOLOR_CNAME", "SCOLOR_DNAME",
+ "SCOLOR_CREF", "SCOLOR_DREF", "SCOLOR_CREFTAIL", "SCOLOR_DREFTAIL"),
+ "seg": ("SCOLOR_SEGNAME",),
+ "cmt": ("SCOLOR_AUTOCMT", "SCOLOR_REGCMT", "SCOLOR_RPTCMT", "SCOLOR_VOIDOP"),
+ "punct": ("SCOLOR_SYMBOL", "SCOLOR_ALTOP", "SCOLOR_HIDNAME"),
+ "err": ("SCOLOR_ERROR",),
+ }
+ tag_kinds = {}
+ for kind, names in span_names.items():
+ for name in names:
+ value = getattr(ida_lines, name, None)
+ if isinstance(value, str) and value:
+ tag_kinds[value[0]] = kind
+ elif isinstance(value, int):
+ tag_kinds[chr(value)] = kind
+
+ def spans(tagged):
+ on, off, esc = "\x01", "\x02", "\x03"
+ addr_tag = chr(getattr(ida_lines, "COLOR_ADDR", 0x28))
+ addr_len = int(getattr(ida_lines, "COLOR_ADDR_SIZE", 16))
+ out, stack, buf = [], [], []
+ def flush():
+ if buf:
+ out.append([stack[-1] if stack else "text", "".join(buf)])
+ buf.clear()
+ i = 0
+ while i < len(tagged):
+ ch = tagged[i]
+ if ch == on and i + 1 < len(tagged):
+ tag = tagged[i + 1]
+ if tag == addr_tag:
+ i += 2 + addr_len
+ continue
+ flush(); stack.append(tag_kinds.get(tag, "text")); i += 2; continue
+ if ch == off and i + 1 < len(tagged):
+ flush()
+ if stack: stack.pop()
+ i += 2; continue
+ if ch == esc and i + 1 < len(tagged):
+ buf.append(tagged[i + 1]); i += 2; continue
+ buf.append(ch); i += 1
+ flush()
+ collapsed, previous_space = [], False
+ for kind, text in out:
+ acc = []
+ for ch in text:
+ if ch.isspace():
+ if previous_space: continue
+ acc.append(" "); previous_space = True
+ else:
+ acc.append(ch); previous_space = False
+ if acc: collapsed.append([kind, "".join(acc)])
+ if collapsed:
+ collapsed[0][1] = collapsed[0][1].lstrip()
+ collapsed[-1][1] = collapsed[-1][1].rstrip()
+ return [[kind, text] for kind, text in collapsed if text]
+
+ def row(ea):
+ flags = ida_bytes.get_flags(ea)
+ kind = "code" if ida_bytes.is_code(flags) else ("data" if ida_bytes.is_data(flags) else "unknown")
+ tagged = ida_lines.generate_disasm_line(ea, 0) or ""
+ text = " ".join(ida_lines.tag_remove(tagged).split()) if tagged else ""
+ item = {"ea": hex(ea), "kind": kind, "size": int(ida_bytes.get_item_size(ea)), "text": text}
+ if tagged:
+ rich = spans(tagged)
+ if " ".join("".join(x[1] for x in rich).split()) == text:
+ item["spans"] = rich
+ name = ida_name.get_ea_name(ea)
+ if name: item["name"] = name
+ return item
+
+ def unknown_row(ea, size):
+ if size <= 1: return row(ea)
+ item = {"ea": hex(ea), "kind": "unknown", "size": int(size), "text": f"db {size} dup(?)"}
+ name = ida_name.get_ea_name(ea)
+ if name: item["name"] = name
+ return item
+
+ def members(ea):
+ tif = db.types.get_at(ea)
+ if tif is None or not tif.is_udt(): return []
+ answer = []
+ for member in db.types.get_udt_members(tif):
+ type_text = member.type.dstr() or ""
+ text = f"+{member.offset:X} {member.name}" + (f" {type_text}" if type_text else "")
+ answer.append({"ea": hex(ea + member.offset), "kind": "member",
+ "size": int(member.size), "text": text})
+ return answer
+
+ def is_unknown(ea):
+ flags = ida_bytes.get_flags(ea)
+ return not (ida_bytes.is_code(flags) or ida_bytes.is_data(flags))
+ def run_end(ea):
+ nxt = ida_bytes.next_head(ea, hi)
+ return nxt if nxt != ida_idaapi.BADADDR and ea < nxt <= hi else hi
+ def advance(ea):
+ if is_unknown(ea): return run_end(ea)
+ nxt = ida_bytes.get_item_end(ea)
+ return nxt if nxt > ea else ea + 1
+ def rows_for(ea):
+ if is_unknown(ea): return [unknown_row(ea, run_end(ea) - ea)]
+ fn = db.functions.get_at(ea) if annotate else None
+ at_start = fn is not None and int(fn.start_ea) == ea
+ answer = []
+ if at_start:
+ name = db.functions.get_name(fn) or f"sub_{ea:X}"
+ answer += [
+ {"ea": hex(ea), "kind": "sep", "size": 0, "text": ""},
+ {"ea": hex(ea), "kind": "sep", "size": 0,
+ "text": "; " + "=" * 15 + " S U B R O U T I N E " + "=" * 15},
+ {"ea": hex(ea), "kind": "funchdr", "size": 0,
+ "text": name + " proc", "name": name},
+ ]
+ item = row(ea)
+ if at_start:
+ item["name"] = None
+ elif annotate and item["kind"] == "code" and item.get("name"):
+ name = item["name"]
+ answer.append({"ea": hex(ea), "kind": "label", "size": 0,
+ "text": name + ":", "name": name})
+ item["name"] = None
+ answer.append(item)
+ if item["kind"] == "data": answer += members(ea)
+ if fn is not None and ida_bytes.get_item_end(ea) >= int(fn.end_ea):
+ name = db.functions.get_name(fn) or f"sub_{int(fn.start_ea):X}"
+ answer += [
+ {"ea": hex(ea), "kind": "funchdr", "size": 0,
+ "text": name + " endp", "name": name},
+ {"ea": hex(ea), "kind": "sep", "size": 0, "text": "; " + "-" * 60},
+ ]
+ return answer
+
+ ea = ida_bytes.get_item_head(start)
+ if ea == ida_idaapi.BADADDR: ea = start
+ for _ in range(offset):
+ if ea >= hi: break
+ ea = advance(ea)
+ rows = []
+ more = False
+ while ea != ida_idaapi.BADADDR and ea < hi:
+ if len(rows) >= count:
+ more = True; break
+ rows += rows_for(ea)
+ ea = advance(ea)
+ result = {"addr": a["addr"], "heads": rows,
+ "cursor": {"next": hex(ea)} if more else {"done": True}}
+result
+'''
+
+
+_DECOMP_MAP_HELPER = r'''
+def line_map(cfunc):
+ import ida_hexrays
+ answer = []
+ for sl in cfunc.get_pseudocode():
+ tagged, eas, seen = sl.line, [], set()
+ for x in range(len(tagged) + 1):
+ head = ida_hexrays.ctree_item_t(); item = ida_hexrays.ctree_item_t(); tail = ida_hexrays.ctree_item_t()
+ if not cfunc.get_line_item(tagged, x, False, head, item, tail): continue
+ text = item.dstr() or ""
+ try: ea = int(text.split(": ", 1)[0], 16)
+ except (ValueError, IndexError): continue
+ if ea not in seen: seen.add(ea); eas.append(ea)
+ answer.append(eas)
+ return answer
+'''
+
+
+_OPERATIONS: dict[str, str] = {
+ "list_funcs": r'''
+import fnmatch
+queries = a.get("queries") or [{}]
+q = queries[0]
+offset, count = max(0, int(q.get("offset", 0))), max(1, int(q.get("count", 500)))
+pattern = str(q.get("filter") or "").lower()
+if pattern and not any(ch in pattern for ch in "*?["): pattern = "*" + pattern + "*"
+rows = []
+for fn in db.functions.get_all():
+ name = db.functions.get_name(fn) or f"sub_{int(fn.start_ea):X}"
+ if pattern and not fnmatch.fnmatchcase(name.lower(), pattern): continue
+ rows.append({"addr": hex(int(fn.start_ea)), "name": name,
+ "size": int(fn.end_ea) - int(fn.start_ea)})
+page = rows[offset:offset + count]
+result = {"result": [{"data": page, "next_offset": offset + len(page), "total": len(rows)}]}
+result
+''',
+ "disasm": r'''
+ea = int(str(a["addr"]), 16)
+fn = db.functions.get_at(ea)
+if fn is None:
+ result = {"instructions": [], "total_instructions": 0, "instruction_count": 0}
+else:
+ instructions = list(db.functions.get_instructions(fn))
+ limit = max(1, int(a.get("max_instructions", len(instructions) or 1)))
+ rows = [{"addr": hex(int(insn.ea)), "instruction": db.instructions.get_disassembly(insn)}
+ for insn in instructions[:limit]]
+ result = {"instructions": rows, "total_instructions": len(instructions),
+ "instruction_count": len(instructions)}
+result
+''',
+ "file_regions": r'''
+import idaapi
+rows = []
+for seg in db.segments.get_all():
+ try: file_off = int(idaapi.get_fileregion_offset(seg.start_ea))
+ except Exception: file_off = -1
+ if file_off < 0 or file_off >= (1 << 48): file_off = -1
+ rows.append({"start": hex(int(seg.start_ea)), "end": hex(int(seg.end_ea)),
+ "file_off": file_off, "name": db.segments.get_name(seg) or ""})
+result = {"regions": rows}
+result
+''',
+ "read_raw": r'''
+import ida_bytes
+ea, size = int(str(a["addr"]), 16), max(0, int(a["size"]))
+raw = ida_bytes.get_bytes(ea, size) or b""
+raw = raw[:size] + b"\xff" * max(0, size - len(raw))
+data = bytearray(raw)
+for index, value in enumerate(data):
+ if value == 0xFF and not ida_bytes.is_loaded(ea + index): data[index] = 0
+result = {"addr": a["addr"], "hex": bytes(data).hex(), "n": len(data)}
+result
+''',
+ "get_bytes": r'''
+rows = []
+for region in a.get("regions", []):
+ ea, size = int(str(region["addr"]), 16), int(region["size"])
+ raw = db.bytes.get_bytes_at(ea, size) or b""
+ rows.append({"addr": region["addr"], "data": " ".join(f"{b:02x}" for b in raw)})
+result = {"result": rows}
+result
+''',
+ "search_structs": r'''
+needle = str(a.get("filter") or "").lower()
+rows = []
+for tif in db.types.get_all():
+ name = tif.get_type_name() or ""
+ if not name or needle not in name.lower() or not tif.is_udt(): continue
+ members = list(db.types.get_udt_members(tif))
+ rows.append({"name": name, "size": int(tif.get_size()), "is_union": bool(tif.is_union()),
+ "cardinality": len(members), "ordinal": int(tif.get_ordinal())})
+result = {"result": rows}
+result
+''',
+ "type_inspect": r'''
+rows = []
+for query in a.get("queries", []):
+ name = str(query.get("name") or "")
+ tif = db.types.get_by_name(name)
+ if tif is None:
+ rows.append({"name": name, "error": "type not found"}); continue
+ members = [{"name": m.name, "type": m.type.dstr() or str(m.type),
+ "offset": int(m.offset), "size": int(m.size)}
+ for m in db.types.get_udt_members(tif)] if tif.is_udt() else []
+ rows.append({"name": name, "size": int(tif.get_size()), "is_union": bool(tif.is_union()),
+ "members": members})
+result = {"result": rows}
+result
+''',
+ "declare_type": r'''
+import ida_typeinf
+decls = a.get("decls", "")
+if isinstance(decls, str): decls = [decls]
+rows = []
+for declaration in decls:
+ try:
+ errors = int(db.types.parse_declarations(ida_typeinf.get_idati(), declaration))
+ rows.append({"ok": errors == 0, **({} if errors == 0 else {"error": f"{errors} parse error(s)"})})
+ except Exception as exc:
+ rows.append({"ok": False, "error": str(exc)})
+result = {"result": rows}
+result
+''',
+ "del_type": r'''
+import ida_typeinf
+name = str(a["name"])
+ok = bool(ida_typeinf.del_named_type(ida_typeinf.get_idati(), name, ida_typeinf.NTF_TYPE))
+result = {"name": name, "deleted": ok, **({} if ok else {"error": f"Type {name!r} not found or could not be deleted"})}
+result
+''',
+ "func_types": r'''
+import ida_typeinf
+ea = int(str(a["addr"]), 16)
+fn = db.functions.get_at(ea)
+if fn is None:
+ result = {"addr": a["addr"], "error": "no function at address"}
+else:
+ pseudo = db.pseudocode.decompile(fn)
+ name = db.functions.get_name(fn) or ""
+ tif = pseudo.get_func_type()
+ try: prototype = ida_typeinf.print_tinfo("", 0, 0, ida_typeinf.PRTYPE_1LINE, tif, name, "") if tif else ""
+ except Exception: prototype = tif.dstr() if tif else ""
+ lvars = [{"name": var.name, "type": var.type_info.dstr() if var.type_info else "",
+ "is_arg": bool(var.is_arg)} for var in pseudo.local_variables]
+ result = {"addr": hex(int(fn.start_ea)), "name": name,
+ "prototype": (prototype or "").strip(), "lvars": lvars}
+result
+''',
+ "set_lvar_type": r'''
+import ida_typeinf
+ea, variable, declaration = int(str(a["addr"]), 16), str(a["variable"]), str(a["type"])
+fn = db.functions.get_at(ea)
+if fn is None:
+ result = {"error": "no function at address"}
+else:
+ pseudo = db.pseudocode.decompile(fn)
+ var = pseudo.find_local_variable(variable)
+ if var is None:
+ result = {"error": f"local variable {variable!r} not found"}
+ else:
+ try:
+ tif = db.types.parse_one_declaration(ida_typeinf.get_idati(), declaration)
+ accepted = bool(var.set_type(tif))
+ saved = bool(pseudo.save_local_variable_info(var, save_type=True)) if accepted else False
+ result = {"addr": hex(int(fn.start_ea)), "variable": variable,
+ "type": declaration, "ok": accepted and saved}
+ except Exception as exc:
+ result = {"error": f"bad type {declaration!r}: {exc}"}
+result
+''',
+ "set_type": r'''
+from ida_domain.types import TypeApplyFlags
+rows = []
+for edit in a.get("edits", []):
+ ea = int(str(edit["addr"]), 16)
+ declaration = str(edit.get("signature") or edit.get("type") or "")
+ try:
+ ok = bool(db.types.apply_declaration_at(ea, declaration, TypeApplyFlags.DEFINITE))
+ rows.append({"addr": hex(ea), "ok": ok, **({} if ok else {"error": "IDA rejected the type"})})
+ except Exception as exc:
+ rows.append({"addr": hex(ea), "ok": False, "error": str(exc)})
+result = {"result": rows}
+result
+''',
+ "data_type": r'''
+ea = int(str(a["addr"]), 16)
+try:
+ tif = db.types.get_at(ea)
+ fn = db.functions.get_at(ea)
+ result = {"addr": hex(ea), "name": db.names.get_at(ea) or "",
+ "type": tif.dstr() if tif else "", "size": int(db.heads.size(ea)) if db.heads.is_head(ea) else 0,
+ "is_func": bool(fn)}
+except Exception as exc:
+ result = {"addr": hex(ea), "error": str(exc)}
+result
+''',
+ "force_recompile": r'''
+import ida_hexrays
+rows = []
+for item in a.get("items", []):
+ ea = int(str(item["addr"]), 16)
+ ida_hexrays.mark_cfunc_dirty(ea, False)
+ rows.append({"addr": hex(ea), "ok": True})
+result = {"result": rows}
+result
+''',
+ "undefine": r'''
+import ida_bytes
+rows = []
+for item in a.get("items", []):
+ ea = int(str(item["addr"]), 16)
+ size = max(1, int(item.get("size") or ida_bytes.get_item_size(ea) or 1))
+ ok = bool(ida_bytes.del_items(ea, ida_bytes.DELIT_SIMPLE, size))
+ rows.append({"addr": hex(ea), "ok": ok, **({} if ok else {"error": "delete items failed"})})
+result = {"result": rows}
+result
+''',
+ "define_code": r'''
+import ida_ua
+rows = []
+for item in a.get("items", []):
+ ea = int(str(item["addr"]), 16); size = int(ida_ua.create_insn(ea))
+ rows.append({"addr": hex(ea), "ok": size > 0, "size": size,
+ **({} if size > 0 else {"error": "instruction did not decode"})})
+result = {"result": rows}
+result
+''',
+ "define_func": r'''
+rows = []
+for item in a.get("items", []):
+ ea = int(str(item["addr"]), 16); ok = bool(db.functions.create(ea))
+ rows.append({"addr": hex(ea), "ok": ok, **({} if ok else {"error": "IDA refused the function"})})
+result = {"result": rows}
+result
+''',
+ "make_data": r'''
+import ida_bytes, ida_idaapi, ida_typeinf
+from ida_domain.types import TypeApplyFlags
+rows = []
+for item in a.get("items", []):
+ ea, declaration = int(str(item["addr"]), 16), str(item["type"])
+ try:
+ tif = db.types.parse_one_declaration(ida_typeinf.get_idati(), declaration)
+ size = max(1, int(tif.get_size()))
+ saved_names = [(addr, name) for addr, name in db.names.get_all()
+ if ea <= int(addr) < ea + size]
+ ida_bytes.del_items(ea, ida_bytes.DELIT_EXPAND | ida_bytes.DELIT_DELNAMES,
+ max(size, int(ida_bytes.get_item_size(ea) or 1)))
+ created = bool(ida_bytes.create_data(ea, ida_bytes.FF_BYTE, size, ida_idaapi.BADADDR))
+ ok = created and bool(db.types.apply_at(tif, ea, TypeApplyFlags.DEFINITE))
+ for address, name in saved_names:
+ db.names.set_name(int(address), name)
+ if ok and item.get("name"): ok = bool(db.names.set_name(ea, str(item["name"])))
+ rows.append({"addr": hex(ea), "ok": ok, "size": size,
+ **({} if ok else {"error": "IDA rejected the data type"})})
+ except Exception as exc:
+ rows.append({"addr": hex(ea), "ok": False, "error": str(exc)})
+result = {"result": rows}
+result
+''',
+ "make_string": r'''
+from ida_domain.strings import StringType
+ea, length = int(str(a["addr"]), 16), max(0, int(a.get("length", 0)))
+kind = {"c": StringType.C, "c16": StringType.C_16, "c32": StringType.C_32,
+ "pascal": StringType.PASCAL}.get(str(a.get("kind", "c")).lower(), StringType.C)
+import ida_bytes
+try:
+ ida_bytes.del_items(ea, ida_bytes.DELIT_SIMPLE, length if length > 0 else 1)
+except Exception:
+ pass
+try:
+ ok = bool(db.bytes.create_string_at(ea, length or None, kind))
+ text = db.bytes.get_string_at(ea) or "" if ok else ""
+ result = {"addr": hex(ea), "ok": ok, "size": int(db.heads.size(ea)) if ok else 0, "text": text}
+except Exception as exc:
+ result = {"addr": hex(ea), "ok": False, "error": str(exc)}
+result
+''',
+ "list_strings": r'''
+from ida_domain.strings import StringListConfig
+offset, count, min_len = max(0, int(a.get("offset", 0))), max(1, int(a.get("count", 2000))), max(1, int(a.get("min_len", 4)))
+if offset == 0 or a.get("refresh"):
+ from ida_domain.strings import StringType
+ db.strings.rebuild(StringListConfig(string_types=list(StringType), min_len=min_len,
+ only_ascii_7bit=False))
+items = list(db.strings.get_all())
+page = items[offset:offset + count]
+rows = []
+for item in page:
+ try: text = str(item)
+ except Exception: text = item.contents.decode("utf-8", "replace") if item.contents else ""
+ rows.append({"addr": hex(int(item.address)), "text": text, "len": int(item.length), "type": item.type.name})
+result = {"strings": rows, "total": len(items), "next_offset": offset + len(rows)}
+result
+''',
+ "list_linkage": r'''
+imports = [{"addr": hex(int(item.address)), "name": item.name, "module": item.module_name}
+ for item in db.imports.get_all_imports() if item.name]
+exports = [{"addr": hex(int(item.address)), "name": item.name, "ordinal": int(item.ordinal)}
+ for item in db.entries.get_all() if item.name]
+result = {"imports": imports, "exports": exports,
+ "n_imports": len(imports), "n_exports": len(exports)}
+result
+''',
+ "lookup_funcs": r'''
+rows = []
+for query in a.get("queries", []):
+ raw = str(query)
+ try: ea = int(raw, 16)
+ except ValueError:
+ fn = db.functions.get_by_name(raw); ea = int(fn.start_ea) if fn else None
+ else: fn = db.functions.get_at(ea)
+ if fn is None:
+ rows.append({"query": raw, "fn": None})
+ else:
+ rows.append({"query": raw, "fn": {"addr": hex(int(fn.start_ea)),
+ "name": db.functions.get_name(fn) or f"sub_{int(fn.start_ea):X}",
+ "size": int(fn.end_ea) - int(fn.start_ea)}})
+result = {"result": rows}
+result
+''',
+ "resolve_names": r'''
+import ida_idaapi, ida_name
+rows = []
+for query in a.get("queries", []):
+ name = str(query).strip(); ea = ida_name.get_name_ea(ida_idaapi.BADADDR, name)
+ rows.append({"query": name, "ea": hex(int(ea)) if ea != ida_idaapi.BADADDR else None})
+result = {"result": rows}
+result
+''',
+ "xref_types": r'''
+queries = a.get("queries") or []
+all_results = []
+for query in queries:
+ ea, direction = int(str(query["addr"]), 16), str(query.get("direction", "both"))
+ refs = []
+ if direction in ("to", "both"): refs += list(db.xrefs.to_ea(ea))
+ if direction in ("from", "both"): refs += list(db.xrefs.from_ea(ea))
+ rows, seen = [], set()
+ for ref in refs:
+ key = (int(ref.from_ea), int(ref.to_ea), int(ref.type))
+ if query.get("dedup") and key in seen: continue
+ seen.add(key)
+ fn = db.functions.get_at(int(ref.from_ea))
+ kind = ("call" if ref.is_call else "jump" if ref.is_jump else "flow" if ref.is_flow
+ else "read" if ref.is_read else "write" if ref.is_write else ref.type.name.lower())
+ row = {"from": hex(int(ref.from_ea)), "to": hex(int(ref.to_ea)),
+ "type": "code" if ref.is_code else "data", "kind": kind}
+ if query.get("include_fn") and fn is not None:
+ row["fn"] = {"addr": hex(int(fn.start_ea)), "name": db.functions.get_name(fn) or ""}
+ rows.append(row)
+ if len(rows) >= int(query.get("count", 2000)): break
+ all_results.append({"data": rows})
+result = {"result": all_results}
+result
+''',
+ "xref_query": r'''
+queries = a.get("queries") or []
+all_results = []
+for query in queries:
+ ea, direction = int(str(query["addr"]), 16), str(query.get("direction", "both"))
+ refs = []
+ if direction in ("to", "both"): refs += list(db.xrefs.to_ea(ea))
+ if direction in ("from", "both"): refs += list(db.xrefs.from_ea(ea))
+ rows = []
+ for ref in refs[:int(query.get("count", 2000))]:
+ fn = db.functions.get_at(int(ref.from_ea))
+ row = {"from": hex(int(ref.from_ea)), "to": hex(int(ref.to_ea)),
+ "type": "code" if ref.is_code else "data"}
+ if query.get("include_fn") and fn is not None:
+ row["fn"] = {"addr": hex(int(fn.start_ea)), "name": db.functions.get_name(fn) or ""}
+ rows.append(row)
+ all_results.append({"data": rows})
+result = {"result": all_results}
+result
+''',
+ "set_comments": r'''
+rows = []
+for item in a.get("items", []):
+ ea, text = int(str(item["addr"]), 16), str(item.get("comment") or "")
+ try:
+ if text: ok = bool(db.comments.set_at(ea, text))
+ else: db.comments.delete_at(ea); ok = True
+ rows.append({"addr": hex(ea), "ok": ok})
+ except Exception as exc:
+ rows.append({"addr": hex(ea), "ok": False, "error": str(exc)})
+result = {"result": rows}
+result
+''',
+ "rename": r'''
+import ida_idaapi, ida_name, ida_typeinf
+batch = a.get("batch") or {}
+out = {}; ok_count = failed = 0
+for category, edit in batch.items():
+ try:
+ if category == "func":
+ ea, new = int(str(edit["addr"]), 16), str(edit["name"])
+ fn = db.functions.get_at(ea); ok = bool(fn and db.functions.set_name(fn, new))
+ elif category == "data":
+ new = str(edit.get("new") or "")
+ if edit.get("addr") is not None: ea = int(str(edit["addr"]), 16)
+ else: ea = int(ida_name.get_name_ea(ida_idaapi.BADADDR, str(edit.get("old") or "")))
+ ok = bool(db.names.set_name(ea, new))
+ elif category in ("local", "stack"):
+ ea, old, new = int(str(edit["func_addr"]), 16), str(edit["old"]), str(edit["new"])
+ pseudo = db.pseudocode.decompile(ea); var = pseudo.find_local_variable(old)
+ if var is None: ok = False
+ else:
+ var.set_user_name(new)
+ ok = bool(pseudo.save_local_variable_info(var, save_name=True))
+ else:
+ raise ValueError(f"unsupported rename category: {category}")
+ row = {"ok": ok, **({} if ok else {"error": "IDA rejected the name"})}
+ except Exception as exc:
+ row = {"ok": False, "error": str(exc)}
+ out[category] = [row]
+ if row["ok"]: ok_count += 1
+ else: failed += 1
+out["summary"] = {"ok": ok_count, "failed": failed}
+result = out
+result
+''',
+}
+
+
+_OPERATIONS["decompile"] = _DECOMP_MAP_HELPER + r'''
+ea = int(str(a["addr"]), 16)
+fn = db.functions.get_at(ea)
+if fn is None:
+ result = {"error": f"no function at {ea:#x}"}
+else:
+ pseudo = db.pseudocode.decompile(fn)
+ mapping = line_map(pseudo.raw_cfunc)
+ plain = pseudo.to_text()
+ marked = [line + (f" /*0x{eas[0]:X}*/" if eas else "")
+ for line, eas in zip(plain, mapping)]
+ import ida_name
+ refs, seen = [], set()
+ for expr in pseudo.find_objects():
+ target = int(expr.obj_ea)
+ if target in seen or not (db.is_valid_ea(target) or db.is_private_ea(target)): continue
+ seen.add(target)
+ name = expr.obj_name or ida_name.get_name(target) or ""
+ try: string = db.bytes.get_string_at(target) if db.is_valid_ea(target) else None
+ except Exception: string = None
+ refs.append({"addr": hex(target), "name": name, "string": string})
+ result = {"addr": hex(int(fn.start_ea)), "code": "\n".join(marked), "refs": refs}
+result
+'''
+
+_OPERATIONS["decomp_map"] = _DECOMP_MAP_HELPER + r'''
+ea = int(str(a["addr"]), 16)
+fn = db.functions.get_at(ea)
+if fn is None:
+ result = {"error": f"no function at {ea:#x}"}
+else:
+ pseudo = db.pseudocode.decompile(fn)
+ mapping = line_map(pseudo.raw_cfunc)
+ result = {"addr": hex(int(fn.start_ea)),
+ "lines": [{"ea": hex(eas[0]) if eas else None,
+ "eas": [hex(item) for item in eas]} for eas in mapping]}
+result
+'''
+
+_OPERATIONS["define_code_run"] = r'''
+import ida_bytes, ida_idp, ida_segment, ida_ua, idaapi
+ea, limit = int(str(a["addr"]), 16), max(1, min(int(a.get("limit", 20000)), 200000))
+seg = ida_segment.getseg(ea)
+if seg is None:
+ result = {"addr": a["addr"], "error": "no segment", "count": 0}
+else:
+ start, count, stopped, hi = ea, 0, "limit", int(seg.end_ea)
+ while count < limit:
+ if ea >= hi: stopped = "segment"; break
+ flags = ida_bytes.get_flags(ea)
+ if ida_bytes.is_code(flags) or ida_bytes.is_data(flags): stopped = "defined"; break
+ size = int(ida_ua.create_insn(ea))
+ if size <= 0: stopped = "undecodable"; break
+ count += 1
+ insn = ida_ua.insn_t()
+ if ida_ua.decode_insn(insn, ea) > 0:
+ try: is_ret = bool(ida_idp.is_ret_insn(insn))
+ except Exception: is_ret = False
+ if is_ret or (insn.get_canon_feature() & idaapi.CF_STOP):
+ ea += size; stopped = "flow"; break
+ ea += size
+ result = {"start": hex(start), "end": hex(ea), "count": count, "stopped": stopped}
+result
+'''
+
+_OPERATIONS["define_func_run"] = r'''
+import ida_bytes, ida_funcs, ida_segment
+ea = int(str(a["addr"]), 16)
+fn = db.functions.get_at(ea)
+if fn is not None and int(fn.start_ea) == ea:
+ result = {"addr": hex(ea), "ok": True, "start": hex(ea), "end": hex(int(fn.end_ea)), "how": "existed"}
+else:
+ automatic = bool(db.functions.create(ea))
+ if not automatic:
+ seg = db.segments.get_at(ea); end = ea; hi = int(seg.end_ea) if seg else ea
+ while end < hi and ida_bytes.is_code(ida_bytes.get_flags(end)):
+ nxt = int(ida_bytes.get_item_end(end))
+ if nxt <= end: break
+ end = nxt
+ ok = bool(end > ea and ida_funcs.add_func(ea, end))
+ else: ok = True
+ fn = db.functions.get_at(ea)
+ result = ({"addr": hex(ea), "ok": True, "start": hex(int(fn.start_ea)),
+ "end": hex(int(fn.end_ea)), "how": "auto" if automatic else "explicit-end"}
+ if ok and fn is not None else
+ {"addr": hex(ea), "ok": False, "error": f"IDA refused a function at {ea:#x}"})
+result
+'''
+
+_OPERATIONS["set_thumb"] = r'''
+import ida_bytes, ida_ida, ida_idp, ida_segment, ida_segregs
+ea = int(str(a["addr"]), 16); treg = ida_idp.str2reg("T")
+seg = ida_segment.getseg(ea)
+if treg is None or treg < 0:
+ result = {"addr": hex(ea), "error": "no T register (not an ARM database)"}
+elif seg is None:
+ result = {"addr": hex(ea), "error": "no segment"}
+else:
+ current = ida_segregs.get_sreg(ea, treg)
+ current = 0 if current in (None, 0xFFFFFFFF, -1) else int(current)
+ want = {"on": 1, "off": 0}.get(str(a.get("mode", "toggle")).lower(), 0 if current else 1)
+ changed = False
+ if want and seg.bitness != 1:
+ ida_segment.set_segm_addressing(seg, 1); changed = True
+ size = max(int(ida_bytes.get_item_size(ea)), 2)
+ ida_bytes.del_items(ea, 0, size)
+ ok = bool(ida_segregs.split_sreg_range(ea, treg, want, ida_segregs.SR_user))
+ now = ida_segregs.get_sreg(ea, treg)
+ result = {"addr": hex(ea), "thumb": bool(now), "was": bool(current), "ok": ok,
+ "bitness": ida_segment.getseg(ea).bitness, "forced_32bit": changed,
+ "db_64bit": bool(ida_ida.inf_get_app_bitness() == 64 and want)}
+result
+'''
+
+_OPERATIONS["thumb_scan"] = r'''
+import ida_bytes, ida_funcs, ida_idp, ida_segment, ida_segregs, ida_ua
+lo, hi = int(str(a["start"]), 16), int(str(a["end"]), 16)
+apply, limit = bool(a.get("apply", True)), int(a.get("limit", 512))
+treg = ida_idp.str2reg("T"); found = []; applied = 0; cursor = lo
+while cursor + 4 <= hi and len(found) < limit:
+ at = cursor; value = int(ida_bytes.get_dword(cursor)); cursor += 4
+ if not value & 1: continue
+ target = value & ~1; seg = ida_segment.getseg(target)
+ if seg is None or not (seg.perm & ida_segment.SEGPERM_EXEC or seg.perm == 0): continue
+ flags = ida_bytes.get_flags(target)
+ if ida_bytes.is_data(flags): continue
+ item = {"at": hex(at), "value": hex(value), "target": hex(target),
+ "was_code": bool(ida_bytes.is_code(flags))}; found.append(item)
+ if not apply: continue
+ if treg is not None and treg >= 0: ida_segregs.split_sreg_range(target, treg, 1, ida_segregs.SR_user)
+ if not ida_bytes.is_code(ida_bytes.get_flags(target)):
+ ida_bytes.del_items(target, 0, 2)
+ if ida_ua.create_insn(target) <= 0: item["decoded"] = False; continue
+ item["decoded"] = True; item["function"] = bool(db.functions.get_at(target) or db.functions.create(target)); applied += 1
+result = {"start": hex(lo), "end": hex(hi), "found": found, "applied": applied, "n": len(found)}
+result
+'''
+
+_OPERATIONS["decomp_error"] = r'''
+import ida_hexrays, ida_ida
+ea = int(str(a["addr"]), 16); fn = db.functions.get_at(ea)
+result = {"addr": hex(ea), "bitness": ida_ida.inf_get_app_bitness()}
+if fn is None:
+ result["reason"] = "no function here"
+else:
+ try:
+ failure = ida_hexrays.hexrays_failure_t(); cfunc = ida_hexrays.decompile_func(fn, failure)
+ if cfunc is not None: result["reason"] = ""
+ else:
+ result.update({"reason": failure.desc() or f"error {failure.code}",
+ "code": int(failure.code), "errea": hex(int(failure.errea))})
+ except Exception as exc: result["reason"] = f"{type(exc).__name__}: {exc}"
+result
+'''
+
+
+class CodeModeClient:
+ """A leased GUI/idalib database accessed through ``ida_codemode``."""
+
+ def __init__(
+ self,
+ binary_path: str,
+ *,
+ ttl: int = 0,
+ load_args: str = "",
+ processor: str | None = None,
+ loading_address: int | None = None,
+ file_type: str | None = None,
+ output_database: str | None = None,
+ spawn: bool = True,
+ new_database: bool = False,
+ ) -> None:
+ del ttl # managed-worker lifetime is lease-based, not idle-TTL based
+ self._path = os.path.abspath(os.path.expanduser(binary_path))
+ parsed_processor, parsed_address, parsed_file_type = _parse_load_args(load_args)
+ self._processor = processor or parsed_processor
+ self._loading_address = loading_address if loading_address is not None else parsed_address
+ self._file_type = file_type or parsed_file_type
+ self._output_database = output_database
+ self._spawn = spawn
+ self._new_database = new_database
+ self._handle: DatabaseHandle | None = None
+ self._last_entry: RegistryEntry | None = None
+ self._connect_lock = threading.Lock()
+
+ def connect(self, timeout: float = 1800.0, progress=None) -> "CodeModeClient":
+ _require_codemode()
+ with self._connect_lock:
+ if self._handle is not None and self._handle.connected:
+ return self
+ if progress:
+ progress(f"discovering Code Mode database for {os.path.basename(self._path)}…")
+ try:
+ # A Ctrl+L reload releases its current managed-worker lease, but
+ # that worker remains registered during Code Mode's final-lease
+ # grace period. Retry only that known handoff window. A GUI or
+ # another long-lived client remains busy and yields a clear
+ # failure rather than being modified underneath its owner.
+ deadline = time.monotonic() + min(timeout, 60.0)
+ while True:
+ try:
+ handle = DatabaseHandle.open(
+ self._path,
+ spawn=self._spawn,
+ timeout=max(0.1, timeout),
+ output_database=self._output_database,
+ processor=self._processor,
+ loading_address=self._loading_address,
+ file_type=self._file_type,
+ new_database=self._new_database,
+ )
+ break
+ except IdbBusy:
+ if not self._new_database or time.monotonic() >= deadline:
+ raise
+ if progress:
+ progress("waiting for the previous Code Mode lease to close…")
+ # Remember the record before managed shutdown withdraws
+ # its JSON. The lifetime lock remains held until IDA has
+ # actually closed the IDB; waiting on it avoids racing a
+ # replacement worker into the old process's file lock.
+ expected = canonical_path(
+ self._output_database or expected_idb_path(self._path)
+ )
+ owners = [item.entry for item in scan_instances(timeout=0.5)
+ if item.entry.idb_key == idb_key(expected)]
+ if owners:
+ self._wait_for_entry_release(
+ owners[0], max(0.0, deadline - time.monotonic())
+ )
+ else:
+ time.sleep(0.2)
+ if progress:
+ backend = handle.entry.backend
+ progress(f"attached to {backend} database; waiting for auto-analysis…")
+ handle.wait_autoanalysis(timeout=timeout)
+ except Exception as exc: # normalize the dependency's transport errors
+ raise self._connection_error(exc) from exc
+ self._handle = handle
+ self._last_entry = handle.entry
+ return self
+
+ @staticmethod
+ def _connection_error(exc: BaseException) -> IDAConnectionError:
+ return IDAConnectionError(str(exc) or type(exc).__name__)
+
+ @property
+ def connected(self) -> bool:
+ return self._handle is not None and self._handle.connected
+
+ @property
+ def pid(self) -> int | None:
+ return self._handle.entry.pid if self._handle is not None else None
+
+ @property
+ def backend(self) -> str | None:
+ return self._handle.entry.backend if self._handle is not None else None
+
+ def execute_python(self, code: str, *, timeout: float | None = None) -> Any:
+ if not self.connected:
+ self.connect()
+ handle = self._handle
+ if handle is None:
+ raise IDAConnectionError("Code Mode database is not connected")
+ try:
+ response = handle.execute_python(code, timeout=timeout)
+ except RemoteError as exc:
+ details = exc.details or {}
+ message = str(exc)
+ if details.get("traceback"):
+ message += f"\n{details['traceback']}"
+ if exc.code == "operation_timeout":
+ raise IDATimeoutError(message) from exc
+ raise IDAToolError("execute_python", message) from exc
+ except (InstanceDisconnectedError, ClientError) as exc:
+ raise self._connection_error(exc) from exc
+ if not isinstance(response, dict) or "result" not in response:
+ raise IDAToolError("execute_python", "Code Mode returned an invalid execution result")
+ return response["result"]
+
+ def invoke(self, operation: str, *, timeout: float | None = None, **args) -> Any:
+ """Execute one TUI domain operation through Code Mode."""
+ if operation in ("idb_save", "save"):
+ return self.save_database()
+ if operation in ("server_health", "ping", "health", "state"):
+ return self.health()
+ body = _HEADS if operation == "heads" else _OPERATIONS.get(operation)
+ if body is None:
+ raise IDAToolError(operation, f"unknown ida-tui Code Mode operation: {operation}")
+ try:
+ return self.execute_python(_script(args, body), timeout=timeout)
+ except IDAToolError as exc:
+ if exc.tool == "execute_python":
+ raise IDAToolError(operation, exc.message) from exc
+ raise
+
+ # Temporary source compatibility for external drivers/tests that used the
+ # old WorkerClient. Application code uses the accurately named invoke().
+ call = invoke
+
+ def save_database(self) -> dict[str, Any]:
+ if not self.connected:
+ self.connect()
+ handle = self._handle
+ if handle is None:
+ raise IDAConnectionError("Code Mode database is not connected")
+ try:
+ return handle.save_database()
+ except RemoteError as exc:
+ raise IDAToolError("save_database", str(exc)) from exc
+ except (InstanceDisconnectedError, ClientError) as exc:
+ raise self._connection_error(exc) from exc
+
+ def health(self) -> dict[str, Any]:
+ if not self.connected:
+ self.connect()
+ assert self._handle is not None
+ entry = self._handle.entry
+ module = os.path.basename(entry.exe_path or entry.idb_path or self._path)
+ return {
+ "ok": self._handle.connected,
+ "module": module,
+ "backend": entry.backend,
+ "record_id": entry.record_id,
+ "input_path": entry.exe_path,
+ "idb_path": entry.idb_path,
+ }
+
+ def keepalive(self, interval: float = 120.0) -> _NoopKeepAlive:
+ del interval
+ return _NoopKeepAlive()
+
+ def resolve_db(self) -> str:
+ if not self.connected:
+ self.connect()
+ assert self._handle is not None
+ return self._handle.entry.record_id
+
+ def set_db(self, db: str | None) -> None:
+ del db # one handle is permanently bound to one registered database
+
+ def list_sessions(self) -> list[Session]:
+ if not self.connected:
+ self.connect()
+ assert self._handle is not None
+ entry = self._handle.entry
+ path = entry.exe_path or entry.idb_path or self._path
+ return [Session(session_id=entry.record_id, filename=os.path.basename(path),
+ input_path=path, is_active=True)]
+
+ def close(self, grace: float = 0.0) -> None:
+ del grace
+ with self._connect_lock:
+ handle, self._handle = self._handle, None
+ if handle is not None:
+ self._last_entry = handle.entry
+ handle.close() # release our lease; never close a GUI/other client's DB
+
+ @staticmethod
+ def _wait_for_entry_release(entry: "RegistryEntry", timeout: float) -> bool:
+ _require_codemode()
+ path = REGISTRY_DIR / f"{entry.record_id}.lock"
+ deadline = time.monotonic() + max(0.0, timeout)
+ while True:
+ lock = FileLock(path)
+ try:
+ if lock.try_acquire():
+ return True
+ except OSError:
+ pass
+ finally:
+ lock.close()
+ if time.monotonic() >= deadline:
+ return False
+ time.sleep(min(0.1, deadline - time.monotonic()))
+
+ def wait_released(self, timeout: float = 45.0) -> bool:
+ """Wait until a managed instance releases its lifetime lock.
+
+ Normal application shutdown must not wait: another client may retain the
+ worker. This is an explicit test/maintenance helper for deleting a
+ temporary IDB safely after this client closes. GUI instances return
+ ``False`` immediately because clients never own their lifetime.
+ """
+ entry = self._last_entry
+ if entry is None or entry.backend != "idalib":
+ return False
+ return self._wait_for_entry_release(entry, timeout)
+
+ def __enter__(self) -> "CodeModeClient":
+ return self.connect()
+
+ def __exit__(self, *exc) -> None:
+ self.close()
diff --git a/idatui/domain.py b/idatui/domain.py
index aedade9..6fe73c2 100644
--- a/idatui/domain.py
+++ b/idatui/domain.py
@@ -1,19 +1,15 @@
-"""Domain / paging layer: address-centric models over the raw MCP client.
+"""Domain / paging layer: address-centric models over IDA Code Mode.
This is where the "millions of lines" problem is solved, so the TUI widgets only
ever see a viewport-sized slice. Every hard-won constraint from
``docs/PAGING_FINDINGS.md`` is encoded here:
-* Per-call caps are silent (over the cap the server returns 10, not a clamp), so
- we clamp page sizes ourselves: ``LIST_PAGE`` / ``DISASM_BLOCK`` <= the caps.
-* ``next_offset`` is unreliable; we paginate by advancing ``len(data)``.
-* ``disasm offset=N`` is O(N) with no resumable cursor, so windowed disassembly
- is **block-cached** (revisits are free) and **prefetches** the next block on a
- background thread (the client is concurrency-safe).
-* ``include_total`` scans the whole function (~200ms on monsters); totals are
- fetched once and cached.
-* ``decompile`` can hard-fail on huge functions as a *soft* error (``code`` is
- null); that is surfaced as data, not an exception.
+* Page sizes remain bounded so remote execution returns viewport-scale JSON.
+* Pagination advances by the number of rows actually returned.
+* Deep head walks are block-cached (revisits are free) and neighboring blocks
+ prefetch through the thread-safe Code Mode client.
+* Expensive function totals are fetched once and cached.
+* Decompilation failures are surfaced as data, not application crashes.
Everything here is synchronous and thread-safe. The TUI runs these calls from
Textual worker threads; the internal prefetch pool is separate and small.
@@ -22,10 +18,8 @@ Textual worker threads; the internal prefetch pool is separate and small.
from __future__ import annotations
import bisect
-import json
import re
import threading
-import urllib.request
from concurrent.futures import ThreadPoolExecutor
from collections.abc import Sequence
from dataclasses import dataclass, field, replace
@@ -36,7 +30,7 @@ from . import diag
from .errors import IDAToolError
if TYPE_CHECKING: # type hint only
- from .worker_client import WorkerClient # noqa: F401
+ from .codemode_client import CodeModeClient
# Clamps derived from measured caps (list ~700, disasm ~500). Margin included.
LIST_PAGE = 500
@@ -66,7 +60,7 @@ class Func:
def from_raw(cls, d: dict) -> "Func":
addr = _as_int(d["addr"])
name = d.get("name")
- # An unnamed function (server returns null/empty) must still have a
+ # An unnamed function must still have a
# usable string name — synthesize IDA's sub_ADDR so every consumer
# (palette, sort, rename prefill) can treat name as a str.
if not name:
@@ -93,7 +87,7 @@ class Line:
class Head(NamedTuple):
- """One flat-listing item (from the ``heads`` server tool): a code
+ """One flat-listing item (from the Code Mode ``heads`` operation): a code
instruction, a data item, or an undefined byte run.
A ``NamedTuple`` rather than a dataclass because this is by far the
@@ -114,7 +108,7 @@ class Head(NamedTuple):
name: str | None = None
raw: bytes | None = None # opcode/item bytes (filled in for code by the model)
#: [(kind, text)] from IDA's own colour tags — mnem/reg/num/name/str/punct/…
- #: None when the worker didn't provide them (older worker, or the spans
+ #: None when Code Mode didn't provide them (or the spans
#: disagreed with the plain text, in which case the text wins).
#:
#: Held exactly as it came off the wire, and **read-only**. The worker
@@ -301,7 +295,7 @@ class FunctionIndex:
"""A lazily-paginated, cached view of the function list.
Loads pages of ``LIST_PAGE`` on demand, advancing by ``len(data)`` (never by
- ``next_offset``). A single index instance corresponds to one server-side
+ ``next_offset``). A single index instance corresponds to one remote
``filter`` glob (``None`` = all functions).
"""
@@ -321,7 +315,7 @@ class FunctionIndex:
query: dict = {"offset": offset, "count": LIST_PAGE}
if self.filter:
query["filter"] = self.filter
- data = _query_data(self._prog.client.call("list_funcs", queries=[query]))
+ data = _query_data(self._prog.client.invoke("list_funcs", queries=[query]))
added = 0
with self._lock:
for d in data:
@@ -431,7 +425,7 @@ class DisasmModel:
code function this equals the heads row count that backs the lines."""
if self._total is not None:
return self._total
- payload = self._prog.client.call(
+ payload = self._prog.client.invoke(
"disasm", addr=hex(self.ea), max_instructions=1, include_total=True
)
total = payload.get("total_instructions")
@@ -492,7 +486,7 @@ class DisasmModel:
# The function disasm view is a listing filtered to the function: fetch a
# block of heads (one per instruction for code). Over-fetch one row so
# the block knows where its last instruction ends (opcode-byte sizing).
- payload = self._prog.client.call(
+ payload = self._prog.client.invoke(
"heads", addr=hex(self.ea), offset=b * self.BLOCK,
count=self.BLOCK + 1, **self._end_kw(),
)
@@ -633,15 +627,15 @@ class ListingModel:
"""A flat, IDA-style disassembly *listing* over one segment: code, data and
undefined heads interleaved, unlike ``DisasmModel`` (one function, code only).
- Backed by the injected ``heads`` server tool, which walks item heads and
- renders each via ``generate_disasm_line``. The segment is walked lazily in
+ Backed by the Code Mode adapter's ``heads`` operation, which walks item heads
+ and renders each via ``generate_disasm_line``. The segment is walked lazily in
forward pages (``FunctionIndex`` style); line index == position in the walked
head list. Random access to an address is O(distance-from-seg-start) the
first time (then cached) — the same tradeoff as ``disasm offset=N``. Grows
on demand as the viewport scrolls. Synchronous + thread-safe.
"""
- PAGE = 500 # heads per server call (well under the tool's 2000 cap)
+ PAGE = 500 # viewport-scale heads per Code Mode execution
def __init__(self, program: "Program", seg_start: int, seg_end: int,
name: str | None = None):
@@ -754,7 +748,7 @@ class ListingModel:
if self._done or self._next is None:
return 0
frm = self._next
- payload = self._prog.client.call(
+ payload = self._prog.client.invoke(
"heads", addr=hex(frm), count=self.PAGE, annotate=True)
rows = payload.get("heads", []) if isinstance(payload, dict) else []
cur = payload.get("cursor", {}) if isinstance(payload, dict) else {}
@@ -1019,7 +1013,7 @@ class ListingModel:
# the expectation rather than asking first means a page that HAS changed
# still costs one round trip.
try:
- payload = self._prog.client.call(
+ payload = self._prog.client.invoke(
"heads", addr=hex(addr), count=self.PAGE, annotate=True,
expect="" if want_digest is None else str(want_digest))
except Exception: # noqa: BLE001 -- keep the old text rather than blank
@@ -1233,7 +1227,7 @@ class HexModel:
class Program:
"""The bound analysis session: models, caches, and a small prefetch pool."""
- def __init__(self, client: "WorkerClient", prefetch_workers: int = 2):
+ def __init__(self, client: "CodeModeClient", prefetch_workers: int = 2):
self.client = client
self._pool = ThreadPoolExecutor(
max_workers=prefetch_workers, thread_name_prefix="idatui-prefetch"
@@ -1257,7 +1251,7 @@ class Program:
self._sections: list[tuple[int, int, str]] | None = None
self._fileregions: list[tuple[int, int, int]] | None = None
self._hexmodel: "HexModel | None" = None
- self._no_read_raw = False # set if the server lacks the read_raw tool
+ self._no_read_raw = False # compatibility fallback for alternate clients
self._lock = threading.Lock()
# -- prefetch plumbing ------------------------------------------------- #
@@ -1284,17 +1278,14 @@ class Program:
"""Sorted raw segment map [(start, end, file_off, name)] — the single
source for sections()/file_regions()/image_range. Cached.
- Uses the injected ``file_regions`` tool (a plain segment walk, ~ms).
- This deliberately AVOIDS ``survey_binary``, which also computes function
- counts / strings / stats and takes *seconds* on a large IDB (it was the
- cause of the multi-second hex-pane open). Falls back to survey_binary
- only if the injected tool is missing.
+ Uses the Code Mode adapter's ``file_regions`` operation (a plain segment
+ walk, ~ms), avoiding broad binary surveys on the hex-pane open path.
"""
if self._segments_cache is not None:
return self._segments_cache
segs: list[tuple[int, int, int, str]] = []
try:
- r = self.client.call("file_regions")
+ r = self.client.invoke("file_regions")
for d in (r.get("regions", []) if isinstance(r, dict) else []):
if isinstance(d, dict) and "start" in d:
segs.append((_as_int(d["start"]), _as_int(d["end"]),
@@ -1303,7 +1294,7 @@ class Program:
segs = []
if not segs: # older server without file_regions -> survey_binary (slow)
try:
- sb = self.client.call("survey_binary")
+ sb = self.client.invoke("survey_binary")
for s in (sb.get("segments", []) if isinstance(sb, dict) else []):
try:
segs.append((_as_int(s["start"]), _as_int(s["end"]), -1,
@@ -1345,8 +1336,7 @@ class Program:
def file_regions(self) -> list[tuple[int, int, int]]:
"""Sorted [(start, end, file_off)] mapping loaded segments to raw file
- offsets (file_off == -1 for non-file-backed, e.g. .bss). Cached; needs
- the injected ``file_regions`` server tool."""
+ offsets (file_off == -1 for non-file-backed, e.g. .bss). Cached."""
if self._fileregions is not None:
return self._fileregions
regions = [(s, e, fo) for s, e, fo, _nm in self._segments()]
@@ -1364,15 +1354,14 @@ class Program:
def read_bytes(self, ea: int, n: int) -> bytes:
"""Raw bytes [ea, ea+n) from IDA (gaps read as zero).
- Fast path: the injected ``read_raw`` tool returns one contiguous hex
- string (C-speed both ends). Falls back to the stock ``get_bytes`` (a
- per-byte '0x..'-with-spaces string) on an older server without it.
+ The Code Mode adapter returns one contiguous hex string (C-speed in IDA).
+ A legacy ``get_bytes`` decoding fallback remains for alternate clients.
"""
if n <= 0:
return b""
if not self._no_read_raw:
try:
- r = self.client.call("read_raw", addr=hex(ea), size=int(n))
+ r = self.client.invoke("read_raw", addr=hex(ea), size=int(n))
h = r.get("hex") if isinstance(r, dict) else None
if isinstance(h, str):
out = bytes.fromhex(h)
@@ -1386,7 +1375,7 @@ class Program:
except (ValueError, KeyError):
pass # malformed hex -> fall through to the legacy decoder
try:
- r = self.client.call("get_bytes", regions=[{"addr": hex(ea), "size": int(n)}])
+ r = self.client.invoke("get_bytes", regions=[{"addr": hex(ea), "size": int(n)}])
except IDAToolError:
return b"\x00" * n
res = r.get("result", []) if isinstance(r, dict) else []
@@ -1437,7 +1426,7 @@ class Program:
def list_structs(self, filter: str = "") -> list[Struct]:
"""All local structs/unions (optionally name-substring filtered), sorted
by name."""
- payload = self.client.call("search_structs", filter=filter)
+ payload = self.client.invoke("search_structs", filter=filter)
res = payload.get("result", []) if isinstance(payload, dict) else []
out = [Struct.from_raw(d) for d in res
if isinstance(d, dict) and d.get("name")
@@ -1447,9 +1436,9 @@ class Program:
def struct_source(self, name: str) -> str:
"""A C definition for ``name`` reconstructed from its member layout
- (the server exposes members, not printable source). Faithful to IDA's
+ (the remote operation exposes members, not printable source). Faithful to IDA's
field names/types; array dims are moved after the field name."""
- payload = self.client.call(
+ payload = self.client.invoke(
"type_inspect", queries=[{"name": name, "include_members": True}])
res = payload.get("result", []) if isinstance(payload, dict) else []
info = res[0] if res and isinstance(res[0], dict) else {}
@@ -1472,7 +1461,7 @@ class Program:
def declare_type(self, decl: str) -> str | None:
"""Create or update a C type. Returns None on success, else the parse
error. (Re-declaring a name updates it in place.)"""
- payload = self.client.call("declare_type", decls=decl)
+ payload = self.client.invoke("declare_type", decls=decl)
res = payload.get("result", []) if isinstance(payload, dict) else []
if res and isinstance(res[0], dict):
return res[0].get("error")
@@ -1481,10 +1470,9 @@ class Program:
# -- function / variable types ---------------------------------------- #
def func_types(self, ea: int) -> FuncTypes | None:
"""Structured decompiler types for the function at ``ea`` (prototype +
- local variables). None if ``ea`` isn't a decompilable function. Requires
- the injected ``func_types`` server tool."""
+ local variables). None if ``ea`` isn't a decompilable function."""
try:
- r = self.client.call("func_types", addr=hex(ea))
+ r = self.client.invoke("func_types", addr=hex(ea))
except IDAToolError:
return None
if not isinstance(r, dict) or r.get("error"):
@@ -1497,7 +1485,7 @@ class Program:
def set_function_type(self, ea: int, signature: str) -> str | None:
"""Set a function's prototype. None on success, else an error string."""
- r = self.client.call("set_type", edits=[{"addr": hex(ea), "signature": signature}])
+ r = self.client.invoke("set_type", edits=[{"addr": hex(ea), "signature": signature}])
res = r.get("result", []) if isinstance(r, dict) else []
row = res[0] if res and isinstance(res[0], dict) else {}
if row.get("ok"):
@@ -1506,9 +1494,9 @@ class Program:
def data_type(self, ea: int) -> dict | None:
"""Current type info for a data item/global: {addr,name,type,size,is_func}.
- None if the tool is unavailable or the address isn't mapped."""
+ None if the operation fails or the address isn't mapped."""
try:
- r = self.client.call("data_type", addr=hex(ea))
+ r = self.client.invoke("data_type", addr=hex(ea))
except IDAToolError:
return None
if not isinstance(r, dict) or r.get("error"):
@@ -1517,7 +1505,7 @@ class Program:
def set_data_type(self, ea: int, decl: str) -> str | None:
"""Set a global/data item's type. None on success, else an error string."""
- r = self.client.call(
+ r = self.client.invoke(
"set_type", edits=[{"kind": "global", "addr": hex(ea), "type": decl}])
res = r.get("result", []) if isinstance(r, dict) else []
row = res[0] if res and isinstance(res[0], dict) else {}
@@ -1526,9 +1514,9 @@ class Program:
return row.get("error") or "failed to set the type"
def set_lvar_type(self, fn_ea: int, var: str, ty: str) -> str | None:
- """Set a decompiler local variable's type (via the injected server tool).
+ """Set a decompiler local variable's type through ida-domain pseudocode.
None on success, else an error string."""
- r = self.client.call("set_lvar_type", addr=hex(fn_ea), variable=var, type=ty)
+ r = self.client.invoke("set_lvar_type", addr=hex(fn_ea), variable=var, type=ty)
if isinstance(r, dict) and r.get("error"):
return r["error"]
if isinstance(r, dict) and not r.get("ok"):
@@ -1537,15 +1525,14 @@ class Program:
def delete_type(self, name: str) -> str | None:
"""Delete a named type. Returns None on success, else an error string.
- Requires a server-side ``del_type`` tool; if absent, a clear message is
- returned instead of raising."""
+ Returns a clear error instead of raising when the runtime cannot do it."""
try:
- self.client.call("del_type", name=name)
+ self.client.invoke("del_type", name=name)
return None
except IDAToolError as e:
msg = e.message
if "not found" in msg.lower() and "del_type" in msg:
- return "delete needs a 'del_type' tool on the ida-pro-mcp server"
+ return "the connected Code Mode runtime cannot delete local types"
return msg
# -- disassembly ------------------------------------------------------- #
@@ -1559,13 +1546,7 @@ class Program:
# -- decompilation ----------------------------------------------------- #
def decompile(self, ea: int, refresh: bool = False) -> Decompilation:
- """Full pseudocode for a function.
-
- The server truncates responses over 50KB (strings clipped to 1000
- chars) but caches the full output and exposes it at
- ``_meta.ida_mcp.download_url``. We transparently fetch that so the view
- always gets the complete body, not a 1KB stub.
- """
+ """Full pseudocode for a function, returned directly by Code Mode."""
if not refresh:
with self._lock:
hit = self._decomp.get(ea)
@@ -1574,10 +1555,10 @@ class Program:
dec, hit_gen = hit
if hit_gen == gen:
return dec
- # Cached before a rename: names may be stale. Drop the server's
- # Hex-Rays cache so the refetch reflects the new names.
+ # Cached before a rename: names may be stale. Drop Hex-Rays'
+ # cache so the refetch reflects the new names.
try:
- self.client.call("force_recompile", items=[{"addr": hex(ea)}])
+ self.client.invoke("force_recompile", items=[{"addr": hex(ea)}])
except Exception: # noqa: BLE001
pass
# Bound the decompile: a function Hex-Rays can't handle tends to stall
@@ -1587,7 +1568,10 @@ class Program:
# rpcclient socket timeout, and cache the failure below so a re-request
# returns instantly instead of re-grinding.
try:
- envelope = self.client.call_envelope(
+ # Code Mode returns the complete JSON result directly; unlike the
+ # old MCP tool transport there is no structured-content envelope or
+ # out-of-band download URL to unwrap.
+ payload = self.client.invoke(
"decompile", addr=hex(ea), timeout=DECOMPILE_TIMEOUT
)
except Exception as e: # noqa: BLE001 -- surface as a failed decompile
@@ -1595,15 +1579,6 @@ class Program:
with self._lock:
self._decomp[ea] = (dec, self._name_gen)
return dec
- result = envelope.get("result", {})
- payload = result.get("structuredContent")
- if payload is None: # fall back to text content
- payload = self.client._extract_payload("decompile", result)
- meta = (result.get("_meta") or {}).get("ida_mcp")
- if isinstance(meta, dict) and meta.get("download_url"):
- full = self._fetch_output(meta["download_url"])
- if isinstance(full, dict) and full.get("code"):
- payload = full
dec = _parse_decompilation(ea, payload)
with self._lock:
self._decomp[ea] = (dec, self._name_gen)
@@ -1681,18 +1656,18 @@ class Program:
Undefine first so it works even when the bytes are currently part of a
data/align item — ``create_insn`` refuses to carve into a live item."""
try:
- self.client.call("undefine", items=[{"addr": hex(ea)}])
+ self.client.invoke("undefine", items=[{"addr": hex(ea)}])
except IDAToolError:
pass # nothing defined here yet -> just try to create the insn
res = self._first_result(
- self.client.call("define_code", items=[{"addr": hex(ea)}]))
+ self.client.invoke("define_code", items=[{"addr": hex(ea)}]))
if res.get("error"):
raise IDAToolError("define_code", f"@ {ea:#x}: {res['error']}")
def decomp_error(self, ea: int) -> str:
"""Hex-Rays' own reason for refusing ``ea``, or "" if it won't say."""
try:
- r = self.client.call("decomp_error", addr=hex(ea))
+ r = self.client.invoke("decomp_error", addr=hex(ea))
except IDAToolError:
return ""
if not isinstance(r, dict):
@@ -1711,7 +1686,7 @@ class Program:
def thumb_scan(self, start: int, end: int, apply: bool = True) -> dict:
"""Find Thumb entry points from odd pointers in ``[start, end)``."""
- r = self.client.call("thumb_scan", start=hex(start), end=hex(end),
+ r = self.client.invoke("thumb_scan", start=hex(start), end=hex(end),
apply=bool(apply))
if not isinstance(r, dict) or r.get("error"):
raise IDAToolError("thumb_scan",
@@ -1720,7 +1695,7 @@ class Program:
def set_thumb(self, ea: int, mode: str = "toggle") -> dict:
"""Switch ARM/Thumb decoding at ``ea``. Returns the resulting state."""
- r = self.client.call("set_thumb", addr=hex(ea), mode=mode)
+ r = self.client.invoke("set_thumb", addr=hex(ea), mode=mode)
if not isinstance(r, dict) or r.get("error"):
raise IDAToolError("set_thumb",
f"@ {ea:#x}: {(r or {}).get('error', 'failed')}")
@@ -1729,11 +1704,11 @@ class Program:
def define_code_run(self, ea: int, limit: int = 20000) -> dict:
"""Disassemble consecutively from ``ea`` until something stops it.
- Falls back to a single instruction when the worker predates the tool, so
- an old worker degrades to the previous behaviour instead of failing.
+ Falls back to a single instruction for alternate clients that do not
+ provide the run operation.
"""
try:
- r = self.client.call("define_code_run", addr=hex(ea), limit=int(limit))
+ r = self.client.invoke("define_code_run", addr=hex(ea), limit=int(limit))
except IDAToolError:
self.define_code(ea)
return {"count": 1, "stopped": "single", "end": hex(ea)}
@@ -1745,14 +1720,14 @@ class Program:
def define_func(self, ea: int) -> dict:
"""Create a function starting at ``ea`` (IDA's 'p').
- Prefers the injected tool, which works out the end when IDA can't;
- falls back to the plain one for an older worker.
+ Prefers the Code Mode operation, which works out the end when IDA can't;
+ falls back to a plain create for alternate clients.
"""
try:
- r = self.client.call("define_func_run", addr=hex(ea))
+ r = self.client.invoke("define_func_run", addr=hex(ea))
except IDAToolError:
res = self._first_result(
- self.client.call("define_func", items=[{"addr": hex(ea)}]))
+ self.client.invoke("define_func", items=[{"addr": hex(ea)}]))
if res.get("error"):
raise IDAToolError("define_func", f"@ {ea:#x}: {res['error']}")
return {"ok": True, "how": "legacy"}
@@ -1766,7 +1741,7 @@ class Program:
item: dict = {"addr": hex(ea)}
if size:
item["size"] = int(size)
- res = self._first_result(self.client.call("undefine", items=[item]))
+ res = self._first_result(self.client.invoke("undefine", items=[item]))
if res.get("error"):
raise IDAToolError("undefine", f"@ {ea:#x}: {res['error']}")
@@ -1776,7 +1751,7 @@ class Program:
item: dict = {"addr": hex(ea), "type": type_decl}
if name:
item["name"] = name
- res = self._first_result(self.client.call("make_data", items=[item]))
+ res = self._first_result(self.client.invoke("make_data", items=[item]))
if res.get("ok") is False or res.get("error"):
raise IDAToolError(
"make_data", f"@ {ea:#x}: {res.get('error') or 'rejected'}")
@@ -1784,7 +1759,7 @@ class Program:
def make_string(self, ea: int, length: int = 0, kind: str = "c") -> str:
"""Create a string literal at ``ea`` (IDA's 'A'); auto-length when 0.
Returns the decoded contents."""
- r = self.client.call("make_string", addr=hex(ea), length=int(length), kind=kind)
+ r = self.client.invoke("make_string", addr=hex(ea), length=int(length), kind=kind)
res = r if isinstance(r, dict) else {}
if not res.get("ok"):
raise IDAToolError(
@@ -1801,7 +1776,7 @@ class Program:
``cycle``/``back`` (step the stops that make sense for this value) or a
format by name. ``show`` reports without changing anything.
"""
- r = self.client.call("op_format", addr=hex(ea), mode=str(mode),
+ r = self.client.invoke("op_format", addr=hex(ea), mode=str(mode),
col=int(col), n=int(n))
res = r if isinstance(r, dict) else {}
if res.get("error"):
@@ -1825,7 +1800,7 @@ class Program:
if hit is not None and hit[1] == gen:
return hit[0]
try:
- r = self.client.call("pc_nums", addr=hex(fn_ea))
+ r = self.client.invoke("pc_nums", addr=hex(fn_ea))
except Exception: # noqa: BLE001 -- an older worker hasn't got the tool
r = {}
out: dict[int, list[tuple[int, int, str, int, int]]] = {}
@@ -1848,7 +1823,7 @@ class Program:
listing's format doesn't reach the pseudocode and vice versa, so this is
a separate call rather than a flag on ``op_format``.
"""
- r = self.client.call("pc_num_format", addr=hex(fn_ea), mode=str(mode),
+ r = self.client.invoke("pc_num_format", addr=hex(fn_ea), mode=str(mode),
line=int(line), col=int(col))
res = r if isinstance(r, dict) else {}
if res.get("error"):
@@ -1865,18 +1840,6 @@ class Program:
sec = None
return f"{sec} @ {ea:#x}" if sec else f"<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/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()
generated by cgit v1.3.1-sl0p (git 2.54.0) at 2026-08-08 14:08:20 +0000 — sl0p fork