diff options
| author | blasty <blasty@local> | 2026-08-07 15:18:56 +0200 |
|---|---|---|
| committer | blasty <blasty@local> | 2026-08-07 15:18:56 +0200 |
| commit | 4f0e17c69c5424f307ecf2ef4ffb278e775c8349 (patch) | |
| tree | 43d561beb58b18a1859beb915fed84ed495e0c0d | |
| parent | Merge the IDA Code Mode port (diff) | |
| download | ida-tui-4f0e17c69c5424f307ecf2ef4ffb278e775c8349.tar.gz ida-tui-4f0e17c69c5424f307ecf2ef4ffb278e775c8349.tar.xz ida-tui-4f0e17c69c5424f307ecf2ef4ffb278e775c8349.zip | |
Stop tracking the upstream port patch
It was swept into the tree by a broad 'git add -A' during the port. The merge
itself plus docs/CODEMODE_PORT.md are the record; a 251KB patch of superseded
content is not. Still on disk, now ignored.
| -rw-r--r-- | .gitignore | 3 | ||||
| -rw-r--r-- | ida-codemode-mcp.patch | 5848 |
2 files changed, 3 insertions, 5848 deletions
@@ -21,3 +21,6 @@ core core.* .fastfeedback/ tests/.synthetic/ + +# the upstream port patch, kept locally for reference only +ida-codemode-mcp.patch diff --git a/ida-codemode-mcp.patch b/ida-codemode-mcp.patch deleted file mode 100644 index 82a9b3b..0000000 --- a/ida-codemode-mcp.patch +++ /dev/null @@ -1,5848 +0,0 @@ -From 0db18f870a45c0355ca1b28db1a2e763b60f6226 Mon Sep 17 00:00:00 2001 -From: Duncan Ogilvie <mr.exodia.tpodt@gmail.com> -Date: Fri, 31 Jul 2026 01:40:21 +0200 -Subject: [PATCH] WIP: vibeslop ida-codemode port - ---- - README.md | 100 +-- - TODO | 17 +- - docs/CODEMODE_PORT.md | 175 +++++ - docs/PAGING_FINDINGS.md | 54 +- - docs/PROJECTS.md | 41 +- - docs/SPLIT_VIEW.md | 12 +- - experiments/worker_smoke.py | 104 ++- - ida-tui | 7 +- - idatui/__init__.py | 5 +- - idatui/app.py | 206 +++--- - idatui/codemode_client.py | 1107 +++++++++++++++++++++++++++++ - idatui/domain.py | 194 +++-- - idatui/drive.py | 3 +- - idatui/errors.py | 10 +- - idatui/launch.py | 92 +-- - idatui/pane.py | 92 +-- - idatui/pool.py | 100 +-- - idatui/project.py | 35 +- - idatui/worker.py | 233 ------ - idatui/worker_client.py | 234 ------- - pyproject.toml | 17 +- - server/patch_server.py | 1248 --------------------------------- - tests/test_codemode_client.py | 137 ++++ - tests/test_pool.py | 61 +- - tests/test_project.py | 2 +- - tests/test_scenarios.py | 18 +- - uv.lock | 65 +- - 27 files changed, 2055 insertions(+), 2314 deletions(-) - create mode 100644 docs/CODEMODE_PORT.md - create mode 100644 idatui/codemode_client.py - delete mode 100644 idatui/worker.py - delete mode 100644 idatui/worker_client.py - delete mode 100644 server/patch_server.py - create mode 100644 tests/test_codemode_client.py - -diff --git a/README.md b/README.md -index 7d7d8d7..47d5366 100644 ---- a/README.md -+++ b/README.md -@@ -1,13 +1,13 @@ - # ida-tui - - A minimal, keyboard-first (mouse-capable) **TUI frontend for IDA Pro**, built with --[Textual](https://textual.textualize.io/) and driving **idalib** (IDA headless). -+[Textual](https://textual.textualize.io/) and using -+[ida-codemode-mcp](../ida-codemode-mcp) as a Python library. - --Opening a binary spawns our own **idalib worker** — a private subprocess talking a --unix socket (`idatui/worker.py` + `WorkerClient`), ~50–100× cheaper per call than --an HTTP transport. It reuses [ida-pro-mcp](https://github.com/mrexodia/ida-pro-mcp)'s --tool implementations in-process; the old ida-pro-mcp HTTP server/supervisor path --has been **removed**. -+ida-tui attaches to databases through `ida_codemode.client.DatabaseHandle`. A -+matching database already open in the IDA GUI is reused; otherwise Code Mode -+reuses or starts a shared managed idalib worker. The TUI owns only a client lease, -+never the GUI or worker process. - - ## ⚠️ Status: not ready for public consumption - -@@ -31,7 +31,7 @@ don't file expectations. **Use at your own risk.** - - A unified **IDA-style listing** (continuous disassembly interleaved with data / - undefined heads) as the default code view; `F5`/`Tab` drops into the - **decompiler (pseudocode)** for the function under the cursor. Both are -- line-virtualized and page lazily over the worker. -+ line-virtualized and page lazily over the Code Mode database. - - A **Ghidra-style split view** (`s`): listing and pseudocode side by side, kept - in cursor sync — the focused pane drives and the other highlights the linked - region (every instruction a C line owns), following you across functions. -@@ -49,46 +49,58 @@ don't file expectations. **Use at your own risk.** - - ## Architecture (three layers, kept separate) - --- **`idatui/worker.py` + `idatui/worker_client.py`** — the backend. `worker.py` -- opens one DB with idalib (on its main thread) and serves ida-pro-mcp's tool -- functions over a unix socket; `WorkerClient` spawns it and is a stdlib-only -- drop-in client (length-prefixed pickle, calls serialized under a lock). Shared -- error types + the `Session` model live in `idatui/errors.py`. --- **`idatui/domain.py`** — paging/caching over the worker client (`FunctionIndex`, -- `DisasmModel`, `ListingModel`, `decompile`, xrefs, resolve). Synchronous, -- thread-safe. Tools ida-pro-mcp lacks (`heads`, `read_raw`, `resolve_names`, -- `xref_types`, …) are injected by `server/patch_server.py`, which the worker -- runs itself on startup. -+- **`idatui/codemode_client.py`** — lifecycle and execution adapter. It leases a -+ registered GUI/idalib instance with `DatabaseHandle`, waits for autoanalysis, -+ normalizes errors, saves, and releases the lease. Address-centric operations -+ are sent through Code Mode's `execute_python` surface and use its preloaded -+ `ida-domain` `db` object. -+- **`idatui/domain.py`** — synchronous, thread-safe paging/caching -+ (`FunctionIndex`, `DisasmModel`, `ListingModel`, decompile, xrefs, resolve). -+ It has no process/database ownership logic. - - **`idatui/app.py`** — the Textual app (virtualized `ScrollView`s, shared cursor/ - search/nav mixins, modals). - --The domain + worker-client layers are intentionally **stdlib-only** (the worker --process links idalib); only the TUI layer pulls in Textual + Pygments. -+`idatui/pool.py` retains LRU project leases. Releasing an entry never kills a GUI -+or another client's worker. See `docs/CODEMODE_PORT.md` for what maps to public -+ida-domain APIs and which remaining features require IDAPython inside the Code -+Mode execution sandbox. - - ## Requirements - - - Python ≥ 3.11 --- A working **IDA Pro** with **idalib** and **ida-pro-mcp** installed (the worker -- reuses ida-pro-mcp's tool implementations in-process — no server runs). --- Textual ≥ 8 and Pygments ≥ 2 for the TUI (`pip install -e '.[tui]'`). -+- IDA Pro 9.4+ with idalib configured -+- `ida-codemode-mcp` installed in the TUI environment (this checkout uses the -+ editable sibling path `../ida-codemode-mcp`) -+- The ida-codemode IDA plugin installed so GUI databases register themselves -+- Textual ≥ 8 and Pygments ≥ 2 (`uv sync` installs both) - --Two python environments are expected: one with **textual + idapro** for the TUI --(`~/ida-venv`, override `$IDATUI_PYTHON`) and one with **idapro + ida_pro_mcp** --for the worker (auto-detected, override `$IDATUI_WORKER_PYTHON`). -+Code Mode's own worker launcher carries the correct Python environment; ida-tui -+no longer searches for a second Python or imports `ida_pro_mcp`. - - ## Running - --One command — it spawns a private idalib worker for the binary (which opens + --auto-analyzes it in its own process over a unix socket) and drops you into the --TUI behind a loading overlay: -+Install the project and its TUI dependencies: - - ```sh --./ida-tui /path/to/binary # open a binary and drive it — that's it -+uv sync - ``` - --It uses `~/ida-venv/bin/python` for the TUI (override with `$IDATUI_PYTHON`) and --resolves binary paths against your real cwd. The binary's directory must be --writable (idalib writes a `.i64` there). -+Pass an executable/IDB path. If the plugin has registered a matching GUI session, -+ida-tui attaches to it; otherwise Code Mode opens a managed idalib database: -+ -+```sh -+./ida-tui /path/to/binary -+``` -+ -+With exactly one registered database, the path may be omitted: -+ -+```sh -+./ida-tui -+``` -+ -+When several databases are registered, the launcher lists their paths and asks -+for one explicitly. A newly managed single-binary database still needs a writable -+output location; projects stage binaries and IDBs in their sidecar directory. - - Headerless blobs need to be told what they are — a raw firmware dump has no - format to detect, and IDA falls back to x86 at address 0, which analyses to -@@ -102,15 +114,16 @@ ARM images that use Thumb need one more thing: press `t` on the listing to switc - ARM/Thumb decoding at the cursor (it sets IDA's `T` register, and the segment to - 32-bit, since Thumb doesn't exist in AArch64). - --`--base` is a real address (IDA's own `-b` is in paragraphs; the conversion is --done for you). In a project the options are recorded per binary, which is what a --multi-image firmware wants. They apply to the first open only — after that the --`.i64` records how the image was loaded. See `docs/PROJECTS.md`. -+`--base` is a real address (Code Mode's typed loading address is also natural, -+so no paragraph conversion crosses the dependency boundary). In a project the -+options are recorded per binary. They apply only when Code Mode must create the -+first database; a registered or existing IDB already records them. Arbitrary -+`--ida-args` are rejected because `DatabaseHandle.open()` has no equivalent; -+processor, base, and loader/file type are the supported import surface. - --> Recovering a wedged database: if a worker was hard-killed it leaves unpacked --> `foo.id0/.id1/.id2/.nam/.til` next to `foo.i64`, and the `.i64` then refuses to --> reopen. Delete those stale files (never the `.i64`) and retry — `ida-tui` does --> this automatically. -+ida-tui never deletes unpacked IDA scratch files during discovery: those files -+may belong to a registered GUI or another Code Mode client. Registry locks and -+health probes are the ownership authority. - - ## Execution traces - -@@ -166,8 +179,8 @@ See `docs/RPC.md` for the full protocol. - - ## Tests - --A headless Textual `Pilot` suite lives in `tests/`; it spawns a worker on the --given binary (default `targets/echo`): -+A headless Textual `Pilot` suite lives in `tests/`; it attaches through Code Mode -+(or starts a managed worker) for the given binary: - - ```sh - python tests/test_scenarios.py targets/echo # full UI suite -@@ -177,6 +190,7 @@ python tests/test_scenarios.py --only hex,rename - ## Docs - - - `docs/RPC.md` — the RPC protocol --- `docs/PAGING_FINDINGS.md` — idalib tool paging/scale quirks -+- `docs/CODEMODE_PORT.md` — port coverage, API gaps, and lifecycle semantics -+- `docs/PAGING_FINDINGS.md` — historical paging/scale findings - - `docs/TEXTUAL_NOTES.md` — Textual pitfalls encountered - - `docs/TUI_DRIVING_BLUEPRINT.md` — generalizing the driving layer -diff --git a/TODO b/TODO -index 58b4693..c578cff 100644 ---- a/TODO -+++ b/TODO -@@ -4,14 +4,15 @@ - bugs: - - current: --[~] DITCH ida-pro-mcp -> our own idalib worker (idatui/worker.py + WorkerClient) -- [x] worker + WorkerClient (drop-in for IDAClient, same tool shapes) -- [x] --backend {worker,mcp}; worker is now the DEFAULT for opening a binary -- [x] worker spawns under the IDA python; {"result":...} wrapping to match MCP -- [ ] run the pilot suite against --backend worker (blocked: idalib reaping here) -- [ ] progress reporting during analysis (worker streams notes to the overlay) -- [ ] once solid: delete client.py, server/patch_server.py, spawn.sh, and -- launch.py's whole supervisor/ensure_server/lock-sweep dance -+[x] PORT TO ida-codemode-mcp as a library dependency -+ [x] DatabaseHandle discovery prefers registered GUI sessions -+ [x] shared managed idalib workers + SSE lease lifecycle -+ [x] domain operations execute against ida-domain through Code Mode -+ [x] delete the private pickle worker and ida-pro-mcp patch injection -+ [x] stop sweeping/reaping resources that may belong to another client -+ [ ] run the full live Pilot suite against both GUI and managed backends -+ [ ] add database revision/change notifications for cross-client cache invalidation -+ [ ] decide how "discard changes" should work (Code Mode final workers save) - [x] RPC endpoint for robot-spectator-ida - -> progressssss - -diff --git a/docs/CODEMODE_PORT.md b/docs/CODEMODE_PORT.md -new file mode 100644 -index 0000000..4b4bc45 ---- /dev/null -+++ b/docs/CODEMODE_PORT.md -@@ -0,0 +1,175 @@ -+# ida-tui → IDA Code Mode port -+ -+This port is an experiment: can ida-tui be implemented as an ordinary client of -+`ida_codemode`, sharing GUI databases and managed idalib workers instead of -+owning a private worker and depending on ida-pro-mcp tool functions? -+ -+## Result -+ -+Yes for the database lifecycle and the complete current TUI feature set, with a -+small number of operations implemented using IDAPython inside Code Mode's -+`execute_python` sandbox because ida-domain does not yet expose the required -+behavior. -+ -+The old components are gone: -+ -+- `idatui/worker.py` (private pickle/socket idalib process) -+- `idatui/worker_client.py` -+- `server/patch_server.py` (ida-pro-mcp tool injection) -+ -+The replacement is `idatui/codemode_client.py`. -+ -+## Lifecycle mapping -+ -+`CodeModeClient.connect()` calls `ida_codemode.client.DatabaseHandle.open()`. -+Resolution is therefore Code Mode's resolution, not ida-tui's: -+ -+1. Match a registered GUI by executable path. -+2. Otherwise match the owner of the expected IDB. -+3. Otherwise serialize creation and start a managed `ida-codemode-worker`. -+4. Establish an authenticated SSE lease. -+5. Wait through the public autoanalysis route. -+ -+The handle's registry entry supplies the backend, PID, executable path, IDB path, -+and record ID used by the status/pool layers. -+ -+Closing ida-tui closes only its lease. It never closes a GUI or kills an idalib -+process. A managed worker saves and exits under Code Mode's own policy after its -+last lease disappears. A second agent or TUI can keep using the same instance. -+ -+This also changes project pooling semantics. `DatabasePool` is an LRU pool of -+leases, not process ownership. Managed-IDB save-on-evict remains; budget eviction -+does not implicitly save a GUI. Eviction cannot force a shared worker to exit, -+and GUI process memory is only advisory. -+ -+## ida-domain coverage -+ -+The remote snippets receive Code Mode's preloaded `db` (`ida_domain.Database`). -+The following TUI needs map to public ida-domain entities: -+ -+| TUI need | ida-domain surface | -+|---|---| -+| Function paging, lookup, names, sizes | `db.functions` | -+| Segments and names | `db.segments` | -+| Instructions and plain disassembly | `db.instructions`, `db.functions.get_instructions()` | -+| Heads and item classification | `db.heads`, `db.bytes` | -+| Bytes and strings | `db.bytes`, `db.strings` | -+| Symbol resolution and rename | `db.names`, `db.functions` | -+| Comments | `db.comments` | -+| Imports and exports | `db.imports`, `db.entries` | -+| Xrefs and fine type predicates | `db.xrefs` / `XrefInfo` | -+| Named types, members, parse/apply | `db.types` | -+| Function prototypes and local variables | `db.pseudocode`, `PseudocodeFunction.local_variables` | -+| Decompilation text and object references | `db.pseudocode` | -+ -+All values are reduced to JSON primitives inside the database process. No SWIG -+or ida-domain object crosses the Code Mode boundary. -+ -+## Remaining IDAPython gaps -+ -+Code Mode intentionally allows regular Python imports, so these features still -+work, but they identify useful additions to ida-domain: -+ -+1. **Rich continuous listing** -+ - ida-domain enumerates defined heads and renders plain disassembly. -+ - ida-tui also needs coalesced undefined runs, IDA colour-tag spans, function -+ banners, code-label rows, file-region offsets, and expanded struct members. -+ - The `heads` operation uses `ida_bytes`, `ida_lines`, and related modules for -+ this presentation model. -+ -+2. **Instruction/function carving** -+ - Creating an instruction and walking a speculative decode run requires -+ `ida_ua.create_insn` and processor flow/return checks. -+ - Function creation exists in ida-domain; the explicit-end fallback still -+ needs lower-level item boundaries. -+ -+3. **ARM/Thumb state** -+ - T-register ranges and segment addressing use `ida_segregs`, `ida_idp`, and -+ `ida_segment`. There is no equivalent ida-domain operation. -+ -+4. **Detailed decompiler diagnostics and line maps** -+ - Pseudocode text, ctree objects, and the address map are available through -+ ida-domain. -+ - Reproducing IDA's per-rendered-line coverage uses -+ `cfunc.get_line_item`; obtaining the exact Hex-Rays failure description -+ uses `hexrays_failure_t`. -+ -+5. **A few type/item primitives** -+ - Deleting a named local type and some exact item-undefinition/data-creation -+ behavior still use `ida_typeinf`/`ida_bytes` directly. -+ -+These uses are isolated in `idatui/codemode_client.py`; the paging and Textual -+layers do not import IDAPython. -+ -+## API limitations exposed by the port -+ -+### No rollback or close-without-save -+ -+A Code Mode lease has no rollback operation. Closing a GUI handle leaves the GUI -+state as-is. A managed idalib worker currently saves when its final lease closes. -+Consequently ida-tui's old “discard & quit” guarantee cannot be implemented. -+The UI now labels this choice “leave as-is & quit” and does not explicitly save, -+but managed-worker policy may still persist the changes. -+ -+A true discard action would need a Code Mode/database API for transaction-like -+rollback, a close policy on a newly-owned worker, or a TUI-managed disposable DB -+copy. -+ -+### Typed loader options only -+ -+`DatabaseHandle.open()` supports processor, natural loading address, file type, -+output database, and fresh-database selection. It does not support ida-tui's -+arbitrary `ida_args` escape hatch. The adapter rejects unsupported switches -+rather than silently loading at the wrong architecture/base. -+ -+### No database-change notification stream -+ -+The lease reports liveness, not mutations. If a GUI user or another Code Mode -+client renames/retypes content while ida-tui is open, already-materialized TUI -+caches are not invalidated automatically. TUI-originated edits invalidate their -+own caches correctly. A database revision counter or change feed would make -+shared interactive editing robust. -+ -+### Discovery requires a path for ambiguity -+ -+`ida-tui` with no path attaches automatically when exactly one database is -+registered. With several registrations it lists them and requires an explicit -+executable/IDB path. There is not yet a pre-connection database picker in the -+Textual UI. -+ -+### `DatabaseHandle` import stability -+ -+The usable library primitive currently lives at -+`ida_codemode.client.DatabaseHandle`; `ida_codemode.__init__` exports nothing. -+The port therefore depends on a submodule path. Exporting the handle and public -+client exceptions from the package root would make the supported library API -+clearer. -+ -+## Safety differences -+ -+ida-tui no longer removes `.id0/.id1/.id2/.nam/.til` files before opening. That -+was only defensible when the TUI exclusively owned a private process; it is -+unsafe when a GUI or another client may own the database. Code Mode registry -+locks, health probes, and IDA itself now arbitrate ownership. -+ -+The old pane “reap private workers” behavior is obsolete. A TUI crash closes its -+lease at the socket/kernel boundary; Code Mode decides whether a managed worker -+still has clients and when it should stop. -+ -+## Verification surfaces -+ -+The non-IDA suite verifies project staging, LRU lease behavior, load-option -+translation, and adapter response/error normalization. The existing live suites -+remain the end-to-end contract: -+ -+```sh -+uv run python tests/test_codemode_client.py -+uv run python tests/test_pool.py -+uv run python tests/test_project.py -+uv run python tests/test_scenarios.py /path/to/binary -+``` -+ -+For GUI reuse, open the same binary in an IDA with the Code Mode plugin, confirm -+it appears in `ida_codemode.registry.discover_instances()`, then launch -+`ida-tui /path/to/binary`. The TUI status/`CodeModeClient.backend` should report -+`gui`, and closing the TUI must leave IDA open. -diff --git a/docs/PAGING_FINDINGS.md b/docs/PAGING_FINDINGS.md -index bcde583..bd6c38f 100644 ---- a/docs/PAGING_FINDINGS.md -+++ b/docs/PAGING_FINDINGS.md -@@ -2,10 +2,10 @@ - - Measured against a real target: `libcrypto.so.3` (5.7 MB, **10,092 functions**, - biggest function **52,120 instructions**). These constraints drive the domain / --paging layer. They describe the ida-pro-mcp *tool functions* (`list_funcs`, --`disasm`, `decompile`, `xref_query`, …) which the idalib worker now calls --in-process (`idatui/worker.py`) — the shapes and caps below are the tools' --behaviour and are unchanged by dropping the HTTP transport. -+paging layer. The measurements below came from the former ida-pro-mcp tool -+backend. The Code Mode port preserves the adapter response shapes and conservative -+page sizes, but executes enumeration through ida-domain; old server caps and RTT -+numbers are historical rather than Code Mode constraints. - - ## Response shape (list_* / *_query tools) - -@@ -93,32 +93,26 @@ disasm totals are **top-level** fields, not under `asm`: - (correct). The pseudocode view must handle "decompilation failed" gracefully — - fall back to the disassembly view or show an error panel. - --Normal decompile bodies are server-truncated with a `[N chars total]` marker --(still to be solved for full-body display — see Phase 2). -- --## Worker lifecycle (idatui's own idalib worker) -- --idatui no longer uses ida-pro-mcp's shared HTTP supervisor. `idatui/worker.py` --opens exactly **one** database with `idapro.open_database(...)` in its own process --and serves tool calls over a unix socket (`WorkerClient`). Consequences vs the --old supervisor model, which several design choices here were built around: -- --* **No `max_workers` cap, no cross-session contention.** Each TUI owns its -- worker; there is no "Maximum idalib worker count reached" and no shared license -- slot to free. --* **No idle self-exit / keepalive dance.** The old per-worker `WorkerLifecycle` -- watchdog (`idle_ttl_sec`, default 600s) and the `KeepAlive` heartbeat that -- fought it are gone with the supervisor. The worker lives as long as the TUI -- holds the socket and dies with it. `WorkerClient.keepalive()` is a no-op kept -- for API parity, and `--ttl` is passed through but the single owned worker does -- not self-reap. --* **A crashed worker drops the socket**, surfacing as `IDAConnectionError`; the -- app's `_reconnect` respawns a fresh worker (re-opening + re-analyzing the -- binary). The hard-kill lock recovery below still applies. -+Code Mode returns the complete execution result directly; ida-tui no longer -+needs MCP structured-content/download-URL recovery for large pseudocode bodies. -+ -+## Code Mode lifecycle -+ -+`CodeModeClient` owns an authenticated SSE lease on a registered database: -+ -+* A matching GUI is preferred and remains open when the TUI exits. -+* Otherwise Code Mode reuses or starts a shared managed idalib worker. -+* Releasing one lease never terminates another client's session. A managed -+ worker saves and exits after its final lease under Code Mode's grace policy. -+* Lease loss surfaces as `IDAConnectionError`; reconnect performs discovery -+ again and may bind a newly-created instance. It does not silently swap the -+ handle underneath an operation. -+* `--ttl` and the old keepalive flag are compatibility no-ops; the lease itself -+ carries heartbeats. - - ## Writable path requirement (operational) - --`idb_open` writes the `.i64` next to the input binary, so the path must be --**writable**. Opening from read-only dirs (e.g. `/usr/lib`) fails with --`"Failed to open database"`. Copy targets into a writable dir first --(`targets/` in this repo). -+Attaching to a registered GUI does not require ida-tui to write beside the input. -+Creating a managed database does require a writable output path. Multi-binary -+projects provide one in their sidecar. ida-tui does not sweep IDA scratch files, -+because another registered session may own them. -diff --git a/docs/PROJECTS.md b/docs/PROJECTS.md -index fe6c2a8..0efee31 100644 ---- a/docs/PROJECTS.md -+++ b/docs/PROJECTS.md -@@ -7,9 +7,11 @@ search across all of them, and (later) follow calls from one into another. - - ## The constraint that shapes everything - --`idatui/worker.py` is `serve(sock, binpath)` — **one worker process holds exactly --one database** (idalib is main-thread-only and single-DB). So N binaries = N --worker processes, each with the analyzed DB resident. -+IDA still exposes one active database per GUI/idalib process. Code Mode makes -+those instances discoverable and shareable: each project entry retains one -+`DatabaseHandle` lease, which may target a registered GUI or a managed idalib -+worker. N resident project databases can therefore mean up to N processes, but -+ida-tui no longer owns or terminates them. - - Measured cost (this box, `targets/`): - -@@ -30,13 +32,13 @@ crypto library. - - Two capabilities that feel like one, but aren't: - --1. **Switching** to a binary needs a *live worker*. -+1. **Switching** to a binary needs a *live Code Mode lease*. - 2. **Searching across** binaries does *not* — if a per-binary index (functions, - strings, imports/exports) is cached on disk. - - That split is the unlock: project-wide search stays instant across every binary, - including ones never opened this session, and only *jumping* to a hit costs a --worker spawn. -+Code Mode attach/open. - - ## Layout - -@@ -83,17 +85,16 @@ basename and must be unique (it names the staged file). - - ## Runtime - --- **`WorkerPool`** — one `WorkerClient` per binary, spawned lazily on first -- switch, kept resident until the memory budget is exceeded, then LRU-evicted. -- Eviction **saves the DB first**, so returning to a binary is a DB load, not a -- re-analysis. Binaries can be pinned to stay resident. -+- **`DatabasePool`** — one `CodeModeClient` lease per resident binary, attached -+ lazily on first switch and LRU-released when the advisory memory budget is -+ exceeded. Eviction explicitly saves managed IDBs but never implicitly saves a -+ GUI. Closing a lease never kills a GUI or another client's managed worker; -+ Code Mode owns final worker shutdown. - - **`BinaryState`** — per binary: `client, program, nav, cur, func_index, - pref/active/split, filter`. Switching snapshots the current state and restores -- the target's. `_after_reconnect` already does exactly this swap (client + -- program, reload the index, re-open the entry) — switching reuses that seam. --- **Clean shutdown** — the worker currently does `close_database(save=False)` and -- is hard-killed on exit, which is why wedge files accumulate. Projects need -- save-on-evict and an orderly close anyway, so that gets fixed here. -+ the target's. `_after_reconnect` provides the client/program swap seam. -+- **Clean shutdown** — release all leases. Managed idalib workers save/close on -+ their own main thread after the final lease; GUI sessions remain open. - - ## UI - -@@ -109,8 +110,8 @@ basename and must be unique (it names the staged file). - ## Phases - - **Phase 1 — project model + switching. DONE.** Project file + staging --(`idatui/project.py`), `WorkerPool` with budget eviction / save-on-evict / --clean shutdown (`idatui/pool.py`), `BinaryState` snapshot+restore and the switch -+(`idatui/project.py`), `DatabasePool` with budgeted lease release and -+save-on-evict (`idatui/pool.py`), `BinaryState` snapshot+restore and the switch - itself, the `Ctrl+O` switcher palette, the active binary in the status line, and - `--project` (which creates the project when given binaries). One active binary; - no cross-binary search yet. -@@ -118,9 +119,9 @@ no cross-binary search yet. - Project mode is **additive**: with no `--project` the app is byte-for-byte the - single-binary tool it was, which is what keeps the 167-check pilot honest. - Switching reuses the `_after_reconnect` shape — swap client+program, rebuild the --index, reopen the entry. A binary whose worker is still resident restores --instantly (its `Program` and index are still in memory); an evicted one comes --back with a fresh worker but keeps its nav history, since that is just addresses. -+index, reopen the entry. A binary whose lease is still resident restores -+instantly (its `Program` and index are still in memory); an evicted one attaches -+again but keeps its nav history, since that is just addresses. - - **Phase 2 — index cache + project-wide search. (symbols done)** - `idatui/index.py` keeps one **SQLite FTS5 trigram** index at -@@ -222,7 +223,7 @@ records nothing — that's not navigation. - *Pre-warm follows the linkage graph, not list order.* When a binary finishes - indexing, `_prewarm_provider` warms the binary that provides the most of its - imports — where a follow is most likely to take you, so its startup is paid --before you ask. `WorkerPool.prewarm()` refuses rather than evicting: spending a -+before you ask. `DatabasePool.prewarm()` refuses rather than evicting: spending a - binary you visited on one you haven't is a straight downgrade, and it would throw - away that binary's caches too. At a tight budget pre-warm simply does nothing. It - estimates the cost of a not-yet-spawned worker from the largest resident one, -diff --git a/docs/SPLIT_VIEW.md b/docs/SPLIT_VIEW.md -index de37dc1..c421656 100644 ---- a/docs/SPLIT_VIEW.md -+++ b/docs/SPLIT_VIEW.md -@@ -29,11 +29,11 @@ Ghidra highlights **all** instructions a C line owns. We have one ea per line - (the marker), not the set. Getting the set is the only real work, and it's a - known technique: - --ida-pro-mcp derives the per-line marker via -+The old ida-pro-mcp backend derived the per-line marker via - `cfunc.get_line_item(line, col=0, …).get_ea()`. To get the **full set**, sweep - every column of the line (`get_line_item(line, x, …).get_ea()` for `x` in --`0..len`) and collect distinct non-`BADADDR` EAs. Same proven API, swept across --the line. A custom `decomp_map(ea)` tool in `server/patch_server.py` returns -+`0..len`) and collect distinct non-`BADADDR` EAs. The Code Mode adapter's -+`decomp_map(ea)` operation returns - `[{line, primary_ea, eas:[…]}, …]`; invert for `ea → line`. - - ## State model -@@ -78,14 +78,14 @@ decomp→listing uses `ListingModel.ensure_ea`. Tab re-links from the new driver - Still single-ea per line (one instruction highlighted); the region comes in - phase 3. - --**Phase 3 — rich highlight. DONE.** `decomp_map` custom tool --(`server/patch_server.py`) sweeps `cfunc.get_line_item` across every column of -+**Phase 3 — rich highlight. DONE.** The Code Mode `decomp_map` operation -+(`idatui/codemode_client.py`) sweeps `cfunc.get_line_item` across every column of - each pseudocode line and collects the EAs from each item's `dstr()` (`'EA: desc'` - — the same source as the `/*ea*/` marker, so it aligns). `Program.decomp_map(ea)` - returns the per-line ea lists (cached by name-gen); the app loads it async into - `_split_eamap` / `_split_ea2line` and `_sync_split` bands the **whole** instruction - region of a C line (and uses the exact ea→line inverse for the reverse). Falls --back to the single marker until the map lands. Verified on the pilot's real worker -+back to the single marker until the map lands. Verified on a real Code Mode database - (alignment + multi-instruction region band). - - **Phase 4 — polish. DONE.** -diff --git a/experiments/worker_smoke.py b/experiments/worker_smoke.py -index b215a57..9b55138 100644 ---- a/experiments/worker_smoke.py -+++ b/experiments/worker_smoke.py -@@ -1,58 +1,56 @@ --"""Runnable read-path smoke: drives the REAL domain.Program through WorkerClient --(our idalib worker over a unix socket). Run when idalib can spawn: -- ~/ida-venv/bin/python experiments/worker_smoke.py --""" --import os, sys, shutil, time --REPO=os.path.expanduser("~/dev/ida-tui-maybe"); sys.path.insert(0, REPO); os.chdir(REPO) --# fresh copy so the worker's idalib doesn't fight any running server --src=f"{REPO}/targets/echo"; tmp="/tmp/echo_worker" --shutil.copy(src, tmp) --for e in ".i64 .id0 .id1 .id2 .nam .til".split(): -- try: os.remove(tmp+e) -- except OSError: pass -- --from idatui.worker_client import WorkerClient --from idatui.domain import Program -- --print("spawning worker + opening echo…", flush=True) --t=time.time() --cl=WorkerClient(tmp) --cl.connect(progress=lambda m: None) --print(f" worker ready in {time.time()-t:.2f}s session={cl.resolve_db()}", flush=True) --prog=Program(cl) -- --# --- drive the REAL domain layer through the worker (read path) --- --main=prog.resolve("main") --print("resolve('main') =", hex(main), flush=True) -+"""Exercise the real domain.Program through an IDA Code Mode lease. - --idx=prog.functions(); idx.load_all() --print("functions() ->", len(idx), "funcs", flush=True) -- --fn=prog.function_of(main) --print("function_of(main) ->", fn.name, hex(fn.addr), "size", fn.size, flush=True) -- --b=prog.read_bytes(main, 16) --print("read_bytes(main,16) ->", b.hex(), flush=True) -- --lm=prog.listing(main) --for _ in range(3): lm.load_next_page() --rows=[lm.get(i) for i in range(min(6,len(lm)))] --print("listing() first rows:", flush=True) --for h in rows: -- if h: print(" ", hex(h.ea), h.kind, repr(h.text[:44]), flush=True) -+A matching registered GUI is reused; otherwise Code Mode starts a managed -+idalib worker. Usage: ``uv run python experiments/worker_smoke.py FILE``. -+""" -+from __future__ import annotations - --d=prog.decompile(main) --print("decompile(main) -> failed?", d.failed, "lines:", len((d.code or '').splitlines()), flush=True) -+import os -+import sys -+import time - --regs=prog.file_regions() --print("file_regions ->", len(regs), "segments", flush=True) -+from idatui.codemode_client import CodeModeClient -+from idatui.domain import Program - --# xrefs to a called function --callee=next((f.addr for f in idx.all_loaded() if f.name.startswith("sub_")), None) --if callee: -- xr=prog.xrefs_to(callee) -- print("xrefs_to(", hex(callee), ") ->", len(xr), "refs", flush=True) - --ok = (fn.name=="main" and len(idx)>100 and b and not d.failed and len(regs)>0) --print("VERDICT:", "OK — domain.Program runs unchanged on the worker" if ok else "FAIL", flush=True) --cl.close() -+def main() -> int: -+ target = os.path.abspath(sys.argv[1] if len(sys.argv) > 1 else "experiments/fibonacci.elf") -+ print(f"attaching Code Mode to {target}…", flush=True) -+ started = time.time() -+ client = CodeModeClient(target) -+ client.connect(progress=lambda message: print(f" {message}", flush=True)) -+ print( -+ f" ready in {time.time() - started:.2f}s; backend={client.backend}; " -+ f"session={client.resolve_db()}", -+ flush=True, -+ ) -+ program = Program(client) -+ try: -+ index = program.functions() -+ index.load_all() -+ print(f"functions() -> {len(index)}", flush=True) -+ first = index.get(0) -+ if first is None: -+ print("VERDICT: FAIL — no functions", flush=True) -+ return 1 -+ fn = program.function_of(first.addr) -+ data = program.read_bytes(first.addr, 16) -+ decompilation = program.decompile(first.addr) -+ print(f"function_of() -> {fn}", flush=True) -+ print(f"read_bytes() -> {data.hex()}", flush=True) -+ print( -+ f"decompile() -> failed={decompilation.failed}; " -+ f"lines={len((decompilation.code or '').splitlines())}", -+ flush=True, -+ ) -+ print(f"file_regions() -> {len(program.file_regions())}", flush=True) -+ ok = fn is not None and bool(data) and bool(program.file_regions()) -+ print(f"VERDICT: {'OK' if ok else 'FAIL'}", flush=True) -+ return 0 if ok else 1 -+ finally: -+ program.close() -+ client.close() -+ -+ -+if __name__ == "__main__": -+ raise SystemExit(main()) -diff --git a/ida-tui b/ida-tui -index 9ceb949..a488cac 100755 ---- a/ida-tui -+++ b/ida-tui -@@ -3,10 +3,9 @@ - # - # ./ida-tui foo.elf # open a binary and drive it — that's it - # --# Opening a binary spins up our own idalib worker (a unix-socket subprocess; --# no HTTP, no supervisor). Uses the venv python that has textual (override with --# $IDATUI_PYTHON); the worker auto-picks the python that has ida_pro_mcp --# (override with $IDATUI_WORKER_PYTHON). -+# The launcher leases a registered IDA GUI or shared managed idalib worker -+# through ida_codemode. The selected Python must have ida-tui's dependencies; -+# override it with $IDATUI_PYTHON. - set -eu - - SCRIPT_DIR=$(CDPATH= cd -- "$(dirname -- "$0")" && pwd) -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 1edc65e..984140f 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 -@@ -30,7 +30,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 -@@ -46,8 +46,8 @@ from textual.widgets.option_list import Option - 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") -@@ -126,8 +126,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 -@@ -812,9 +812,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] -@@ -2228,12 +2228,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"), - ] - -@@ -2250,7 +2254,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: -@@ -3392,15 +3396,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 "") - self._trace_path = trace_path or "" # Tenet execution trace to explore - self._trace = None # the loaded Trace, once analysed -@@ -3414,7 +3419,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 -@@ -3512,8 +3517,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: -@@ -3544,6 +3549,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: -@@ -3557,7 +3567,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" -@@ -3576,10 +3590,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: -@@ -3591,7 +3608,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) -@@ -3605,23 +3624,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 -@@ -3665,6 +3667,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 - -@@ -3718,10 +3725,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 -@@ -3786,9 +3789,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): -@@ -3820,15 +3824,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 -@@ -3836,7 +3840,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 -@@ -3856,13 +3860,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 -@@ -3877,14 +3881,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)) -@@ -3895,14 +3899,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 -@@ -3974,7 +3979,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) -@@ -3992,7 +3997,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) -@@ -4077,7 +4082,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) -@@ -4267,6 +4278,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": -@@ -4283,7 +4297,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) -@@ -4323,7 +4337,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, -@@ -4344,8 +4358,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) -@@ -4382,7 +4396,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 -@@ -4439,28 +4453,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).""" -@@ -4807,7 +4799,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: -@@ -4980,7 +4972,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. -@@ -5422,7 +5414,7 @@ class IdaTui(App): - 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: - self.app.call_from_thread(self._status, f"rename failed: {e.message}") - return -@@ -5438,7 +5430,6 @@ class IdaTui(App): - self.app.call_from_thread(self._after_rename, kind, addr, old, new) - - def _after_rename(self, kind: str, addr: int | None, old: str, new: str) -> None: -- cur = self._cur - # A renamed symbol can appear in many functions, so invalidate globally; - # each function refreshes its names the next time it's viewed. - self.program.bump_names() -@@ -5465,7 +5456,7 @@ class IdaTui(App): - — unlike the symbol-by-name path, this names the address directly.""" - assert self.program is not None - try: -- res = self.program.client.call( -+ res = self.program.client.invoke( - "rename", batch={"data": {"addr": hex(addr), "new": name}}) - except IDAToolError as e: - self.app.call_from_thread(self._status, f"name failed: {e.message}") -@@ -5729,7 +5720,6 @@ class IdaTui(App): - screen row, so the eye tracks straight across. - """ - lst = self.query_one(ListingView) -- dec = self.query_one(DecompView) - if lst.model is None: - return False - row = lst.model.ensure_ea(pc) -@@ -6022,7 +6012,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 -@@ -6951,7 +6941,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: -@@ -6963,7 +6955,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..ab4e698 ---- /dev/null -+++ b/idatui/codemode_client.py -@@ -0,0 +1,1107 @@ -+"""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 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 -+ -+from .errors import IDAConnectionError, IDATimeoutError, IDAToolError, Session -+ -+ -+def registered_database(path: str, output_database: str | None = None) -> bool: -+ """Whether a live/lock-held Code Mode instance owns this target.""" -+ 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": -+ 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: -+ 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 97b042d..2d6500e 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 dataclasses import dataclass, field, replace - from typing import Callable, TYPE_CHECKING -@@ -33,7 +27,7 @@ from typing import Callable, TYPE_CHECKING - 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 -@@ -63,7 +57,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: -@@ -91,7 +85,7 @@ class Line: - - @dataclass(frozen=True) - class Head: -- """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.""" - - ea: int -@@ -101,7 +95,7 @@ class Head: - 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). - spans: tuple[tuple[str, str], ...] | None = None - -@@ -240,7 +234,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). - """ - -@@ -260,7 +254,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: -@@ -370,7 +364,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") -@@ -431,7 +425,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(), - ) -@@ -572,15 +566,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): -@@ -654,7 +648,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 {} -@@ -956,7 +950,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" -@@ -973,7 +967,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 ------------------------------------------------- # -@@ -1000,17 +994,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"]), -@@ -1019,7 +1010,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, -@@ -1061,8 +1052,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()] -@@ -1080,15 +1070,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) -@@ -1102,7 +1091,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 [] -@@ -1151,7 +1140,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") -@@ -1161,9 +1150,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 {} -@@ -1186,7 +1175,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") -@@ -1195,10 +1184,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"): -@@ -1211,7 +1199,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"): -@@ -1220,9 +1208,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"): -@@ -1231,7 +1219,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 {} -@@ -1240,9 +1228,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"): -@@ -1251,15 +1239,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 ------------------------------------------------------- # -@@ -1273,13 +1260,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) -@@ -1288,10 +1269,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 -@@ -1301,7 +1282,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 -@@ -1309,15 +1293,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) -@@ -1365,18 +1340,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): -@@ -1395,7 +1370,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", -@@ -1404,7 +1379,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')}") -@@ -1413,11 +1388,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)} -@@ -1429,14 +1404,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"} -@@ -1450,7 +1425,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']}") - -@@ -1460,7 +1435,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'}") -@@ -1468,7 +1443,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( -@@ -1483,15 +1458,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: # noqa: BLE001 -- fall back to the truncated preview -- 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.""" -@@ -1504,7 +1470,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: -@@ -1529,13 +1495,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): -@@ -1566,7 +1532,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 [] -@@ -1579,13 +1545,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}], - ) -@@ -1597,9 +1563,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 ------------------------------------------------ # -@@ -1617,17 +1583,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 [] -@@ -1665,7 +1631,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 b6fa641..0d865b5 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/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 64e1546..f78213c 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).""" - from .formats import load_args -@@ -34,24 +26,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.""" -- 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: -- try: -- os.remove(base + suf) -- 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)") -@@ -59,9 +46,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", -@@ -76,7 +63,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 = {} -@@ -134,27 +121,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 - rpc_path = os.path.abspath(os.path.expanduser(args.rpc)) if args.rpc else None - IdaTui(open_path=binary, keepalive=not args.no_keepalive, -diff --git a/idatui/pane.py b/idatui/pane.py -index 37a5f0c..0f8fa57 100644 ---- a/idatui/pane.py -+++ b/idatui/pane.py -@@ -15,9 +15,9 @@ then close it — all without a human touching the keyboard. - python -m idatui.pane list - python -m idatui.pane stop --sock <sock> # graceful quit + kill pane - --Requires: running inside tmux. 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. Each pane leases a registered GUI or shared -+managed idalib database through Code Mode. Uses ~/ida-venv/bin/python for the -+TUI unless --python / IDATUI_PYTHON says otherwise. - """ - from __future__ import annotations - -@@ -25,7 +25,6 @@ import argparse - import json - import os - import secrets --import signal - import subprocess - import sys - import time -@@ -72,56 +71,14 @@ def _tmux(*args: str) -> str: - check=True).stdout.strip() - - --# --------------------------------------------------------------------------- # --# 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 -- -- - def _count_live_panes() -> int: - return sum(1 for r in _load_registry() if _pane_alive(r.get("pane", ""))) - - - 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 - - - # --------------------------------------------------------------------------- # -@@ -146,16 +103,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). -@@ -202,9 +151,8 @@ def _wait_ready(sock: str, timeout: float, pane: str, - stuck_after: float = 45.0) -> 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 -@@ -225,9 +173,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 -@@ -317,13 +264,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 - - -@@ -361,9 +304,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) - - args = p.parse_args(argv) -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..6afd8a1 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,32 @@ 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 ida_codemode.registry import canonical_path, idb_key, scan_instances -+ expected_key = idb_key(ref.db) -+ staged_path = canonical_path(ref.staged) -+ owner = next( -+ (item.entry for item in scan_instances(timeout=0.5) -+ if item.entry.idb_key == expected_key -+ or (item.entry.exe_path and canonical_path(item.entry.exe_path) == staged_path)), -+ None, -+ ) -+ 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 +358,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/worker.py b/idatui/worker.py -deleted file mode 100644 -index a4e3509..0000000 ---- a/idatui/worker.py -+++ /dev/null -@@ -1,233 +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 uuid -- -- --# --------------------------------------------------------------------------- # --# 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] -- -- 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}") -- result = fn(**args) -- # 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 79a6db9..0000000 ---- a/idatui/worker_client.py -+++ /dev/null -@@ -1,234 +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() -- -- # -- 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: -- 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() -- 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(0.2) -- 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 -- 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._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() -diff --git a/pyproject.toml b/pyproject.toml -index 30f8cc2..72f1ff3 100644 ---- a/pyproject.toml -+++ b/pyproject.toml -@@ -1,15 +1,17 @@ - [project] - name = "idatui" - version = "0.0.1" --description = "A minimal keyboard-first TUI frontend for IDA Pro over the ida-pro-mcp (idalib) server." -+description = "A keyboard-first TUI frontend for shared IDA Code Mode databases." - requires-python = ">=3.11" --# The client layer is intentionally stdlib-only (urllib/http.client), matching the --# ida-mcp skill philosophy: no install needed to talk to the server. --dependencies = [] -+# ida-codemode supplies GUI discovery, shared idalib workers, leases, and the -+# execute_python/ida-domain database surface. -+dependencies = [ -+ "ida-codemode-mcp", -+ "textual>=8", -+ "pygments>=2", # Used directly for pseudocode highlighting. -+] - - [project.optional-dependencies] --# The TUI layer pulls in Textual; the client/domain layers are stdlib-only. --tui = ["textual>=8", "pygments>=2"] # pygments ships with rich; explicit for the C lexer - dev = ["pytest>=8"] - - [project.scripts] -@@ -22,3 +24,6 @@ build-backend = "hatchling.build" - - [tool.hatch.build.targets.wheel] - packages = ["idatui"] -+ -+[tool.uv.sources] -+ida-codemode-mcp = { path = "../ida-codemode-mcp", editable = true } -diff --git a/server/patch_server.py b/server/patch_server.py -deleted file mode 100644 -index 160b15b..0000000 ---- a/server/patch_server.py -+++ /dev/null -@@ -1,1248 +0,0 @@ --#!/usr/bin/env python3 --"""Inject idatui's extra ida-pro-mcp tools into the installed server package. -- --DEPRECATED along with the ida-pro-mcp transport: the default backend is now the --idalib worker (idatui/worker.py), which registers these same tools in-process and --needs no patching. Kept only for `--backend mcp`; slated for removal. -- --ida-pro-mcp lacks a few tools idatui needs. Rather than vendor/fork the server, --we keep the tool source here and inject it (idempotently) into the installed --``api_types.py``. That module is imported by every worker --(``python -m ida_pro_mcp.idalib_server``), so the tools register themselves via --``@tool`` on the shared ``MCP_SERVER`` — no server code is forked, and re-running --this (spawn.sh does, on every start) re-applies it after a reinstall/upgrade. -- --Injected tools: -- * ``del_type`` — delete a named local type (struct editor CRUD). -- * ``func_types`` — structured decompiler types for a function (prototype + -- local variables), so clients don't parse pseudocode text. -- * ``set_lvar_type`` — set a decompiler local variable's type; works on auto/ -- register vars too (the stock set_type only updates lvars -- that already have user-saved info). -- --The block between the BEGIN/END markers is *replaced* on each run, so editing --BODY here and restarting the supervisor updates the tools. -- --Run with the *same* interpreter the server uses (the idalib-mcp entry point's --``/usr/bin/python``), so it patches the file the workers actually import. --Changing a tool needs a supervisor restart so workers respawn. --""" --from __future__ import annotations -- --import importlib.util --import pathlib --import sys -- --BEGIN = "# >>> idatui-ext: begin (auto-injected by server/patch_server.py) >>>" --END = "# <<< idatui-ext: end <<<" -- --# Appended to ida_pro_mcp/ida_mcp/api_types.py, which already imports --# ``Annotated``, ``tool``, ``idasync``, ``ida_typeinf``, ``parse_address`` and --# ``_parse_type_tinfo``. --BODY = ''' --def _idatui_lv_get(x): -- return x() if callable(x) else x -- -- --@tool --@idasync --def resolve_names( -- queries: Annotated[list, "Symbol name(s) to resolve to their OWN address"], --) -> list: -- """Resolve named locations (functions, labels like loc_/locret_, data) to the -- exact address the NAME denotes, via get_name_ea. Unlike lookup_funcs, a -- mid-function label resolves to the label's address, not the containing -- function's entry.""" -- import idaapi -- qs = queries if isinstance(queries, list) else [queries] -- out = [] -- for q in qs: -- q = str(q).strip() -- ea = idaapi.get_name_ea(idaapi.BADADDR, q) -- out.append({"query": q, "ea": (hex(ea) if ea != idaapi.BADADDR else None)}) -- return out -- -- --@tool --@idasync --def del_type( -- name: Annotated[str, "Local type name to delete (struct/union/enum/typedef)"], --) -> dict: -- """Delete a named local type from the local type library.""" -- til = ida_typeinf.get_idati() -- ok = ida_typeinf.del_named_type(til, name, ida_typeinf.NTF_TYPE) -- if not ok: -- return {"name": name, "error": f"Type '{name}' not found or could not be deleted"} -- return {"name": name, "deleted": True} -- -- --@tool --@idasync --def func_types( -- addr: Annotated[str, "Function address or name"], --) -> dict: -- """Structured decompiler types for a function: its prototype plus each local -- variable (name/type/is_arg). Lets clients read/edit types without parsing -- pseudocode text.""" -- import ida_hexrays -- import idaapi -- -- def _tstr(tif): -- try: -- s = tif.dstr() -- if s: -- return s -- except Exception: -- pass -- return str(tif) -- -- ea = parse_address(addr) -- f = idaapi.get_func(ea) -- if not f: -- return {"addr": str(addr), "error": "no function at address"} -- try: -- cf = ida_hexrays.decompile(f.start_ea) -- except Exception as e: -- return {"addr": hex(f.start_ea), "error": f"decompile failed: {e}"} -- if cf is None: -- return {"addr": hex(f.start_ea), "error": "decompilation failed"} -- name = idaapi.get_func_name(f.start_ea) or "" -- try: -- proto = ida_typeinf.print_tinfo( -- "", 0, 0, ida_typeinf.PRTYPE_1LINE, cf.type, name, "") -- except Exception: -- proto = "" -- lvars = [] -- for lv in cf.get_lvars(): -- try: -- ty = _tstr(_idatui_lv_get(lv.type)) -- except Exception: -- ty = "" -- lvars.append({ -- "name": _idatui_lv_get(lv.name), -- "type": ty, -- "is_arg": bool(_idatui_lv_get(lv.is_arg_var)), -- }) -- return { -- "addr": hex(f.start_ea), -- "name": name, -- "prototype": (proto or "").strip(), -- "lvars": lvars, -- } -- -- --@tool --@idasync --def set_lvar_type( -- addr: Annotated[str, "Function address or name"], -- variable: Annotated[str, "Local variable name"], -- type: Annotated[str, "New C type for the variable"], --) -> dict: -- """Set a decompiler local variable's type. Handles auto/register vars (unlike -- set_type, which only updates lvars that already have user-saved info).""" -- import ida_hexrays -- import idaapi -- -- ea = parse_address(addr) -- f = idaapi.get_func(ea) -- if not f: -- return {"error": "no function at address"} -- try: -- cf = ida_hexrays.decompile(f.start_ea) -- except Exception as e: -- return {"error": f"decompile failed: {e}"} -- if cf is None: -- return {"error": "decompilation failed"} -- target = None -- for lv in cf.get_lvars(): -- if _idatui_lv_get(lv.name) == variable: -- target = lv -- break -- if target is None: -- return {"error": f"local variable {variable!r} not found"} -- try: -- tif = _parse_type_tinfo(type) -- except Exception as e: -- return {"error": f"bad type {type!r}: {e}"} -- lsi = ida_hexrays.lvar_saved_info_t() -- try: -- lsi.ll = target -- except Exception: -- try: -- lsi.ll.location = _idatui_lv_get(target.location) -- lsi.ll.defea = target.defea -- except Exception as e: -- return {"error": f"could not locate variable: {e}"} -- lsi.type = tif -- ok = bool(ida_hexrays.modify_user_lvar_info( -- f.start_ea, ida_hexrays.MLI_TYPE, lsi)) -- return {"addr": hex(f.start_ea), "variable": variable, "type": type, "ok": ok} -- -- --@tool --@idasync --def file_regions() -> dict: -- """Loaded segments mapped to their raw file offsets (get_fileregion_offset), -- so clients can convert a virtual address to an on-disk file offset without a -- format-specific header parser. file_off is -1 for non-file-backed segments -- (e.g. .bss).""" -- import ida_segment -- import idaapi -- -- out = [] -- seg = ida_segment.get_first_seg() -- while seg is not None: -- try: -- fo = int(idaapi.get_fileregion_offset(seg.start_ea)) -- except Exception: -- fo = -1 -- if fo < 0 or fo >= (1 << 48): -- fo = -1 -- try: -- nm = ida_segment.get_segm_name(seg) or "" -- except Exception: -- nm = "" -- out.append({"start": hex(seg.start_ea), "end": hex(seg.end_ea), -- "file_off": fo, "name": nm}) -- seg = ida_segment.get_next_seg(seg.start_ea) -- return {"regions": out} -- -- --@tool --@idasync --def make_string( -- addr: Annotated[str, "Address of the string start"], -- length: Annotated[int, "Length in bytes (0 = auto-detect to the terminator)"] = 0, -- kind: Annotated[str, "String kind: c | c16 | c32 | pascal"] = "c", --) -> dict: -- """Create a string literal at ``addr`` (IDA's 'A'). ``length`` 0 auto-detects -- to the terminator. Undefines any items in the way first, like the UI does. -- Returns the created byte size and the decoded contents.""" -- import ida_bytes -- import ida_nalt -- -- ea = parse_address(addr) -- strtype = { -- "c": ida_nalt.STRTYPE_C, -- "c16": ida_nalt.STRTYPE_C_16, -- "c32": ida_nalt.STRTYPE_C_32, -- "pascal": ida_nalt.STRTYPE_PASCAL, -- }.get(str(kind).lower(), ida_nalt.STRTYPE_C) -- n = max(int(length), 0) -- # Free any existing item(s) so create_strlit can carve the literal. -- ida_bytes.del_items(ea, ida_bytes.DELIT_SIMPLE, n if n > 0 else 1) -- ok = bool(ida_bytes.create_strlit(ea, n, strtype)) -- if not ok: -- return {"addr": addr, "ok": False, "error": "create_strlit failed"} -- size = int(ida_bytes.get_item_size(ea)) -- try: -- raw = ida_bytes.get_strlit_contents(ea, -1, strtype) -- text = raw.decode("utf-8", "replace") if raw else "" -- except Exception: -- text = "" -- return {"addr": addr, "ok": True, "size": size, "text": text} -- -- --@tool --@idasync --def read_raw( -- addr: Annotated[str, "Start address (hex or name)"], -- size: Annotated[int, "Number of bytes to read"], --) -> dict: -- """Read ``size`` bytes at ``addr`` as ONE contiguous lowercase hex string -- (no per-byte '0x'/spaces). The hot path for the hex view and disasm opcode -- bytes. -- -- Fast: does a single bulk ``ida_bytes.get_bytes`` (C-speed) instead of the -- per-byte read_bytes_bss_safe loop (2 IDA calls/byte). Unloaded bytes come -- back from IDA as the 0xFF sentinel, so we only re-check is_loaded for the -- (usually sparse) 0xFF bytes and zero the genuinely-unloaded ones — matching -- get_bytes' bss semantics without paying per-byte for the whole range. -- -- Encoding is compact hex (~2.5x smaller than get_bytes' '0x..'-with-spaces) -- and, unlike get_bytes, does not truncate on large reads.""" -- import ida_bytes -- -- ea = parse_address(addr) -- n = max(int(size), 0) -- if n == 0: -- return {"addr": addr, "hex": "", "n": 0} -- raw = ida_bytes.get_bytes(ea, n) -- if raw is None or len(raw) < n: # nothing (or not all) mapped -- base = bytearray(raw or b"") -- base.extend(b"\\xff" * (n - len(base))) -- raw = bytes(base) -- ba = bytearray(raw) -- # Only unloaded bytes read as 0xFF; correct just those to 0 (bss => zero). -- i = ba.find(0xFF) -- while i != -1: -- if not ida_bytes.is_loaded(ea + i): -- ba[i] = 0 -- i = ba.find(0xFF, i + 1) -- return {"addr": addr, "hex": bytes(ba).hex(), "n": len(ba)} -- -- --def _idatui_head_row(ea): -- """One flat-listing row for the head at ``ea``: kind (code/data/unknown), -- byte size, rendered text, and any symbol name.""" -- import ida_bytes -- import ida_lines -- import ida_name -- -- f = ida_bytes.get_flags(ea) -- if ida_bytes.is_code(f): -- kind = "code" -- elif ida_bytes.is_data(f): -- kind = "data" -- else: -- kind = "unknown" -- line = ida_lines.generate_disasm_line(ea, 0) -- text = ida_lines.tag_remove(line) if line else "" -- text = " ".join(text.split()) # collapse IDA's column padding -- row = { -- "ea": hex(ea), -- "kind": kind, -- "size": int(ida_bytes.get_item_size(ea)), -- "text": text, -- } -- if line: -- # Keep IDA's own token classification for syntax highlighting. Built from -- # the SAME line as `text`, then whitespace-collapsed identically so the -- # two never disagree about what the row says. -- spans = _idatui_spans(line) -- joined = "".join(t for _k, t in spans) -- if " ".join(joined.split()) == text: -- row["spans"] = spans -- nm = ida_name.get_ea_name(ea) -- if nm: -- row["name"] = nm -- return row -- -- --#: IDA colour tag -> the semantic kind the TUI styles. IDA already classifies --#: every token in a disassembly line, for every processor it supports, so there --#: is nothing to lex: generate_disasm_line emits \x01<tag>text\x02<tag> and the --#: tag says what the text IS. A pygments assembly lexer would be a worse guess at --#: this and would need one dialect per architecture. --_IDATUI_SPAN_KINDS = { -- "insn": ("SCOLOR_INSN", "SCOLOR_KEYWORD", "SCOLOR_ASMDIR", "SCOLOR_MACRO"), -- "reg": ("SCOLOR_REG",), -- "num": ("SCOLOR_NUMBER", "SCOLOR_CHAR", "SCOLOR_BINPREF"), -- "str": ("SCOLOR_STRING",), -- # NB the real constant names: DATNAME/CODNAME, not "DNAME". Guessing here -- # fails silently — an unmapped tag renders as plain body text, so symbols -- # just quietly aren't blue and nothing tells you why. -- "name": ("SCOLOR_DATNAME", "SCOLOR_CODNAME", "SCOLOR_LOCNAME", -- "SCOLOR_IMPNAME", "SCOLOR_DEMNAME", "SCOLOR_LIBNAME", -- "SCOLOR_CNAME", "SCOLOR_DNAME", -- "SCOLOR_CREF", "SCOLOR_DREF", "SCOLOR_CREFTAIL", "SCOLOR_DREFTAIL"), -- "seg": ("SCOLOR_SEGNAME",), -- "cmt": ("SCOLOR_AUTOCMT", "SCOLOR_REGCMT", "SCOLOR_RPTCMT", "SCOLOR_VOIDOP"), -- "punct": ("SCOLOR_SYMBOL", "SCOLOR_ALTOP", "SCOLOR_HIDNAME"), -- "err": ("SCOLOR_ERROR",), --} -- -- --def _idatui_tag_map(): -- """{tag character: kind}, built once from whatever this IDA actually has.""" -- import ida_lines -- out = {} -- for kind, names in _IDATUI_SPAN_KINDS.items(): -- for n in names: -- v = getattr(ida_lines, n, None) -- if isinstance(v, str) and v: -- out[v[0]] = kind -- elif isinstance(v, int): -- out[chr(v)] = kind -- return out -- -- --_IDATUI_TAGS = None -- -- --def _idatui_spans(line): -- """A tagged disasm line as [[kind, text], ...], colour tags resolved. -- -- Unknown tags become 'text' rather than being dropped: a processor module can -- emit a colour we don't classify, and losing the characters would corrupt the -- line.""" -- global _IDATUI_TAGS -- import ida_lines -- if _IDATUI_TAGS is None: -- _IDATUI_TAGS = _idatui_tag_map() -- 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)) -- spans, stack, buf = [], [], [] -- i, n = 0, len(line) -- -- def flush(): -- if buf: -- spans.append([stack[-1] if stack else "text", "".join(buf)]) -- del buf[:] -- -- while i < n: -- ch = line[i] -- if ch == on and i + 1 < n: -- tag = line[i + 1] -- if tag == addr_tag: -- # An embedded target address, not display text: 16 hex digits -- # that must not reach the screen. -- i += 2 + addr_len -- continue -- flush() -- stack.append(_IDATUI_TAGS.get(tag, "text")) -- i += 2 -- continue -- if ch == off and i + 1 < n: -- flush() -- if stack: -- stack.pop() -- i += 2 -- continue -- if ch == esc and i + 1 < n: # escaped literal -- buf.append(line[i + 1]) -- i += 2 -- continue -- buf.append(ch) -- i += 1 -- flush() -- # Collapse IDA's column padding EXACTLY as the plain text does. A run of -- # spaces can straddle two spans, so this walks characters rather than -- # collapsing each span on its own — otherwise the spans and `text` disagree -- # about the line and the row silently loses its highlighting. -- out, prev_space = [], False -- for kind, txt in spans: -- acc = [] -- for ch in txt: -- if ch.isspace(): -- if prev_space: -- continue -- acc.append(" ") -- prev_space = True -- else: -- acc.append(ch) -- prev_space = False -- if acc: -- out.append([kind, "".join(acc)]) -- while out and out[0][1] == " ": -- out.pop(0) -- while out and out[-1][1] == " ": -- out.pop() -- if out and out[0][1].startswith(" "): -- out[0][1] = out[0][1].lstrip() -- if out and out[-1][1].endswith(" "): -- out[-1][1] = out[-1][1].rstrip() -- return [[k, t] for k, t in out if t] -- -- --def _idatui_unknown_row(ea, size): -- """One collapsed row for a run of ``size`` undefined bytes starting at -- ``ea``. A single byte is rendered normally (shows its value); a longer run -- collapses to ``db N dup(?)`` so a big .bss/gap doesn't explode into millions -- of one-byte rows.""" -- import ida_name -- -- if size <= 1: -- return _idatui_head_row(ea) -- row = {"ea": hex(ea), "kind": "unknown", "size": int(size), -- "text": f"db {size} dup(?)"} -- nm = ida_name.get_ea_name(ea) -- if nm: -- row["name"] = nm -- return row -- -- --def _idatui_struct_member_rows(ea): -- """Indented member rows for a struct-typed data item at ``ea`` (expansion), -- or [] if it isn't a struct. Top-level fields only.""" -- import ida_nalt -- import ida_typeinf -- import idaapi -- -- tif = ida_typeinf.tinfo_t() -- if not (ida_nalt.get_tinfo(tif, ea) and tif.is_udt()): -- return [] -- udt = ida_typeinf.udt_type_data_t() -- if not tif.get_udt_details(udt): -- return [] -- rows = [] -- for m in udt: -- off = m.begin() // 8 -- try: -- mtype = m.type._print() or "" -- except Exception: -- mtype = "" -- try: -- sz = int(m.type.get_size()) -- if sz == idaapi.BADSIZE: -- sz = 0 -- except Exception: -- sz = 0 -- name = m.name or "" -- text = f"+{off:X} {name}" + (f" {mtype}" if mtype else "") -- rows.append({"ea": hex(ea + off), "kind": "member", "size": sz, -- "text": text}) -- return rows -- -- --def _idatui_func_header_rows(ea): -- """IDA-style subroutine banner rows shown just before a function's entry.""" -- import ida_funcs -- -- name = ida_funcs.get_func_name(ea) or "sub_%X" % ea -- bar = "=" * 15 + " S U B R O U T I N E " + "=" * 15 -- return [ -- {"ea": hex(ea), "kind": "sep", "size": 0, "text": ""}, -- {"ea": hex(ea), "kind": "sep", "size": 0, "text": "; " + bar}, -- {"ea": hex(ea), "kind": "funchdr", "size": 0, -- "text": name + " proc", "name": name}, -- ] -- -- --def _idatui_func_footer_rows(ea, func): -- """End-of-function marker shown just after a function's last item.""" -- import ida_funcs -- -- name = ida_funcs.get_func_name(func.start_ea) or "sub_%X" % func.start_ea -- return [ -- {"ea": hex(ea), "kind": "funchdr", "size": 0, -- "text": name + " endp", "name": name}, -- {"ea": hex(ea), "kind": "sep", "size": 0, "text": "; " + "-" * 60}, -- ] -- -- --@tool --@idasync --def heads( -- addr: Annotated[str, "Start address or name to walk from"], -- count: Annotated[int, "Max heads to return (default 200, max 2000)"] = 200, -- offset: Annotated[int, "Skip first N heads from addr (default 0)"] = 0, -- end: Annotated[str, "Optional exclusive end address; default = segment end"] = "", -- back: Annotated[bool, "Walk backwards: return the count heads ENDING just before addr, in forward order"] = False, -- annotate: Annotated[bool, "Emit IDA-style function boundary banner rows (kind sep/funchdr)"] = False, --) -> dict: -- """Walk item heads from ``addr`` as a flat listing: every head is rendered -- (code OR data OR undefined) via generate_disasm_line and stepped with -- next_head/prev_head. Unlike ``disasm`` (code-only, bails at the first data -- byte) this shows db/dw/dd/... lines for data and undefined regions — IDA's -- real disassembly view. Address-paged: page forward by re-calling with -- ``addr`` = the returned cursor.next; page up with ``back=true``.""" -- import ida_bytes -- import ida_segment -- import idaapi -- -- count = 2000 if count > 2000 else (1 if count < 1 else count) -- offset = max(int(offset), 0) -- try: -- start = parse_address(addr) -- except Exception as e: -- return {"addr": str(addr), "error": str(e), "heads": [], "cursor": {"done": True}} -- seg = ida_segment.getseg(start) -- if not seg: -- return {"addr": str(addr), "error": "no segment", "heads": [], "cursor": {"done": True}} -- lo, hi = seg.start_ea, seg.end_ea -- if end: -- try: -- hi = min(hi, parse_address(end)) -- except Exception: -- pass -- -- rows = [] -- if back: -- # Collect up to (count+offset) heads strictly before `start`, then take -- # the window closest to `start`, returned in forward order. -- walk = [] -- cur = ida_bytes.prev_head(start, lo) -- while cur != idaapi.BADADDR and cur >= lo and len(walk) < count + offset: -- walk.append(cur) -- cur = ida_bytes.prev_head(cur, lo) -- walk.reverse() -- chosen = walk[: len(walk) - offset] if offset else walk -- chosen = chosen[-count:] -- rows = [_idatui_head_row(e) for e in chosen] -- first = chosen[0] if chosen else start -- pea = ida_bytes.prev_head(first, lo) -- cursor = {"done": True} if pea == idaapi.BADADDR or pea < lo else {"prev": hex(pea)} -- return {"addr": str(addr), "heads": rows, "cursor": cursor} -- -- # Walk by item END (not next_head): next_head SKIPS undefined bytes, but a -- # flat listing must show them (IDA renders undefined as `db ?` lines, and -- # navigating to an unmarked address must land ON it). Defined items advance -- # by get_item_end; a run of undefined bytes is COLLAPSED into one row (its -- # end found in O(1) via next_head, which skips undefined) so a large .bss or -- # gap doesn't explode into millions of one-byte rows. -- def _is_unknown(e): -- f = ida_bytes.get_flags(e) -- return not (ida_bytes.is_code(f) or ida_bytes.is_data(f)) -- -- def _run_end(e): -- """End (exclusive) of the undefined run starting at ``e``.""" -- nh = ida_bytes.next_head(e, hi) -- return nh if (nh != idaapi.BADADDR and e < nh <= hi) else hi -- -- def _advance(e): -- if _is_unknown(e): -- return _run_end(e) -- nxt = ida_bytes.get_item_end(e) -- return nxt if nxt > e else e + 1 -- -- def _rows_for(e): -- if _is_unknown(e): -- return [_idatui_unknown_row(e, _run_end(e) - e)] -- func = idaapi.get_func(e) if annotate else None -- at_start = func is not None and func.start_ea == e -- out = [] -- if at_start: -- out.extend(_idatui_func_header_rows(e)) -- row = _idatui_head_row(e) -- if at_start: -- row = dict(row) -- row["name"] = None # the name is shown on the proc header line -- elif annotate and row.get("kind") == "code" and row.get("name"): -- # A code label (loc_XXX/jump target) gets its OWN line at depth 0, -- # like IDA; strip it from the instruction row below. -- nm = row["name"] -- out.append({"ea": hex(e), "kind": "label", "size": 0, -- "text": nm + ":", "name": nm}) -- row = dict(row) -- row["name"] = None -- out.append(row) -- if row.get("kind") == "data": -- out.extend(_idatui_struct_member_rows(e)) # expand struct fields -- if func is not None and ida_bytes.get_item_end(e) >= func.end_ea: -- out.extend(_idatui_func_footer_rows(e, func)) -- return out -- -- ea = ida_bytes.get_item_head(start) -- for _ in range(offset): -- if ea >= hi or ea == idaapi.BADADDR: -- break -- ea = _advance(ea) -- more = False -- while ea != idaapi.BADADDR and ea < hi: -- if len(rows) >= count: -- more = True -- break -- rows.extend(_rows_for(ea)) # a struct head expands into member rows -- ea = _advance(ea) -- cursor = {"next": hex(ea)} if more else {"done": True} -- return {"addr": str(addr), "heads": rows, "cursor": cursor} -- -- --@tool --@idasync --def xref_types( -- queries: Annotated[list, "[{addr, direction:'to'|'from'|'both', include_fn, dedup, count}]"], --) -> dict: -- """Like xref_query, but every row carries a fine-grained ``kind`` derived from -- the IDA xref type \u2014 call/jump/flow for code, read/write/offset/text/info for -- data \u2014 alongside the coarse ``type`` (code/data). Feeds the xref dialog's -- r/w/call badges. Same query/envelope shape as xref_query.""" -- import idaapi, idautils, ida_funcs, ida_bytes, ida_xref -- code_kind = {ida_xref.fl_CF: "call", ida_xref.fl_CN: "call", -- ida_xref.fl_JF: "jump", ida_xref.fl_JN: "jump", -- ida_xref.fl_F: "flow"} -- data_kind = {ida_xref.dr_O: "offset", ida_xref.dr_W: "write", -- ida_xref.dr_R: "read", ida_xref.dr_T: "text", ida_xref.dr_I: "info"} -- -- def _kind(xr): -- table = code_kind if xr.iscode else data_kind -- return table.get(xr.type, "code" if xr.iscode else "data") -- -- def _fn(ea): -- f = ida_funcs.get_func(ea) -- if not f: -- return None -- return {"addr": hex(f.start_ea), "name": ida_funcs.get_func_name(f.start_ea)} -- -- def _resolve(raw): -- raw = str(raw).strip() -- try: -- return int(raw, 16) # handles '0x2490' and '2490' -- except ValueError: -- return idaapi.get_name_ea(idaapi.BADADDR, raw) -- -- qs = queries if isinstance(queries, list) else [queries] -- result = [] -- for q in qs: -- q = q if isinstance(q, dict) else {"addr": q} -- raw = str(q.get("addr", "")).strip() -- direction = str(q.get("direction", "to") or "to").lower() -- include_fn = bool(q.get("include_fn", True)) -- dedup = bool(q.get("dedup", True)) -- try: -- count = int(q.get("count", 2000) or 2000) -- except (TypeError, ValueError): -- count = 2000 -- target = _resolve(raw) -- rows = [] -- if target is not None and target != idaapi.BADADDR and ida_bytes.is_mapped(target): -- if direction in ("to", "both"): -- for xr in idautils.XrefsTo(target, 0): -- row = {"direction": "to", "addr": hex(xr.frm), "from": hex(xr.frm), -- "to": hex(target), "type": "code" if xr.iscode else "data", -- "kind": _kind(xr)} -- if include_fn: -- row["fn"] = _fn(xr.frm) -- rows.append(row) -- if direction in ("from", "both"): -- for xr in idautils.XrefsFrom(target, 0): -- row = {"direction": "from", "addr": hex(xr.to), "from": hex(target), -- "to": hex(xr.to), "type": "code" if xr.iscode else "data", -- "kind": _kind(xr)} -- if include_fn: -- row["fn"] = _fn(xr.to) -- rows.append(row) -- if dedup: -- seen = set() -- deduped = [] -- for r in rows: -- k = (r["direction"], r["from"], r["to"], r["kind"]) -- if k in seen: -- continue -- seen.add(k) -- deduped.append(r) -- rows = deduped -- rows = rows[:count] -- result.append({"query": raw, "data": rows, "next_offset": None}) -- return {"result": result} -- -- --@tool --@idasync --def data_type( -- addr: Annotated[str, "Address or name of a data item / global"], --) -> dict: -- """The current C type of a data item, for prefilling a retype prompt: -- {addr, name, type, size, is_func}. ``type`` is empty when the item is -- untyped; ``is_func`` distinguishes a global from a function so the caller -- knows which flavour of set_type to use.""" -- import idaapi -- import ida_bytes -- import ida_name -- import idc -- raw = str(addr).strip() -- try: -- ea = int(raw, 16) -- except ValueError: -- ea = idaapi.get_name_ea(idaapi.BADADDR, raw) -- if ea == idaapi.BADADDR or not ida_bytes.is_mapped(ea): -- return {"addr": raw, "error": f"not a mapped address: {raw}"} -- return { -- "addr": hex(ea), -- "name": ida_name.get_name(ea) or "", -- "type": idc.get_type(ea) or "", -- "size": int(ida_bytes.get_item_size(ea) or 0), -- "is_func": bool(idaapi.get_func(ea)), -- } -- -- --@tool --@idasync --def decomp_map( -- addr: Annotated[str, "Function address or name"], --) -> dict: -- """Per-pseudocode-line instruction coverage for the split view's region -- highlight: for each line, the set of EAs the decompiler attributes to it, -- swept across the line's columns via get_line_item. Shape: -- {addr, lines:[{ea: primary|None, eas:[hex,...]}, ...]}.""" -- import ida_hexrays -- import idaapi -- try: -- ea = int(str(addr), 16) -- except ValueError: -- ea = idaapi.get_name_ea(idaapi.BADADDR, str(addr).strip()) -- func = idaapi.get_func(ea) -- if not func: -- return {"error": f"no function at {addr}"} -- try: -- cfunc = ida_hexrays.decompile(func.start_ea) -- except Exception as e: # noqa: BLE001 -- return {"error": f"decompile failed: {e}"} -- if cfunc is None: -- return {"error": "decompile failed"} -- lines = [] -- for sl in cfunc.get_pseudocode(): -- line = sl.line -- eas, seen = [], set() -- for x in range(len(line) + 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(line, x, False, head, item, tail): -- continue -- # Match the /*ea*/ marker's source (decompile_function_safe): the -- # item's dstr() is 'EA: description'; get_ea() reports a different ea. -- dstr = item.dstr() -- if not dstr: -- continue -- parts = dstr.split(": ", 1) -- if len(parts) != 2: -- continue -- try: -- e = int(parts[0], 16) -- except ValueError: -- continue -- if e not in seen: -- seen.add(e) -- eas.append(hex(e)) -- lines.append({"ea": eas[0] if eas else None, "eas": eas}) -- return {"addr": hex(func.start_ea), "lines": lines} -- -- --_idatui_strings_cache = {} -- -- --def _idatui_build_strings(min_len): -- """[(ea, text, length, typename)] for every string IDA found, cached by -- min_len (rebuilding the list is O(n) and the browser pages through it).""" -- import idautils -- import ida_nalt -- hit = _idatui_strings_cache.get(min_len) -- if hit is not None: -- return hit -- tnames = {} -- for nm, lbl in (("STRTYPE_C", "C"), ("STRTYPE_C_16", "utf16"), -- ("STRTYPE_C_32", "utf32"), ("STRTYPE_PASCAL", "pascal")): -- v = getattr(ida_nalt, nm, None) -- if v is not None: -- tnames[v & 0xFF] = lbl -- items = [] -- for s in idautils.Strings(): -- if s is None: -- continue -- try: -- text = str(s) -- except Exception: # noqa: BLE001 -- undecodable literal -- continue -- if len(text) < min_len: -- continue -- st = getattr(s, "strtype", 0) & 0xFF -- items.append((s.ea, text, getattr(s, "length", len(text)), -- tnames.get(st, "t%d" % st))) -- _idatui_strings_cache[min_len] = items -- return items -- -- --@tool --@idasync --def list_strings( -- offset: Annotated[int, "Start index into the strings list"] = 0, -- count: Annotated[int, "Max strings to return (page size)"] = 2000, -- min_len: Annotated[int, "Minimum string length to include"] = 4, -- refresh: Annotated[bool, "Rebuild the cached strings list"] = False, --) -> dict: -- """Every string literal IDA found in the binary (IDA's Shift+F12 window), -- paginated: {strings:[{addr,text,len,type}], total, next_offset}. Feeds the -- TUI's strings browser.""" -- try: -- min_len = max(int(min_len), 1) -- except (TypeError, ValueError): -- min_len = 4 -- try: -- offset = max(int(offset), 0) -- except (TypeError, ValueError): -- offset = 0 -- try: -- count = max(int(count), 1) -- except (TypeError, ValueError): -- count = 2000 -- if refresh: -- _idatui_strings_cache.pop(min_len, None) -- items = _idatui_build_strings(min_len) -- page = items[offset:offset + count] -- return { -- "strings": [{"addr": hex(ea), "text": text, "len": ln, "type": ty} -- for (ea, text, ln, ty) in page], -- "total": len(items), -- "next_offset": offset + len(page), -- } -- --@tool --@idasync --def list_linkage( -- kind: Annotated[str, "'import', 'export' or 'both'"] = "both", --) -> dict: -- """What this binary imports from, and exports to, other modules: -- {imports:[{addr,name,module}], exports:[{addr,name,ordinal}]}. Feeds the -- project-wide import/export join, which resolves a PLT stub in one binary to -- the real implementation in another.""" -- import idaapi -- import idautils -- import ida_nalt -- want = str(kind or "both").lower() -- imports = [] -- exports = [] -- if want in ("import", "both"): -- n = ida_nalt.get_import_module_qty() -- for i in range(n): -- mod = ida_nalt.get_import_module_name(i) or "" -- -- def _cb(ea, name, ordinal, _mod=mod): -- # An ordinal-only import has no name; skip rather than invent one. -- if name: -- imports.append({"addr": hex(ea), "name": name, "module": _mod}) -- return True -- -- ida_nalt.enum_import_names(i, _cb) -- if want in ("export", "both"): -- for index, ordinal, ea, name in idautils.Entries(): -- if name: -- exports.append({"addr": hex(ea), "name": name, -- "ordinal": int(ordinal)}) -- return {"imports": imports, "exports": exports, -- "n_imports": len(imports), "n_exports": len(exports)} -- --@tool --@idasync --def define_code_run( -- addr: Annotated[str, "Address to start disassembling from"], -- limit: Annotated[int, "Max instructions to create (safety stop)"] = 20000, --) -> dict: -- """Disassemble CONSECUTIVELY from ``addr`` until something stops it, the way -- IDA's 'c' does — one instruction is rarely what you want when carving a raw -- image. Returns {start,end,count,stopped} where ``stopped`` says why: -- 'undecodable' (bytes aren't an instruction), 'flow' (the last instruction -- doesn't fall through, e.g. RET/B), 'defined' (ran into existing code/data), -- 'segment' (hit the end) or 'limit'. -- -- Runs in-process: doing this from the client would be one round trip per -- instruction, which is minutes on a real firmware image.""" -- import ida_bytes -- import ida_idp -- import ida_segment -- import ida_ua -- import idaapi -- -- try: -- ea = parse_address(addr) -- except Exception as e: -- return {"addr": str(addr), "error": str(e), "count": 0} -- -- seg = ida_segment.getseg(ea) -- if not seg: -- return {"addr": str(addr), "error": "no segment", "count": 0} -- hi = seg.end_ea -- try: -- limit = max(1, min(int(limit), 200000)) -- except (TypeError, ValueError): -- limit = 20000 -- -- start, count, stopped = ea, 0, "limit" -- 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): -- # Already defined: stop rather than clobber. Undefining someone's -- # existing work to keep a speculative run going is not a trade the -- # user asked for. -- stopped = "defined" -- break -- n = ida_ua.create_insn(ea) -- if n <= 0: -- stopped = "undecodable" -- break -- count += 1 -- # Stop where control flow stops. Past a RET the next bytes are usually -- # padding or a new function's data, and running on turns a clean carve -- # into a mess that has to be undone by hand. -- # -- # Ask ida_idp.is_ret_insn, NOT the canonical feature bits: on AArch64 -- # get_canon_feature() returns 0 for RET, so a CF_STOP test silently never -- # fires and the run walks straight through the end of the routine. -- insn = ida_ua.insn_t() -- if ida_ua.decode_insn(insn, ea) > 0: -- try: -- is_ret = ida_idp.is_ret_insn(insn) -- except Exception: -- is_ret = False -- if is_ret or (insn.get_canon_feature() & idaapi.CF_STOP): -- ea += n -- stopped = "flow" -- break -- ea += n -- -- return {"start": hex(start), "end": hex(ea), "count": count, -- "stopped": stopped} -- --@tool --@idasync --def set_thumb( -- addr: Annotated[str, "Address to change the ARM decoding mode at"], -- mode: Annotated[str, "'toggle', 'on' (Thumb) or 'off' (ARM)"] = "toggle", -- end: Annotated[str, "Optional exclusive end address (default: this item)"] = "", --) -> dict: -- """Switch ARM/Thumb decoding at ``addr`` (IDA's T segment register). -- -- Thumb is not a property of the bytes, it's a mode the CPU is in, so a raw -- image gives IDA no way to know: at a Thumb entry point it decodes 16-bit -- instructions as 32-bit ARM and produces confident nonsense -- (``push {r3,lr}`` reads as ``SVCLT 0xBF00``). -- -- Also forces the segment to 32-bit when turning Thumb ON. Thumb does not -- exist in AArch64, and a headerless blob loaded with -parm defaults to -- 64-bit — so setting T alone changes nothing and looks broken. Asking for -- Thumb IS asking for ARM32.""" -- import ida_bytes -- import ida_idp -- import ida_segment -- import ida_segregs -- -- try: -- ea = parse_address(addr) -- except Exception as e: -- return {"addr": str(addr), "error": str(e)} -- treg = ida_idp.str2reg("T") -- if treg is None or treg < 0: -- return {"addr": hex(ea), "error": "no T register (not an ARM database)"} -- seg = ida_segment.getseg(ea) -- if not seg: -- return {"addr": hex(ea), "error": "no segment"} -- -- import ida_ida -- db64 = ida_ida.inf_get_app_bitness() == 64 -- cur = ida_segregs.get_sreg(ea, treg) -- cur = 0 if cur in (None, 0xFFFFFFFF, -1) else int(cur) -- want = {"on": 1, "off": 0}.get(str(mode).lower(), 0 if cur else 1) -- -- changed_bits = False -- if want and seg.bitness != 1: -- ida_segment.set_segm_addressing(seg, 1) -- changed_bits = True -- -- try: -- stop = parse_address(end) if end else 0 -- except Exception: -- stop = 0 -- size = max(int(stop) - ea, 0) or max(ida_bytes.get_item_size(ea), 2) -- # The bytes are currently decoded in the OLD mode; leaving that item defined -- # pins the wrong instruction length and the new mode has nothing to apply to. -- 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) -- return {"addr": hex(ea), "thumb": bool(now), "was": bool(cur), "ok": ok, -- "bitness": ida_segment.getseg(ea).bitness, -- "forced_32bit": changed_bits, -- # The DATABASE's bitness is fixed at load and can't be corrected -- # here (setting it post-hoc makes the decompiler INTERR). In a -- # 64-bit database a 32-bit function disassembles but Hex-Rays -- # refuses it outright, so say so instead of leaving the user to -- # discover that F5 does nothing. -- "db_64bit": bool(db64 and want)} -- --def _idatui_add_func(ea): -- """add_func at ``ea``, falling back to an explicit end. -- -- ida_funcs.add_func(ea) asks IDA to find the end and on carved or -- freshly-marked code it often can't, failing with no reason given.""" -- import ida_bytes -- import ida_funcs -- import ida_segment -- import idaapi -- -- if idaapi.get_func(ea) is not None: -- return True -- if ida_funcs.add_func(ea): -- return True -- seg = ida_segment.getseg(ea) -- hi = seg.end_ea if seg else ea -- end = ea -- while end < hi and ida_bytes.is_code(ida_bytes.get_flags(end)): -- nxt = ida_bytes.get_item_end(end) -- if nxt <= end: -- break -- end = nxt -- return bool(end > ea and ida_funcs.add_func(ea, end)) -- -- --@tool --@idasync --def define_func_run( -- addr: Annotated[str, "Entry point of the function to create"], --) -> dict: -- """Create a function at ``addr``, working out its end if IDA can't. -- -- ida_funcs.add_func(ea) asks IDA to find the end itself, and on hand-carved -- code it often can't — a run that ends in a tail call, or whose last -- instruction isn't recognised as a return, simply fails with no reason given. -- You then have a disassembled routine that refuses to become a function, and -- F5 has nothing to work with. -- -- So: try IDA's way, and if that fails, use the end of the contiguous -- instruction run starting at ``addr``.""" -- import ida_bytes -- import ida_funcs -- import ida_segment -- import idaapi -- -- try: -- ea = parse_address(addr) -- except Exception as e: -- return {"addr": str(addr), "error": str(e), "ok": False} -- fn = idaapi.get_func(ea) -- if fn is not None and fn.start_ea == ea: -- return {"addr": hex(ea), "ok": True, "start": hex(fn.start_ea), -- "end": hex(fn.end_ea), "how": "existed"} -- auto = ida_funcs.add_func(ea) -- if not auto and not _idatui_add_func(ea): -- return {"addr": hex(ea), "ok": False, -- "error": f"IDA refused a function at {ea:#x}"} -- f = idaapi.get_func(ea) -- if f is None: -- return {"addr": hex(ea), "ok": False, "error": "function did not stick"} -- return {"addr": hex(ea), "ok": True, "start": hex(f.start_ea), -- "end": hex(f.end_ea), "how": "auto" if auto else "explicit-end"} -- --@tool --@idasync --def decomp_error( -- addr: Annotated[str, "Address of the function that failed to decompile"], --) -> dict: -- """Why Hex-Rays refused this function, in its own words. -- -- The plain decompile tool reports "Decompilation failed at 0x0" and drops the -- reason, which is the only useful part. Hex-Rays fills in a hexrays_failure_t -- saying things like "only 64-bit functions can be decompiled in the current -- database" — that one is unfixable in place (the database's bitness is set at -- load), so a user who can't see it has no way to know they must reload.""" -- import ida_funcs -- import ida_hexrays -- import ida_ida -- -- try: -- ea = parse_address(addr) -- except Exception as e: -- return {"addr": str(addr), "error": str(e)} -- out = {"addr": hex(ea), "bitness": ida_ida.inf_get_app_bitness()} -- fn = ida_funcs.get_func(ea) -- if fn is None: -- out["reason"] = "no function here" -- return out -- try: -- if not ida_hexrays.init_hexrays_plugin(): -- out["reason"] = "the decompiler is not available for this processor" -- return out -- hf = ida_hexrays.hexrays_failure_t() -- cf = ida_hexrays.decompile_func(fn, hf) -- if cf is not None: -- out["reason"] = "" # it decompiles now -- return out -- out["reason"] = hf.desc() or f"error {hf.code}" -- out["code"] = int(hf.code) -- out["errea"] = hex(hf.errea) -- except Exception as e: # noqa: BLE001 -- out["reason"] = f"{type(e).__name__}: {e}" -- return out -- --@tool --@idasync --def thumb_scan( -- start: Annotated[str, "Start of the range to scan for entry pointers"] = "", -- end: Annotated[str, "Exclusive end of the range (default: 1KB from start)"] = "", -- apply: Annotated[bool, "Mark the targets as Thumb and disassemble them"] = True, -- limit: Annotated[int, "Max entries to act on"] = 512, --) -> dict: -- """Find Thumb entry points from ODD pointers, e.g. a Cortex-M vector table. -- -- An ARM function pointer carries the mode in bit 0: odd means Thumb. A vector -- table is therefore a list of Thumb entry points that IDA won't follow on a -- headerless image, because nothing tells it those words are pointers at all. -- -- Being wrong here is expensive — marking a data word as code corrupts the -- listing — so a word only counts when it is odd, lands inside a loaded -- segment, and its target is EXECUTABLE and not already defined as data. The -- even words in a vector table (the initial stack pointer) fail the first test, -- which is the point.""" -- import ida_bytes -- import ida_funcs -- import ida_idp -- import ida_segment -- import ida_segregs -- import ida_ua -- -- seg0 = ida_segment.getseg(parse_address(start)) if start else None -- if seg0 is None: -- seg0 = ida_segment.getnseg(0) -- if seg0 is None: -- return {"error": "no segments", "found": [], "applied": 0} -- try: -- lo = parse_address(start) if start else seg0.start_ea -- hi = parse_address(end) if end else min(lo + 0x400, seg0.end_ea) -- except Exception as e: -- return {"error": str(e), "found": [], "applied": 0} -- -- treg = ida_idp.str2reg("T") -- found, applied = [], 0 -- ea = lo -- while ea + 4 <= hi and len(found) < limit: -- w = ida_bytes.get_dword(ea) -- ea += 4 -- if not (w & 1): -- continue # even: not a Thumb pointer -- tgt = w & ~1 -- seg = ida_segment.getseg(tgt) -- if seg is None or not (seg.perm & ida_segment.SEGPERM_EXEC or seg.perm == 0): -- continue # points outside the image, or at data -- f = ida_bytes.get_flags(tgt) -- if ida_bytes.is_data(f): -- continue # already something else; don't fight it -- rec = {"at": hex(ea - 4), "value": hex(w), "target": hex(tgt), -- "was_code": bool(ida_bytes.is_code(f))} -- found.append(rec) -- if not apply: -- continue -- if treg is not None and treg >= 0: -- ida_segregs.split_sreg_range(tgt, treg, 1, ida_segregs.SR_user) -- if not ida_bytes.is_code(ida_bytes.get_flags(tgt)): -- ida_bytes.del_items(tgt, 0, 2) -- if ida_ua.create_insn(tgt) <= 0: -- rec["decoded"] = False -- continue -- rec["decoded"] = True -- rec["function"] = _idatui_add_func(tgt) -- applied += 1 -- return {"start": hex(lo), "end": hex(hi), "found": found, -- "applied": applied, "n": len(found)} --''' -- --SNIPPET = f"{BEGIN}\n{BODY.strip()}\n{END}\n" -- -- --def api_types_path() -> pathlib.Path | None: -- """Locate ida_pro_mcp/ida_mcp/api_types.py without importing it (importing the -- submodule would pull in IDA, which isn't available outside a worker).""" -- spec = importlib.util.find_spec("ida_pro_mcp") # top-level pkg is IDA-free -- if spec is None or not spec.submodule_search_locations: -- return None -- p = pathlib.Path(spec.submodule_search_locations[0]) / "ida_mcp" / "api_types.py" -- return p if p.exists() else None -- -- --def main() -> int: -- path = api_types_path() -- if path is None: -- print("idatui: ida_pro_mcp not found; skipping tool injection", file=sys.stderr) -- return 0 -- text = path.read_text() -- if BEGIN in text and END in text: # replace the existing block in place -- pre = text[: text.index(BEGIN)].rstrip() -- post = text[text.index(END) + len(END):].lstrip("\n") -- new = pre + "\n\n" + SNIPPET + ("\n" + post if post else "") -- else: -- new = text.rstrip() + "\n\n" + SNIPPET -- if new == text: -- return 0 -- try: -- path.write_text(new) -- except OSError as e: -- print(f"idatui: could not patch {path}: {e}", file=sys.stderr) -- return 1 -- print(f"idatui: injected/updated idatui-ext tools in {path}", file=sys.stderr) -- return 0 -- -- --if __name__ == "__main__": -- raise SystemExit(main()) -diff --git a/tests/test_codemode_client.py b/tests/test_codemode_client.py -new file mode 100644 -index 0000000..6303eb9 ---- /dev/null -+++ b/tests/test_codemode_client.py -@@ -0,0 +1,137 @@ -+"""IDA-free contract tests for the Code Mode client adapter.""" -+from __future__ import annotations -+ -+import os -+import sys -+import tempfile -+from dataclasses import dataclass -+ -+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) -+ -+import idatui.codemode_client as module # noqa: E402 -+from idatui.codemode_client import CodeModeClient, _parse_load_args # noqa: E402 -+from idatui.errors import IDAToolError # noqa: E402 -+ -+PASS = FAIL = 0 -+ -+ -+def check(name: str, condition: bool, detail="") -> None: -+ global PASS, FAIL -+ if condition: -+ PASS += 1 -+ print(f" ok {name}") -+ else: -+ FAIL += 1 -+ print(f" FAIL {name} {detail}") -+ -+ -+@dataclass(frozen=True) -+class FakeEntry: -+ pid: int = 123 -+ backend: str = "gui" -+ record_id: str = "123-abcdef" -+ exe_path: str = "" -+ idb_path: str = "" -+ -+ -+class FakeHandle: -+ def __init__(self, path: str) -> None: -+ self.connected = True -+ self.entry = FakeEntry(exe_path=path, idb_path=path + ".i64") -+ self.waited = None -+ self.saved = 0 -+ self.closed = False -+ self.code = "" -+ self.code_timeout = None -+ -+ def wait_autoanalysis(self, timeout=None): -+ self.waited = timeout -+ return {"complete": True, "status": "complete"} -+ -+ def execute_python(self, code, timeout=None): -+ self.code = code -+ self.code_timeout = timeout -+ return {"result": {"sentinel": 7}, "stdout": "", "stderr": ""} -+ -+ def save_database(self): -+ self.saved += 1 -+ return {"saved": True, "idb_path": self.entry.idb_path} -+ -+ def close(self): -+ self.connected = False -+ self.closed = True -+ -+ -+class FakeDatabaseHandle: -+ opened = None -+ kwargs = None -+ -+ @classmethod -+ def open(cls, path, **kwargs): -+ cls.opened = path -+ cls.kwargs = kwargs -+ return FakeHandle(path) -+ -+ -+def main() -> int: -+ proc, base, file_type = _parse_load_args("-parm:ARMv7-M -b800000 -TRaw") -+ check("legacy switches map to typed Code Mode options", -+ (proc, base, file_type) == ("arm:ARMv7-M", 0x8000000, "Raw"), -+ (proc, base, file_type)) -+ try: -+ _parse_load_args("-parm -zcustom") -+ except ValueError as exc: -+ check("arbitrary IDA switches fail loudly", "cannot represent" in str(exc), exc) -+ else: -+ check("arbitrary IDA switches fail loudly", False) -+ -+ original = module.DatabaseHandle -+ module.DatabaseHandle = FakeDatabaseHandle -+ try: -+ with tempfile.TemporaryDirectory() as tmp: -+ path = os.path.join(tmp, "sample.bin") -+ with open(path, "wb") as file: -+ file.write(b"sample") -+ client = CodeModeClient(path, load_args="-parm:ARMv7-A -b100") -+ notes = [] -+ client.connect(timeout=42, progress=notes.append) -+ handle = client._handle -+ check("connect delegates database discovery to DatabaseHandle.open", -+ FakeDatabaseHandle.opened == path and handle is not None) -+ check("typed loader options cross the dependency boundary", -+ FakeDatabaseHandle.kwargs["processor"] == "arm:ARMv7-A" -+ and FakeDatabaseHandle.kwargs["loading_address"] == 0x1000, -+ FakeDatabaseHandle.kwargs) -+ check("connect waits for Code Mode autoanalysis", -+ handle.waited == 42, getattr(handle, "waited", None)) -+ check("progress distinguishes discovery and backend attachment", -+ len(notes) == 2 and "gui" in notes[-1], notes) -+ result = client.invoke("list_funcs", queries=[{"offset": 0, "count": 2}]) -+ check("invoke returns execute_python's result", result == {"sentinel": 7}, result) -+ check("operation scripts use the preloaded ida-domain database", -+ "db.functions.get_all()" in handle.code, handle.code[:200]) -+ check("health exposes registry identity", -+ client.health()["record_id"] == "123-abcdef") -+ client.save_database() -+ check("save uses the public Code Mode save route", handle.saved == 1) -+ client.close() -+ check("close releases only the handle lease", handle.closed) -+ check("GUI lifetime is never claimed by the client", -+ client.wait_released(0) is False) -+ finally: -+ module.DatabaseHandle = original -+ -+ client = CodeModeClient(__file__) -+ try: -+ client.invoke("not-an-operation") -+ except IDAToolError as exc: -+ check("unknown adapter operations are explicit", exc.tool == "not-an-operation") -+ else: -+ check("unknown adapter operations are explicit", False) -+ -+ print(f"\n{PASS} passed, {FAIL} failed") -+ return 1 if FAIL else 0 -+ -+ -+if __name__ == "__main__": -+ raise SystemExit(main()) -diff --git a/tests/test_pool.py b/tests/test_pool.py -index 5c6e2c4..ff0c016 100644 ---- a/tests/test_pool.py -+++ b/tests/test_pool.py -@@ -1,8 +1,7 @@ - #!/usr/bin/env python3 --"""Unit tests for idatui.pool (worker residency: LRU + memory budget). -+"""Unit tests for idatui.pool (Code Mode lease residency and LRU budget). - --Pure stdlib with a fake client injected, so the eviction policy is testable --without spawning real idalib workers. -+A fake client keeps the policy testable without IDA or Textual. - - python tests/test_pool.py - """ -@@ -11,7 +10,7 @@ import sys - import tempfile - - sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) --from idatui.pool import WorkerPool # noqa: E402 -+from idatui.pool import DatabasePool # noqa: E402 - from idatui.project import Project # noqa: E402 - - PASS = FAIL = 0 -@@ -28,11 +27,12 @@ def check(name, cond, detail=""): - - - class FakeClient: -- """Stands in for a WorkerClient: records saves/closes, reports fixed memory.""" -+ """Stands in for a CodeModeClient lease and records saves/closes.""" - -- def __init__(self, ref, mem=100): -+ def __init__(self, ref, mem=100, backend="idalib"): - self.ref = ref - self.mem = mem -+ self.backend = backend - self.saved = 0 - self.closed = False - self.connected = False -@@ -41,10 +41,9 @@ class FakeClient: - self.connected = True - return self - -- def call(self, tool, **kw): -- if tool == "idb_save": -- self.saved += 1 -- return {} -+ def save_database(self): -+ self.saved += 1 -+ return {"saved": True} - - def close(self, grace=None): - self.closed = True -@@ -72,15 +71,15 @@ def main() -> int: - made[ref.label] = c - return c - -- pool = WorkerPool(proj, budget_mb=350, spawn=spawn, -+ pool = DatabasePool(proj, budget_mb=350, spawn=spawn, - mem_fn=lambda c: c.mem) - - # -- lazy spawn + reuse -------------------------------------------- # - a = pool.get("bin0") -- check("get() spawns a worker on first use", a is made["bin0"] and a.connected) -+ check("get() spawns a database lease on first use", a is made["bin0"] and a.connected) - check("get() stages the binary first", - os.path.isfile(proj.by_label("bin0").staged)) -- check("get() reuses the resident worker", pool.get("bin0") is a) -+ check("get() reuses the resident lease", pool.get("bin0") is a) - check("resident() reports it", pool.resident() == ["bin0"], pool.resident()) - - # -- LRU ordering ---------------------------------------------------- # -@@ -96,9 +95,9 @@ def main() -> int: - check("exceeding the budget evicts the least-recently-used", - pool.evicted == ["bin1"] and not pool.is_resident("bin1"), - f"evicted={pool.evicted} resident={pool.resident()}") -- check("the just-spawned worker is never the victim", pool.is_resident("bin3")) -+ check("the just-attached lease is never the victim", pool.is_resident("bin3")) - check("eviction saves the database first", made["bin1"].saved == 1) -- check("eviction closes the worker", made["bin1"].closed) -+ check("eviction closes the lease", made["bin1"].closed) - check("pool is back within budget", pool.memory_mb() <= pool.budget_mb, - f"{pool.memory_mb()}/{pool.budget_mb}") - -@@ -112,7 +111,7 @@ def main() -> int: - - # -- pinning ---------------------------------------------------------- # - pool.close_all() -- pool2 = WorkerPool(proj, budget_mb=250, spawn=spawn, mem_fn=lambda c: c.mem) -+ pool2 = DatabasePool(proj, budget_mb=250, spawn=spawn, mem_fn=lambda c: c.mem) - pool2.get("bin0") - pool2.pin("bin0") - pool2.get("bin1") -@@ -139,7 +138,7 @@ def main() -> int: - - # -- teardown ----------------------------------------------------------- # - pool2.close_all() -- check("close_all() closes every worker", -+ check("close_all() closes every lease", - not pool2.resident() and all(c.closed for c in made.values())) - check("close_all() clears the active binary", pool2.active is None) - -@@ -151,8 +150,8 @@ def main() -> int: - check("an unknown label raises KeyError", True) - - # -- default budget comes from the project's memory_pct ------------------- # -- pool3 = WorkerPool(proj, spawn=spawn, mem_fn=lambda c: c.mem) -- check("default budget is derived, not a fixed worker count", -+ pool3 = DatabasePool(proj, spawn=spawn, mem_fn=lambda c: c.mem) -+ check("default budget is derived, not a fixed lease count", - pool3.budget_mb >= 256, pool3.budget_mb) - - # -- prewarm: speculative, and never at the cost of a real binary ------ # -@@ -165,7 +164,7 @@ def main() -> int: - made2[ref.label] = c - return c - -- pool = WorkerPool(proj, budget_mb=250, spawn=spawn2, -+ pool = DatabasePool(proj, budget_mb=250, spawn=spawn2, - mem_fn=lambda c: c.mem) - labels = [r.label for r in proj.refs] - a, b, c_ = labels[0], labels[1], labels[2] -@@ -184,6 +183,28 @@ def main() -> int: - check("prewarm ignores a label outside the project", - pool.prewarm("nope") is False) - -+ # Budget eviction releases GUI leases but must not save somebody's open IDA -+ # implicitly. An explicit save-and-close remains authoritative. -+ with tempfile.TemporaryDirectory() as tmp: -+ proj = _mkproject(tmp, n=1) -+ made_gui = [] -+ -+ def spawn_gui(ref, ttl): -+ client = FakeClient(ref, backend="gui") -+ made_gui.append(client) -+ return client -+ -+ pool = DatabasePool(proj, spawn=spawn_gui, mem_fn=lambda c: c.mem) -+ label = proj.refs[0].label -+ pool.get(label) -+ pool.evict(label) -+ check("LRU release does not implicitly save a GUI database", -+ made_gui[-1].saved == 0) -+ pool.get(label) -+ pool.close_all(save=True) -+ check("explicit close_all(save=True) does save a GUI database", -+ made_gui[-1].saved == 1) -+ - print(f"\n{PASS} passed, {FAIL} failed") - return 1 if FAIL else 0 - -diff --git a/tests/test_project.py b/tests/test_project.py -index 91fd250..7c0dcba 100644 ---- a/tests/test_project.py -+++ b/tests/test_project.py -@@ -1,7 +1,7 @@ - #!/usr/bin/env python3 - """Unit tests for idatui.project (the multi-binary project model + staging). - --Pure stdlib: no IDA, no textual, no worker — runs anywhere in under a second. -+IDA-free: exercises staging plus Code Mode ownership checks without opening a database. - - python tests/test_project.py - """ -diff --git a/tests/test_scenarios.py b/tests/test_scenarios.py -index 9d77680..ab510f4 100644 ---- a/tests/test_scenarios.py -+++ b/tests/test_scenarios.py -@@ -2316,7 +2316,13 @@ async def _build_pristine(binary, cache): - await pilot.pause(0.05) - if app._func_index is not None and app._func_index.complete: - break -- app.program.client.call("idb_save", timeout=600.0) -+ app.program.client.save_database() -+ # Textual's headless run_test context does not reliably emit App.Unmount on -+ # every platform/version; release the Code Mode lease explicitly. -+ if app.program is not None: -+ app.program.close() -+ if app.client is not None: -+ app.client.close() - db = binary + ".i64" - if os.path.exists(db): - shutil.copy2(db, cache) -@@ -2346,7 +2352,7 @@ async def run(binary, only=None): - - - async def _run_on(binary, only=None): -- # Own idalib worker: opens the binary in-process over a unix socket. -+ # Code Mode attaches a registered GUI or starts/reuses a managed worker. - app = IdaTui(open_path=binary, keepalive=False) - async with app.run_test(size=(140, 44)) as pilot: - c = Ctx(app, pilot) -@@ -2366,6 +2372,14 @@ async def _run_on(binary, only=None): - print(f"── {name} ({asyncio.get_event_loop().time() - _t0:.1f}s) CRASHED") - c.check("scenario did not crash", False, f"{type(e).__name__}: {e}") - traceback.print_exc() -+ # Headless run_test does not reliably emit App.Unmount; explicitly release -+ # the lease. Then wait through the managed worker's final-lease grace and -+ # IDB close so Windows can remove this suite's TemporaryDirectory safely. -+ if app.program is not None: -+ app.program.close() -+ if app.client is not None: -+ app.client.close() -+ await asyncio.to_thread(app.client.wait_released, 45.0) - - - def main(argv): -diff --git a/uv.lock b/uv.lock -index 91414b5..392ebf0 100644 ---- a/uv.lock -+++ b/uv.lock -@@ -11,27 +11,69 @@ wheels = [ - { url = "https://files.pythonhosted.org/packages/d1/d6/3965ed04c63042e047cb6a3e6ed1a63a35087b6a609aa3a15ed8ac56c221/colorama-0.4.6-py2.py3-none-any.whl", hash = "sha256:4f1d9991f5acc0ca119f9d443620b77f9d6b33703e51011c16baf57afb285fc6", size = 25335, upload-time = "2022-10-25T02:36:20.889Z" }, - ] - -+[[package]] -+name = "ida-codemode-mcp" -+version = "0.2.0" -+source = { editable = "../ida-codemode-mcp" } -+dependencies = [ -+ { name = "ida-domain" }, -+ { name = "zeromcp" }, -+] -+ -+[package.metadata] -+requires-dist = [ -+ { name = "ida-domain", git = "https://github.com/HexRaysSA/ida-domain?branch=main" }, -+ { name = "zeromcp", specifier = ">=1.5.0" }, -+] -+ -+[package.metadata.requires-dev] -+dev = [ -+ { name = "pytest", specifier = ">=9.0.3" }, -+ { name = "ruff", specifier = ">=0.12.0" }, -+] -+ -+[[package]] -+name = "ida-domain" -+version = "0.5.1.dev1" -+source = { git = "https://github.com/HexRaysSA/ida-domain?branch=main#8f36bbce94f0dd55e4ad5f7c8b5f0ef59b9c557a" } -+dependencies = [ -+ { name = "idapro" }, -+ { name = "packaging" }, -+ { name = "typing-extensions" }, -+] -+ -+[[package]] -+name = "idapro" -+version = "0.0.10" -+source = { registry = "https://pypi.org/simple" } -+sdist = { url = "https://files.pythonhosted.org/packages/f8/75/249c605cc144a6b3778c48381d31ff9242f3e0b7ae23a9ca9c27224e641a/idapro-0.0.10.tar.gz", hash = "sha256:417c03c4605d18417e470f6a748e397b39d6d5829ebd3bbdedd92ff5b9092d11", size = 1060989, upload-time = "2026-07-15T12:55:22.313Z" } -+wheels = [ -+ { url = "https://files.pythonhosted.org/packages/e8/83/7b02832cc8b057f686cccdb771fe80282f801a9b798961f8070bb468c73c/idapro-0.0.10-py3-none-any.whl", hash = "sha256:43f227953a0e348ced21c050d277b7ce34103e2ce05fc739b7d8c186ef0e1542", size = 2194897, upload-time = "2026-07-15T12:55:20.88Z" }, -+] -+ - [[package]] - name = "idatui" - version = "0.0.1" - source = { editable = "." } -+dependencies = [ -+ { name = "ida-codemode-mcp" }, -+ { name = "pygments" }, -+ { name = "textual" }, -+] - - [package.optional-dependencies] - dev = [ - { name = "pytest" }, - ] --tui = [ -- { name = "pygments" }, -- { name = "textual" }, --] - - [package.metadata] - requires-dist = [ -- { name = "pygments", marker = "extra == 'tui'", specifier = ">=2" }, -+ { name = "ida-codemode-mcp", editable = "../ida-codemode-mcp" }, -+ { name = "pygments", specifier = ">=2" }, - { name = "pytest", marker = "extra == 'dev'", specifier = ">=8" }, -- { name = "textual", marker = "extra == 'tui'", specifier = ">=8" }, -+ { name = "textual", specifier = ">=8" }, - ] --provides-extras = ["tui", "dev"] -+provides-extras = ["dev"] - - [[package]] - name = "iniconfig" -@@ -191,3 +233,12 @@ sdist = { url = "https://files.pythonhosted.org/packages/78/67/9a363818028526e2d - wheels = [ - { url = "https://files.pythonhosted.org/packages/61/73/d21edf5b204d1467e06500080a50f79d49ef2b997c79123a536d4a17d97c/uc_micro_py-2.0.0-py3-none-any.whl", hash = "sha256:3603a3859af53e5a39bc7677713c78ea6589ff188d70f4fee165db88e22b242c", size = 6383, upload-time = "2026-03-01T06:31:26.257Z" }, - ] -+ -+[[package]] -+name = "zeromcp" -+version = "1.5.0" -+source = { registry = "https://pypi.org/simple" } -+sdist = { url = "https://files.pythonhosted.org/packages/95/10/0c5018221766413c808b62f229a3b6b2cd0e4b10bc9ac25fee6152c22938/zeromcp-1.5.0.tar.gz", hash = "sha256:ef4e590ddb20a30a2ceaee86dbf893c9edb5d3e583c22a0ea7025e94763e59d2", size = 95257, upload-time = "2026-07-22T13:39:29.535Z" } -+wheels = [ -+ { url = "https://files.pythonhosted.org/packages/8c/46/aa0e0941b511969a3eb70ea19add43b22c7336a4bf65a4fe4ea2176faf25/zeromcp-1.5.0-py3-none-any.whl", hash = "sha256:ca3b67687850ed463a255c180a286901ea69343612dc84ef279ec75b460f77ee", size = 21875, upload-time = "2026-07-22T13:39:28.44Z" }, -+] --- -2.53.0.windows.2 - |