diff options
Diffstat (limited to '')
34 files changed, 7999 insertions, 3992 deletions
@@ -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. - The startup splash draws the **real logo image** on terminals that speak the kitty graphics protocol (~10× the resolution of the block art), and falls back to `logo.ans` everywhere else. Support is detected by *asking the terminal*, @@ -77,46 +77,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 @@ -130,15 +142,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 @@ -195,7 +208,9 @@ See `docs/RPC.md` for the full protocol. ## Tests -`tests/run.py` is the front door — it runs every suite and prints one table: +`tests/run.py` is the front door — it runs every suite and prints one table. +The live suites attach through Code Mode (a registered GUI database, or a managed +idalib worker started on demand) for the given binary: ```sh python3 tests/run.py --fast # 257 checks, ~0.5s, any python3 — between edits @@ -227,6 +242,7 @@ The individual suites still run standalone, which is how you iterate on one: ## 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 @@ -1,6 +1,16 @@ RESTART. TODO: +[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] add support for toggling literal types, ala `o` in IDA. (decimal to hex to reference etc.) -> `o` / `O` on either code view cycle the literal under the cursor; 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). +Code Mode returns the complete execution result directly; ida-tui no longer +needs MCP structured-content/download-URL recovery for large pseudocode bodies. -## Worker lifecycle (idatui's own idalib worker) +## Code Mode lifecycle -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: +`CodeModeClient` owns an authenticated SSE lease on a registered database: -* **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. +* 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) +"""Exercise the real domain.Program through an IDA Code Mode lease. -# --- drive the REAL domain layer through the worker (read path) --- -main=prog.resolve("main") -print("resolve('main') =", hex(main), flush=True) - -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) +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 -b=prog.read_bytes(main, 16) -print("read_bytes(main,16) ->", b.hex(), flush=True) +import os +import sys +import time -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) +from idatui.codemode_client import CodeModeClient +from idatui.domain import Program -d=prog.decompile(main) -print("decompile(main) -> failed?", d.failed, "lines:", len((d.code or '').splitlines()), flush=True) -regs=prog.file_regions() -print("file_regions ->", len(regs), "segments", flush=True) +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() -# 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() +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/ida-codemode-mcp.patch b/ida-codemode-mcp.patch new file mode 100644 index 0000000..82a9b3b --- /dev/null +++ b/ida-codemode-mcp.patch @@ -0,0 +1,5848 @@ +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 + @@ -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 4a58b16..a4b03a1 100644 --- a/idatui/app.py +++ b/idatui/app.py @@ -10,8 +10,8 @@ Design notes: without ever materializing 52k lines in a widget. * All network/domain work runs in Textual worker threads; the UI never blocks. * An address-history stack backs Enter (follow) / Esc (back), IDA-style. -* On startup we bump the worker idle-TTL and run a keepalive heartbeat so the - session never gets reaped while we chill. +* Database lifecycle is lease-based through ida_codemode: matching GUI sessions + are reused, otherwise a shared managed idalib worker is opened on demand. """ from __future__ import annotations @@ -32,7 +32,7 @@ from textual import work from textual.app import App, ComposeResult from textual.binding import Binding from textual.command import DiscoveryHit, Hit, Provider -from textual.containers import Grid, Horizontal, Vertical, VerticalScroll +from textual.containers import Horizontal, Vertical, VerticalScroll from textual.geometry import Region, Size from textual.message import Message from textual.reactive import reactive @@ -53,8 +53,8 @@ from .trace_ctl import TraceController from .highlight import highlight_c from .errors import IDAToolError, IDAConnectionError -from .worker_client import WorkerClient -from .domain import DisasmModel, Func, Head, ListingModel, Program, Struct +from .codemode_client import CodeModeClient, registered_database +from .domain import Func, Head, ListingModel, Program, Struct # Styles for the disassembly listing. _S_ADDR = Style(color="#6b7684") @@ -173,8 +173,8 @@ _ADDR_MARK_STRIP_RE = re.compile(r"\s*/\*\s*0x[0-9A-Fa-f]+\s*\*/") @dataclass class BinaryState: """Everything that makes one project binary's session resumable across a - switch. Addresses outlive the worker, so nav history survives eviction; the - Program/index only survive while that worker is still resident.""" + switch. Addresses outlive a database lease, so nav history survives eviction; + the Program/index only survive while that lease remains resident.""" label: str program: object | None = None @@ -1085,9 +1085,9 @@ class ListingView(SearchMixin, NavMixin, ColumnCursor, ScrollView, can_focus=Tru def _span_segments(h: Head, fallback: Style): """Segments for a row's disassembly text. - Uses IDA's own token classification when the worker supplied it; falls - back to the old mnemonic/rest split so an older worker (or a row whose - spans didn't match the text) still renders. + Uses IDA's own token classification when Code Mode supplies it; falls + back to the mnemonic/rest split when spans are absent or disagree with + the plain text. """ if h.spans: return [Segment(t, _S_SPAN.get(k, fallback)) for k, t in h.spans] @@ -3493,12 +3493,16 @@ _HELP = ( class QuitScreen(ModalScreen): - """Asked before exiting with unsaved database changes. Dismisses with - "save", "discard" or None (stay).""" + """Asked before exiting with unsaved database changes. + + Code Mode clients cannot roll a shared database back. The ``d`` choice means + "do not explicitly save": a GUI keeps the changes dirty, while a managed + idalib worker may persist them when its final lease closes. + """ BINDINGS = [ Binding("s", "save", "Save & quit"), - Binding("d", "discard", "Discard & quit"), + Binding("d", "discard", "Leave & quit"), Binding("escape,c", "cancel", "Cancel"), ] @@ -3515,7 +3519,7 @@ class QuitScreen(ModalScreen): for label in self._labels: body.append(f" \u2022 {label}\n", _S_LABEL) yield Static(body, id="quit-list") - yield Static("s save & quit d discard & quit Esc cancel", + yield Static("s save & quit d leave as-is & quit Esc cancel", id="quit-help") def action_save(self) -> None: @@ -4780,15 +4784,16 @@ class IdaTui(App): self._index = None # project-wide symbol/string index if project is not None: from .index import ProjectIndex - from .pool import WorkerPool - self._pool = WorkerPool(project, ttl=ttl) + from .pool import DatabasePool + self._pool = DatabasePool(project, ttl=ttl) self._index = ProjectIndex( os.path.join(project.index_dir, "project.db")) self._binary = project.refs[0].label open_path = project.refs[0].staged self._open_path = open_path self._ttl = ttl - self._load_args = load_args or "" # IDA switches for a headerless blob + self._load_args = load_args or "" # first-open options for a headerless blob + self._new_database = False # Ctrl+L asks Code Mode for a fresh IDB self._title = (os.path.basename(open_path) if open_path else "") #: Where we are in the execution trace, and everything that moves us. #: Owns the trace state; the _trace/_t/_trail_* properties below @@ -4797,7 +4802,7 @@ class IdaTui(App): self._do_keepalive = keepalive self._rpc_path = rpc_path self._rpc = None - self.client: WorkerClient | None = None + self.client: CodeModeClient | None = None self.program: Program | None = None self._loading_screen: LoadingScreen | None = None self._ka = None @@ -4900,8 +4905,8 @@ class IdaTui(App): if self._rpc_path: self._start_rpc() # A file no loader recognises has to be described before it can be - # opened, so ask BEFORE the worker starts — once IDA has made a database - # the answer is baked in and changing it means deleting the .i64. + # opened, so ask BEFORE Code Mode creates it — once IDA has made a database + # the answer is baked in and changing it requires a fresh-IDB reopen. if self._project is not None: ref = self._pending_load_ref() if ref is not None: @@ -4932,6 +4937,11 @@ class IdaTui(App): if os.path.exists(self._open_path + ".i64") or os.path.exists( os.path.splitext(self._open_path)[0] + ".i64"): return False + try: + if registered_database(self._open_path): + return False + except Exception: + pass # connect() will surface registry failures with full diagnostics return needs_load_options(self._open_path) def action_load_options(self) -> None: @@ -4945,7 +4955,11 @@ class IdaTui(App): forward. """ if not self._can_reload(): - self._status("nothing to reload") + if self.client is not None and self.client.backend == "gui": + self._status( + "reload unavailable for a GUI-owned database — reopen it in IDA") + else: + self._status("nothing to reload") return n = len(self._func_index) if self._func_index else 0 note = ("this image has no functions, so nothing is lost" @@ -4964,10 +4978,13 @@ class IdaTui(App): ref = self._project.by_label(self._binary) if ref is not None: path, label = ref.source, ref.label - # Drop the worker first: it holds the database open, and the .i64 can't - # be removed (or rebuilt) underneath a live one. - self._release_worker() - self._drop_database() + # Release our lease first. Code Mode waits for a managed worker's final + # lease grace, then creates the replacement IDB atomically. A GUI-backed + # database is rejected by _can_reload(): the TUI must never close it. + self._release_database() + self._new_database = True + if label is not None and self._pool is not None: + self._pool.recreate_on_next_open(label) self._reset_for_reload() self._load_args = "" if label is not None and self._project is not None: @@ -4979,7 +4996,9 @@ class IdaTui(App): self._pending_switch = None self._ask_load_options(path, label=label) - def _release_worker(self) -> None: + def _release_database(self) -> None: + if self.program is not None: + self.program.close() if self._pool is not None and self._binary is not None: try: self._pool.evict(self._binary, save=False) @@ -4993,23 +5012,6 @@ class IdaTui(App): self.client = None self.program = None - def _drop_database(self) -> None: - """Remove the .i64 (and any unpacked scratch) so the next open re-reads - the raw image with new options.""" - base = self._open_path - if self._project is not None and self._binary is not None: - ref = self._project.by_label(self._binary) - if ref is not None: - base = ref.staged - if not base: - return - for suffix in (".i64", ".id0", ".id1", ".id2", ".nam", ".til"): - for cand in (base + suffix, os.path.splitext(base)[0] + suffix): - try: - os.remove(cand) - except OSError: - pass - def _reset_for_reload(self) -> None: self._no_functions = False self._func_index = None @@ -5053,6 +5055,11 @@ class IdaTui(App): if os.path.exists(ref.db) or os.path.exists( os.path.splitext(ref.staged)[0] + ".i64"): return None # already analysed: the .i64 records how + try: + if registered_database(ref.staged, output_database=ref.db): + return None + except Exception: + pass from .formats import needs_load_options return ref if needs_load_options(ref.source) else None @@ -5106,10 +5113,6 @@ class IdaTui(App): asyncio.get_running_loop().create_task(_serve()) - async def on_unmount(self) -> None: - if self._rpc is not None: - await self._rpc.stop() - # -- status helper ----------------------------------------------------- # def _status(self, text: str, priority: bool = False) -> None: """Write the status bar. ``priority`` marks the RESULT of something the @@ -5174,9 +5177,10 @@ class IdaTui(App): # -- connection loss / recovery --------------------------------------- # def _handle_exception(self, error: BaseException) -> None: - """Intercept a lost-connection error from any worker so the whole app - doesn't die when the analysis server goes away (it can idle out, be - killed, or the box can sleep). Everything else crashes as usual.""" + """Intercept a lost Code Mode lease so the app can rediscover the DB. + + Everything unrelated to database connectivity crashes as usual. + """ from textual.worker import WorkerFailed orig = error.error if isinstance(error, WorkerFailed) else error if isinstance(orig, IDAConnectionError): @@ -5208,15 +5212,15 @@ class IdaTui(App): @work(thread=True, exclusive=True, group="reconnect") def _reconnect(self) -> None: - # The worker died (segfault -> dropped socket). Respawn it: it re-opens - # and re-analyzes the binary in a fresh process, then we rebuild. + # The registered instance disappeared. Rediscover it; Code Mode may find + # a GUI/replacement worker, then we rebuild caches against the new handle. try: if self._open_path is None: self.app.call_from_thread(self._reconnect_failed, "no binary to reopen") return - client = WorkerClient(self._open_path, ttl=self._ttl, - load_args=self._load_args) + client = CodeModeClient(self._open_path, ttl=self._ttl, + load_args=self._load_args) client.connect(progress=lambda m: self.app.call_from_thread( self._conn_note, m)) except Exception as e: # noqa: BLE001 @@ -5224,7 +5228,7 @@ class IdaTui(App): return self.app.call_from_thread(self._after_reconnect, client, Program(client)) - def _after_reconnect(self, client: "WorkerClient", program: "Program") -> None: + def _after_reconnect(self, client: "CodeModeClient", program: "Program") -> None: self.client = client self.program = program self._reconnecting = False @@ -5244,13 +5248,13 @@ class IdaTui(App): @work(thread=True, exclusive=True, group="connect") def _connect(self) -> None: try: - client = self._open_worker_client() + client = self._open_database_client() if client is None: return # the opener already reported + dismissed the overlay module = client.health().get("module", "?") if self._do_keepalive: - # Keep the session warm while we run; don't make it immortal, so - # it's reclaimed after the TUI closes. (No-op for the worker.) + # Compatibility shim: DatabaseHandle's SSE lease already owns + # liveness and heartbeat behavior. self._ka = client.keepalive(interval=120.0).start() program = Program(client) except Exception as e: # noqa: BLE001 @@ -5265,14 +5269,14 @@ class IdaTui(App): return self.client = client self.program = program - self.app.call_from_thread(self._status, f"{module} — loading functions…") + self._new_database = False + self.app.call_from_thread( + self._status, f"{module} [{client.backend}] — loading functions…") self._load_functions() - def _open_worker_client(self): # type: ignore[no-untyped-def] - """Our idalib-worker path: spawn the worker (it opens + analyzes the - binary in its own process) and connect. Returns the client, or None.""" - from .worker_client import WorkerClient - if self._pool is not None: # project mode: the pool owns the workers + def _open_database_client(self): # type: ignore[no-untyped-def] + """Attach through Code Mode, reusing a GUI or managed idalib database.""" + if self._pool is not None: # project mode: the pool owns the leases label = self._binary or self._project.refs[0].label client = self._pool.get(label, progress=lambda m: self.app.call_from_thread(self._status, m)) @@ -5283,14 +5287,15 @@ class IdaTui(App): return client if not self._open_path: self.app.call_from_thread( - self._status, "the worker backend needs a binary path") + self._status, "Code Mode needs a database or executable path") self.app.call_from_thread(self._dismiss_loading) return None base = os.path.basename(self._open_path) self.app.call_from_thread( - self._status, f"starting worker — initial auto-analysis of {base}…") - client = WorkerClient(self._open_path, ttl=self._ttl, - load_args=self._load_args) + self._status, f"discovering Code Mode database for {base}…") + client = CodeModeClient(self._open_path, ttl=self._ttl, + load_args=self._load_args, + new_database=self._new_database) client.connect(progress=lambda m: self.app.call_from_thread( self._status, m)) return client @@ -5362,7 +5367,7 @@ class IdaTui(App): @work(thread=True, exclusive=True, group="index") def _index_binary(self) -> None: """Fold this binary's symbols + strings into the project index, so it can - be searched later even when its worker is gone.""" + be searched later even when its Code Mode lease is gone.""" if self._index is None or self._project is None or self._binary is None: return ref = self._project.by_label(self._binary) @@ -5380,7 +5385,7 @@ class IdaTui(App): imps, exps = self.program.linkage() entries += [(KIND_IMPORT, i.addr, i.name) for i in imps] entries += [(KIND_EXPORT, e.addr, e.name) for e in exps] - except Exception: # noqa: BLE001 -- an old worker has no list_linkage + except Exception: # noqa: BLE001 -- indexing is best-effort pass try: n = self._index.reindex(self._binary, entries, source=ref.source) @@ -5465,7 +5470,13 @@ class IdaTui(App): cursor=0, push=True, is_region=True) def _can_reload(self) -> bool: - """Whether we're able to re-open this binary with different options.""" + """Whether Code Mode can replace this IDB with different options. + + A GUI database is owned by the user and has no remote close/rollback + route. Managed idalib databases can be released and reopened fresh. + """ + if self.client is not None and self.client.backend == "gui": + return False if self._project is not None and self._binary is not None: return True return bool(self._open_path) @@ -5654,6 +5665,9 @@ class IdaTui(App): def _on_quit_choice(self, choice: str | None) -> None: if choice == "discard": + # Code Mode has no rollback/close-without-save operation. For GUI + # sessions this leaves changes dirty in IDA; a managed worker owns + # its final save policy and may persist them on final lease release. self._save_on_exit = False self.exit() elif choice == "save": @@ -5670,7 +5684,7 @@ class IdaTui(App): if self._pool is not None: self._pool.close_all(save=True) # saves each resident worker elif self.program is not None: - self.program.client.call("idb_save", timeout=600.0) + self.program.client.save_database() except Exception as e: # noqa: BLE001 -- still exit, but say so self.app.call_from_thread(self._status, f"save failed: {e}") self.app.call_from_thread(self._finish_exit) @@ -5710,7 +5724,7 @@ class IdaTui(App): self._ask_load_options(ref.source, label=label) return # Snapshot what we're leaving so coming back restores the view, then let - # the pool hand us a worker (spawning + evicting as the budget dictates). + # the pool hand us a lease (attaching + evicting as the budget dictates). if self._binary is not None: self._states[self._binary] = BinaryState( label=self._binary, program=self.program, @@ -5731,8 +5745,8 @@ class IdaTui(App): self.app.call_from_thread(self._switch_failed, label, str(e)) return st = self._states.get(label) - # The Program (and its caches) only survive while that worker does; a - # binary that was evicted comes back with a fresh one. Either way the nav + # The Program (and its caches) only survive while that lease does; an + # evicted binary reattaches. Either way the nav # history is just addresses, so it always survives. reuse = (st is not None and st.program is not None and getattr(st.program, "client", None) is client) @@ -5769,7 +5783,7 @@ class IdaTui(App): self._did_auto_land = False self._auto_land() return - # Cold (first visit, or the worker was evicted): rebuild the index, then + # Cold (first visit, or the lease was evicted): rebuild the index, then # land back where we were via _pending_restore. self._cur = None self._func_index = None @@ -5826,28 +5840,6 @@ class IdaTui(App): return self._goto_ea(addr, push=True) # land on the literal in the listing - def on_descendant_focus(self, event) -> None: # type: ignore[no-untyped-def] - """Keep ``_active`` in step with focus while split. - - Tab moves both together, but focus also moves on its own — a click, or a - pane focusing itself after a load — and then ``_active`` still names the - pane you're NOT in. Everything downstream trusts ``_active``: follow - resolves the word under that pane's cursor and pushes history for it, so - Enter in the pseudocode would follow something from the listing and the - next Esc got spent undoing it. - """ - if not self._split: - return - w = self.focused - mode = ("decomp" if isinstance(w, DecompView) - else "listing" if isinstance(w, ListingView) else None) - if mode is None or mode == self._active: - return - self._active = mode - self._sync_split(mode) # re-link the band from the new driver - if not self.query_one(DecompView).loading: - self._status_for_cur("split") # never clobber "decompiling…" - def action_toggle_view(self) -> None: """Tab: switch the code pane between disassembly and pseudocode (or leave the hex view back to the preferred code view).""" @@ -6356,7 +6348,7 @@ class IdaTui(App): def _cross_binary_impl(self, name: str) -> tuple[str, int] | None: """``(binary, addr)`` of a project binary that EXPORTS ``name``. - Reads the on-disk index, so a provider resolves even when its worker was + Reads the on-disk index, so a provider resolves even when its lease was evicted — the whole reason the index exists. """ if self._index is None or self._project is None or not name: @@ -6529,7 +6521,7 @@ class IdaTui(App): def _foreign_importers(self, subj: int, subj_name, fn): # type: ignore[no-untyped-def] """Project binaries that IMPORT the symbol at ``subj`` — the other half of the phase-3 join, read from the on-disk index so a caller shows up - whether or not its worker is resident. + whether or not its database lease is resident. Only for a symbol this binary actually exports: a local name that happens to collide with another binary's import isn't a caller of ours. @@ -6742,7 +6734,7 @@ class IdaTui(App): def _save(self) -> None: assert self.program is not None try: - self.program.client.call("idb_save", timeout=300.0) + self.program.client.save_database() except Exception as e: # noqa: BLE001 self.app.call_from_thread(self._status, f"save failed: {e}") return @@ -7727,7 +7719,9 @@ class IdaTui(App): "(c code · p func · u undefine · Enter follow)") # -- teardown ---------------------------------------------------------- # - def on_unmount(self) -> None: + async def on_unmount(self) -> None: + if self._rpc is not None: + await self._rpc.stop() if self._ka is not None: self._ka.stop() if self.program is not None: @@ -7739,7 +7733,7 @@ class IdaTui(App): elif self.client is not None: if self._save_on_exit is None and self._dirty: try: # unexpected teardown with edits: don't drop them - self.client.call("idb_save", timeout=600.0) + self.client.save_database() except Exception: # noqa: BLE001 pass self.client.close() diff --git a/idatui/codemode_client.py b/idatui/codemode_client.py new file mode 100644 index 0000000..77a0227 --- /dev/null +++ b/idatui/codemode_client.py @@ -0,0 +1,1164 @@ +"""Client adapter from ida-tui's domain operations to IDA Code Mode. + +``DatabaseHandle`` is the lifecycle boundary: it discovers an already-registered +GUI database, reuses a shared managed idalib worker, or starts one when needed. +The TUI never owns or terminates an IDA process. Closing this client releases +only its lease. + +The Code Mode transport intentionally exposes one broad operation, +``execute_python``. ``CodeModeClient.invoke`` turns the small, address-centric +operations needed by the paging layer into self-contained snippets. The +snippets prefer the public ``ida-domain`` ``db`` object. A handful of features +that ida-domain does not currently expose (IDA-coloured listing rows, creating +instructions, ARM T-state, and detailed Hex-Rays line maps/failures) use the +IDAPython modules that Code Mode deliberately makes importable. +""" +from __future__ import annotations + +import json +import os +import shlex +import threading +import time +from textwrap import dedent +from typing import Any + +from .errors import IDAConnectionError, IDATimeoutError, IDAToolError, Session + +# ida_codemode is imported EAGERLY-IF-PRESENT but never at hard import cost. +# +# The paging/graph/trace layers and their offline test suites must keep importing +# `idatui` on a machine with no IDA and no Code Mode installed -- that is the +# house rule the stdlib-only worker client used to satisfy for free, and +# `tests/run.py --fast` (257 checks, any python3) depends on it. A hard top-level +# import here makes the whole package unimportable, so the failure is deferred to +# the first operation that genuinely needs the library. +_CODEMODE_ERROR: Exception | None = None +try: + from ida_codemode.client import ( + ClientError, + DatabaseHandle, + InstanceDisconnectedError, + RemoteError, + ) + from ida_codemode.registry import ( + REGISTRY_DIR, + FileLock, + RegistryEntry, + canonical_path, + idb_key, + scan_instances, + ) + from ida_codemode.resolver import IdbBusy, expected_idb_path +except ImportError as _exc: # library absent: usable only for offline layers + _CODEMODE_ERROR = _exc + # Bound to None rather than left undefined so the names stay patchable: the + # offline contract tests inject a fake DatabaseHandle here. + ClientError = InstanceDisconnectedError = RemoteError = None # type: ignore[assignment,misc] + DatabaseHandle = RegistryEntry = FileLock = None # type: ignore[assignment,misc] + REGISTRY_DIR = canonical_path = idb_key = scan_instances = None # type: ignore[assignment] + IdbBusy = expected_idb_path = None # type: ignore[assignment] + + +def _require_codemode() -> None: + """Raise an actionable error when the Code Mode library is missing. + + Gated on the binding, not on the original import result, so a test that + injects a fake ``DatabaseHandle`` exercises the real adapter logic. + """ + if DatabaseHandle is None: + raise IDAConnectionError( + "ida-codemode-mcp is not installed in this environment " + f"({_CODEMODE_ERROR}). Install it (e.g. `uv sync`, or " + "`pip install -e ../ida-codemode-mcp`) so ida-tui can lease a " + "database.") from _CODEMODE_ERROR + + +def database_owner(idb_path: str, staged_path: str | None = None): + """The registry entry that owns ``idb_path``/``staged_path``, else None. + + Returns None when the Code Mode library is absent: with no library there is + no client in this environment that could be holding the database, and the + IDA-free layers (project staging) must keep working. Registry errors that + happen WITH the library installed still propagate -- those mean "we could + not determine ownership", which is not the same as "nobody owns it". + """ + if DatabaseHandle is None: + return None + expected_key = idb_key(idb_path) + staged = canonical_path(staged_path) if staged_path else None + for item in scan_instances(timeout=0.5): + entry = item.entry + if entry.idb_key == expected_key: + return entry + if staged and entry.exe_path and canonical_path(entry.exe_path) == staged: + return entry + return None + + +def registered_database(path: str, output_database: str | None = None) -> bool: + """Whether a live/lock-held Code Mode instance owns this target.""" + _require_codemode() + source = canonical_path(path) + expected = canonical_path(output_database) if output_database else expected_idb_path(source) + expected_key = idb_key(expected) + for instance in scan_instances(timeout=0.5): + entry = instance.entry + if entry.idb_key == expected_key: + return True + if not output_database and entry.backend == "gui" and entry.exe_path: + if canonical_path(entry.exe_path) == source: + return True + return False + + +class _NoopKeepAlive: + """Compatibility shim: the DatabaseHandle's SSE lease is the heartbeat.""" + + def __init__(self) -> None: + self.beats = self.failures = 0 + + def start(self) -> "_NoopKeepAlive": + return self + + def stop(self) -> None: + pass + + +def _parse_load_args(value: str) -> tuple[str | None, int | None, str | None]: + """Translate ida-tui's legacy first-open switches to Code Mode options. + + Code Mode has typed options for processor, natural loading address and file + type. It deliberately has no arbitrary command-line escape hatch; reject + switches we cannot represent instead of silently loading a blob wrongly. + """ + processor: str | None = None + loading_address: int | None = None + file_type: str | None = None + unsupported: list[str] = [] + try: + words = shlex.split(value or "", posix=os.name != "nt") + except ValueError as exc: + raise ValueError(f"invalid IDA load options: {exc}") from exc + for word in words: + if word.startswith("-p") and len(word) > 2: + processor = word[2:] + elif word.startswith("-b") and len(word) > 2: + try: + # IDA's -b is in 16-byte paragraphs. DatabaseHandle expects the + # natural address, which is the safer public API. + loading_address = int(word[2:], 16) << 4 + except ValueError as exc: + raise ValueError(f"invalid IDA loading address: {word!r}") from exc + elif word.startswith("-T") and len(word) > 2: + file_type = word[2:] + else: + unsupported.append(word) + if unsupported: + joined = " ".join(unsupported) + raise ValueError( + "ida-codemode cannot represent arbitrary IDA load options: " + f"{joined!r}; use processor/base/file type options instead" + ) + return processor, loading_address, file_type + + +def _script(args: dict[str, Any], body: str) -> str: + """Bind JSON arguments without interpolating user text into Python code.""" + encoded = json.dumps(args, ensure_ascii=False, separators=(",", ":")) + return f"import json\na = json.loads({encoded!r})\n{dedent(body).strip()}\n" + + +# Rich flat-listing generation is the largest ida-domain gap in this port. +# ida-domain can enumerate heads and render plain disassembly, but it does not +# expose undefined runs, IDA colour spans, function banners, or expanded UDT +# members. Keep that IDAPython-only logic isolated in this one operation. +_HEADS = r''' +import ida_bytes, ida_funcs, ida_idaapi, ida_lines, ida_name, ida_nalt, ida_segment, ida_typeinf +start = int(str(a["addr"]), 16) +count = max(1, min(int(a.get("count", 200)), 2000)) +offset = max(0, int(a.get("offset", 0))) +annotate = bool(a.get("annotate", False)) +seg = db.segments.get_at(start) +if seg is None: + result = {"addr": a["addr"], "error": "no segment", "heads": [], "cursor": {"done": True}} +else: + lo, hi = int(seg.start_ea), int(seg.end_ea) + if a.get("end"): + hi = min(hi, int(str(a["end"]), 16)) + + span_names = { + "insn": ("SCOLOR_INSN", "SCOLOR_KEYWORD", "SCOLOR_ASMDIR", "SCOLOR_MACRO"), + "reg": ("SCOLOR_REG",), + "num": ("SCOLOR_NUMBER", "SCOLOR_CHAR", "SCOLOR_BINPREF"), + "str": ("SCOLOR_STRING",), + "name": ("SCOLOR_DATNAME", "SCOLOR_CODNAME", "SCOLOR_LOCNAME", "SCOLOR_IMPNAME", + "SCOLOR_DEMNAME", "SCOLOR_LIBNAME", "SCOLOR_CNAME", "SCOLOR_DNAME", + "SCOLOR_CREF", "SCOLOR_DREF", "SCOLOR_CREFTAIL", "SCOLOR_DREFTAIL"), + "seg": ("SCOLOR_SEGNAME",), + "cmt": ("SCOLOR_AUTOCMT", "SCOLOR_REGCMT", "SCOLOR_RPTCMT", "SCOLOR_VOIDOP"), + "punct": ("SCOLOR_SYMBOL", "SCOLOR_ALTOP", "SCOLOR_HIDNAME"), + "err": ("SCOLOR_ERROR",), + } + tag_kinds = {} + for kind, names in span_names.items(): + for name in names: + value = getattr(ida_lines, name, None) + if isinstance(value, str) and value: + tag_kinds[value[0]] = kind + elif isinstance(value, int): + tag_kinds[chr(value)] = kind + + def spans(tagged): + on, off, esc = "\x01", "\x02", "\x03" + addr_tag = chr(getattr(ida_lines, "COLOR_ADDR", 0x28)) + addr_len = int(getattr(ida_lines, "COLOR_ADDR_SIZE", 16)) + out, stack, buf = [], [], [] + def flush(): + if buf: + out.append([stack[-1] if stack else "text", "".join(buf)]) + buf.clear() + i = 0 + while i < len(tagged): + ch = tagged[i] + if ch == on and i + 1 < len(tagged): + tag = tagged[i + 1] + if tag == addr_tag: + i += 2 + addr_len + continue + flush(); stack.append(tag_kinds.get(tag, "text")); i += 2; continue + if ch == off and i + 1 < len(tagged): + flush() + if stack: stack.pop() + i += 2; continue + if ch == esc and i + 1 < len(tagged): + buf.append(tagged[i + 1]); i += 2; continue + buf.append(ch); i += 1 + flush() + collapsed, previous_space = [], False + for kind, text in out: + acc = [] + for ch in text: + if ch.isspace(): + if previous_space: continue + acc.append(" "); previous_space = True + else: + acc.append(ch); previous_space = False + if acc: collapsed.append([kind, "".join(acc)]) + if collapsed: + collapsed[0][1] = collapsed[0][1].lstrip() + collapsed[-1][1] = collapsed[-1][1].rstrip() + return [[kind, text] for kind, text in collapsed if text] + + def row(ea): + flags = ida_bytes.get_flags(ea) + kind = "code" if ida_bytes.is_code(flags) else ("data" if ida_bytes.is_data(flags) else "unknown") + tagged = ida_lines.generate_disasm_line(ea, 0) or "" + text = " ".join(ida_lines.tag_remove(tagged).split()) if tagged else "" + item = {"ea": hex(ea), "kind": kind, "size": int(ida_bytes.get_item_size(ea)), "text": text} + if tagged: + rich = spans(tagged) + if " ".join("".join(x[1] for x in rich).split()) == text: + item["spans"] = rich + name = ida_name.get_ea_name(ea) + if name: item["name"] = name + return item + + def unknown_row(ea, size): + if size <= 1: return row(ea) + item = {"ea": hex(ea), "kind": "unknown", "size": int(size), "text": f"db {size} dup(?)"} + name = ida_name.get_ea_name(ea) + if name: item["name"] = name + return item + + def members(ea): + tif = db.types.get_at(ea) + if tif is None or not tif.is_udt(): return [] + answer = [] + for member in db.types.get_udt_members(tif): + type_text = member.type.dstr() or "" + text = f"+{member.offset:X} {member.name}" + (f" {type_text}" if type_text else "") + answer.append({"ea": hex(ea + member.offset), "kind": "member", + "size": int(member.size), "text": text}) + return answer + + def is_unknown(ea): + flags = ida_bytes.get_flags(ea) + return not (ida_bytes.is_code(flags) or ida_bytes.is_data(flags)) + def run_end(ea): + nxt = ida_bytes.next_head(ea, hi) + return nxt if nxt != ida_idaapi.BADADDR and ea < nxt <= hi else hi + def advance(ea): + if is_unknown(ea): return run_end(ea) + nxt = ida_bytes.get_item_end(ea) + return nxt if nxt > ea else ea + 1 + def rows_for(ea): + if is_unknown(ea): return [unknown_row(ea, run_end(ea) - ea)] + fn = db.functions.get_at(ea) if annotate else None + at_start = fn is not None and int(fn.start_ea) == ea + answer = [] + if at_start: + name = db.functions.get_name(fn) or f"sub_{ea:X}" + answer += [ + {"ea": hex(ea), "kind": "sep", "size": 0, "text": ""}, + {"ea": hex(ea), "kind": "sep", "size": 0, + "text": "; " + "=" * 15 + " S U B R O U T I N E " + "=" * 15}, + {"ea": hex(ea), "kind": "funchdr", "size": 0, + "text": name + " proc", "name": name}, + ] + item = row(ea) + if at_start: + item["name"] = None + elif annotate and item["kind"] == "code" and item.get("name"): + name = item["name"] + answer.append({"ea": hex(ea), "kind": "label", "size": 0, + "text": name + ":", "name": name}) + item["name"] = None + answer.append(item) + if item["kind"] == "data": answer += members(ea) + if fn is not None and ida_bytes.get_item_end(ea) >= int(fn.end_ea): + name = db.functions.get_name(fn) or f"sub_{int(fn.start_ea):X}" + answer += [ + {"ea": hex(ea), "kind": "funchdr", "size": 0, + "text": name + " endp", "name": name}, + {"ea": hex(ea), "kind": "sep", "size": 0, "text": "; " + "-" * 60}, + ] + return answer + + ea = ida_bytes.get_item_head(start) + if ea == ida_idaapi.BADADDR: ea = start + for _ in range(offset): + if ea >= hi: break + ea = advance(ea) + rows = [] + more = False + while ea != ida_idaapi.BADADDR and ea < hi: + if len(rows) >= count: + more = True; break + rows += rows_for(ea) + ea = advance(ea) + result = {"addr": a["addr"], "heads": rows, + "cursor": {"next": hex(ea)} if more else {"done": True}} +result +''' + + +_DECOMP_MAP_HELPER = r''' +def line_map(cfunc): + import ida_hexrays + answer = [] + for sl in cfunc.get_pseudocode(): + tagged, eas, seen = sl.line, [], set() + for x in range(len(tagged) + 1): + head = ida_hexrays.ctree_item_t(); item = ida_hexrays.ctree_item_t(); tail = ida_hexrays.ctree_item_t() + if not cfunc.get_line_item(tagged, x, False, head, item, tail): continue + text = item.dstr() or "" + try: ea = int(text.split(": ", 1)[0], 16) + except (ValueError, IndexError): continue + if ea not in seen: seen.add(ea); eas.append(ea) + answer.append(eas) + return answer +''' + + +_OPERATIONS: dict[str, str] = { + "list_funcs": r''' +import fnmatch +queries = a.get("queries") or [{}] +q = queries[0] +offset, count = max(0, int(q.get("offset", 0))), max(1, int(q.get("count", 500))) +pattern = str(q.get("filter") or "").lower() +if pattern and not any(ch in pattern for ch in "*?["): pattern = "*" + pattern + "*" +rows = [] +for fn in db.functions.get_all(): + name = db.functions.get_name(fn) or f"sub_{int(fn.start_ea):X}" + if pattern and not fnmatch.fnmatchcase(name.lower(), pattern): continue + rows.append({"addr": hex(int(fn.start_ea)), "name": name, + "size": int(fn.end_ea) - int(fn.start_ea)}) +page = rows[offset:offset + count] +result = {"result": [{"data": page, "next_offset": offset + len(page), "total": len(rows)}]} +result +''', + "disasm": r''' +ea = int(str(a["addr"]), 16) +fn = db.functions.get_at(ea) +if fn is None: + result = {"instructions": [], "total_instructions": 0, "instruction_count": 0} +else: + instructions = list(db.functions.get_instructions(fn)) + limit = max(1, int(a.get("max_instructions", len(instructions) or 1))) + rows = [{"addr": hex(int(insn.ea)), "instruction": db.instructions.get_disassembly(insn)} + for insn in instructions[:limit]] + result = {"instructions": rows, "total_instructions": len(instructions), + "instruction_count": len(instructions)} +result +''', + "file_regions": r''' +import idaapi +rows = [] +for seg in db.segments.get_all(): + try: file_off = int(idaapi.get_fileregion_offset(seg.start_ea)) + except Exception: file_off = -1 + if file_off < 0 or file_off >= (1 << 48): file_off = -1 + rows.append({"start": hex(int(seg.start_ea)), "end": hex(int(seg.end_ea)), + "file_off": file_off, "name": db.segments.get_name(seg) or ""}) +result = {"regions": rows} +result +''', + "read_raw": r''' +import ida_bytes +ea, size = int(str(a["addr"]), 16), max(0, int(a["size"])) +raw = ida_bytes.get_bytes(ea, size) or b"" +raw = raw[:size] + b"\xff" * max(0, size - len(raw)) +data = bytearray(raw) +for index, value in enumerate(data): + if value == 0xFF and not ida_bytes.is_loaded(ea + index): data[index] = 0 +result = {"addr": a["addr"], "hex": bytes(data).hex(), "n": len(data)} +result +''', + "get_bytes": r''' +rows = [] +for region in a.get("regions", []): + ea, size = int(str(region["addr"]), 16), int(region["size"]) + raw = db.bytes.get_bytes_at(ea, size) or b"" + rows.append({"addr": region["addr"], "data": " ".join(f"{b:02x}" for b in raw)}) +result = {"result": rows} +result +''', + "search_structs": r''' +needle = str(a.get("filter") or "").lower() +rows = [] +for tif in db.types.get_all(): + name = tif.get_type_name() or "" + if not name or needle not in name.lower() or not tif.is_udt(): continue + members = list(db.types.get_udt_members(tif)) + rows.append({"name": name, "size": int(tif.get_size()), "is_union": bool(tif.is_union()), + "cardinality": len(members), "ordinal": int(tif.get_ordinal())}) +result = {"result": rows} +result +''', + "type_inspect": r''' +rows = [] +for query in a.get("queries", []): + name = str(query.get("name") or "") + tif = db.types.get_by_name(name) + if tif is None: + rows.append({"name": name, "error": "type not found"}); continue + members = [{"name": m.name, "type": m.type.dstr() or str(m.type), + "offset": int(m.offset), "size": int(m.size)} + for m in db.types.get_udt_members(tif)] if tif.is_udt() else [] + rows.append({"name": name, "size": int(tif.get_size()), "is_union": bool(tif.is_union()), + "members": members}) +result = {"result": rows} +result +''', + "declare_type": r''' +import ida_typeinf +decls = a.get("decls", "") +if isinstance(decls, str): decls = [decls] +rows = [] +for declaration in decls: + try: + errors = int(db.types.parse_declarations(ida_typeinf.get_idati(), declaration)) + rows.append({"ok": errors == 0, **({} if errors == 0 else {"error": f"{errors} parse error(s)"})}) + except Exception as exc: + rows.append({"ok": False, "error": str(exc)}) +result = {"result": rows} +result +''', + "del_type": r''' +import ida_typeinf +name = str(a["name"]) +ok = bool(ida_typeinf.del_named_type(ida_typeinf.get_idati(), name, ida_typeinf.NTF_TYPE)) +result = {"name": name, "deleted": ok, **({} if ok else {"error": f"Type {name!r} not found or could not be deleted"})} +result +''', + "func_types": r''' +import ida_typeinf +ea = int(str(a["addr"]), 16) +fn = db.functions.get_at(ea) +if fn is None: + result = {"addr": a["addr"], "error": "no function at address"} +else: + pseudo = db.pseudocode.decompile(fn) + name = db.functions.get_name(fn) or "" + tif = pseudo.get_func_type() + try: prototype = ida_typeinf.print_tinfo("", 0, 0, ida_typeinf.PRTYPE_1LINE, tif, name, "") if tif else "" + except Exception: prototype = tif.dstr() if tif else "" + lvars = [{"name": var.name, "type": var.type_info.dstr() if var.type_info else "", + "is_arg": bool(var.is_arg)} for var in pseudo.local_variables] + result = {"addr": hex(int(fn.start_ea)), "name": name, + "prototype": (prototype or "").strip(), "lvars": lvars} +result +''', + "set_lvar_type": r''' +import ida_typeinf +ea, variable, declaration = int(str(a["addr"]), 16), str(a["variable"]), str(a["type"]) +fn = db.functions.get_at(ea) +if fn is None: + result = {"error": "no function at address"} +else: + pseudo = db.pseudocode.decompile(fn) + var = pseudo.find_local_variable(variable) + if var is None: + result = {"error": f"local variable {variable!r} not found"} + else: + try: + tif = db.types.parse_one_declaration(ida_typeinf.get_idati(), declaration) + accepted = bool(var.set_type(tif)) + saved = bool(pseudo.save_local_variable_info(var, save_type=True)) if accepted else False + result = {"addr": hex(int(fn.start_ea)), "variable": variable, + "type": declaration, "ok": accepted and saved} + except Exception as exc: + result = {"error": f"bad type {declaration!r}: {exc}"} +result +''', + "set_type": r''' +from ida_domain.types import TypeApplyFlags +rows = [] +for edit in a.get("edits", []): + ea = int(str(edit["addr"]), 16) + declaration = str(edit.get("signature") or edit.get("type") or "") + try: + ok = bool(db.types.apply_declaration_at(ea, declaration, TypeApplyFlags.DEFINITE)) + rows.append({"addr": hex(ea), "ok": ok, **({} if ok else {"error": "IDA rejected the type"})}) + except Exception as exc: + rows.append({"addr": hex(ea), "ok": False, "error": str(exc)}) +result = {"result": rows} +result +''', + "data_type": r''' +ea = int(str(a["addr"]), 16) +try: + tif = db.types.get_at(ea) + fn = db.functions.get_at(ea) + result = {"addr": hex(ea), "name": db.names.get_at(ea) or "", + "type": tif.dstr() if tif else "", "size": int(db.heads.size(ea)) if db.heads.is_head(ea) else 0, + "is_func": bool(fn)} +except Exception as exc: + result = {"addr": hex(ea), "error": str(exc)} +result +''', + "force_recompile": r''' +import ida_hexrays +rows = [] +for item in a.get("items", []): + ea = int(str(item["addr"]), 16) + ida_hexrays.mark_cfunc_dirty(ea, False) + rows.append({"addr": hex(ea), "ok": True}) +result = {"result": rows} +result +''', + "undefine": r''' +import ida_bytes +rows = [] +for item in a.get("items", []): + ea = int(str(item["addr"]), 16) + size = max(1, int(item.get("size") or ida_bytes.get_item_size(ea) or 1)) + ok = bool(ida_bytes.del_items(ea, ida_bytes.DELIT_SIMPLE, size)) + rows.append({"addr": hex(ea), "ok": ok, **({} if ok else {"error": "delete items failed"})}) +result = {"result": rows} +result +''', + "define_code": r''' +import ida_ua +rows = [] +for item in a.get("items", []): + ea = int(str(item["addr"]), 16); size = int(ida_ua.create_insn(ea)) + rows.append({"addr": hex(ea), "ok": size > 0, "size": size, + **({} if size > 0 else {"error": "instruction did not decode"})}) +result = {"result": rows} +result +''', + "define_func": r''' +rows = [] +for item in a.get("items", []): + ea = int(str(item["addr"]), 16); ok = bool(db.functions.create(ea)) + rows.append({"addr": hex(ea), "ok": ok, **({} if ok else {"error": "IDA refused the function"})}) +result = {"result": rows} +result +''', + "make_data": r''' +import ida_bytes, ida_idaapi, ida_typeinf +from ida_domain.types import TypeApplyFlags +rows = [] +for item in a.get("items", []): + ea, declaration = int(str(item["addr"]), 16), str(item["type"]) + try: + tif = db.types.parse_one_declaration(ida_typeinf.get_idati(), declaration) + size = max(1, int(tif.get_size())) + saved_names = [(addr, name) for addr, name in db.names.get_all() + if ea <= int(addr) < ea + size] + ida_bytes.del_items(ea, ida_bytes.DELIT_EXPAND | ida_bytes.DELIT_DELNAMES, + max(size, int(ida_bytes.get_item_size(ea) or 1))) + created = bool(ida_bytes.create_data(ea, ida_bytes.FF_BYTE, size, ida_idaapi.BADADDR)) + ok = created and bool(db.types.apply_at(tif, ea, TypeApplyFlags.DEFINITE)) + for address, name in saved_names: + db.names.set_name(int(address), name) + if ok and item.get("name"): ok = bool(db.names.set_name(ea, str(item["name"]))) + rows.append({"addr": hex(ea), "ok": ok, "size": size, + **({} if ok else {"error": "IDA rejected the data type"})}) + except Exception as exc: + rows.append({"addr": hex(ea), "ok": False, "error": str(exc)}) +result = {"result": rows} +result +''', + "make_string": r''' +from ida_domain.strings import StringType +ea, length = int(str(a["addr"]), 16), max(0, int(a.get("length", 0))) +kind = {"c": StringType.C, "c16": StringType.C_16, "c32": StringType.C_32, + "pascal": StringType.PASCAL}.get(str(a.get("kind", "c")).lower(), StringType.C) +import ida_bytes +try: + ida_bytes.del_items(ea, ida_bytes.DELIT_SIMPLE, length if length > 0 else 1) +except Exception: + pass +try: + ok = bool(db.bytes.create_string_at(ea, length or None, kind)) + text = db.bytes.get_string_at(ea) or "" if ok else "" + result = {"addr": hex(ea), "ok": ok, "size": int(db.heads.size(ea)) if ok else 0, "text": text} +except Exception as exc: + result = {"addr": hex(ea), "ok": False, "error": str(exc)} +result +''', + "list_strings": r''' +from ida_domain.strings import StringListConfig +offset, count, min_len = max(0, int(a.get("offset", 0))), max(1, int(a.get("count", 2000))), max(1, int(a.get("min_len", 4))) +if offset == 0 or a.get("refresh"): + from ida_domain.strings import StringType + db.strings.rebuild(StringListConfig(string_types=list(StringType), min_len=min_len, + only_ascii_7bit=False)) +items = list(db.strings.get_all()) +page = items[offset:offset + count] +rows = [] +for item in page: + try: text = str(item) + except Exception: text = item.contents.decode("utf-8", "replace") if item.contents else "" + rows.append({"addr": hex(int(item.address)), "text": text, "len": int(item.length), "type": item.type.name}) +result = {"strings": rows, "total": len(items), "next_offset": offset + len(rows)} +result +''', + "list_linkage": r''' +imports = [{"addr": hex(int(item.address)), "name": item.name, "module": item.module_name} + for item in db.imports.get_all_imports() if item.name] +exports = [{"addr": hex(int(item.address)), "name": item.name, "ordinal": int(item.ordinal)} + for item in db.entries.get_all() if item.name] +result = {"imports": imports, "exports": exports, + "n_imports": len(imports), "n_exports": len(exports)} +result +''', + "lookup_funcs": r''' +rows = [] +for query in a.get("queries", []): + raw = str(query) + try: ea = int(raw, 16) + except ValueError: + fn = db.functions.get_by_name(raw); ea = int(fn.start_ea) if fn else None + else: fn = db.functions.get_at(ea) + if fn is None: + rows.append({"query": raw, "fn": None}) + else: + rows.append({"query": raw, "fn": {"addr": hex(int(fn.start_ea)), + "name": db.functions.get_name(fn) or f"sub_{int(fn.start_ea):X}", + "size": int(fn.end_ea) - int(fn.start_ea)}}) +result = {"result": rows} +result +''', + "resolve_names": r''' +import ida_idaapi, ida_name +rows = [] +for query in a.get("queries", []): + name = str(query).strip(); ea = ida_name.get_name_ea(ida_idaapi.BADADDR, name) + rows.append({"query": name, "ea": hex(int(ea)) if ea != ida_idaapi.BADADDR else None}) +result = {"result": rows} +result +''', + "xref_types": r''' +queries = a.get("queries") or [] +all_results = [] +for query in queries: + ea, direction = int(str(query["addr"]), 16), str(query.get("direction", "both")) + refs = [] + if direction in ("to", "both"): refs += list(db.xrefs.to_ea(ea)) + if direction in ("from", "both"): refs += list(db.xrefs.from_ea(ea)) + rows, seen = [], set() + for ref in refs: + key = (int(ref.from_ea), int(ref.to_ea), int(ref.type)) + if query.get("dedup") and key in seen: continue + seen.add(key) + fn = db.functions.get_at(int(ref.from_ea)) + kind = ("call" if ref.is_call else "jump" if ref.is_jump else "flow" if ref.is_flow + else "read" if ref.is_read else "write" if ref.is_write else ref.type.name.lower()) + row = {"from": hex(int(ref.from_ea)), "to": hex(int(ref.to_ea)), + "type": "code" if ref.is_code else "data", "kind": kind} + if query.get("include_fn") and fn is not None: + row["fn"] = {"addr": hex(int(fn.start_ea)), "name": db.functions.get_name(fn) or ""} + rows.append(row) + if len(rows) >= int(query.get("count", 2000)): break + all_results.append({"data": rows}) +result = {"result": all_results} +result +''', + "xref_query": r''' +queries = a.get("queries") or [] +all_results = [] +for query in queries: + ea, direction = int(str(query["addr"]), 16), str(query.get("direction", "both")) + refs = [] + if direction in ("to", "both"): refs += list(db.xrefs.to_ea(ea)) + if direction in ("from", "both"): refs += list(db.xrefs.from_ea(ea)) + rows = [] + for ref in refs[:int(query.get("count", 2000))]: + fn = db.functions.get_at(int(ref.from_ea)) + row = {"from": hex(int(ref.from_ea)), "to": hex(int(ref.to_ea)), + "type": "code" if ref.is_code else "data"} + if query.get("include_fn") and fn is not None: + row["fn"] = {"addr": hex(int(fn.start_ea)), "name": db.functions.get_name(fn) or ""} + rows.append(row) + all_results.append({"data": rows}) +result = {"result": all_results} +result +''', + "set_comments": r''' +rows = [] +for item in a.get("items", []): + ea, text = int(str(item["addr"]), 16), str(item.get("comment") or "") + try: + if text: ok = bool(db.comments.set_at(ea, text)) + else: db.comments.delete_at(ea); ok = True + rows.append({"addr": hex(ea), "ok": ok}) + except Exception as exc: + rows.append({"addr": hex(ea), "ok": False, "error": str(exc)}) +result = {"result": rows} +result +''', + "rename": r''' +import ida_idaapi, ida_name, ida_typeinf +batch = a.get("batch") or {} +out = {}; ok_count = failed = 0 +for category, edit in batch.items(): + try: + if category == "func": + ea, new = int(str(edit["addr"]), 16), str(edit["name"]) + fn = db.functions.get_at(ea); ok = bool(fn and db.functions.set_name(fn, new)) + elif category == "data": + new = str(edit.get("new") or "") + if edit.get("addr") is not None: ea = int(str(edit["addr"]), 16) + else: ea = int(ida_name.get_name_ea(ida_idaapi.BADADDR, str(edit.get("old") or ""))) + ok = bool(db.names.set_name(ea, new)) + elif category in ("local", "stack"): + ea, old, new = int(str(edit["func_addr"]), 16), str(edit["old"]), str(edit["new"]) + pseudo = db.pseudocode.decompile(ea); var = pseudo.find_local_variable(old) + if var is None: ok = False + else: + var.set_user_name(new) + ok = bool(pseudo.save_local_variable_info(var, save_name=True)) + else: + raise ValueError(f"unsupported rename category: {category}") + row = {"ok": ok, **({} if ok else {"error": "IDA rejected the name"})} + except Exception as exc: + row = {"ok": False, "error": str(exc)} + out[category] = [row] + if row["ok"]: ok_count += 1 + else: failed += 1 +out["summary"] = {"ok": ok_count, "failed": failed} +result = out +result +''', +} + + +_OPERATIONS["decompile"] = _DECOMP_MAP_HELPER + r''' +ea = int(str(a["addr"]), 16) +fn = db.functions.get_at(ea) +if fn is None: + result = {"error": f"no function at {ea:#x}"} +else: + pseudo = db.pseudocode.decompile(fn) + mapping = line_map(pseudo.raw_cfunc) + plain = pseudo.to_text() + marked = [line + (f" /*0x{eas[0]:X}*/" if eas else "") + for line, eas in zip(plain, mapping)] + import ida_name + refs, seen = [], set() + for expr in pseudo.find_objects(): + target = int(expr.obj_ea) + if target in seen or not (db.is_valid_ea(target) or db.is_private_ea(target)): continue + seen.add(target) + name = expr.obj_name or ida_name.get_name(target) or "" + try: string = db.bytes.get_string_at(target) if db.is_valid_ea(target) else None + except Exception: string = None + refs.append({"addr": hex(target), "name": name, "string": string}) + result = {"addr": hex(int(fn.start_ea)), "code": "\n".join(marked), "refs": refs} +result +''' + +_OPERATIONS["decomp_map"] = _DECOMP_MAP_HELPER + r''' +ea = int(str(a["addr"]), 16) +fn = db.functions.get_at(ea) +if fn is None: + result = {"error": f"no function at {ea:#x}"} +else: + pseudo = db.pseudocode.decompile(fn) + mapping = line_map(pseudo.raw_cfunc) + result = {"addr": hex(int(fn.start_ea)), + "lines": [{"ea": hex(eas[0]) if eas else None, + "eas": [hex(item) for item in eas]} for eas in mapping]} +result +''' + +_OPERATIONS["define_code_run"] = r''' +import ida_bytes, ida_idp, ida_segment, ida_ua, idaapi +ea, limit = int(str(a["addr"]), 16), max(1, min(int(a.get("limit", 20000)), 200000)) +seg = ida_segment.getseg(ea) +if seg is None: + result = {"addr": a["addr"], "error": "no segment", "count": 0} +else: + start, count, stopped, hi = ea, 0, "limit", int(seg.end_ea) + while count < limit: + if ea >= hi: stopped = "segment"; break + flags = ida_bytes.get_flags(ea) + if ida_bytes.is_code(flags) or ida_bytes.is_data(flags): stopped = "defined"; break + size = int(ida_ua.create_insn(ea)) + if size <= 0: stopped = "undecodable"; break + count += 1 + insn = ida_ua.insn_t() + if ida_ua.decode_insn(insn, ea) > 0: + try: is_ret = bool(ida_idp.is_ret_insn(insn)) + except Exception: is_ret = False + if is_ret or (insn.get_canon_feature() & idaapi.CF_STOP): + ea += size; stopped = "flow"; break + ea += size + result = {"start": hex(start), "end": hex(ea), "count": count, "stopped": stopped} +result +''' + +_OPERATIONS["define_func_run"] = r''' +import ida_bytes, ida_funcs, ida_segment +ea = int(str(a["addr"]), 16) +fn = db.functions.get_at(ea) +if fn is not None and int(fn.start_ea) == ea: + result = {"addr": hex(ea), "ok": True, "start": hex(ea), "end": hex(int(fn.end_ea)), "how": "existed"} +else: + automatic = bool(db.functions.create(ea)) + if not automatic: + seg = db.segments.get_at(ea); end = ea; hi = int(seg.end_ea) if seg else ea + while end < hi and ida_bytes.is_code(ida_bytes.get_flags(end)): + nxt = int(ida_bytes.get_item_end(end)) + if nxt <= end: break + end = nxt + ok = bool(end > ea and ida_funcs.add_func(ea, end)) + else: ok = True + fn = db.functions.get_at(ea) + result = ({"addr": hex(ea), "ok": True, "start": hex(int(fn.start_ea)), + "end": hex(int(fn.end_ea)), "how": "auto" if automatic else "explicit-end"} + if ok and fn is not None else + {"addr": hex(ea), "ok": False, "error": f"IDA refused a function at {ea:#x}"}) +result +''' + +_OPERATIONS["set_thumb"] = r''' +import ida_bytes, ida_ida, ida_idp, ida_segment, ida_segregs +ea = int(str(a["addr"]), 16); treg = ida_idp.str2reg("T") +seg = ida_segment.getseg(ea) +if treg is None or treg < 0: + result = {"addr": hex(ea), "error": "no T register (not an ARM database)"} +elif seg is None: + result = {"addr": hex(ea), "error": "no segment"} +else: + current = ida_segregs.get_sreg(ea, treg) + current = 0 if current in (None, 0xFFFFFFFF, -1) else int(current) + want = {"on": 1, "off": 0}.get(str(a.get("mode", "toggle")).lower(), 0 if current else 1) + changed = False + if want and seg.bitness != 1: + ida_segment.set_segm_addressing(seg, 1); changed = True + size = max(int(ida_bytes.get_item_size(ea)), 2) + ida_bytes.del_items(ea, 0, size) + ok = bool(ida_segregs.split_sreg_range(ea, treg, want, ida_segregs.SR_user)) + now = ida_segregs.get_sreg(ea, treg) + result = {"addr": hex(ea), "thumb": bool(now), "was": bool(current), "ok": ok, + "bitness": ida_segment.getseg(ea).bitness, "forced_32bit": changed, + "db_64bit": bool(ida_ida.inf_get_app_bitness() == 64 and want)} +result +''' + +_OPERATIONS["thumb_scan"] = r''' +import ida_bytes, ida_funcs, ida_idp, ida_segment, ida_segregs, ida_ua +lo, hi = int(str(a["start"]), 16), int(str(a["end"]), 16) +apply, limit = bool(a.get("apply", True)), int(a.get("limit", 512)) +treg = ida_idp.str2reg("T"); found = []; applied = 0; cursor = lo +while cursor + 4 <= hi and len(found) < limit: + at = cursor; value = int(ida_bytes.get_dword(cursor)); cursor += 4 + if not value & 1: continue + target = value & ~1; seg = ida_segment.getseg(target) + if seg is None or not (seg.perm & ida_segment.SEGPERM_EXEC or seg.perm == 0): continue + flags = ida_bytes.get_flags(target) + if ida_bytes.is_data(flags): continue + item = {"at": hex(at), "value": hex(value), "target": hex(target), + "was_code": bool(ida_bytes.is_code(flags))}; found.append(item) + if not apply: continue + if treg is not None and treg >= 0: ida_segregs.split_sreg_range(target, treg, 1, ida_segregs.SR_user) + if not ida_bytes.is_code(ida_bytes.get_flags(target)): + ida_bytes.del_items(target, 0, 2) + if ida_ua.create_insn(target) <= 0: item["decoded"] = False; continue + item["decoded"] = True; item["function"] = bool(db.functions.get_at(target) or db.functions.create(target)); applied += 1 +result = {"start": hex(lo), "end": hex(hi), "found": found, "applied": applied, "n": len(found)} +result +''' + +_OPERATIONS["decomp_error"] = r''' +import ida_hexrays, ida_ida +ea = int(str(a["addr"]), 16); fn = db.functions.get_at(ea) +result = {"addr": hex(ea), "bitness": ida_ida.inf_get_app_bitness()} +if fn is None: + result["reason"] = "no function here" +else: + try: + failure = ida_hexrays.hexrays_failure_t(); cfunc = ida_hexrays.decompile_func(fn, failure) + if cfunc is not None: result["reason"] = "" + else: + result.update({"reason": failure.desc() or f"error {failure.code}", + "code": int(failure.code), "errea": hex(int(failure.errea))}) + except Exception as exc: result["reason"] = f"{type(exc).__name__}: {exc}" +result +''' + + +class CodeModeClient: + """A leased GUI/idalib database accessed through ``ida_codemode``.""" + + def __init__( + self, + binary_path: str, + *, + ttl: int = 0, + load_args: str = "", + processor: str | None = None, + loading_address: int | None = None, + file_type: str | None = None, + output_database: str | None = None, + spawn: bool = True, + new_database: bool = False, + ) -> None: + del ttl # managed-worker lifetime is lease-based, not idle-TTL based + self._path = os.path.abspath(os.path.expanduser(binary_path)) + parsed_processor, parsed_address, parsed_file_type = _parse_load_args(load_args) + self._processor = processor or parsed_processor + self._loading_address = loading_address if loading_address is not None else parsed_address + self._file_type = file_type or parsed_file_type + self._output_database = output_database + self._spawn = spawn + self._new_database = new_database + self._handle: DatabaseHandle | None = None + self._last_entry: RegistryEntry | None = None + self._connect_lock = threading.Lock() + + def connect(self, timeout: float = 1800.0, progress=None) -> "CodeModeClient": + _require_codemode() + with self._connect_lock: + if self._handle is not None and self._handle.connected: + return self + if progress: + progress(f"discovering Code Mode database for {os.path.basename(self._path)}…") + try: + # A Ctrl+L reload releases its current managed-worker lease, but + # that worker remains registered during Code Mode's final-lease + # grace period. Retry only that known handoff window. A GUI or + # another long-lived client remains busy and yields a clear + # failure rather than being modified underneath its owner. + deadline = time.monotonic() + min(timeout, 60.0) + while True: + try: + handle = DatabaseHandle.open( + self._path, + spawn=self._spawn, + timeout=max(0.1, timeout), + output_database=self._output_database, + processor=self._processor, + loading_address=self._loading_address, + file_type=self._file_type, + new_database=self._new_database, + ) + break + except IdbBusy: + if not self._new_database or time.monotonic() >= deadline: + raise + if progress: + progress("waiting for the previous Code Mode lease to close…") + # Remember the record before managed shutdown withdraws + # its JSON. The lifetime lock remains held until IDA has + # actually closed the IDB; waiting on it avoids racing a + # replacement worker into the old process's file lock. + expected = canonical_path( + self._output_database or expected_idb_path(self._path) + ) + owners = [item.entry for item in scan_instances(timeout=0.5) + if item.entry.idb_key == idb_key(expected)] + if owners: + self._wait_for_entry_release( + owners[0], max(0.0, deadline - time.monotonic()) + ) + else: + time.sleep(0.2) + if progress: + backend = handle.entry.backend + progress(f"attached to {backend} database; waiting for auto-analysis…") + handle.wait_autoanalysis(timeout=timeout) + except Exception as exc: # normalize the dependency's transport errors + raise self._connection_error(exc) from exc + self._handle = handle + self._last_entry = handle.entry + return self + + @staticmethod + def _connection_error(exc: BaseException) -> IDAConnectionError: + return IDAConnectionError(str(exc) or type(exc).__name__) + + @property + def connected(self) -> bool: + return self._handle is not None and self._handle.connected + + @property + def pid(self) -> int | None: + return self._handle.entry.pid if self._handle is not None else None + + @property + def backend(self) -> str | None: + return self._handle.entry.backend if self._handle is not None else None + + def execute_python(self, code: str, *, timeout: float | None = None) -> Any: + if not self.connected: + self.connect() + handle = self._handle + if handle is None: + raise IDAConnectionError("Code Mode database is not connected") + try: + response = handle.execute_python(code, timeout=timeout) + except RemoteError as exc: + details = exc.details or {} + message = str(exc) + if details.get("traceback"): + message += f"\n{details['traceback']}" + if exc.code == "operation_timeout": + raise IDATimeoutError(message) from exc + raise IDAToolError("execute_python", message) from exc + except (InstanceDisconnectedError, ClientError) as exc: + raise self._connection_error(exc) from exc + if not isinstance(response, dict) or "result" not in response: + raise IDAToolError("execute_python", "Code Mode returned an invalid execution result") + return response["result"] + + def invoke(self, operation: str, *, timeout: float | None = None, **args) -> Any: + """Execute one TUI domain operation through Code Mode.""" + if operation in ("idb_save", "save"): + return self.save_database() + if operation in ("server_health", "ping", "health", "state"): + return self.health() + body = _HEADS if operation == "heads" else _OPERATIONS.get(operation) + if body is None: + raise IDAToolError(operation, f"unknown ida-tui Code Mode operation: {operation}") + try: + return self.execute_python(_script(args, body), timeout=timeout) + except IDAToolError as exc: + if exc.tool == "execute_python": + raise IDAToolError(operation, exc.message) from exc + raise + + # Temporary source compatibility for external drivers/tests that used the + # old WorkerClient. Application code uses the accurately named invoke(). + call = invoke + + def save_database(self) -> dict[str, Any]: + if not self.connected: + self.connect() + handle = self._handle + if handle is None: + raise IDAConnectionError("Code Mode database is not connected") + try: + return handle.save_database() + except RemoteError as exc: + raise IDAToolError("save_database", str(exc)) from exc + except (InstanceDisconnectedError, ClientError) as exc: + raise self._connection_error(exc) from exc + + def health(self) -> dict[str, Any]: + if not self.connected: + self.connect() + assert self._handle is not None + entry = self._handle.entry + module = os.path.basename(entry.exe_path or entry.idb_path or self._path) + return { + "ok": self._handle.connected, + "module": module, + "backend": entry.backend, + "record_id": entry.record_id, + "input_path": entry.exe_path, + "idb_path": entry.idb_path, + } + + def keepalive(self, interval: float = 120.0) -> _NoopKeepAlive: + del interval + return _NoopKeepAlive() + + def resolve_db(self) -> str: + if not self.connected: + self.connect() + assert self._handle is not None + return self._handle.entry.record_id + + def set_db(self, db: str | None) -> None: + del db # one handle is permanently bound to one registered database + + def list_sessions(self) -> list[Session]: + if not self.connected: + self.connect() + assert self._handle is not None + entry = self._handle.entry + path = entry.exe_path or entry.idb_path or self._path + return [Session(session_id=entry.record_id, filename=os.path.basename(path), + input_path=path, is_active=True)] + + def close(self, grace: float = 0.0) -> None: + del grace + with self._connect_lock: + handle, self._handle = self._handle, None + if handle is not None: + self._last_entry = handle.entry + handle.close() # release our lease; never close a GUI/other client's DB + + @staticmethod + def _wait_for_entry_release(entry: "RegistryEntry", timeout: float) -> bool: + _require_codemode() + path = REGISTRY_DIR / f"{entry.record_id}.lock" + deadline = time.monotonic() + max(0.0, timeout) + while True: + lock = FileLock(path) + try: + if lock.try_acquire(): + return True + except OSError: + pass + finally: + lock.close() + if time.monotonic() >= deadline: + return False + time.sleep(min(0.1, deadline - time.monotonic())) + + def wait_released(self, timeout: float = 45.0) -> bool: + """Wait until a managed instance releases its lifetime lock. + + Normal application shutdown must not wait: another client may retain the + worker. This is an explicit test/maintenance helper for deleting a + temporary IDB safely after this client closes. GUI instances return + ``False`` immediately because clients never own their lifetime. + """ + entry = self._last_entry + if entry is None or entry.backend != "idalib": + return False + return self._wait_for_entry_release(entry, timeout) + + def __enter__(self) -> "CodeModeClient": + return self.connect() + + def __exit__(self, *exc) -> None: + self.close() diff --git a/idatui/domain.py b/idatui/domain.py index aedade9..6fe73c2 100644 --- a/idatui/domain.py +++ b/idatui/domain.py @@ -1,19 +1,15 @@ -"""Domain / paging layer: address-centric models over the raw MCP client. +"""Domain / paging layer: address-centric models over IDA Code Mode. This is where the "millions of lines" problem is solved, so the TUI widgets only ever see a viewport-sized slice. Every hard-won constraint from ``docs/PAGING_FINDINGS.md`` is encoded here: -* Per-call caps are silent (over the cap the server returns 10, not a clamp), so - we clamp page sizes ourselves: ``LIST_PAGE`` / ``DISASM_BLOCK`` <= the caps. -* ``next_offset`` is unreliable; we paginate by advancing ``len(data)``. -* ``disasm offset=N`` is O(N) with no resumable cursor, so windowed disassembly - is **block-cached** (revisits are free) and **prefetches** the next block on a - background thread (the client is concurrency-safe). -* ``include_total`` scans the whole function (~200ms on monsters); totals are - fetched once and cached. -* ``decompile`` can hard-fail on huge functions as a *soft* error (``code`` is - null); that is surfaced as data, not an exception. +* Page sizes remain bounded so remote execution returns viewport-scale JSON. +* Pagination advances by the number of rows actually returned. +* Deep head walks are block-cached (revisits are free) and neighboring blocks + prefetch through the thread-safe Code Mode client. +* Expensive function totals are fetched once and cached. +* Decompilation failures are surfaced as data, not application crashes. Everything here is synchronous and thread-safe. The TUI runs these calls from Textual worker threads; the internal prefetch pool is separate and small. @@ -22,10 +18,8 @@ Textual worker threads; the internal prefetch pool is separate and small. from __future__ import annotations import bisect -import json import re import threading -import urllib.request from concurrent.futures import ThreadPoolExecutor from collections.abc import Sequence from dataclasses import dataclass, field, replace @@ -36,7 +30,7 @@ from . import diag from .errors import IDAToolError if TYPE_CHECKING: # type hint only - from .worker_client import WorkerClient # noqa: F401 + from .codemode_client import CodeModeClient # Clamps derived from measured caps (list ~700, disasm ~500). Margin included. LIST_PAGE = 500 @@ -66,7 +60,7 @@ class Func: def from_raw(cls, d: dict) -> "Func": addr = _as_int(d["addr"]) name = d.get("name") - # An unnamed function (server returns null/empty) must still have a + # An unnamed function must still have a # usable string name — synthesize IDA's sub_ADDR so every consumer # (palette, sort, rename prefill) can treat name as a str. if not name: @@ -93,7 +87,7 @@ class Line: class Head(NamedTuple): - """One flat-listing item (from the ``heads`` server tool): a code + """One flat-listing item (from the Code Mode ``heads`` operation): a code instruction, a data item, or an undefined byte run. A ``NamedTuple`` rather than a dataclass because this is by far the @@ -114,7 +108,7 @@ class Head(NamedTuple): name: str | None = None raw: bytes | None = None # opcode/item bytes (filled in for code by the model) #: [(kind, text)] from IDA's own colour tags — mnem/reg/num/name/str/punct/… - #: None when the worker didn't provide them (older worker, or the spans + #: None when Code Mode didn't provide them (or the spans #: disagreed with the plain text, in which case the text wins). #: #: Held exactly as it came off the wire, and **read-only**. The worker @@ -301,7 +295,7 @@ class FunctionIndex: """A lazily-paginated, cached view of the function list. Loads pages of ``LIST_PAGE`` on demand, advancing by ``len(data)`` (never by - ``next_offset``). A single index instance corresponds to one server-side + ``next_offset``). A single index instance corresponds to one remote ``filter`` glob (``None`` = all functions). """ @@ -321,7 +315,7 @@ class FunctionIndex: query: dict = {"offset": offset, "count": LIST_PAGE} if self.filter: query["filter"] = self.filter - data = _query_data(self._prog.client.call("list_funcs", queries=[query])) + data = _query_data(self._prog.client.invoke("list_funcs", queries=[query])) added = 0 with self._lock: for d in data: @@ -431,7 +425,7 @@ class DisasmModel: code function this equals the heads row count that backs the lines.""" if self._total is not None: return self._total - payload = self._prog.client.call( + payload = self._prog.client.invoke( "disasm", addr=hex(self.ea), max_instructions=1, include_total=True ) total = payload.get("total_instructions") @@ -492,7 +486,7 @@ class DisasmModel: # The function disasm view is a listing filtered to the function: fetch a # block of heads (one per instruction for code). Over-fetch one row so # the block knows where its last instruction ends (opcode-byte sizing). - payload = self._prog.client.call( + payload = self._prog.client.invoke( "heads", addr=hex(self.ea), offset=b * self.BLOCK, count=self.BLOCK + 1, **self._end_kw(), ) @@ -633,15 +627,15 @@ class ListingModel: """A flat, IDA-style disassembly *listing* over one segment: code, data and undefined heads interleaved, unlike ``DisasmModel`` (one function, code only). - Backed by the injected ``heads`` server tool, which walks item heads and - renders each via ``generate_disasm_line``. The segment is walked lazily in + Backed by the Code Mode adapter's ``heads`` operation, which walks item heads + and renders each via ``generate_disasm_line``. The segment is walked lazily in forward pages (``FunctionIndex`` style); line index == position in the walked head list. Random access to an address is O(distance-from-seg-start) the first time (then cached) — the same tradeoff as ``disasm offset=N``. Grows on demand as the viewport scrolls. Synchronous + thread-safe. """ - PAGE = 500 # heads per server call (well under the tool's 2000 cap) + PAGE = 500 # viewport-scale heads per Code Mode execution def __init__(self, program: "Program", seg_start: int, seg_end: int, name: str | None = None): @@ -754,7 +748,7 @@ class ListingModel: if self._done or self._next is None: return 0 frm = self._next - payload = self._prog.client.call( + payload = self._prog.client.invoke( "heads", addr=hex(frm), count=self.PAGE, annotate=True) rows = payload.get("heads", []) if isinstance(payload, dict) else [] cur = payload.get("cursor", {}) if isinstance(payload, dict) else {} @@ -1019,7 +1013,7 @@ class ListingModel: # the expectation rather than asking first means a page that HAS changed # still costs one round trip. try: - payload = self._prog.client.call( + payload = self._prog.client.invoke( "heads", addr=hex(addr), count=self.PAGE, annotate=True, expect="" if want_digest is None else str(want_digest)) except Exception: # noqa: BLE001 -- keep the old text rather than blank @@ -1233,7 +1227,7 @@ class HexModel: class Program: """The bound analysis session: models, caches, and a small prefetch pool.""" - def __init__(self, client: "WorkerClient", prefetch_workers: int = 2): + def __init__(self, client: "CodeModeClient", prefetch_workers: int = 2): self.client = client self._pool = ThreadPoolExecutor( max_workers=prefetch_workers, thread_name_prefix="idatui-prefetch" @@ -1257,7 +1251,7 @@ class Program: self._sections: list[tuple[int, int, str]] | None = None self._fileregions: list[tuple[int, int, int]] | None = None self._hexmodel: "HexModel | None" = None - self._no_read_raw = False # set if the server lacks the read_raw tool + self._no_read_raw = False # compatibility fallback for alternate clients self._lock = threading.Lock() # -- prefetch plumbing ------------------------------------------------- # @@ -1284,17 +1278,14 @@ class Program: """Sorted raw segment map [(start, end, file_off, name)] — the single source for sections()/file_regions()/image_range. Cached. - Uses the injected ``file_regions`` tool (a plain segment walk, ~ms). - This deliberately AVOIDS ``survey_binary``, which also computes function - counts / strings / stats and takes *seconds* on a large IDB (it was the - cause of the multi-second hex-pane open). Falls back to survey_binary - only if the injected tool is missing. + Uses the Code Mode adapter's ``file_regions`` operation (a plain segment + walk, ~ms), avoiding broad binary surveys on the hex-pane open path. """ if self._segments_cache is not None: return self._segments_cache segs: list[tuple[int, int, int, str]] = [] try: - r = self.client.call("file_regions") + r = self.client.invoke("file_regions") for d in (r.get("regions", []) if isinstance(r, dict) else []): if isinstance(d, dict) and "start" in d: segs.append((_as_int(d["start"]), _as_int(d["end"]), @@ -1303,7 +1294,7 @@ class Program: segs = [] if not segs: # older server without file_regions -> survey_binary (slow) try: - sb = self.client.call("survey_binary") + sb = self.client.invoke("survey_binary") for s in (sb.get("segments", []) if isinstance(sb, dict) else []): try: segs.append((_as_int(s["start"]), _as_int(s["end"]), -1, @@ -1345,8 +1336,7 @@ class Program: def file_regions(self) -> list[tuple[int, int, int]]: """Sorted [(start, end, file_off)] mapping loaded segments to raw file - offsets (file_off == -1 for non-file-backed, e.g. .bss). Cached; needs - the injected ``file_regions`` server tool.""" + offsets (file_off == -1 for non-file-backed, e.g. .bss). Cached.""" if self._fileregions is not None: return self._fileregions regions = [(s, e, fo) for s, e, fo, _nm in self._segments()] @@ -1364,15 +1354,14 @@ class Program: def read_bytes(self, ea: int, n: int) -> bytes: """Raw bytes [ea, ea+n) from IDA (gaps read as zero). - Fast path: the injected ``read_raw`` tool returns one contiguous hex - string (C-speed both ends). Falls back to the stock ``get_bytes`` (a - per-byte '0x..'-with-spaces string) on an older server without it. + The Code Mode adapter returns one contiguous hex string (C-speed in IDA). + A legacy ``get_bytes`` decoding fallback remains for alternate clients. """ if n <= 0: return b"" if not self._no_read_raw: try: - r = self.client.call("read_raw", addr=hex(ea), size=int(n)) + r = self.client.invoke("read_raw", addr=hex(ea), size=int(n)) h = r.get("hex") if isinstance(r, dict) else None if isinstance(h, str): out = bytes.fromhex(h) @@ -1386,7 +1375,7 @@ class Program: except (ValueError, KeyError): pass # malformed hex -> fall through to the legacy decoder try: - r = self.client.call("get_bytes", regions=[{"addr": hex(ea), "size": int(n)}]) + r = self.client.invoke("get_bytes", regions=[{"addr": hex(ea), "size": int(n)}]) except IDAToolError: return b"\x00" * n res = r.get("result", []) if isinstance(r, dict) else [] @@ -1437,7 +1426,7 @@ class Program: def list_structs(self, filter: str = "") -> list[Struct]: """All local structs/unions (optionally name-substring filtered), sorted by name.""" - payload = self.client.call("search_structs", filter=filter) + payload = self.client.invoke("search_structs", filter=filter) res = payload.get("result", []) if isinstance(payload, dict) else [] out = [Struct.from_raw(d) for d in res if isinstance(d, dict) and d.get("name") @@ -1447,9 +1436,9 @@ class Program: def struct_source(self, name: str) -> str: """A C definition for ``name`` reconstructed from its member layout - (the server exposes members, not printable source). Faithful to IDA's + (the remote operation exposes members, not printable source). Faithful to IDA's field names/types; array dims are moved after the field name.""" - payload = self.client.call( + payload = self.client.invoke( "type_inspect", queries=[{"name": name, "include_members": True}]) res = payload.get("result", []) if isinstance(payload, dict) else [] info = res[0] if res and isinstance(res[0], dict) else {} @@ -1472,7 +1461,7 @@ class Program: def declare_type(self, decl: str) -> str | None: """Create or update a C type. Returns None on success, else the parse error. (Re-declaring a name updates it in place.)""" - payload = self.client.call("declare_type", decls=decl) + payload = self.client.invoke("declare_type", decls=decl) res = payload.get("result", []) if isinstance(payload, dict) else [] if res and isinstance(res[0], dict): return res[0].get("error") @@ -1481,10 +1470,9 @@ class Program: # -- function / variable types ---------------------------------------- # def func_types(self, ea: int) -> FuncTypes | None: """Structured decompiler types for the function at ``ea`` (prototype + - local variables). None if ``ea`` isn't a decompilable function. Requires - the injected ``func_types`` server tool.""" + local variables). None if ``ea`` isn't a decompilable function.""" try: - r = self.client.call("func_types", addr=hex(ea)) + r = self.client.invoke("func_types", addr=hex(ea)) except IDAToolError: return None if not isinstance(r, dict) or r.get("error"): @@ -1497,7 +1485,7 @@ class Program: def set_function_type(self, ea: int, signature: str) -> str | None: """Set a function's prototype. None on success, else an error string.""" - r = self.client.call("set_type", edits=[{"addr": hex(ea), "signature": signature}]) + r = self.client.invoke("set_type", edits=[{"addr": hex(ea), "signature": signature}]) res = r.get("result", []) if isinstance(r, dict) else [] row = res[0] if res and isinstance(res[0], dict) else {} if row.get("ok"): @@ -1506,9 +1494,9 @@ class Program: def data_type(self, ea: int) -> dict | None: """Current type info for a data item/global: {addr,name,type,size,is_func}. - None if the tool is unavailable or the address isn't mapped.""" + None if the operation fails or the address isn't mapped.""" try: - r = self.client.call("data_type", addr=hex(ea)) + r = self.client.invoke("data_type", addr=hex(ea)) except IDAToolError: return None if not isinstance(r, dict) or r.get("error"): @@ -1517,7 +1505,7 @@ class Program: def set_data_type(self, ea: int, decl: str) -> str | None: """Set a global/data item's type. None on success, else an error string.""" - r = self.client.call( + r = self.client.invoke( "set_type", edits=[{"kind": "global", "addr": hex(ea), "type": decl}]) res = r.get("result", []) if isinstance(r, dict) else [] row = res[0] if res and isinstance(res[0], dict) else {} @@ -1526,9 +1514,9 @@ class Program: return row.get("error") or "failed to set the type" def set_lvar_type(self, fn_ea: int, var: str, ty: str) -> str | None: - """Set a decompiler local variable's type (via the injected server tool). + """Set a decompiler local variable's type through ida-domain pseudocode. None on success, else an error string.""" - r = self.client.call("set_lvar_type", addr=hex(fn_ea), variable=var, type=ty) + r = self.client.invoke("set_lvar_type", addr=hex(fn_ea), variable=var, type=ty) if isinstance(r, dict) and r.get("error"): return r["error"] if isinstance(r, dict) and not r.get("ok"): @@ -1537,15 +1525,14 @@ class Program: def delete_type(self, name: str) -> str | None: """Delete a named type. Returns None on success, else an error string. - Requires a server-side ``del_type`` tool; if absent, a clear message is - returned instead of raising.""" + Returns a clear error instead of raising when the runtime cannot do it.""" try: - self.client.call("del_type", name=name) + self.client.invoke("del_type", name=name) return None except IDAToolError as e: msg = e.message if "not found" in msg.lower() and "del_type" in msg: - return "delete needs a 'del_type' tool on the ida-pro-mcp server" + return "the connected Code Mode runtime cannot delete local types" return msg # -- disassembly ------------------------------------------------------- # @@ -1559,13 +1546,7 @@ class Program: # -- decompilation ----------------------------------------------------- # def decompile(self, ea: int, refresh: bool = False) -> Decompilation: - """Full pseudocode for a function. - - The server truncates responses over 50KB (strings clipped to 1000 - chars) but caches the full output and exposes it at - ``_meta.ida_mcp.download_url``. We transparently fetch that so the view - always gets the complete body, not a 1KB stub. - """ + """Full pseudocode for a function, returned directly by Code Mode.""" if not refresh: with self._lock: hit = self._decomp.get(ea) @@ -1574,10 +1555,10 @@ class Program: dec, hit_gen = hit if hit_gen == gen: return dec - # Cached before a rename: names may be stale. Drop the server's - # Hex-Rays cache so the refetch reflects the new names. + # Cached before a rename: names may be stale. Drop Hex-Rays' + # cache so the refetch reflects the new names. try: - self.client.call("force_recompile", items=[{"addr": hex(ea)}]) + self.client.invoke("force_recompile", items=[{"addr": hex(ea)}]) except Exception: # noqa: BLE001 pass # Bound the decompile: a function Hex-Rays can't handle tends to stall @@ -1587,7 +1568,10 @@ class Program: # rpcclient socket timeout, and cache the failure below so a re-request # returns instantly instead of re-grinding. try: - envelope = self.client.call_envelope( + # Code Mode returns the complete JSON result directly; unlike the + # old MCP tool transport there is no structured-content envelope or + # out-of-band download URL to unwrap. + payload = self.client.invoke( "decompile", addr=hex(ea), timeout=DECOMPILE_TIMEOUT ) except Exception as e: # noqa: BLE001 -- surface as a failed decompile @@ -1595,15 +1579,6 @@ class Program: with self._lock: self._decomp[ea] = (dec, self._name_gen) return dec - result = envelope.get("result", {}) - payload = result.get("structuredContent") - if payload is None: # fall back to text content - payload = self.client._extract_payload("decompile", result) - meta = (result.get("_meta") or {}).get("ida_mcp") - if isinstance(meta, dict) and meta.get("download_url"): - full = self._fetch_output(meta["download_url"]) - if isinstance(full, dict) and full.get("code"): - payload = full dec = _parse_decompilation(ea, payload) with self._lock: self._decomp[ea] = (dec, self._name_gen) @@ -1681,18 +1656,18 @@ class Program: Undefine first so it works even when the bytes are currently part of a data/align item — ``create_insn`` refuses to carve into a live item.""" try: - self.client.call("undefine", items=[{"addr": hex(ea)}]) + self.client.invoke("undefine", items=[{"addr": hex(ea)}]) except IDAToolError: pass # nothing defined here yet -> just try to create the insn res = self._first_result( - self.client.call("define_code", items=[{"addr": hex(ea)}])) + self.client.invoke("define_code", items=[{"addr": hex(ea)}])) if res.get("error"): raise IDAToolError("define_code", f"@ {ea:#x}: {res['error']}") def decomp_error(self, ea: int) -> str: """Hex-Rays' own reason for refusing ``ea``, or "" if it won't say.""" try: - r = self.client.call("decomp_error", addr=hex(ea)) + r = self.client.invoke("decomp_error", addr=hex(ea)) except IDAToolError: return "" if not isinstance(r, dict): @@ -1711,7 +1686,7 @@ class Program: def thumb_scan(self, start: int, end: int, apply: bool = True) -> dict: """Find Thumb entry points from odd pointers in ``[start, end)``.""" - r = self.client.call("thumb_scan", start=hex(start), end=hex(end), + r = self.client.invoke("thumb_scan", start=hex(start), end=hex(end), apply=bool(apply)) if not isinstance(r, dict) or r.get("error"): raise IDAToolError("thumb_scan", @@ -1720,7 +1695,7 @@ class Program: def set_thumb(self, ea: int, mode: str = "toggle") -> dict: """Switch ARM/Thumb decoding at ``ea``. Returns the resulting state.""" - r = self.client.call("set_thumb", addr=hex(ea), mode=mode) + r = self.client.invoke("set_thumb", addr=hex(ea), mode=mode) if not isinstance(r, dict) or r.get("error"): raise IDAToolError("set_thumb", f"@ {ea:#x}: {(r or {}).get('error', 'failed')}") @@ -1729,11 +1704,11 @@ class Program: def define_code_run(self, ea: int, limit: int = 20000) -> dict: """Disassemble consecutively from ``ea`` until something stops it. - Falls back to a single instruction when the worker predates the tool, so - an old worker degrades to the previous behaviour instead of failing. + Falls back to a single instruction for alternate clients that do not + provide the run operation. """ try: - r = self.client.call("define_code_run", addr=hex(ea), limit=int(limit)) + r = self.client.invoke("define_code_run", addr=hex(ea), limit=int(limit)) except IDAToolError: self.define_code(ea) return {"count": 1, "stopped": "single", "end": hex(ea)} @@ -1745,14 +1720,14 @@ class Program: def define_func(self, ea: int) -> dict: """Create a function starting at ``ea`` (IDA's 'p'). - Prefers the injected tool, which works out the end when IDA can't; - falls back to the plain one for an older worker. + Prefers the Code Mode operation, which works out the end when IDA can't; + falls back to a plain create for alternate clients. """ try: - r = self.client.call("define_func_run", addr=hex(ea)) + r = self.client.invoke("define_func_run", addr=hex(ea)) except IDAToolError: res = self._first_result( - self.client.call("define_func", items=[{"addr": hex(ea)}])) + self.client.invoke("define_func", items=[{"addr": hex(ea)}])) if res.get("error"): raise IDAToolError("define_func", f"@ {ea:#x}: {res['error']}") return {"ok": True, "how": "legacy"} @@ -1766,7 +1741,7 @@ class Program: item: dict = {"addr": hex(ea)} if size: item["size"] = int(size) - res = self._first_result(self.client.call("undefine", items=[item])) + res = self._first_result(self.client.invoke("undefine", items=[item])) if res.get("error"): raise IDAToolError("undefine", f"@ {ea:#x}: {res['error']}") @@ -1776,7 +1751,7 @@ class Program: item: dict = {"addr": hex(ea), "type": type_decl} if name: item["name"] = name - res = self._first_result(self.client.call("make_data", items=[item])) + res = self._first_result(self.client.invoke("make_data", items=[item])) if res.get("ok") is False or res.get("error"): raise IDAToolError( "make_data", f"@ {ea:#x}: {res.get('error') or 'rejected'}") @@ -1784,7 +1759,7 @@ class Program: def make_string(self, ea: int, length: int = 0, kind: str = "c") -> str: """Create a string literal at ``ea`` (IDA's 'A'); auto-length when 0. Returns the decoded contents.""" - r = self.client.call("make_string", addr=hex(ea), length=int(length), kind=kind) + r = self.client.invoke("make_string", addr=hex(ea), length=int(length), kind=kind) res = r if isinstance(r, dict) else {} if not res.get("ok"): raise IDAToolError( @@ -1801,7 +1776,7 @@ class Program: ``cycle``/``back`` (step the stops that make sense for this value) or a format by name. ``show`` reports without changing anything. """ - r = self.client.call("op_format", addr=hex(ea), mode=str(mode), + r = self.client.invoke("op_format", addr=hex(ea), mode=str(mode), col=int(col), n=int(n)) res = r if isinstance(r, dict) else {} if res.get("error"): @@ -1825,7 +1800,7 @@ class Program: if hit is not None and hit[1] == gen: return hit[0] try: - r = self.client.call("pc_nums", addr=hex(fn_ea)) + r = self.client.invoke("pc_nums", addr=hex(fn_ea)) except Exception: # noqa: BLE001 -- an older worker hasn't got the tool r = {} out: dict[int, list[tuple[int, int, str, int, int]]] = {} @@ -1848,7 +1823,7 @@ class Program: listing's format doesn't reach the pseudocode and vice versa, so this is a separate call rather than a flag on ``op_format``. """ - r = self.client.call("pc_num_format", addr=hex(fn_ea), mode=str(mode), + r = self.client.invoke("pc_num_format", addr=hex(fn_ea), mode=str(mode), line=int(line), col=int(col)) res = r if isinstance(r, dict) else {} if res.get("error"): @@ -1865,18 +1840,6 @@ class Program: sec = None return f"{sec} @ {ea:#x}" if sec else f"<no function> @ {ea:#x}" - @staticmethod - def _fetch_output(url: str, timeout: float = 15.0): - """GET the server's cached full-output blob (plain HTTP, not MCP).""" - try: - with urllib.request.urlopen(url, timeout=timeout) as r: - return json.loads(r.read().decode("utf-8", "replace")) - except Exception as e: # noqa: BLE001 -- fall back to the truncated preview - # The user gets CLIPPED pseudocode with no indication that a fetch - # failed rather than the function genuinely being that short. - diag.note(f"decompile: full-body fetch {url}", e) - return None - def strings(self, min_len: int = 4, refresh: bool = False) -> list[StrLit]: """Every string literal in the binary (IDA's Shift+F12 list), paged in full and cached. ``[]`` if the tool is unavailable.""" @@ -1889,7 +1852,7 @@ class Program: offset, page = 0, 2000 while True: try: - payload = self.client.call( + payload = self.client.invoke( "list_strings", offset=offset, count=page, min_len=min_len, refresh=(refresh and offset == 0)) except IDAToolError: @@ -1914,13 +1877,13 @@ class Program: def linkage(self) -> tuple[list[Linkage], list[Linkage]]: """``(imports, exports)`` for this binary, cached. ``([], [])`` if the - tool is unavailable — an old worker must not break the caller.""" + operation is unavailable — an alternate client must not break the caller.""" with self._lock: hit = self._linkage if hit is not None: return hit try: - payload = self.client.call("list_linkage", kind="both") + payload = self.client.invoke("list_linkage", kind="both") except IDAToolError: return ([], []) if not isinstance(payload, dict): @@ -1951,7 +1914,7 @@ class Program: if hit is not None and hit[1] == gen: return hit[0] try: - payload = self.client.call("decomp_map", addr=hex(ea)) + payload = self.client.invoke("decomp_map", addr=hex(ea)) except IDAToolError: return [] lines = payload.get("lines", []) if isinstance(payload, dict) else [] @@ -1981,7 +1944,7 @@ class Program: if hit is not None and hit[1] == gen: return hit[0] try: - payload = self.client.call("flowchart", addr=hex(ea)) + payload = self.client.invoke("flowchart", addr=hex(ea)) except IDAToolError: return None if not isinstance(payload, dict) or payload.get("error"): @@ -2054,7 +2017,7 @@ class Program: for _ in range(64): # bounded: ~128k heads if addr >= hi: break - payload = self.client.call("heads", addr=hex(addr), end=hex(hi), + payload = self.client.invoke("heads", addr=hex(addr), end=hex(hi), count=2000) rows = payload.get("heads", []) if isinstance(payload, dict) else [] if not rows: @@ -2083,13 +2046,13 @@ class Program: # -- cross-references & containing function --------------------------- # def function_of(self, ea: int) -> Func | None: """Return the function containing ``ea`` (resolves mid-function addrs).""" - payload = self.client.call("lookup_funcs", queries=[hex(ea)]) + payload = self.client.invoke("lookup_funcs", queries=[hex(ea)]) res = payload.get("result", []) if isinstance(payload, dict) else [] fn = res[0].get("fn") if res and isinstance(res[0], dict) else None return Func.from_raw(fn) if fn else None def xrefs_from(self, ea: int) -> list[Xref]: - payload = self.client.call( + payload = self.client.invoke( "xref_query", queries=[{"addr": hex(ea), "direction": "from", "include_fn": True}], ) @@ -2101,9 +2064,9 @@ class Program: try: # xref_types adds a fine-grained `kind` (call/read/write/...) for the # xref dialog; fall back to xref_query (code/data only) if absent. - payload = self.client.call("xref_types", queries=q) + payload = self.client.invoke("xref_types", queries=q) except IDAToolError: - payload = self.client.call("xref_query", queries=q) + payload = self.client.invoke("xref_query", queries=q) return _parse_xrefs(payload) # -- address resolution ------------------------------------------------ # @@ -2121,17 +2084,17 @@ class Program: # (loc_/locret_): lookup_funcs would map a label to its *containing* # function's entry, so double-clicking a label jumped to the wrong place. try: - payload = self.client.call("resolve_names", queries=[s]) + payload = self.client.invoke("resolve_names", queries=[s]) res = payload.get("result", []) if isinstance(payload, dict) else [] ea = res[0].get("ea") if res and isinstance(res[0], dict) else None if ea: return _as_int(ea) except IDAToolError: - pass # older server without resolve_names -> fall back below + pass # alternate client without resolve_names -> fall back below # Fall back to function-name resolution (also drives the 'did you mean' # suggestion when the name is unknown). try: - payload = self.client.call("lookup_funcs", queries=[s]) + payload = self.client.invoke("lookup_funcs", queries=[s]) except IDAToolError as e: raise KeyError(f"cannot resolve {target!r}: {e}") from e res = payload.get("result", []) if isinstance(payload, dict) else [] @@ -2169,7 +2132,7 @@ class Program: """Set (empty text clears) the comment at ``ea``; affects both the disasm and decompiler views. Returns the raw payload so the caller can surface a soft per-item error. The caller must invalidate/recompile to see it.""" - return self.client.call("set_comments", items=[{"addr": hex(ea), "comment": text}]) + return self.client.invoke("set_comments", items=[{"addr": hex(ea), "comment": text}]) # -- invalidation (after edits) --------------------------------------- # def invalidate(self, ea: int) -> None: diff --git a/idatui/drive.py b/idatui/drive.py index 1c7d8c2..825111c 100644 --- a/idatui/drive.py +++ b/idatui/drive.py @@ -121,7 +121,8 @@ def cmd_pc(c, args): lines = d["code"].splitlines() if needle: nlow = needle.lower() - lines = [f"{i:4} {l}" for i, l in enumerate(lines) if nlow in l.lower()] + lines = [f"{i:4} {line}" for i, line in enumerate(lines) + if nlow in line.lower()] return "\n".join(lines) or f"(no line matches {needle!r})" return d["code"] diff --git a/idatui/edit_ctl.py b/idatui/edit_ctl.py index 435fab9..80e1109 100644 --- a/idatui/edit_ctl.py +++ b/idatui/edit_ctl.py @@ -242,7 +242,7 @@ class EditController: kind = "stack" batch = {"stack": {"func_addr": hex(cur.ea), "old": old, "new": new}} try: - res = prog.client.call("rename", batch=batch) + res = prog.client.invoke("rename", batch=batch) except IDAToolError as e: app.call_from_thread(app._status, f"rename failed: {e.message}") return @@ -274,7 +274,7 @@ class EditController: app = self.app assert app.program is not None try: - res = app.program.client.call( + res = app.program.client.invoke( "rename", batch={"data": {"addr": hex(addr), "new": name}}) except IDAToolError as e: app.call_from_thread(app._status, f"name failed: {e.message}") diff --git a/idatui/errors.py b/idatui/errors.py index 29b09ae..aaf2dc5 100644 --- a/idatui/errors.py +++ b/idatui/errors.py @@ -1,10 +1,8 @@ -"""Transport-agnostic error hierarchy and the Session model. +"""TUI-facing error hierarchy and lightweight database session model. -These were originally defined in client.py (the ida-pro-mcp HTTP client), but the -idalib worker path (worker_client / domain / app) needs the same exception types -and Session dataclass without dragging in the HTTP transport. They live here so -both backends share one definition; client.py re-exports them for backwards -compatibility with the (deprecated) mcp tooling and the stress tests. +The Code Mode adapter normalizes ``ida_codemode.client`` transport and execution +errors into these types so the domain and Textual layers do not depend on HTTP or +registry implementation details. """ from __future__ import annotations diff --git a/idatui/launch.py b/idatui/launch.py index 35221a1..0c53987 100644 --- a/idatui/launch.py +++ b/idatui/launch.py @@ -1,15 +1,13 @@ -"""One-shot launcher: ``ida-tui foo.elf`` and you're in the TUI. +"""One-shot launcher for the IDA Code Mode-backed TUI. -Spawns a private idalib worker (``idatui.worker``) that opens + auto-analyzes -THIS binary in its own process, talking to the TUI over a unix socket. No shared -supervisor, no HTTP: everything slow (open + analysis) happens behind the TUI's -loading overlay. +A path first resolves to a registered GUI database; when none matches, Code Mode +reuses or starts a managed idalib worker. With no path, a single registered +database is selected automatically. -Usage: +Usage:: - ida-tui /path/to/binary # open a binary and drive it - -Extras: --ttl, --no-keepalive, --rpc (all forwarded to the TUI). + ida-tui /path/to/binary + ida-tui # attach when exactly one database is registered """ from __future__ import annotations @@ -17,12 +15,6 @@ import argparse import os import sys -# The unpacked working-copy files IDA writes next to a `.i64` while a database is -# open. A hard-killed worker leaves them behind and the `.i64` then refuses to -# reopen ("Failed to open database"). Safe to delete when nothing holds the DB. -_LOCK_SUFFIXES = (".id0", ".id1", ".id2", ".nam", ".til") - - def _load_args(load: dict) -> str: """``load`` as IDA switches, for the single-binary path (no project ref). @@ -41,35 +33,19 @@ def _log(msg: str) -> None: print(f"ida-tui: {msg}", file=sys.stderr) -def _sweep_locks(binary: str) -> int: - """Remove stale unpacked DB files next to ``binary``. Returns how many. - - Never touches the ``.i64`` -- that is the real database, and nothing is - saved unless ``idb_save`` was called -- and never the input file itself. - The second guard is not theoretical: ``.til`` is both an unpacked-DB suffix - and the extension of an IDA type library, so ``ida-tui mylib.til`` swept its - own argument out of existence. Same for anything named ``*.id0``/``*.nam``. - """ - keep = os.path.abspath(binary) - stem = os.path.splitext(binary)[0] - n = 0 - for base in (binary, stem): # IDA may key on the full name or the stem - for suf in _LOCK_SUFFIXES: - victim = base + suf - if os.path.abspath(victim) == keep: - continue # that's what the user asked us to open - try: - os.remove(victim) - n += 1 - except OSError: - pass - return n +def _registered_databases() -> tuple[list[dict], list[dict]]: + """Ready and blocked Code Mode registrations, with normalized errors.""" + try: + from ida_codemode.registry import discover_instances + return discover_instances() + except Exception as exc: # discovery diagnostics belong at the CLI boundary + return [], [{"error": str(exc)}] def main(argv: list[str] | None = None) -> int: p = argparse.ArgumentParser( prog="ida-tui", - description="Open a binary in the IDA TUI (private idalib worker).") + description="Open a registered GUI or managed idalib database in the IDA TUI.") p.add_argument("binary", nargs="*", help="binary to open and analyze (several with --project " "creates/extends that project)") @@ -77,9 +53,9 @@ def main(argv: list[str] | None = None) -> int: help="open a multi-binary project (created from the given " "binaries if FILE doesn't exist)") p.add_argument("--ttl", type=int, default=1800, - help="worker idle-TTL seconds (default 1800)") + help="deprecated compatibility option (Code Mode uses leases)") p.add_argument("--no-keepalive", action="store_true", - help="do not run the keepalive heartbeat") + help="deprecated compatibility option (the lease is the heartbeat)") p.add_argument("--rpc", metavar="PATH", help="listen for RPC on this unix socket (puppeteer the TUI)") p.add_argument("--trace", metavar="FILE", @@ -94,7 +70,7 @@ def main(argv: list[str] | None = None) -> int: g.add_argument("--base", metavar="ADDR", help="load address, e.g. 0x8000000 (any base; NOT paragraphs)") g.add_argument("--ida-args", metavar="STR", dest="ida_args", - help="extra IDA command-line switches, passed through as-is") + help="legacy switches; only Code Mode-representable -p/-b/-T are accepted") args = p.parse_args(argv) load: dict = {} @@ -152,27 +128,42 @@ def main(argv: list[str] | None = None) -> int: _log(str(e)) return 2 else: - if len(args.binary) != 1: - _log("give exactly one binary, or use --project for several") + ready, blocked = _registered_databases() + if len(args.binary) > 1: + _log("give at most one binary, or use --project for several") return 2 - binary = os.path.abspath(os.path.expanduser(args.binary[0])) - if not os.path.isfile(binary): - _log(f"no such file: {binary}") + if args.binary: + binary = os.path.abspath(os.path.expanduser(args.binary[0])) + key = os.path.normcase(os.path.realpath(binary)) + registered = any( + key == os.path.normcase(os.path.realpath(str(item.get(field) or ""))) + for item in ready for field in ("exe_path", "idb_path") + if item.get(field) + ) + if not os.path.isfile(binary) and not registered: + _log(f"no such file or registered database: {binary}") + return 2 + elif len(ready) == 1: + item = ready[0] + binary = str(item.get("exe_path") or item.get("idb_path") or "") + _log(f"attaching to registered {item.get('backend')} database: {binary}") + elif not ready: + detail = f" ({blocked[0].get('error')})" if blocked else "" + _log(f"no registered Code Mode database; pass a binary path{detail}") return 2 - if not os.access(os.path.dirname(binary), os.W_OK): - _log(f"directory not writable (IDA writes a .i64 there): " - f"{os.path.dirname(binary)}") + else: + _log("several Code Mode databases are registered; pass one of these paths:") + for item in ready: + _log(f" {item.get('exe_path') or item.get('idb_path')} " + f"[{item.get('backend')}, {item.get('record_id')}]") return 2 - swept = _sweep_locks(binary) # a crashed worker can leave the DB wedged - if swept: - _log(f"cleared {swept} stale lock file(s) from a crashed worker") - # Hand off to the TUI (imported late so --help works without textual). It - # spawns the worker behind its loading overlay while auto-analysis runs. + # Hand off to the TUI (imported late so --help works without Textual). Code + # Mode discovery/opening happens behind its loading overlay. try: from .app import IdaTui except ImportError as e: - _log(f"the TUI needs textual; run with ~/ida-venv/bin/python ({e})") + _log(f"TUI dependencies are missing; run `uv sync` ({e})") return 1 # Ask the terminal about graphics support NOW: the query needs a reply from # stdin, and once Textual starts it reads stdin on its own thread and would diff --git a/idatui/pane.py b/idatui/pane.py index 2592581..c31a93c 100644 --- a/idatui/pane.py +++ b/idatui/pane.py @@ -23,9 +23,9 @@ per pane in the registry, so stop/list/capture/keys keep working across both python -m idatui.pane capture --pane <pane> python -m idatui.pane keys --pane <pane> Escape -Requires: running inside tmux or zellij. Each pane spawns its own private idalib -worker (no shared supervisor). Uses ~/ida-venv/bin/python for the TUI (needs -textual) unless --python / IDATUI_PYTHON says otherwise. +Requires: running inside tmux or zellij. Each pane leases a registered GUI or +shared managed idalib database through Code Mode. Uses ~/ida-venv/bin/python for +the TUI (needs textual) unless --python / IDATUI_PYTHON says otherwise. """ from __future__ import annotations @@ -33,7 +33,6 @@ import argparse import json import os import secrets -import signal import subprocess import sys import time @@ -251,35 +250,9 @@ def _pane_keys(pane: str, keys: list[str], mux: str | None = None) -> None: subprocess.run(["tmux", "send-keys", "-t", pane, *keys], check=True) -# --------------------------------------------------------------------------- # -# idalib worker reaping -# -# ``pane stop`` kills the TUI pane, but a hard-killed pane can leave its private -# idalib worker (idatui/worker.py) running. A worker is only *safe* to reap when -# no idatui pane is live (then every worker is orphaned), which avoids killing an -# in-use analyser. -# --------------------------------------------------------------------------- # -_WORKER_PATTERN = r"idatui/worker\.py" - - -def _worker_pids() -> list[int]: - """PIDs of our private per-pane idalib worker processes (idatui/worker.py), - never our own PID.""" - try: - out = subprocess.run(["pgrep", "-f", _WORKER_PATTERN], - capture_output=True, text=True) - except OSError: - return [] - me = os.getpid() - pids: list[int] = [] - for tok in out.stdout.split(): - try: - pid = int(tok) - except ValueError: - continue - if pid != me: - pids.append(pid) - return pids +# Code Mode owns database process lifetime: a closed pane drops its lease at the +# socket/kernel boundary and Code Mode decides whether a managed worker still +# has clients. There is nothing for the pane layer to reap. def _count_live_panes() -> int: @@ -288,20 +261,9 @@ def _count_live_panes() -> int: def _reap_orphan_workers(force: bool = False) -> int: - """Kill leaked idalib workers when it is safe (no live pane) or ``force``. - - Returns the number of workers signalled. Best-effort; never raises. - """ - if not force and _count_live_panes() > 0: - return 0 - reaped = 0 - for pid in _worker_pids(): - try: - os.kill(pid, signal.SIGKILL) - reaped += 1 - except OSError: - pass - return reaped + """Compatibility no-op: Code Mode workers are shared and lease-managed.""" + del force + return 0 # --------------------------------------------------------------------------- # @@ -332,16 +294,8 @@ def spawn(args) -> int: print(f"error: no such project: {project}", file=sys.stderr) return 2 - # Reap workers leaked by previously-stopped/crashed panes so we don't spawn - # into a full IDA_MCP_MAX_WORKERS (which makes the new TUI hang forever, - # never reaching ready). No-op while any pane is live. - reaped = _reap_orphan_workers() - if reaped: - print(f"reaped {reaped} orphaned idalib worker(s) before spawn", - file=sys.stderr) - - # the command the pane runs: the launcher spawns a private idalib worker for - # this binary and becomes the TUI, so kill-pane tears the whole thing down. + # The pane owns only the TUI. Code Mode's lease cleanup handles crashes; + # kill-pane must never reap a shared GUI/idalib database. if project is not None: # launch takes: --project FILE [binaries...]; extra binaries are added to # the project (and a missing project file is created from them). @@ -393,9 +347,8 @@ def _wait_ready(sock: str, timeout: float, pane: str, stuck_after: float = 45.0, mux: str | None = None) -> dict[str, Any]: """Poll the socket + ping until the TUI reports ready (or timeout). - Emits a one-time hint to stderr if it's still not ready after ``stuck_after`` - seconds, so a wedged idalib worker / full worker pool surfaces a diagnostic - instead of an unexplained silent hang. + Emits a one-time hint if Code Mode discovery/opening is still not ready after + ``stuck_after`` seconds. """ start = time.time() deadline = start + timeout @@ -416,9 +369,8 @@ def _wait_ready(sock: str, timeout: float, pane: str, warned = True why = ("RPC socket not created yet" if not os.path.exists(sock) else "TUI up but analysis not ready") - print(f"still waiting ({int(time.time() - start)}s): {why}. If this " - f"hangs, the idalib worker may be stuck — try " - f"`python -m idatui.pane reap`.", file=sys.stderr) + print(f"still waiting ({int(time.time() - start)}s): {why}. " + f"Check Code Mode registrations and worker logs.", file=sys.stderr) time.sleep(0.4) last = dict(last) last["ready"] = False @@ -515,13 +467,9 @@ def list_panes(args) -> int: def reap(args) -> int: - """Kill leaked idalib workers (safe when no pane is live; --force overrides).""" - live = _count_live_panes() - n = _reap_orphan_workers(force=args.force) - print(json.dumps({"reaped_workers": n, "live_panes": live, "forced": args.force})) - if n == 0 and not args.force and live > 0: - print(f"note: {live} live pane(s) — not reaping in-use workers; pass " - f"--force to reap anyway", file=sys.stderr) + """Deprecated no-op; shared Code Mode workers are managed by leases.""" + print(json.dumps({"reaped_workers": 0, "live_panes": _count_live_panes(), + "forced": args.force, "deprecated": True})) return 0 @@ -624,9 +572,8 @@ def main(argv: list[str]) -> int: ls.add_argument("--prune", action="store_true", help="drop dead panes (and their sockets)") ls.set_defaults(fn=list_panes) - rp = sub.add_parser("reap", help="kill leaked idalib workers (frees worker slots)") - rp.add_argument("--force", action="store_true", - help="reap even while panes are live (may kill an in-use analyser)") + rp = sub.add_parser("reap", help="deprecated no-op (Code Mode uses shared leases)") + rp.add_argument("--force", action="store_true", help=argparse.SUPPRESS) rp.set_defaults(fn=reap) cp = sub.add_parser("capture", help="print a pane's visible screen") diff --git a/idatui/pool.py b/idatui/pool.py index ae37c25..465dff2 100644 --- a/idatui/pool.py +++ b/idatui/pool.py @@ -1,23 +1,16 @@ -"""WorkerPool — keeps a live idalib worker per project binary, within a budget. +"""DatabasePool — LRU leases on Code Mode databases for a project. -One worker process holds exactly one database (idalib is single-DB and -main-thread-only), so a project with N binaries means up to N processes. They are -not cheap and they do not share: a worker on ``bash`` measures ~126 MB RSS / -117 MB PSS, and the database working set dominates for anything larger -(``libcrypto.so.3``'s ``.i64`` alone is 72 MB). +Code Mode may bind a lease to an existing IDA GUI or to a shared managed idalib +worker. The pool therefore owns *client interest*, never an IDA process. Releasing +an LRU entry persists managed IDBs but does not implicitly save a GUI, then closes +only this TUI's lease; other clients and GUI sessions remain alive. Managed workers exit themselves after their final lease. -Residency is therefore bounded by a **memory budget**, not a worker count — a -count is the wrong knob when one project holds both a 50 KB helper and a 6 MB -crypto library. Workers are spawned lazily on first use, kept resident while they -fit, and least-recently-used ones evicted when they don't. Eviction **saves the -database first**, so coming back is a load rather than a re-analysis. - -The pool never evicts the active binary, nor anything pinned. +The historical memory budget remains useful for managed idalib instances, while +GUI process memory is only advisory. The active and pinned databases are never +released to satisfy it. """ from __future__ import annotations -import os - from .project import BinaryRef, Project #: Fallback budget if /proc/meminfo can't be read (MB). @@ -36,11 +29,11 @@ def _total_ram_mb() -> int: def _pss_mb(pid: int | None) -> int: - """Proportional set size of a worker, in MB. + """Proportional set size of the leased instance process, in MB. - PSS (not RSS) is the honest per-worker cost: it splits shared pages between - the processes mapping them. In practice workers share very little, so the two - are close, but PSS is what makes summing across workers meaningful. + PSS is useful for managed idalib workers. For GUI/shared processes it is only + advisory because the TUI neither owns all that memory nor controls process + exit. """ if not pid: return 0 @@ -54,13 +47,19 @@ def _pss_mb(pid: int | None) -> int: return 0 -def _default_spawn(ref: BinaryRef, ttl: int): # pragma: no cover - needs idalib - from .worker_client import WorkerClient - return WorkerClient(ref.staged, ttl=ttl, load_args=ref.load_args) +def _default_spawn(ref: BinaryRef, ttl: int, *, new_database: bool = False): # pragma: no cover - needs IDA + from .codemode_client import CodeModeClient + return CodeModeClient( + ref.staged, + ttl=ttl, + load_args=ref.load_args, + output_database=ref.db, + new_database=new_database, + ) -class WorkerPool: - """Live workers for a project's binaries, keyed by label.""" +class DatabasePool: + """Live Code Mode database leases, keyed by project label.""" def __init__(self, project: Project, *, budget_mb: int | None = None, ttl: int = 1800, spawn=None, mem_fn=None) -> None: @@ -71,6 +70,7 @@ class WorkerPool: self._clients: dict[str, object] = {} self._lru: list[str] = [] # least-recently-used first self._pinned: set[str] = set() + self._recreate: set[str] = set() # Ctrl+L: next attachment creates a fresh IDB self.active: str | None = None # never evicted if budget_mb is None: ram = _total_ram_mb() @@ -98,11 +98,11 @@ class WorkerPool: # -- acquire ----------------------------------------------------------- # def get(self, label: str, progress=None): - """A live client for ``label``, spawning it (and making room) if needed. + """A live client for ``label``, attaching or spawning as needed. - Staging and the scratch sweep happen here: a worker killed hard last time - leaves unpacked ``.id0/.id1/...`` behind, and the database then refuses to - reopen. Nothing else holds this DB (one worker per label), so it is safe. + Do not sweep IDA scratch files here: a registered GUI or another Code + Mode client may own the database. Code Mode's registry locks and health + probes are the authority for safe discovery and stale-record cleanup. """ client = self._clients.get(label) if client is not None: @@ -118,28 +118,29 @@ class WorkerPool: note(f"staging {ref.label}\u2026") self.project.stage(ref) - self.project.sweep_scratch(ref) note(f"opening {ref.label}\u2026") - client = self._spawn(ref, self._ttl) + fresh = label in self._recreate + client = (_default_spawn(ref, self._ttl, new_database=fresh) + if self._spawn is _default_spawn else self._spawn(ref, self._ttl)) connect = getattr(client, "connect", None) if connect is not None: connect(progress=progress) if progress is not None else connect() self._clients[label] = client + self._recreate.discard(label) self._lru.append(label) self._enforce_budget(protect=label) return client def prewarm(self, label: str, progress=None) -> bool: - """Spawn a worker for ``label`` only if it fits the budget AS IT STANDS. + """Attach a database for ``label`` only if it fits the current budget. Pre-warming must never cost residency: evicting a binary the user actually visited to speculatively load one they haven't is a straight downgrade, and the eviction would also throw away that binary's caches. So this refuses rather than making room, and returns False. - The cost of a worker that doesn't exist yet can only be estimated; the - largest resident one is the best evidence available (they are all the - same program with a different database). With nothing resident we have + The cost of a database not attached yet can only be estimated; the + largest resident instance is the best evidence available. With nothing resident we have no evidence at all, so we allow one — that is the case where the budget is certainly free. """ @@ -153,13 +154,19 @@ class WorkerPool: return False self.get(label, progress=progress) # get() enforces the budget protecting the NEW label; if that had to - # evict, our estimate was wrong and the speculative worker is the one + # evict, our estimate was wrong and the speculative lease is the one # that should go — never a binary the user chose. if self.memory_mb() > self.budget_mb and label != self.active: self.evict(label) return False return True + def recreate_on_next_open(self, label: str) -> None: + """Request a fresh IDB after the current lease has been released.""" + if self.project.by_label(label) is None: + raise KeyError(f"no such binary in the project: {label}") + self._recreate.add(label) + def _touch(self, label: str) -> None: if label in self._lru: self._lru.remove(label) @@ -171,16 +178,21 @@ class WorkerPool: self._touch(label) # -- release ----------------------------------------------------------- # - def evict(self, label: str, save: bool = True) -> bool: - """Drop a resident worker, persisting its database first.""" + def evict(self, label: str, save: bool = True, + save_gui: bool = False) -> bool: + """Release a resident lease, persisting a managed database first. + + A budget-driven eviction must not save somebody's GUI implicitly. GUI + saves are reserved for an explicit/defensive ``close_all(save=True)``. + """ client = self._clients.pop(label, None) if client is None: return False if label in self._lru: self._lru.remove(label) - if save: + if save and (save_gui or getattr(client, "backend", None) != "gui"): try: # persist analysis + edits so the next open is a load - client.call("idb_save") + client.save_database() except Exception: # noqa: BLE001 -- evict regardless pass try: @@ -198,7 +210,7 @@ class WorkerPool: return None def _enforce_budget(self, protect: str | None = None) -> int: - """Evict LRU workers until the pool fits its budget. Returns how many.""" + """Release LRU leases until the pool fits its budget. Returns how many.""" n = 0 while self.memory_mb() > self.budget_mb: victim = self._evictable(protect) @@ -210,7 +222,7 @@ class WorkerPool: def close_all(self, save: bool = True) -> None: for label in list(self._clients): - self.evict(label, save=save) + self.evict(label, save=save, save_gui=save) self.active = None # -- introspection ------------------------------------------------------ # @@ -231,5 +243,9 @@ class WorkerPool: return out def __repr__(self) -> str: # pragma: no cover - debug aid - return (f"<WorkerPool {len(self._clients)}/{len(self.project.refs)} resident " + return (f"<DatabasePool {len(self._clients)}/{len(self.project.refs)} resident " f"{self.memory_mb()}/{self.budget_mb}MB active={self.active}>") + + +# Source compatibility for callers that imported the pre-Code-Mode name. +WorkerPool = DatabasePool diff --git a/idatui/project.py b/idatui/project.py index e2fc542..53f3b0a 100644 --- a/idatui/project.py +++ b/idatui/project.py @@ -23,7 +23,8 @@ firmware image, a cleaned build tree). A source whose size/mtime no longer matches the staged copy is re-staged, and its now-stale database is dropped (the DB describes the old bytes). -stdlib-only, like the domain/worker layers — the TUI is the only Textual consumer. +The model has no IDA imports. Staging consults ida_codemode's registry before +replacing files so it never mutates a database owned by a GUI/shared worker. """ from __future__ import annotations @@ -56,7 +57,7 @@ class BinaryRef: #: it is, a raw firmware image doesn't, and IDA defaults to metapc at 0. processor: str = "" # IDA processor name: arm, armb, mipsb, metapc, … base: int = 0 # load address (natural, e.g. 0x8000000) - ida_args: str = "" # escape hatch: extra IDA command-line switches + ida_args: str = "" # legacy -p/-b/-T switches accepted by Code Mode adapter @property def db(self) -> str: @@ -310,13 +311,25 @@ class Project: """Ensure ``ref`` is staged in the sidecar; returns the staged path. Re-staging a changed source drops its database: the DB describes the old - bytes, so keeping it would silently mismatch the disassembly (any renames - in it are lost, which is why callers should say so out loud). + bytes. Refuse while Code Mode reports a GUI/idalib owner; replacing a + staged executable or IDB underneath a shared live instance is corruption. """ if not os.path.isfile(ref.source): raise ProjectError(f"no such binary: {ref.source}") if not self.is_stale(ref): return ref.staged + try: + from .codemode_client import database_owner + owner = database_owner(ref.db, ref.staged) + except Exception as exc: + raise ProjectError( + f"cannot verify Code Mode ownership before staging {ref.label}: {exc}" + ) from exc + if owner is not None: + raise ProjectError( + f"cannot restage {ref.label}: Code Mode instance {owner.record_id} " + f"still owns {owner.idb_path}; close/release it first" + ) os.makedirs(self.bin_dir, exist_ok=True) tmp = ref.staged + ".staging" _unlink(tmp) @@ -338,10 +351,11 @@ class Project: return out def sweep_scratch(self, ref: BinaryRef) -> int: - """Delete IDA's unpacked working files (never the ``.i64``) for ``ref``. + """Delete unpacked working files (never the ``.i64``) for maintenance. - A hard-killed worker leaves them behind and the database then refuses to - reopen. Only safe when no worker holds it. + Runtime paths no longer call this: Code Mode instances are shared, so a + registry owner may still be using these files. Callers must independently + prove that no GUI/idalib instance owns the database. """ return sum(1 for suf in SCRATCH_SUFFIXES if _unlink(ref.staged + suf)) diff --git a/idatui/rpc.py b/idatui/rpc.py index 698e4cf..225aeb2 100644 --- a/idatui/rpc.py +++ b/idatui/rpc.py @@ -757,7 +757,7 @@ class RpcServer: batch = {"func": ops, "allow_overwrite": bool(overwrite)} # The worker is blocking and single-threaded; off the event loop it goes, # or the TUI freezes for the length of the batch. - res = await asyncio.to_thread(app.program.client.call, "rename", batch=batch) + res = await asyncio.to_thread(app.program.client.invoke, "rename", batch=batch) summary = res.get("summary", {}) if isinstance(res, dict) else {} failed = [r for r in (res.get("func") or []) if isinstance(r, dict) and r.get("error")] if isinstance(res, dict) else [] @@ -768,7 +768,7 @@ class RpcServer: # this a batch import leaves pseudocode calling sub_98C0 forever while # the listing (and every readback) says memset. try: - await asyncio.to_thread(app.program.client.call, "force_recompile") + await asyncio.to_thread(app.program.client.invoke, "force_recompile") except Exception: # noqa: BLE001 -- older worker without the tool pass app.program.bump_names() diff --git a/idatui/worker.py b/idatui/worker.py deleted file mode 100644 index 556e69a..0000000 --- a/idatui/worker.py +++ /dev/null @@ -1,301 +0,0 @@ -"""idatui's own idalib worker — the replacement for the ida-pro-mcp supervisor. - -Opens ONE database in-process (on the main thread, as idalib requires) and -serves ida-pro-mcp's *tool functions* over a unix socket with length-prefixed -pickle. Same tool implementations as the MCP path (we call -``MCP_SERVER.tools.methods[name](**args)`` directly), so return shapes are -byte-identical — but with ~50us/call instead of the HTTP path's ~5ms, and no -supervisor / HTTP / JSON / 50KB-truncation machinery. - - python -m idatui.worker <sock_path> <binary_path> - -The socket only appears once the database is open + analyzed, so a client can -poll ``connect()`` to know when the worker is ready. Requests are served -serially on the main thread (idalib is single-threaded; every tool runs inline -through its own execute_sync, which is a no-op on the main thread). - -Protocol (both directions length-prefixed: 4-byte big-endian len + pickle): - request = (tool_name: str, kwargs: dict) - response = (ok: bool, result_or_error) - tool_name == "__shutdown__" ends the worker. -""" -from __future__ import annotations - -import os -import pickle -import socket -import struct -import sys -import threading -import time -import uuid - -#: Seconds a single tool call may run before it is cancelled. 0 disables the -#: deadline entirely. -TOOL_TIMEOUT_SEC = float(os.environ.get("IDATUI_TOOL_TIMEOUT_SEC") or 60) - -# ida-pro-mcp enforces its own tool deadline by installing a `sys.setprofile` -# hook for the duration of every call, so that a pure-python loop inside a tool -# body can be interrupted. That hook runs a python function on EVERY python call -# and return -- and our tools are exactly the call-heavy kind: `heads` renders -# hundreds of items per request and measured 92us/row with the hook against -# 28us/row without it. A 3.3x tax on the whole backend to bound loops that are -# already bounded by their `count` argument. -# -# So: turn the upstream mechanism off and re-arm the half that does the real -# work ourselves (see _Deadline). ida_kernwin.set_cancelled() is what actually -# frees the IDA main thread -- decompile, auto_wait, find_bytes and friends poll -# user_cancelled() and bail within a poll cycle -- and it costs nothing until it -# fires. -os.environ["IDA_MCP_TOOL_TIMEOUT_SEC"] = "0" - - -class _Deadline: - """A single watchdog thread that cancels a tool call which overruns. - - Arming is two attribute writes, because it is on the path of every call the - TUI makes (a scroll is dozens of them). The watchdog polls instead of being - signalled for the same reason: waking a thread per call costs more than the - 0.25s of granularity it buys on a 60s deadline. - """ - - TICK = 0.25 - - def __init__(self, seconds: float) -> None: - import ida_kernwin - self._kernwin = ida_kernwin - self.seconds = seconds - self._until: float | None = None - t = threading.Thread(target=self._run, name="idatui-deadline", - daemon=True) - t.start() - - def _run(self) -> None: - while True: - time.sleep(self.TICK) - until = self._until - if until is not None and time.monotonic() >= until: - self._until = None - # THREAD_SAFE in the IDA SDK; upstream fires it off a Timer too. - self._kernwin.set_cancelled() - - def arm(self) -> None: - # Clear unconditionally: the flag is sticky, and one left set would make - # every later user_cancelled() true forever. - self._kernwin.clr_cancelled() - self._until = time.monotonic() + self.seconds - - def disarm(self) -> None: - self._until = None - - -# --------------------------------------------------------------------------- # -# framing -# --------------------------------------------------------------------------- # -def _recvn(sock: socket.socket, n: int) -> bytes | None: - buf = bytearray() - while len(buf) < n: - chunk = sock.recv(n - len(buf)) - if not chunk: - return None - buf += chunk - return bytes(buf) - - -def send(sock: socket.socket, obj) -> None: - data = pickle.dumps(obj, protocol=pickle.HIGHEST_PROTOCOL) - sock.sendall(struct.pack(">I", len(data)) + data) - - -def recv(sock: socket.socket): - hdr = _recvn(sock, 4) - if hdr is None: - return None - (n,) = struct.unpack(">I", hdr) - body = _recvn(sock, n) - return None if body is None else pickle.loads(body) - - -# --------------------------------------------------------------------------- # -# worker -# --------------------------------------------------------------------------- # -def _ensure_tools_injected() -> None: - """Inject idatui's custom tools (heads/read_raw/resolve_names/func_types/...) - into the installed ida_pro_mcp, idempotently, so the worker is self-sufficient - (nothing else has to inject these tools first). Must run BEFORE - ida_pro_mcp.ida_mcp is imported (the injected code lives in api_types.py).""" - import importlib.util - repo = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) - patch = os.path.join(repo, "server", "patch_server.py") - if not os.path.exists(patch): - return - try: - spec = importlib.util.spec_from_file_location("_idatui_patch", patch) - mod = importlib.util.module_from_spec(spec) - spec.loader.exec_module(mod) # IDA-free; just defines + patches api_types - mod.main() - except Exception as e: # noqa: BLE001 -- tools may already be present - sys.stderr.write(f"idatui: tool injection skipped: {e}\n") - - -def _has_database(binpath: str) -> bool: - """Whether IDA already has a database for ``binpath``. - - IDA names it ``<file>.i64`` (keeping the extension), but a database made - from ``foo.bin`` can also appear as ``foo.i64`` depending on how it was - created — check both, because guessing wrong here means re-passing load - switches to an existing database, which fails the open. - """ - return (os.path.exists(binpath + ".i64") - or os.path.exists(os.path.splitext(binpath)[0] + ".i64")) - - -def _open_and_register(binpath: str, load_args: str = ""): - """Open the DB (main thread) then import ida-pro-mcp so every @tool registers - against this live database. Returns (tools_dict, module_name, save_fn). - - ``load_args`` is passed to IDA as command-line switches, which is the only - way to tell it how to read a headerless blob: a raw firmware image has no - format to detect, so without ``-p<processor>`` it loads as metapc at 0 and - finds nothing. Ignored once a database exists — the .i64 already records how - it was loaded, and re-passing conflicting switches is how you corrupt one. - """ - _ensure_tools_injected() # before any ida_pro_mcp import - import idapro - idapro.enable_console_messages(False) - args = load_args or None - if args and _has_database(binpath): - # The .i64 already records how this image was loaded. Passing the - # switches again on reopen makes IDA fail outright (rc != 0) — the load - # options belong to the FIRST open only. - args = None - if idapro.open_database(binpath, run_auto_analysis=True, - args=args): # nonzero == failure - if args: - # With load switches in play they are the likeliest culprit by far: - # IDA refuses an unknown -p name with no diagnostic of its own, so - # saying "the database is locked" here sends people hunting a - # problem they don't have. - raise RuntimeError( - f"failed to open {binpath} with load options {args!r}: IDA " - f"rejected them \u2014 an unknown processor name is the usual " - f"cause (see tools/verify_procs.py for the valid ones)") - raise RuntimeError( - f"failed to open {binpath}: the .i64 is likely held by a running " - f"ida-mcp worker (try: pkill -f idalib) or wedged from a crash " - f"(delete its .id0/.id1/.id2/.nam/.til next to the binary)") - import ida_auto - ida_auto.auto_wait() # block until auto-analysis settles (match ida-mcp) - - # importing the package registers all api_*/patched tools against MCP_SERVER - from ida_pro_mcp.ida_mcp import MCP_SERVER # noqa: WPS433 - - import ida_nalt - module = os.path.basename(ida_nalt.get_root_filename() or binpath) - - def save(): - import idc - try: - idc.save_database(idc.get_idb_path(), 0) - except Exception: # noqa: BLE001 - import ida_loader, ida_pro # noqa: WPS433 - ida_loader.save_database(idc.get_idb_path(), 0) - - return MCP_SERVER.tools.methods, module, save - - -def serve(sockpath: str, binpath: str, load_args: str = "") -> None: - tools, module, save = _open_and_register(binpath, load_args) - sid = uuid.uuid4().hex[:8] - deadline = _Deadline(TOOL_TIMEOUT_SEC) if TOOL_TIMEOUT_SEC > 0 else None - - def dispatch(name: str, args: dict): - args = dict(args) - args.pop("database", None) # single-DB worker: no session routing - # session-management shims (were the supervisor's job): - if name in ("idb_open",): - return {"success": True, - "session": {"session_id": sid, "module": module, - "input_path": binpath}} - if name in ("idb_save", "save"): - save() - return {"success": True} - if name in ("server_health", "ping", "health", "state"): - return {"module": module, "ok": True, "session_id": sid} - if name in ("idb_list",): - return {"sessions": [{"session_id": sid, "module": module, - "input_path": binpath}]} - fn = tools.get(name) - if fn is None: - raise KeyError(f"unknown tool: {name!r}") - if deadline is None: - result = fn(**args) - else: - deadline.arm() - try: - result = fn(**args) - finally: - deadline.disarm() - # Match the MCP server's structuredContent: a dict passes through, any - # other return (list/scalar) is wrapped as {"result": ...}. domain.py - # parses that exact shape (e.g. lookup_funcs -> payload["result"]). - return result if isinstance(result, dict) else {"result": result} - - try: - os.unlink(sockpath) - except OSError: - pass - srv = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) - srv.bind(sockpath) - srv.listen(8) - try: - while True: - conn, _ = srv.accept() - try: - while True: - req = recv(conn) - if req is None: - break - name, args = req - if name == "__shutdown__": - return - try: - send(conn, (True, dispatch(name, args))) - except Exception as e: # noqa: BLE001 -- report, keep serving - send(conn, (False, f"{type(e).__name__}: {e}")) - except (ConnectionError, OSError): - pass - finally: - conn.close() - finally: - try: - import idapro - idapro.close_database(save=False) - except Exception: # noqa: BLE001 - pass - try: - os.unlink(sockpath) - except OSError: - pass - - -def main(argv=None) -> None: - argv = argv if argv is not None else sys.argv[1:] - if len(argv) < 2: - sys.stderr.write( - "usage: python -m idatui.worker <sock> <binary> [ida-load-args]\n") - raise SystemExit(2) - try: - serve(argv[0], argv[1], argv[2] if len(argv) > 2 else "") - except SystemExit: - raise - except BaseException as e: # noqa: BLE001 -- surface a clean cause + code 1 - import traceback - sys.stderr.write(f"\nWORKER-FATAL: {type(e).__name__}: {e}\n") - traceback.print_exc() - sys.stderr.flush() - raise SystemExit(1) - - -if __name__ == "__main__": - main() diff --git a/idatui/worker_client.py b/idatui/worker_client.py deleted file mode 100644 index 173cd9d..0000000 --- a/idatui/worker_client.py +++ /dev/null @@ -1,266 +0,0 @@ -"""WorkerClient — a drop-in replacement for ``IDAClient`` backed by our own -idalib worker (``idatui.worker``) over a unix socket instead of ida-pro-mcp's -HTTP/JSON transport. - -It exposes exactly the surface the app/domain use on the client -(``call``/``call_envelope``/``connect``/``set_db``/``resolve_db``/ -``list_sessions``/``health``/``keepalive``/``close``) and returns byte-identical -payloads (the worker calls the same tool functions), so ``domain.py`` and the -app are unchanged — you just construct a WorkerClient instead of an IDAClient. - -Concurrency: the app fires calls from several worker threads over one client; -the worker is single-threaded, so calls are serialized under a lock (the worker -processes one tool at a time anyway — and at ~50us/call that's free). -""" -from __future__ import annotations - -import os -import socket -import subprocess -import sys -import threading -import time -import uuid -from typing import Any - -from .errors import IDAToolError, IDAConnectionError, Session -from .worker import recv as _recv -from .worker import send as _send - -_WORKER_PY = os.path.join(os.path.dirname(os.path.abspath(__file__)), "worker.py") -_worker_python_cache: str | None = None - - -def _find_worker_python() -> str: - """A python that can import ``ida_pro_mcp`` (and thus idalib) — NOT necessarily - the TUI's python. On a typical box the TUI runs under a venv that has textual - + idalib but not ida_pro_mcp, while the system python has idalib + - ida_pro_mcp. Override with IDATUI_WORKER_PYTHON.""" - global _worker_python_cache - if _worker_python_cache: - return _worker_python_cache - override = os.environ.get("IDATUI_WORKER_PYTHON") - candidates = [override] if override else [] - candidates += ["/usr/bin/python", "/usr/bin/python3", sys.executable] - for py in candidates: - if not py or not os.path.exists(py): - continue - try: - r = subprocess.run([py, "-c", "import ida_pro_mcp"], - capture_output=True, timeout=30) - if r.returncode == 0: - _worker_python_cache = py - return py - except Exception: # noqa: BLE001 - continue - return sys.executable # last resort; the worker will report the real error - - -class _NoopKeepAlive: - """The worker is ours and never idles out, so keepalive is a no-op.""" - - def __init__(self) -> None: - self.beats = self.failures = 0 - - def start(self): - return self - - def stop(self) -> None: - pass - - -class WorkerClient: - def __init__(self, binary_path: str, *, ttl: int = 0, - python: str | None = None, load_args: str = "") -> None: - self._bin = os.path.abspath(os.path.expanduser(binary_path)) - self._load_args = load_args or "" # IDA switches for a headerless blob - self._python = python or _find_worker_python() - tag = f"{os.getpid()}-{uuid.uuid4().hex[:8]}" - self._sock_path = f"/tmp/idatui-worker-{tag}.sock" - self._log_path = f"/tmp/idatui-worker-{tag}.log" - self._proc: subprocess.Popen | None = None - self._sock: socket.socket | None = None - self._sid = uuid.uuid4().hex[:8] - self._lock = threading.Lock() # serialize socket use - self._spawn_lock = threading.Lock() - #: Set by close(). A dropped socket is respawned on the next call (the - #: worker segfaulted and we want it back); a CLOSED one must not be. - #: Teardown and binary-switch both close while @work threads are still - #: in flight, so without this, quitting during a decompile spawned a - #: fresh idalib worker that re-opened the database nobody was looking - #: at any more -- a stray process holding the .i64 we just released. - self._closed = False - - # -- lifecycle --------------------------------------------------------- # - def connect(self, timeout: float = 1800.0, progress=None) -> "WorkerClient": - """Spawn the worker (opens + analyzes the DB) and connect once ready.""" - with self._spawn_lock: - self._closed = False # an explicit reconnect revives this client - if self._sock is not None: - return self - if self._proc is None or self._proc.poll() is not None: - # run worker.py as a SCRIPT (not -m idatui.worker) so we don't - # import the textual-dependent idatui package __init__ under the - # IDA python, which usually has no textual. - argv = [self._python, _WORKER_PY, self._sock_path, self._bin] - if self._load_args: - argv.append(self._load_args) - self._proc = subprocess.Popen( - argv, - stdout=open(self._log_path, "wb"), - stderr=subprocess.STDOUT, stdin=subprocess.DEVNULL, - ) - deadline = time.time() + timeout - t0 = time.time() - # Poll fast at first, then back off. A flat 0.2s cost every caller a - # fifth of a second even when the worker was ready in milliseconds - # (a small binary, or a seeded .i64), which is most of the time in - # the tests and noticeable on a re-open. - # - # Backing off all the way to 0.2s was too eager: a seeded database - # is ready at ~250ms, by which point the delay has grown to 134ms, so - # every open waited ~350ms whatever the binary -- the same number for - # a 47KB `echo` and a 1.2MB `bash`, which is what gives a polling - # artefact away. Cap the backoff at 25ms instead: the overshoot on a - # fast open is bounded by that, and 40 probes a second is nothing - # next to an auto-analysis that runs for minutes. - # - # Do NOT be tempted to hold the 5ms rate instead. This poll runs on a - # background thread while the UI thread is drawing, and 200 wakeups a - # second through a cold analysis cost enough GIL time to delay the - # app's own startup -- it left the loading overlay up long enough for - # project mode's first keypress to land on it. - delay = 0.005 - while time.time() < deadline: - try: - s = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) - s.connect(self._sock_path) - self._sock = s - return self - except OSError: - if self._proc.poll() is not None: - raise IDAConnectionError( - f"worker exited (code {self._proc.returncode}): " - f"{self._log_tail()} [full log: {self._log_path}]") - if progress: - progress(f"auto-analyzing {os.path.basename(self._bin)}… " - f"({int(time.time() - t0)}s)") - time.sleep(delay) - delay = min(delay * 1.6, 0.025) - raise IDAConnectionError("worker did not become ready in time") - - @property - def pid(self) -> int | None: - """The worker process id (for memory accounting), or None if not spawned.""" - return self._proc.pid if self._proc is not None else None - - def close(self, grace: float = 20.0) -> None: - """Shut the worker down cleanly. - - After ``__shutdown__`` the worker still has to ``close_database()``, which - re-packs the ``.i64`` and removes the unpacked ``.id0/.id1/...`` scratch. - Signalling it before that finishes is what leaves databases wedged, so - wait out the grace period first and only escalate if it really is stuck. - """ - with self._lock: - s = self._sock - self._sock = None - self._closed = True - if s is not None: - try: - _send(s, ("__shutdown__", {})) - except Exception: # noqa: BLE001 - pass - try: - s.close() - except Exception: # noqa: BLE001 - pass - if self._proc is not None: - try: - self._proc.wait(timeout=grace) # let it close the DB properly - except Exception: # noqa: BLE001 -- TimeoutExpired: it's stuck - try: - self._proc.terminate() - self._proc.wait(timeout=5) - except Exception: # noqa: BLE001 - try: - self._proc.kill() - except Exception: # noqa: BLE001 - pass - - # -- the call surface -------------------------------------------------- # - def call(self, tool: str, *, timeout: float | None = None, **args) -> Any: - if self._closed: - raise IDAConnectionError( - f"{tool}: this worker was closed (call connect() to revive it)") - if self._sock is None: - self.connect() - with self._lock: - s = self._sock - if s is None: - raise IDAConnectionError("worker connection is closed") - try: - _send(s, (tool, args)) - reply = _recv(s) - except (OSError, ConnectionError) as e: - self._sock = None - raise IDAConnectionError(f"worker transport failed: {e}") from e - if reply is None: - self._sock = None - raise IDAConnectionError("worker closed the connection") - ok, payload = reply - if not ok: - raise IDAToolError(tool, str(payload)) - return payload - - def call_envelope(self, tool: str, *, timeout: float | None = None, - **args) -> dict: - # domain.decompile() reads result.structuredContent — mirror that shape. - return {"result": {"structuredContent": self.call(tool, timeout=timeout, - **args)}} - - # -- session shims (single-DB worker) --------------------------------- # - def set_db(self, db: str | None) -> None: - if db: - self._sid = db - - def resolve_db(self) -> str: - return self._sid - - def list_sessions(self) -> list[Session]: - return [Session(session_id=self._sid, - filename=os.path.basename(self._bin), - input_path=self._bin, is_active=True)] - - def health(self) -> dict: - try: - return self.call("server_health") - except IDAToolError: - return {"module": os.path.basename(self._bin), "ok": True} - - def keepalive(self, interval: float = 120.0) -> _NoopKeepAlive: - return _NoopKeepAlive() - - def _log_tail(self, n: int = 400) -> str: - """Last meaningful line(s) of the worker log (skip IDA's licence banner), - so a startup crash surfaces the real cause instead of just 'code 1'.""" - try: - with open(self._log_path, encoding="utf-8", errors="replace") as f: - lines = [ln.strip() for ln in f if ln.strip()] - except OSError: - return "(no worker log)" - # the worker prints a clean 'WORKER-FATAL: ...' line on a startup crash - for ln in reversed(lines): - if ln.startswith("WORKER-FATAL:"): - return ln[len("WORKER-FATAL:"):].strip()[-n:] - skip = ("thank you", "licensed to", "[mcp]", "ida ", "hex-rays") - meaningful = [ln for ln in lines - if not any(s in ln.lower() for s in skip)] - return " | ".join((meaningful or lines)[-3:])[-n:] - - # context manager parity with IDAClient - def __enter__(self) -> "WorkerClient": - return self.connect() - - def __exit__(self, *exc) -> None: - self.close() 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 6667e12..0000000 --- a/server/patch_server.py +++ /dev/null @@ -1,2258 +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, flags=None): - """One flat-listing row for the head at ``ea``: kind (code/data/unknown), - byte size, rendered text, and any symbol name. - - ``flags`` lets a caller that already asked for them say so -- the walk in - ``heads`` used to fetch them three times per head (here, in _is_unknown from - _advance, and again from _rows_for). - """ - import ida_bytes - import ida_lines - import ida_name - - f = ida_bytes.get_flags(ea) if flags is None else flags - if ida_bytes.is_code(f): - kind = "code" - elif ida_bytes.is_data(f): - kind = "data" - else: - kind = "unknown" - line = ida_lines.generate_disasm_line(ea, 0) - text, spans, ops = _idatui_line_parts(line) if line else ("", None, None) - row = { - "ea": hex(ea), - "kind": kind, - "size": int(ida_bytes.get_item_size(ea)), - "text": text, - } - if spans is not None: - row["spans"] = spans - # Where each operand sits in `text`. Comes out of the same tag walk - # (free), and is what lets the client show WHICH literal a keypress - # would reformat before you press it. - if ops: - row["ops"] = ops - nm = ida_name.get_ea_name(ea) - if nm: - row["name"] = nm - return row - - -import functools as _idatui_functools -import os as _idatui_os - -#: Entries in the per-line render cache. Sized to hold a whole segment's -#: DISTINCT lines rather than a working set, because the listing gets rendered -#: TWICE: once when it is first walked, and again after a rename, which restates -#: every row's text. bash's .text is 228 659 rows but only 53 363 distinct -#: lines, and the difference between thrashing and not is the whole win: -#: -#: maxsize first sweep second sweep worker RSS -#: 16 384 17.2 us/row 16.9 us/row +29 MB -#: 32 768 17.0 17.2 +52 MB -#: 65 536 17.0 11.1 +75 MB -#: 131 072 16.9 11.2 +75 MB (working set fits) -#: -#: It is a bound, not a proportion: a bigger binary fills it and stops, so the -#: cost is capped at ~56 MB whatever is open. Lower it with IDATUI_LINE_CACHE if -#: a pool of workers is competing for memory. -_IDATUI_LINE_CACHE = int(_idatui_os.environ.get("IDATUI_LINE_CACHE") or 65536) - - -@_idatui_functools.lru_cache(maxsize=_IDATUI_LINE_CACHE) -def _idatui_line_parts(line): - """``(text, spans, ops)`` for one tagged disassembly line -- memoised. - - A function of the tagged line and nothing else, so the same line always - gives the same answer: a rename changes the line, which changes the key. - And listings repeat themselves hard -- 196k lines of bash are 53k distinct - ones, so a 16k-entry cache serves ~70% of them and takes the per-line cost - from 10.4us to 3.9us. This is the most expensive thing the backend does per - listing row, and a jump to an address near the end of a big binary walks - hundreds of thousands of them. - - ``spans`` is None when the tag walk and the plain text disagree about what - the line says (then the text wins and the row renders unhighlighted). - - The returned lists are SHARED between every row that has the same line; - treat them as read-only. Pickle notices the sharing too, so a page of - repetitive disassembly also serialises smaller. - """ - import ida_lines - text = " ".join(ida_lines.tag_remove(line).split()) # collapse the padding - spans, ops = _idatui_spans(line) - # Built from the SAME line as `text`, then whitespace-collapsed identically, - # so the two can never disagree about what the row says. - joined = "".join([t for _k, t in spans]) - if " ".join(joined.split()) != text: - return (text, None, None) - return (text, spans, ops) - - -#: 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 -_IDATUI_OPND_TAGS = None -_IDATUI_CTL = None # re: a tag = one of three control chars plus its argument -#: {tag character: (kind, operand index or None)} -- the two maps above merged, -#: because the span walker wants both for the same tag and a dict lookup per -#: tag per line is one of the few things it does often enough to matter. -_IDATUI_TAGINFO = None - - -def _idatui_opnd_tag_map(): - """{tag character: operand index}. IDA wraps each operand of a disassembly - line in COLOR_OPND1..8, so the line already says where operand N starts and - ends -- no need to re-render operands with print_operand to find out (and - the two agree exactly; checked over thousands of instructions).""" - import ida_lines - out = {} - for i in range(1, 9): - v = getattr(ida_lines, "COLOR_OPND%d" % i, None) - if isinstance(v, int): - out[chr(v)] = i - 1 - elif isinstance(v, str) and v: - out[v[0]] = i - 1 - return out - - -def _idatui_spans(line): - """(spans, ops) for a tagged disasm line. - - ``spans`` is [[kind, text], ...] with colour tags resolved; ``ops`` is - [[start, end, n], ...], the extent of each operand in the SAME (collapsed) - coordinates the row's ``text`` uses -- which is what lets a cursor column - name the operand it is standing on. - - Unknown tags become 'text' rather than being dropped: a processor module can - emit a colour we don't classify, and losing the characters would corrupt the - line.""" - global _IDATUI_TAGS, _IDATUI_OPND_TAGS, _IDATUI_CTL, _IDATUI_TAGINFO - import ida_lines - if _IDATUI_TAGS is None: - _IDATUI_TAGS = _idatui_tag_map() - if _IDATUI_OPND_TAGS is None: - _IDATUI_OPND_TAGS = _idatui_opnd_tag_map() - if _IDATUI_CTL is None: - import re as _re - # One capturing split gives [text, tag, text, tag, ..., text] in a - # single C pass. A per-character python loop over the line used to be - # the most expensive thing the `heads` tool did, and a line is ~54 - # characters but only ~13 tags -- everything between two tags is already - # exactly one span's worth of text. - _IDATUI_CTL = _re.compile("([\\x01\\x02\\x03](?s:.))") - if _IDATUI_TAGINFO is None: - _IDATUI_TAGINFO = { - tag: (_IDATUI_TAGS.get(tag, "text"), _IDATUI_OPND_TAGS.get(tag)) - for tag in set(_IDATUI_TAGS) | set(_IDATUI_OPND_TAGS)} - taginfo = _IDATUI_TAGINFO - plain_tag = ("text", None) - on, off, esc = "\x01", "\x02", "\x03" - addr_tag = chr(getattr(ida_lines, "COLOR_ADDR", 0x28)) - addr_len = int(getattr(ida_lines, "COLOR_ADDR_SIZE", 16)) - parts = _IDATUI_CTL.split(line) - spans, stack = [], [] # stack entries: (kind, operand index|None) - kind, opnd = "text", None # state the current run of text belongs to - pend = "" - skip = 0 # characters of an address payload still due - i, n = 0, len(parts) - while i < n: - txt = parts[i] - i += 1 - if skip: - if len(txt) <= skip: - skip -= len(txt) - txt = "" - else: - txt = txt[skip:] - skip = 0 - if txt: - pend += txt - if i >= n: - break - pair = parts[i] - i += 1 - if skip: # a tag INSIDE an address payload: 2 chars - skip = skip - 2 if skip > 2 else 0 - continue - ch = pair[0] - if ch == esc: # escaped literal: keep the char it guards - pend += pair[1] - continue - tag = pair[1] - if ch == on and tag == addr_tag: - # An embedded target address, not display text: 16 hex digits that - # must not reach the screen. Deliberately NOT a span boundary. - skip = addr_len - continue - if pend: - spans.append([kind, pend, opnd]) - pend = "" - if ch == on: - stack.append((kind, opnd)) - kind, o = taginfo.get(tag, plain_tag) - if o is not None: - opnd = o # operands nest: an inner colour keeps the operand - elif stack: - kind, opnd = stack.pop() - else: - kind, opnd = "text", None - if pend: - spans.append([kind, pend, opnd]) - # Collapse IDA's column padding EXACTLY as the plain text does. A run of - # spaces can straddle two spans, so the leading space of a span is dropped - # when the previous one ended in space — otherwise the spans and `text` - # disagree about the line and the row silently loses its highlighting. - # ``" ".join(txt.split())`` splits on exactly what str.isspace() calls - # whitespace, which is what the character walk this replaces tested. - out, prev_space = [], False - for kind, txt, opnd in spans: - core = " ".join(txt.split()) - if core == txt: - # Nothing to collapse and no edge whitespace -- which is the common - # case ("mov", "rax", ", ") and skips both isspace() probes below. - prev_space = False - out.append([kind, txt, opnd]) - continue - if not core: # the span is nothing but padding - if not prev_space: - prev_space = True - out.append([kind, " ", opnd]) - continue - acc = core - if txt[0].isspace() and not prev_space: - acc = " " + acc - if txt[-1].isspace(): - acc += " " - prev_space = acc[-1] == " " - out.append([kind, acc, opnd]) - while out and out[0][1] == " ": - out.pop(0) - while out and out[-1][1] == " ": - out.pop() - if out and out[0][1].startswith(" "): - out[0][1] = out[0][1].lstrip() - if out and out[-1][1].endswith(" "): - out[-1][1] = out[-1][1].rstrip() - out = [s for s in out if s[1]] - # Operand extents, in the coordinates of the collapsed text these spans - # spell out. Adjacent spans of the same operand merge, so an operand like - # ``[rbp+var_40]`` (five differently-coloured tokens) comes back as ONE - # range -- which is the thing a cursor is inside of, and the thing a format - # change applies to. - ops, pos, cur, start = [], 0, None, 0 - for _kind, txt, opnd in out: - if opnd != cur: - if cur is not None and pos > start: - ops.append([start, pos, cur]) - cur, start = opnd, pos - pos += len(txt) - if cur is not None and pos > start: - ops.append([start, pos, cur]) - text = "".join(t for _k, t, _o in out) - trimmed = [] - for lo, hi, k in ops: # don't let a range own trailing space - while hi > lo and text[hi - 1].isspace(): - hi -= 1 - while lo < hi and text[lo].isspace(): - lo += 1 - if hi > lo: - trimmed.append([lo, hi, k]) - return [[k, t] for k, t, _o in out], trimmed - - -def _idatui_rows_digest(rows): - """A value that changes whenever any of ``rows`` would render differently. - - Covers everything a client keeps off a row: address, kind, size, the plain - text, the symbol name and the colour spans (which is what makes it exact - rather than a heuristic -- two lines can collapse to the same text and still - be coloured differently). - - Uses the interpreter's own ``hash``, deliberately. It never has to mean - anything outside this process: the client stores what a page hashed to when - it loaded it and hands the same number back to ask whether the page still - hashes to that. One worker, one process, one hash seed. - """ - acc = 0 - # The per-line render is memoised, so one spans list is shared by every row - # that says the same thing -- about 45% of them within a page. Hash each - # distinct list once and key that by identity, rather than rebuilding a - # tuple of tuples per row (which is the exact cost that was measured and - # removed from the client side for the same reason). - seen = {} - for r in rows: - sp = r.get("spans") - if sp is None: - sh = None - else: - key = id(sp) - sh = seen.get(key) - if sh is None: - sh = seen[key] = hash(tuple(map(tuple, sp))) - acc = hash((acc, r.get("ea"), r.get("kind"), r.get("size"), - r.get("text"), r.get("name"), sh)) - return acc - - -def _idatui_unknown_row(ea, size): - """One collapsed row for a run of ``size`` undefined bytes starting at - ``ea``. A single byte is rendered normally (shows its value); a longer run - collapses to ``db N dup(?)`` so a big .bss/gap doesn't explode into millions - of one-byte rows.""" - import ida_name - - if size <= 1: - return _idatui_head_row(ea) - row = {"ea": hex(ea), "kind": "unknown", "size": int(size), - "text": f"db {size} dup(?)"} - nm = ida_name.get_ea_name(ea) - if nm: - row["name"] = nm - return row - - -def _idatui_struct_member_rows(ea): - """Indented member rows for a struct-typed data item at ``ea`` (expansion), - or [] if it isn't a struct. Top-level fields only.""" - import ida_nalt - import ida_typeinf - import idaapi - - tif = ida_typeinf.tinfo_t() - if not (ida_nalt.get_tinfo(tif, ea) and tif.is_udt()): - return [] - udt = ida_typeinf.udt_type_data_t() - if not tif.get_udt_details(udt): - return [] - rows = [] - for m in udt: - off = m.begin() // 8 - try: - mtype = m.type._print() or "" - except Exception: - mtype = "" - try: - sz = int(m.type.get_size()) - if sz == idaapi.BADSIZE: - sz = 0 - except Exception: - sz = 0 - name = m.name or "" - text = f"+{off:X} {name}" + (f" {mtype}" if mtype else "") - rows.append({"ea": hex(ea + off), "kind": "member", "size": sz, - "text": text}) - return rows - - -def _idatui_func_header_rows(ea): - """IDA-style subroutine banner rows shown just before a function's entry.""" - import ida_funcs - - name = ida_funcs.get_func_name(ea) or "sub_%X" % ea - bar = "=" * 15 + " S U B R O U T I N E " + "=" * 15 - return [ - {"ea": hex(ea), "kind": "sep", "size": 0, "text": ""}, - {"ea": hex(ea), "kind": "sep", "size": 0, "text": "; " + bar}, - {"ea": hex(ea), "kind": "funchdr", "size": 0, - "text": name + " proc", "name": name}, - ] - - -def _idatui_func_footer_rows(ea, func): - """End-of-function marker shown just after a function's last item.""" - import ida_funcs - - name = ida_funcs.get_func_name(func.start_ea) or "sub_%X" % func.start_ea - return [ - {"ea": hex(ea), "kind": "funchdr", "size": 0, - "text": name + " endp", "name": name}, - {"ea": hex(ea), "kind": "sep", "size": 0, "text": "; " + "-" * 60}, - ] - - -@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, - expect: Annotated[str, "Digest a caller already holds: the rows are omitted when they still hash to it"] = "", -) -> dict: - """Walk item heads from ``addr`` as a flat listing: every head is rendered - (code OR data OR undefined) via generate_disasm_line and stepped with - next_head/prev_head. Unlike ``disasm`` (code-only, bails at the first data - byte) this shows db/dw/dd/... lines for data and undefined regions — IDA's - real disassembly view. Address-paged: page forward by re-calling with - ``addr`` = the returned cursor.next; page up with ``back=true``.""" - import ida_bytes - import ida_segment - import idaapi - - count = 2000 if count > 2000 else (1 if count < 1 else count) - offset = max(int(offset), 0) - try: - start = parse_address(addr) - except Exception as e: - return {"addr": str(addr), "error": str(e), "heads": [], "cursor": {"done": True}} - seg = ida_segment.getseg(start) - if not seg: - return {"addr": str(addr), "error": "no segment", "heads": [], "cursor": {"done": True}} - lo, hi = seg.start_ea, seg.end_ea - if end: - try: - hi = min(hi, parse_address(end)) - except Exception: - pass - - rows = [] - if back: - # Collect up to (count+offset) heads strictly before `start`, then take - # the window closest to `start`, returned in forward order. - walk = [] - cur = ida_bytes.prev_head(start, lo) - while cur != idaapi.BADADDR and cur >= lo and len(walk) < count + offset: - walk.append(cur) - cur = ida_bytes.prev_head(cur, lo) - walk.reverse() - chosen = walk[: len(walk) - offset] if offset else walk - chosen = chosen[-count:] - rows = [_idatui_head_row(e) for e in chosen] - first = chosen[0] if chosen else start - pea = ida_bytes.prev_head(first, lo) - cursor = {"done": True} if pea == idaapi.BADADDR or pea < lo else {"prev": hex(pea)} - return {"addr": str(addr), "heads": rows, "cursor": cursor} - - # Walk by item END (not next_head): next_head SKIPS undefined bytes, but a - # flat listing must show them (IDA renders undefined as `db ?` lines, and - # navigating to an unmarked address must land ON it). Defined items advance - # by get_item_end; a run of undefined bytes is COLLAPSED into one row (its - # end found in O(1) via next_head, which skips undefined) so a large .bss or - # gap doesn't explode into millions of one-byte rows. - def _is_unknown_f(f): - return not (ida_bytes.is_code(f) or ida_bytes.is_data(f)) - - def _run_end(e): - """End (exclusive) of the undefined run starting at ``e``.""" - nh = ida_bytes.next_head(e, hi) - return nh if (nh != idaapi.BADADDR and e < nh <= hi) else hi - - def _advance(e, f): - if _is_unknown_f(f): - return _run_end(e) - nxt = ida_bytes.get_item_end(e) - return nxt if nxt > e else e + 1 - - # The function the walk is currently inside, reused while it stays inside. - # get_func is ~0.5us and the walk asks per head; a head is nearly always in - # the same function as the one before it. Only ever consulted when ``e`` - # falls in [start_ea, end_ea), so a tail chunk elsewhere cannot be - # misattributed -- checked against get_func over 437k heads of - # bash/ls_ttl/echo with zero disagreements. - fn_cache = [None] - - def _func_at(e): - cur = fn_cache[0] - if cur is not None and cur.start_ea <= e < cur.end_ea: - return cur - cur = idaapi.get_func(e) - fn_cache[0] = cur - return cur - - def _rows_for(e, f): - if _is_unknown_f(f): - return [_idatui_unknown_row(e, _run_end(e) - e)] - func = _func_at(e) if annotate else None - at_start = func is not None and func.start_ea == e - out = [] - if at_start: - out.extend(_idatui_func_header_rows(e)) - row = _idatui_head_row(e, f) - if at_start: - row = dict(row) - row["name"] = None # the name is shown on the proc header line - elif annotate and row.get("kind") == "code" and row.get("name"): - # A code label (loc_XXX/jump target) gets its OWN line at depth 0, - # like IDA; strip it from the instruction row below. - nm = row["name"] - out.append({"ea": hex(e), "kind": "label", "size": 0, - "text": nm + ":", "name": nm}) - row = dict(row) - row["name"] = None - out.append(row) - if row.get("kind") == "data": - out.extend(_idatui_struct_member_rows(e)) # expand struct fields - if func is not None and ida_bytes.get_item_end(e) >= func.end_ea: - out.extend(_idatui_func_footer_rows(e, func)) - return out - - ea = ida_bytes.get_item_head(start) - get_flags = ida_bytes.get_flags - for _ in range(offset): - if ea >= hi or ea == idaapi.BADADDR: - break - ea = _advance(ea, get_flags(ea)) - more = False - while ea != idaapi.BADADDR and ea < hi: - if len(rows) >= count: - more = True - break - f = get_flags(ea) # once per head, not once per consumer - rows.extend(_rows_for(ea, f)) # a struct head expands into member rows - ea = _advance(ea, f) - cursor = {"next": hex(ea)} if more else {"done": True} - dig = _idatui_rows_digest(rows) - out = {"addr": str(addr), "cursor": cursor, "digest": dig, "count": len(rows)} - # ``expect`` says "I already hold a page that hashed to this". The rows are - # built either way -- generate_disasm_line is the floor and there is no way - # to know a line is unchanged without rendering it -- but pickling several - # hundred rows with their colour spans, unpickling them and rebuilding Heads - # is about 40% of what a page costs, and after a rename almost every page - # comes back identical. - # - # It carries the expected value rather than being a yes/no "digest mode" so - # that a page which HAS changed still costs one round trip: asking first and - # fetching afterwards made every changed page two. - if not (expect and str(dig) == expect): - out["heads"] = rows - return out - - -@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"} - import ida_lines - # Three things this loop must not do, each measured on real functions (the 25 - # largest of bash went 68.3s -> 6.5s; echo's 60 largest 5.4s -> 0.6s, with - # byte-identical output): - # - # * allocate ctree_item_t's per COLUMN. They are SWIG objects and this is - # the innermost loop; one per call is enough, and head/tail are never - # read, so don't ask for them at all. - # * sweep the TAGGED length. ``x`` is a screen column but ``sl.line`` still - # carries IDA's colour tags, so a 23-column line was swept 124 times. - # * call dstr() per column. It formats a whole 'EA: description' string -- - # 24us a call, which is 79% of this tool. Comparing against the PREVIOUS - # column's item id is not enough: items interleave, so `foo(a, b)` flips - # call -> arg -> call -> arg and every flip re-formats an item already - # seen (106 594 calls for 15 417 lines of bash). Memoise id -> ea for the - # whole function instead: obj_id is unique within a cfunc, so the same id - # always yields the same string, and the result is deduped by ``seen`` - # anyway. Items with no ctree node (it is None) have no id to key on and - # still pay per occurrence. - item = ida_hexrays.ctree_item_t() - tag_remove = ida_lines.tag_remove - get_line_item = cfunc.get_line_item - ea_of_id = {} - lines = [] - for sl in cfunc.get_pseudocode(): - line = sl.line - eas, seen = [], set() - prev_id = None - for x in range(len(tag_remove(line)) + 1): - if not get_line_item(line, x, False, None, item, None): - continue - it = item.it - if it is not None: - oid = it.obj_id - if oid == prev_id: - continue - prev_id = oid - if oid in ea_of_id: - e = ea_of_id[oid] - if e is not None and e not in seen: - seen.add(e) - eas.append(hex(e)) - continue - else: - oid = None - prev_id = None - # Match the /*ea*/ marker's source (decompile_function_safe): the - # item's dstr() is 'EA: description'; get_ea() reports a different ea. - e = None - dstr = item.dstr() - if dstr: - parts = dstr.split(": ", 1) - if len(parts) == 2: - try: - e = int(parts[0], 16) - except ValueError: - e = None - if oid is not None: - ea_of_id[oid] = e - if e is not None and e not in seen: - seen.add(e) - eas.append(hex(e)) - lines.append({"ea": eas[0] if eas else None, "eas": eas}) - return {"addr": hex(func.start_ea), "lines": lines} - - -_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)} - - -# --------------------------------------------------------------------------- # -# operand display formats (IDA's 'o' family: hex / dec / char / offset / ...) -# --------------------------------------------------------------------------- # -#: The stops a cycle walks, in order, before filtering to the ones that make -#: sense for the operand in hand. Octal is deliberately NOT one of them -- every -#: extra stop is another keypress and nobody reads octal -- but it is still -#: reachable by name. "default" hands the operand back to IDA's own choice, -#: which for data is how you get an auto-detected offset/string back. -_IDATUI_FMT_CYCLE = ("hex", "dec", "bin", "char", "offset", "default") - -#: Formats we can re-apply from a name alone. enum/stroff/custom carry an id -#: (which enum, which struct) that a nibble doesn't record, so they are never -#: cycled INTO -- and cycling out of one is called out in ``warn``. -_IDATUI_FMT_SETTABLE = ("hex", "dec", "oct", "bin", "char", "offset", "seg", - "float", "stack", "default") - - -def _idatui_fmt_nibbles(): - """{format name: IDA operand-type nibble}. Built on call, not at import: - this module is injected into a file that is imported before a database is - open.""" - import ida_bytes - return { - "default": ida_bytes.FF_N_VOID, "hex": ida_bytes.FF_N_NUMH, - "dec": ida_bytes.FF_N_NUMD, "char": ida_bytes.FF_N_CHAR, - "seg": ida_bytes.FF_N_SEG, "offset": ida_bytes.FF_N_OFF, - "bin": ida_bytes.FF_N_NUMB, "oct": ida_bytes.FF_N_NUMO, - "enum": ida_bytes.FF_N_ENUM, "forced": ida_bytes.FF_N_FOP, - "stroff": ida_bytes.FF_N_STRO, "stack": ida_bytes.FF_N_STK, - "float": ida_bytes.FF_N_FLT, "custom": ida_bytes.FF_N_CUST, - } - - -def _idatui_fmt_name(nib): - for name, v in _idatui_fmt_nibbles().items(): - if v == nib: - return name - return "default" - - -def _idatui_op_fmt(ea, n): - """The format operand ``n`` of the item at ``ea`` is currently displayed in. - - Reads the nibble IDA keeps per operand rather than guessing from the text -- - ``1`` renders identically in hex and decimal, so the rendered line cannot - answer this.""" - import ida_bytes - F = ida_bytes.get_flags(ea) - nib = (F >> ida_bytes.get_operand_type_shift(int(n))) & 0xF - return _idatui_fmt_name(nib) - - -def _idatui_op_value(ea, n): - """(value, byte width) of operand ``n``, or (None, 0) if it hasn't got one. - - The value is what decides which formats are OFFERED: a character constant - for 0x38A9 or an offset to an unmapped address are stops worth skipping.""" - import ida_bytes - import ida_ua - - F = ida_bytes.get_flags(ea) - if ida_bytes.is_code(F): - insn = ida_ua.insn_t() - if ida_ua.decode_insn(insn, ea) <= 0: - return None, 0 - try: - op = insn.ops[int(n)] - except Exception: - return None, 0 - if op.type == ida_ua.o_void: - return None, 0 - v = op.value if op.type == ida_ua.o_imm else op.addr - try: - size = int(ida_ua.get_dtype_size(op.dtype)) - except Exception: - size = 0 - return int(v), size - size = int(ida_bytes.get_item_size(ea)) - read = {1: ida_bytes.get_byte, 2: ida_bytes.get_word, - 4: ida_bytes.get_dword, 8: ida_bytes.get_qword}.get(size) - if read is None: - return None, size - try: - return int(read(ea)), size - except Exception: - return None, size - - -def _idatui_printable(v): - """Whether ``v`` would actually render as a character constant. IDA accepts - op_chr on anything and then prints the number anyway, so a cycle that offers - 'char' for 0x18 has a stop where nothing visibly happens.""" - if v is None or v < 0 or v > 0xFFFFFFFF: - return False - bs, x = [], int(v) - while True: - bs.append(x & 0xFF) - x >>= 8 - if not x: - break - return all(0x20 <= b <= 0x7E or b in (9, 10, 13) for b in bs) - - -def _idatui_offset_worth(v): - """Whether 'offset' is worth OFFERING as a cycle stop for value ``v``. - - Making an offset is not free: IDA invents a dummy name at the target - (``off_18``) and that name STAYS once you cycle past it. So the ring only - stops there when the target is already something you could name -- a symbol, - a function, or an item something else references. In a PIE at base 0 half - the small constants in a function are 'mapped' (they land in the ELF - header); ``sub rsp, 18h`` is not a reference and must not offer to become - one on the way past. - - An explicit request still converts anything mapped: that's a decision, not a - keypress that happened to land here. After it, the target HAS a name, so the - ring includes the stop from then on.""" - import ida_bytes - import ida_name - - return bool(v and ida_bytes.is_mapped(v) and ida_name.get_ea_name(v)) - - -def _idatui_op_candidates(ea): - """Operand indices at ``ea`` whose display format is worth changing. - - Immediates and displacements -- the literals. Deliberately NOT: - - * branch targets (o_near/o_far), or every jump on the listing would offer to - become a bare number, on a view you navigate by label; - * memory references (o_mem), e.g. x86-64's RIP-relative ``lea rdi, name``. - IDA prints those from the reference, not from the operand's number format, - so setting one is accepted and changes nothing on screen -- a keypress - that appears to do nothing is worse than one that says it can't. - - An explicit ``n`` still reaches them; this is what a bare cursor picks.""" - import ida_bytes - import ida_ua - - F = ida_bytes.get_flags(ea) - if ida_bytes.is_data(F): - return [0] # a data item's value is operand 0 - if not ida_bytes.is_code(F): - return [] # undefined bytes: IDA refuses a format outright - insn = ida_ua.insn_t() - if ida_ua.decode_insn(insn, ea) <= 0: - return [] - want = (ida_ua.o_imm, ida_ua.o_displ) - out = [] - for i in range(len(insn.ops)): - op = insn.ops[i] - if op.type == ida_ua.o_void: - break - if op.type in want: - out.append(i) - return out - - -def _idatui_op_spans(ea, text): - """[(start, end, n)] -- where each operand sits inside ``text`` (the - whitespace-collapsed line the TUI shows), so a cursor column can name the - operand it is standing on. - - Read out of IDA's own COLOR_OPND markers on the line, which is both free - (the line is generated anyway) and exact. print_operand is kept as a - fallback for a processor module that emits no operand markers -- it agrees - with the tags where both exist, but it re-renders every operand to say so. - """ - import ida_lines - import ida_ua - - line = ida_lines.generate_disasm_line(ea, 0) - if line: - _spans, ops = _idatui_spans(line) - if ops: - return [tuple(o) for o in ops] - - out, pos = [], 0 - for n in range(8): - try: - raw = ida_ua.print_operand(ea, n) - except Exception: - raw = None - if not raw: - continue - op = " ".join(ida_lines.tag_remove(raw).split()) - if not op: - continue - i = text.find(op, pos) - if i < 0: # duplicated operand text (mov eax, eax) - i = text.find(op) - if i < 0: - continue - out.append((i, i + len(op), n)) - pos = i + len(op) - return out - - -def _idatui_line_text(ea): - import ida_lines - line = ida_lines.generate_disasm_line(ea, 0) - return " ".join(ida_lines.tag_remove(line).split()) if line else "" - - -def _idatui_op_text(ea, text, n): - """How operand ``n`` reads on the line, for a message that names it.""" - for lo, hi, i in _idatui_op_spans(ea, text): - if i == int(n): - return text[lo:hi].strip() - return "" - - -def _idatui_apply_fmt(ea, n, fmt): - """Set operand ``n``'s display format. Returns (ok, error).""" - import ida_bytes - import ida_offset - import idaapi - - n = int(n) - if fmt == "default": - return bool(ida_bytes.clr_op_type(ea, n)), "" - if fmt == "offset": - base = ida_offset.calc_offset_base(ea, n) - if base in (idaapi.BADADDR, None) or base < 0: - base = 0 - return bool(ida_offset.op_plain_offset(ea, n, base)), "" - fn = {"hex": ida_bytes.op_hex, "dec": ida_bytes.op_dec, - "oct": ida_bytes.op_oct, "bin": ida_bytes.op_bin, - "char": ida_bytes.op_chr, "seg": ida_bytes.op_seg, - "float": ida_bytes.op_flt, "stack": ida_bytes.op_stkvar}.get(fmt) - if fn is None: - return False, (f"can't set {fmt!r} from a name alone" - if fmt in _idatui_fmt_nibbles() else - f"unknown format {fmt!r}") - return bool(fn(ea, n)), "" - - -@tool -@idasync -def op_format( - addr: Annotated[str, "Address of the instruction or data item"], - mode: Annotated[str, "cycle | back | show | hex | dec | oct | bin | char | offset | stack | default"] = "cycle", - col: Annotated[int, "Cursor column inside the rendered line (-1: first literal)"] = -1, - n: Annotated[int, "Operand index; -1 derives it from ``col``"] = -1, -) -> dict: - """Change how a literal is DISPLAYED (IDA's 'o' family): hex, decimal, - binary, character, or an offset to the address it names. - - The value in the bytes never changes -- only the representation IDA renders - and remembers. ``cycle``/``back`` step the stops that make sense for THIS - operand: 'char' is skipped unless the value prints as one, 'offset' unless - the target is already named, so no press is ever a no-op you have to press - again. ``show`` reports without changing anything. - - A format the ring can't hold (a stack variable, an enum) is reported in - ``warn`` on the way out, with what to do about it -- ``mode`` takes any of - the names above outright, which is also how you put one back. - - Which operand: ``n`` if given, else the one under ``col`` (a column in the - whitespace-collapsed line, as ``heads`` renders it), else the first literal - on the line.""" - import ida_bytes - - try: - ea = ida_bytes.get_item_head(parse_address(addr)) - except Exception as e: - return {"addr": str(addr), "error": str(e)} - - before = _idatui_line_text(ea) - cands = _idatui_op_candidates(ea) - n = int(n) - if n < 0: - n = -1 - if int(col) >= 0: - for lo, hi, i in _idatui_op_spans(ea, before): - if not (lo <= int(col) < hi): - continue - if i in cands: - n = i - break - # The cursor IS on an operand, just not one with a format. The - # client highlights what the cursor is on, so quietly moving to - # a different operand would make that highlight a lie -- say - # which one can be changed instead. - where = before[lo:hi].strip() - alt = (f"; the literal on this line is operand {cands[0]} " - f"({_idatui_op_text(ea, before, cands[0])})" - if cands else "") - return {"addr": hex(ea), "n": i, "text": before, - "error": f"operand {i} ({where}) has no format to " - f"change{alt}"} - if n < 0: - if not cands: - F = ida_bytes.get_flags(ea) - why = ("no literal on this line to reformat" - if ida_bytes.is_code(F) or ida_bytes.is_data(F) else - "undefined bytes have no format to change -- define " - "them first ('d' makes data, 'c' makes code)") - return {"addr": hex(ea), "text": before, "error": why} - n = cands[0] - - cur = _idatui_op_fmt(ea, n) - value, width = _idatui_op_value(ea, n) - mapped = value is not None and value != 0 and ida_bytes.is_mapped(value) - # The ring is a property of the OPERAND, not of what you last pressed: every - # stop is one that changes what you see for this value, and it is the same - # ring at every step, so a lap always comes home. - choices = [f for f in _IDATUI_FMT_CYCLE - if (f != "char" or _idatui_printable(value)) - and (f != "offset" or _idatui_offset_worth(value))] - # A stack variable is deliberately NOT a stop: ``[rbp+var_40]`` is a frame - # member, not a way of writing a number, and IDA's own "is this a stack - # variable" test isn't exposed to Python here (calc_stkvar_struc_offset - # happily answers for ``[r14+8]`` too, which would put a bogus stop in the - # ring). Leaving one is reported instead, with the command that undoes it. - lossy = cur not in choices and cur != "default" - - mode = str(mode or "cycle").lower() - if mode == "show": - return {"addr": hex(ea), "n": n, "format": cur, "prev": cur, - "choices": choices, "text": before, "before": before, - "value": None if value is None else hex(value), - "width": width, "applied": False} - if mode in ("cycle", "back"): - step = 1 if mode == "cycle" else -1 - if cur in choices: - want = choices[(choices.index(cur) + step) % len(choices)] - else: - # Standing on a format the ring can't hold (an enum names a type a - # nibble doesn't record): enter the ring at its end, don't skip a - # stop working out where we "would have" been. - want = choices[0] if step > 0 else choices[-1] - else: - want = mode - if want not in _idatui_fmt_nibbles(): - return {"addr": hex(ea), "n": n, "text": before, - "error": f"unknown format {mode!r}; one of " - + ", ".join(_IDATUI_FMT_SETTABLE)} - if want == "offset" and not mapped: - return {"addr": hex(ea), "n": n, "text": before, "format": cur, - "error": (f"{'0x%x' % value if value is not None else 'this operand'}" - " isn't a mapped address -- an offset to it would" - " invent a name for nothing")} - - ok, err = _idatui_apply_fmt(ea, n, want) - if err: - return {"addr": hex(ea), "n": n, "text": before, "format": cur, - "error": err} - got = _idatui_op_fmt(ea, n) - out = {"addr": hex(ea), "n": n, "prev": cur, "format": got, - "requested": want, "applied": bool(ok), "choices": choices, - "before": before, "text": _idatui_line_text(ea), - "value": None if value is None else hex(value), "width": width} - if not ok: - out["error"] = f"IDA refused {want} on operand {n}" - elif lossy: - out["warn"] = ( - f"operand {n} was {cur} and the ring has no stop there -- " - + (f"'{cur}' sets it again" if cur in _IDATUI_FMT_SETTABLE else - f"{cur} names a type this can't put back, reassign it by hand")) - return out - - -# --------------------------------------------------------------------------- # -# the same thing in the decompiler (Hex-Rays keeps its own number formats) -# --------------------------------------------------------------------------- # -#: Hex-Rays prints C, so two of the listing's stops are missing here: binary -#: (C has no binary literal -- the format takes and then renders decimal, which -#: would be a lie on screen) and offset (it makes the function fail to -#: decompile outright). -_IDATUI_PC_FMT_CYCLE = ("hex", "dec", "oct", "char", "default") - - -def _idatui_compact(line): - """The ida-pro-mcp whitespace collapse the pseudocode is served through, so - a column in what the client SHOWS can be mapped back to Hex-Rays' line.""" - try: - from ida_pro_mcp.ida_mcp.utils import compact_whitespace - return compact_whitespace(line) - except Exception: - import re as _re - stripped = line.lstrip(" \t") - lead = line[: len(line) - len(stripped)] - return lead + _re.sub(r"[ \t]{2,}", " ", stripped) - - -def _idatui_compact_col(plain, compact, col): - """The inverse of ``_idatui_uncompact_col``: a column in Hex-Rays' own line, - expressed in the collapsed line the client shows.""" - j = 0 - for i in range(min(int(col), len(plain))): - if j < len(compact) and plain[i] == compact[j]: - j += 1 - return j - - -def _idatui_uncompact_col(plain, compact, col): - """Map a column in the collapsed line back to the same character in the - original. The transform only ever DELETES spaces, so walking both in step - and skipping what vanished is exact.""" - i = 0 - for j in range(min(int(col), len(compact))): - c = compact[j] - while i < len(plain) and plain[i] != c: - i += 1 - i += 1 - return min(i, max(len(plain) - 1, 0)) - - -#: Characters that can be part of a C number literal as Hex-Rays prints one -#: (digits, hex letters, the 0x prefix, u/L suffixes). -_IDATUI_LIT_CHARS = frozenset("0123456789abcdefABCDEFxXuUlL") - - -def _idatui_lit_extent(plain, x): - """The [start, end) of the literal token containing column ``x``. - - Hex-Rays says WHICH item a column belongs to, but not how wide the printed - literal is -- and it attributes neighbouring punctuation to the same item, - so ``if ( a1 > 1 )`` reports the closing paren as part of the number. The - identity comes from the ctree; the extent is the run of literal characters - around the column, which cannot reach a ``)`` or a space.""" - if x >= len(plain): - return None - if plain[x] == "'": # a character constant: '-' - end = plain.find("'", x + 1) - return (x, end + 1) if end > x else None - lo = plain.rfind("'", 0, x) - if lo >= 0 and plain.find("'", x) > x and "'" in plain[lo:x] and \ - plain[lo:x].count("'") == 1 and " " not in plain[lo:x]: - return (lo, plain.find("'", x) + 1) # inside 'c' - if plain[x] not in _IDATUI_LIT_CHARS: - return None - lo = x - while lo > 0 and plain[lo - 1] in _IDATUI_LIT_CHARS: - lo -= 1 - hi = x - while hi < len(plain) and plain[hi] in _IDATUI_LIT_CHARS: - hi += 1 - if lo > 0 and plain[lo - 1] == "-": # a unary minus is part of it - lo -= 1 - return (lo, hi) - - -def _idatui_pc_nums(cf, sl): - """Every number literal on one pseudocode line, as - [{x0, x1, ea, opnum, value, nbytes, fmt}]. - - Asks Hex-Rays what each column belongs to rather than pattern-matching the - text: a regex over ``v6 = a1 - 1;`` has to guess which of those characters - are a literal, and ``v11`` looks like one.""" - import ida_bytes - import ida_hexrays - import ida_lines - import idaapi - - plain = ida_lines.tag_remove(sl.line) - out = [] - x = 0 - while x < len(plain): - ch = plain[x] - if ch not in _IDATUI_LIT_CHARS and ch != "'": - x += 1 - continue - head, item, tail = (ida_hexrays.ctree_item_t() for _ in range(3)) - if not cf.get_line_item(sl.line, x, True, head, item, tail): - x += 1 - continue - if item.citype != ida_hexrays.VDI_EXPR: - x += 1 - continue - e = item.e - if e.op != ida_hexrays.cot_num: - x += 1 - continue - extent = _idatui_lit_extent(plain, x) - if extent is None: - x += 1 - continue - nf = e.n.nf - opnum = ord(nf.opnum) if isinstance(nf.opnum, str) else int(nf.opnum) - nbytes = (ord(nf.org_nbytes) if isinstance(nf.org_nbytes, str) - else int(nf.org_nbytes)) - ea = int(e.ea) - if ea == idaapi.BADADDR: - x = extent[1] - continue # synthesised: nothing to key on - nib = (nf.flags >> ida_bytes.get_operand_type_shift(opnum)) & 0xF - # Whether this format is the USER's or Hex-Rays' own guess. The nibble - # can't say: an untouched number reads back as whatever it happens to - # be printed as, and cycling from there would skip that stop forever - # (default already looks like it) and never come back to it. - loc = ida_hexrays.operand_locator_t(ea, opnum) - user = (ida_hexrays.user_numforms_find(cf.numforms, loc) - != ida_hexrays.user_numforms_end(cf.numforms)) - out.append({"x0": extent[0], "x1": extent[1], "ea": ea, - "opnum": opnum, "value": int(e.n._value), - "nbytes": nbytes, "user": user, - "fmt": _idatui_fmt_name(nib) if user else "default", - "shown": _idatui_fmt_name(nib)}) - x = extent[1] # past this literal, not into it - return out - - -@tool -@idasync -def pc_nums( - addr: Annotated[str, "Function address (or any address inside it)"], -) -> dict: - """Every number literal in a function's pseudocode, as - [{line, x0, x1, ea, opnum, value, fmt, user}]. - - One call per decompilation, so a client can show WHICH literal the cursor is - on (and reformat exactly that one) without a round trip per cursor move. - Columns are in the same collapsed coordinates the decompile tool serves its - text in, i.e. what the client actually displays.""" - import ida_hexrays - import ida_lines - import idaapi - - if not ida_hexrays.init_hexrays_plugin(): - return {"addr": str(addr), "error": "no decompiler", "nums": []} - try: - f = idaapi.get_func(parse_address(addr)) - except Exception as e: - return {"addr": str(addr), "error": str(e), "nums": []} - if f is None: - return {"addr": str(addr), "error": "no function here", "nums": []} - try: - cf = ida_hexrays.decompile(f.start_ea) - except Exception as e: - return {"addr": hex(f.start_ea), "error": f"decompile failed: {e}", - "nums": []} - if cf is None: - return {"addr": hex(f.start_ea), "error": "decompilation failed", - "nums": []} - sv = cf.get_pseudocode() - out = [] - for i in range(len(sv)): - plain = ida_lines.tag_remove(sv[i].line) - compact = _idatui_compact(plain) - for rec in _idatui_pc_nums(cf, sv[i]): - out.append({ - "line": i, - "x0": _idatui_compact_col(plain, compact, rec["x0"]), - "x1": _idatui_compact_col(plain, compact, rec["x1"]), - "ea": hex(rec["ea"]), "opnum": rec["opnum"], - "value": hex(rec["value"]), "fmt": rec["fmt"], - "shown": rec["shown"], "user": bool(rec["user"]), - }) - return {"addr": hex(f.start_ea), "nums": out, "lines": len(sv)} - - -@tool -@idasync -def pc_num_format( - addr: Annotated[str, "Function address (or any address inside it)"], - mode: Annotated[str, "cycle | back | show | hex | dec | oct | char | default"] = "cycle", - line: Annotated[int, "0-based pseudocode line index"] = -1, - col: Annotated[int, "Cursor column in the DISPLAYED line (-1: first literal)"] = -1, - ea: Annotated[str, "Address of the number instead of line/col"] = "", - opnum: Annotated[int, "Operand number, with ``ea``"] = -1, -) -> dict: - """Change how a number is displayed in the DECOMPILATION (Hex-Rays keeps its - own number formats, per (address, operand), independent of the listing). - - Same stops as ``op_format`` minus the two C can't express: binary (no such - literal -- IDA takes the format and prints decimal anyway) and offset (it - makes the function stop decompiling). Returns the re-rendered line, and - marks the function dirty so the next decompile is the new text.""" - import ida_hexrays - import ida_lines - import idaapi - - if not ida_hexrays.init_hexrays_plugin(): - return {"addr": str(addr), "error": "no decompiler"} - try: - f = idaapi.get_func(parse_address(addr)) - except Exception as e: - return {"addr": str(addr), "error": str(e)} - if f is None: - return {"addr": str(addr), "error": "no function here"} - try: - cf = ida_hexrays.decompile(f.start_ea) - except Exception as e: - return {"addr": hex(f.start_ea), "error": f"decompile failed: {e}"} - if cf is None: - return {"addr": hex(f.start_ea), "error": "decompilation failed"} - - sv = cf.get_pseudocode() - line = int(line) - target = None - if ea: - try: - want_ea = parse_address(ea) - except Exception as e: - return {"addr": hex(f.start_ea), "error": str(e)} - for i in range(len(sv)): - for rec in _idatui_pc_nums(cf, sv[i]): - if rec["ea"] == want_ea and (int(opnum) < 0 - or rec["opnum"] == int(opnum)): - target, line = rec, i - break - if target: - break - elif 0 <= line < len(sv): - nums = _idatui_pc_nums(cf, sv[line]) - if nums: - if int(col) >= 0: - plain = ida_lines.tag_remove(sv[line].line) - x = _idatui_uncompact_col(plain, _idatui_compact(plain), int(col)) - target = next((r for r in nums if r["x0"] <= x < r["x1"]), None) - target = target or nums[0] - else: - return {"addr": hex(f.start_ea), - "error": f"line {line} is outside the {len(sv)}-line decompilation"} - if target is None: - return {"addr": hex(f.start_ea), "line": line, - "text": (ida_lines.tag_remove(sv[line].line).strip() - if 0 <= line < len(sv) else ""), - "error": "no number literal on this line"} - - cur, value = target["fmt"], target["value"] - choices = [c for c in _IDATUI_PC_FMT_CYCLE - if c != "char" or _idatui_printable(value)] - # Same rule as the listing: one ring per literal, every step. A format the - # ring can't hold (an enum set in the GUI) is reported on the way out - # instead of being kept for one lap and then lost. - lossy = cur not in choices and cur != "default" - out = {"addr": hex(f.start_ea), "ea": hex(target["ea"]), - "opnum": target["opnum"], "line": line, "prev": cur, - "format": cur, "shown": target["shown"], "choices": choices, - "value": hex(value), - "before": ida_lines.tag_remove(sv[line].line).strip()} - - mode = str(mode or "cycle").lower() - if mode == "show": - out["text"] = out["before"] - out["applied"] = False - return out - if mode in ("cycle", "back"): - step = 1 if mode == "cycle" else -1 - if cur in choices: - want = choices[(choices.index(cur) + step) % len(choices)] - else: - want = choices[0] if step > 0 else choices[-1] - else: - want = mode - if want in ("bin", "offset", "stack", "seg", "float"): - out["error"] = (f"Hex-Rays has no {want} format for a number " - f"-- set it on the listing instead") - out["text"] = out["before"] - return out - if want not in ("hex", "dec", "oct", "char", "default"): - out["error"] = (f"unknown format {mode!r}; one of hex, dec, oct, " - f"char, default") - out["text"] = out["before"] - return out - - loc = ida_hexrays.operand_locator_t(target["ea"], target["opnum"]) - it = ida_hexrays.user_numforms_find(cf.numforms, loc) - if it != ida_hexrays.user_numforms_end(cf.numforms): - # std::map::insert is a no-op on an existing key, so a format already - # set here would silently win over the new one. - ida_hexrays.user_numforms_erase(cf.numforms, it) - if want != "default": - import ida_bytes - nf = ida_hexrays.number_format_t(target["opnum"]) - nf.flags = ida_bytes.get_operand_flag(_idatui_fmt_nibbles()[want], - target["opnum"]) - try: - nf.org_nbytes = target["nbytes"] - except Exception: - pass - ida_hexrays.user_numforms_insert(cf.numforms, loc, nf) - cf.save_user_numforms() - try: - ida_hexrays.mark_cfunc_dirty(f.start_ea) - except Exception: - pass - - out["format"] = want - out["applied"] = True - if lossy: - out["warn"] = (f"this number was {cur}, which names a type a radix " - f"can't put back -- reassign it in IDA") - try: - cf2 = ida_hexrays.decompile(f.start_ea, - flags=ida_hexrays.DECOMP_NO_CACHE) - sv2 = cf2.get_pseudocode() if cf2 is not None else None - out["text"] = (ida_lines.tag_remove(sv2[line].line).strip() - if sv2 is not None and line < len(sv2) else out["before"]) - except Exception as e: - out["text"] = out["before"] - out["warn"] = f"re-render failed: {e}" - return out - - -@tool -@idasync -def flowchart( - addr: Annotated[str, "Address or name inside the function to chart"], -) -> dict: - """Basic-block control-flow graph of the function containing ``addr``. - - Returns the blocks and the edges between them -- NOT their text: the block - body is just an address range, which the client already knows how to render - with ``heads``. Keeping text out means the graph view reuses the exact same - listing rows (colours, operand marks and all) instead of growing a second - disassembly renderer. - - Edge ``kind`` is what the graph view colours by: - * ``fall`` -- control falls through to the next address (IDA draws red) - * ``jump`` -- a taken conditional branch (green) - * ``uncond`` -- the block's only successor (blue) - * ``switch`` -- one of an n-way dispatch - """ - import ida_funcs - import ida_gdl - - try: - ea = parse_address(addr) - except Exception as e: - return {"addr": str(addr), "error": str(e), "blocks": []} - fn = ida_funcs.get_func(ea) - if fn is None: - return {"addr": str(addr), "error": "no function at that address", - "blocks": []} - - fc = ida_gdl.FlowChart(fn, flags=ida_gdl.FC_PREDS) - index = {} - order = [] - for bb in fc: - index[bb.start_ea] = len(order) - order.append(bb) - blocks = [] - for bb in order: - sl = [s for s in bb.succs() if s.start_ea in index] - succs = [] - for s in sl: - if len(sl) > 2: - kind = "switch" - elif s.start_ea == bb.end_ea: - kind = "fall" - else: - kind = "jump" - succs.append([index[s.start_ea], kind]) - blocks.append({ - "id": index[bb.start_ea], - "start": hex(bb.start_ea), - "end": hex(bb.end_ea), - "succs": succs, - }) - return { - "addr": hex(ea), - "func": {"addr": hex(fn.start_ea), "end": hex(fn.end_ea), - "name": ida_funcs.get_func_name(fn.start_ea)}, - "entry": index.get(fn.start_ea, 0), - "blocks": blocks, - } -''' - -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/_fixtures.py b/tests/_fixtures.py index 766ecd5..c7a45d3 100644 --- a/tests/_fixtures.py +++ b/tests/_fixtures.py @@ -50,7 +50,13 @@ async def build_pristine(binary: str, cache: str, app_factory) -> None: 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) diff --git a/tests/run.py b/tests/run.py index 4920e54..b38a707 100755 --- a/tests/run.py +++ b/tests/run.py @@ -51,8 +51,8 @@ import time ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) TESTS = os.path.join(ROOT, "tests") -#: The IDA-capable interpreter. The pilot tests need textual AND idapro in one -#: python; the worker python is auto-detected separately by WorkerClient. +#: The IDA-capable interpreter. The pilot tests need textual AND the Code Mode +#: library in one python; the database process is Code Mode's to place. DEFAULT_PY = os.path.expanduser("~/ida-venv/bin/python") #: Both shapes the suites print: "N passed, M failed" and "N checks, M failed". diff --git a/tests/test_codemode_client.py b/tests/test_codemode_client.py new file mode 100644 index 0000000..0954918 --- /dev/null +++ b/tests/test_codemode_client.py @@ -0,0 +1,140 @@ +"""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 + +#: Pure: fakes the DatabaseHandle, never touches IDA or the Code Mode library. +NEEDS_IDA = False + +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_launch.py b/tests/test_launch.py index c6ed8f1..d57ee5f 100644 --- a/tests/test_launch.py +++ b/tests/test_launch.py @@ -1,12 +1,14 @@ #!/usr/bin/env python3 -"""The launcher's file handling -- the part that deletes things. +"""The launcher's option handling, and the file handling it must NOT do. -`_sweep_locks` runs automatically when a database fails to open, and it removes -files next to the user's binary. That is exactly the kind of code that must not -be tested by trying it, so it is tested here: which files it takes, which it -must never take, and what it reports. +The old `_sweep_locks` deleted `.id0/.id1/.id2/.nam/.til` next to the user's +binary when a database failed to open. That was only defensible while the TUI +exclusively owned a private worker; under Code Mode a GUI or another client may +own the database, so the sweep is gone. Its tests are replaced by one that keeps +it gone -- deleting a shared database's working files is unrecoverable, and this +is the cheapest guard against someone reintroducing the "helpful" cleanup. -Pure: no IDA, no worker, no Textual. +Pure: no IDA, no Code Mode library, no Textual. """ from __future__ import annotations @@ -20,7 +22,8 @@ sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) #: Read by tests/run.py (--fast skips every NEEDS_IDA file). NEEDS_IDA = False -from idatui.launch import _LOCK_SUFFIXES, _load_args, _sweep_locks # noqa: E402 +import idatui.launch as launch # noqa: E402 +from idatui.launch import _load_args # noqa: E402 PASS = FAIL = 0 @@ -41,100 +44,18 @@ def touch(*paths): fh.write(b"x") -def t_sweeps_the_scratch_files(): - """IDA unpacks a .i64 into .id0/.id1/.id2/.nam/.til while it is open; a - hard-killed worker leaves them and the .i64 then refuses to reopen.""" - with tempfile.TemporaryDirectory() as d: - binary = os.path.join(d, "echo") - touch(binary, *[binary + s for s in _LOCK_SUFFIXES]) - n = _sweep_locks(binary) - check("every unpacked scratch file is swept", n == len(_LOCK_SUFFIXES), - f"swept {n} of {len(_LOCK_SUFFIXES)}") - check("none of them survive", - not any(os.path.exists(binary + s) for s in _LOCK_SUFFIXES)) - check("the binary itself is untouched", os.path.exists(binary)) +def t_no_lock_sweeping(): + """The launcher must not delete database working files any more. - -def t_sweeps_by_stem_too(): - """IDA keys the scratch on the full name or the stem depending on how the - database was created, so both are swept.""" - with tempfile.TemporaryDirectory() as d: - binary = os.path.join(d, "prog.elf") - stem = os.path.join(d, "prog") - touch(binary, stem + ".id0", stem + ".nam", binary + ".id1") - n = _sweep_locks(binary) - check("scratch named after the stem is swept too", n == 3, f"n={n}") - check("stem-keyed files are gone", - not os.path.exists(stem + ".id0") - and not os.path.exists(stem + ".nam")) - check("full-name-keyed files are gone", not os.path.exists(binary + ".id1")) - check("the binary itself is untouched", os.path.exists(binary)) - - -def t_never_the_database(): - """The .i64 IS the database. Nothing is saved unless idb_save was called, so - deleting it throws away every rename and comment in the session.""" - with tempfile.TemporaryDirectory() as d: - binary = os.path.join(d, "echo") - db = binary + ".i64" - stem_db = os.path.join(d, "echo.i64") - touch(binary, db, binary + ".id0") - _sweep_locks(binary) - check("the .i64 is never swept", os.path.exists(db)) - check("nor the stem-keyed .i64", os.path.exists(stem_db)) - check(".i64 is not in the suffix list", ".i64" not in _LOCK_SUFFIXES, - str(_LOCK_SUFFIXES)) - - -def t_never_the_input_itself(): - """`.til` is both an unpacked-DB suffix and the extension of an IDA type - library, so `ida-tui mylib.til` used to sweep its own argument out of - existence -- irreversibly, on a path that runs automatically when an open - fails. Same for anything named *.id0/*.id1/*.id2/*.nam. + Code Mode's registry locks, health probes and IDA itself arbitrate database + ownership now. A sweep here would delete files out from under a live GUI. """ - for suf in _LOCK_SUFFIXES: - with tempfile.TemporaryDirectory() as d: - binary = os.path.join(d, "mylib" + suf) - touch(binary) - _sweep_locks(binary) - check(f"a binary named *{suf} is not deleted by its own sweep", - os.path.exists(binary), f"{binary} was removed") - - -def t_relative_path_is_still_the_input(): - """The guard compares absolute paths -- a relative argument names the same - file and must be protected the same way.""" - with tempfile.TemporaryDirectory() as d: - cwd = os.getcwd() - try: - os.chdir(d) - touch("mylib.til") - _sweep_locks("mylib.til") - check("a relative path to the input is protected too", - os.path.exists("mylib.til")) - finally: - os.chdir(cwd) - - -def t_missing_files_are_fine(): - with tempfile.TemporaryDirectory() as d: - binary = os.path.join(d, "nothing-here") - touch(binary) - n = _sweep_locks(binary) - check("sweeping with nothing to sweep reports 0", n == 0, f"n={n}") - check("and does not raise", True) - - -def t_leaves_the_neighbours_alone(): - with tempfile.TemporaryDirectory() as d: - binary = os.path.join(d, "echo") - other = os.path.join(d, "other.id0") # another binary's scratch - src = os.path.join(d, "echo.c") - touch(binary, other, src, binary + ".id0") - _sweep_locks(binary) - check("another binary's scratch is left alone", os.path.exists(other)) - check("unrelated neighbours are left alone", os.path.exists(src)) - check("our own scratch is still swept", not os.path.exists(binary + ".id0")) + check("_sweep_locks is gone", not hasattr(launch, "_sweep_locks")) + check("the scratch-suffix list is gone", not hasattr(launch, "_LOCK_SUFFIXES")) + src = open(launch.__file__, encoding="utf-8").read() + check("the launcher does not remove files at all", + "os.remove" not in src and "shutil.rmtree" not in src, + "launch.py deletes something again") def t_load_args(): @@ -154,10 +75,7 @@ def t_load_args(): def main() -> int: - for fn in (t_sweeps_the_scratch_files, t_sweeps_by_stem_too, - t_never_the_database, t_never_the_input_itself, - t_relative_path_is_still_the_input, t_missing_files_are_fine, - t_leaves_the_neighbours_alone, t_load_args): + for fn in (t_no_lock_sweeping, t_load_args): print(f"\n{fn.__name__}") try: fn() diff --git a/tests/test_pool.py b/tests/test_pool.py index 94e197b..9fc1446 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 """ @@ -15,7 +14,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 @@ -32,11 +31,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 @@ -45,10 +45,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 @@ -76,15 +75,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 ---------------------------------------------------- # @@ -100,9 +99,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}") @@ -116,7 +115,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") @@ -143,7 +142,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) @@ -155,8 +154,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 ------ # @@ -169,7 +168,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] @@ -188,6 +187,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 822002f..690f360 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 9b2cc7e..e19ea11 100644 --- a/tests/test_scenarios.py +++ b/tests/test_scenarios.py @@ -694,18 +694,18 @@ async def s_split_view(c: Ctx): # to count. The bound is loose because the bug was three orders of magnitude # out, not a near miss. _lookups = {"n": 0} - _orig_call = c.prog.client.call + _orig_call = c.prog.client.invoke def _counting(name, *a, **kw): if name == "lookup_funcs": _lookups["n"] += 1 return _orig_call(name, *a, **kw) - c.prog.client.call = _counting + c.prog.client.invoke = _counting try: await _split_view_body(c, app, lst, dec) finally: - c.prog.client.call = _orig_call + c.prog.client.invoke = _orig_call c.check("split view doesn't storm the worker with function lookups", _lookups["n"] < 500, f"{_lookups['n']} lookup_funcs calls") @@ -1586,14 +1586,14 @@ async def s_decomp_nav(c: Ctx): old_ea = dec._line_ea(drow) if old_ea is not None and dsym in old_line: tmp = f"stale_{os.getpid()}" - app.program.client.call("rename", batch={"func": {"addr": hex(dstale), "name": tmp}}) + app.program.client.invoke("rename", batch={"func": {"addr": hex(dstale), "name": tmp}}) app.program.bump_names() d2 = len(app._nav) app._follow_decomp(old_line, dsym, old_ea) await c.wait(lambda: len(app._nav) > d2, 25) c.check("decomp follow works with a stale name (ea-marker fallback)", app._cur.ea == dstale, f"cur={app._cur.ea:#x} want={dstale:#x}") - app.program.client.call("rename", batch={"func": {"addr": hex(dstale), "name": dsym}}) + app.program.client.invoke("rename", batch={"func": {"addr": hex(dstale), "name": dsym}}) app.program.bump_names() @@ -1674,7 +1674,7 @@ async def s_rename(c: Ctx): c.check("rename updates the function name", app._func_index.by_addr(dtarget).name == newname, app._func_index.by_addr(dtarget).name) - rr = app.program.client.call("rename", batch={"func": {"addr": hex(dtarget), "name": dsym}}) + rr = app.program.client.invoke("rename", batch={"func": {"addr": hex(dtarget), "name": dsym}}) c.check("rename reverted cleanly", rr.get("summary", {}).get("ok", 0) == 1, str(rr.get("summary"))) # goto label refuse @@ -1724,7 +1724,7 @@ async def s_rename(c: Ctx): and any(cnote in t for t in dec._texts), 25) c.check("comment appears in the pseudocode after ';'", any(cnote in t for t in dec._texts), "comment not shown") - app.program.client.call("set_comments", items=[{"addr": hex(cea), "comment": ""}]) + app.program.client.invoke("set_comments", items=[{"addr": hex(cea), "comment": ""}]) else: c.check("found a pseudocode line to comment", False, "no marker line") @@ -1766,7 +1766,7 @@ async def s_comment_func(c: Ctx): la is not None and lb is not None and lb > la and a not in dec._texts[lb], f"la={la} lb={lb}") - app.program.client.call("set_comments", items=[{"addr": hex(fn.addr), "comment": ""}]) + app.program.client.invoke("set_comments", items=[{"addr": hex(fn.addr), "comment": ""}]) @scenario("retype") @@ -2031,7 +2031,7 @@ async def s_rename_history(c: Ctx): await c.pause(0.15) c.check("caller pseudocode shows renamed callee after 'back'", any(hnew in tx for tx in dec._texts), "pseudocode still stale") - app.program.client.call("rename", batch={"func": {"addr": hex(htarget), "name": hsym}}) + app.program.client.invoke("rename", batch={"func": {"addr": hex(htarget), "name": hsym}}) @scenario("region_define") @@ -2250,7 +2250,7 @@ async def s_listing_name_addr(c: Ctx): finally: # revert: drop the label and restore raw bytes at A try: - c.prog.client.call("rename", batch={"data": {"addr": hex(A + 1), "new": ""}}) + c.prog.client.invoke("rename", batch={"data": {"addr": hex(A + 1), "new": ""}}) except Exception: # noqa: BLE001 pass c.prog.undefine(A, size=8) @@ -2327,7 +2327,7 @@ async def s_listing_struct_expand(c: Ctx): c.check("found a data address for the struct test", False) return try: - c.prog.client.call( + c.prog.client.invoke( "declare_type", decls=["struct TuiExpandS { int a; char b[4]; short c; };"]) c.prog.make_data(A, "TuiExpandS") @@ -3206,7 +3206,7 @@ async def s_graph_rename(c: Ctx): f"resolve({new}) -> {got if got is None else hex(got)} want {ea:#x}") # revert, so the suite stays idempotent if got is not None: - app.program.client.call( + app.program.client.invoke( "rename", batch={"data": {"addr": hex(ea), "new": ""}}) app.program.bump_names() @@ -3259,7 +3259,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) @@ -3279,6 +3279,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/tests/test_worker_client.py b/tests/test_worker_client.py deleted file mode 100644 index cd99377..0000000 --- a/tests/test_worker_client.py +++ /dev/null @@ -1,447 +0,0 @@ -#!/usr/bin/env python3 -"""WorkerClient: spawn, transport, failure reporting, shutdown. - -This is the layer between the app and idalib, and it had no tests -- which is -awkward, because it is where the failures are silent and expensive. A worker -that dies during startup, a socket that drops mid-call, two UI threads sharing -one socket: none of those look like a bug from the outside, they look like the -TUI hanging or showing stale data. - -None of it needs IDA. The client spawns whatever ``_WORKER_PY`` points at, so -these tests point it at a fake that speaks the same length-prefixed pickle -protocol and can be told to misbehave on demand. -""" -from __future__ import annotations - -import os -import sys -import threading -import time - -sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) - -#: pure: a fake worker over a unix socket, no idalib anywhere. -#: Read by tests/run.py (--fast skips every NEEDS_IDA file). -NEEDS_IDA = False - -from idatui import worker_client as wc # noqa: E402 -from idatui.errors import IDAConnectionError, IDAToolError # noqa: E402 - -PASS = FAIL = 0 - - -def check(name, ok, detail=""): - global PASS, FAIL - if ok: - PASS += 1 - print(f" ok {name}") - else: - FAIL += 1 - print(f" FAIL {name} {detail}") - - -# --------------------------------------------------------------------------- # -# A worker that isn't IDA -# --------------------------------------------------------------------------- # -#: Speaks the real protocol (idatui.worker.send/recv) and implements a handful -#: of tools whose only job is to be predictable, plus the misbehaviours we need: -#: dying at startup, dropping the socket mid-conversation, taking its time. -FAKE_WORKER = r''' -import os, socket, sys, time -sys.path.insert(0, %(repo)r) -from idatui.worker import send, recv - -sock_path, binary = sys.argv[1], sys.argv[2] -mode = os.environ.get("FAKE_MODE", "ok") - -if mode == "die": - # A startup crash, the way the real worker reports one. - print("IDA Pro: thank you for using it") # banner noise, must be skipped - print("WORKER-FATAL: could not open database: it is wedged") - sys.stdout.flush() - sys.exit(3) -if mode == "hang": - time.sleep(60) # never binds: connect() must time out - sys.exit(0) - -srv = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) -if os.path.exists(sock_path): - os.unlink(sock_path) -srv.bind(sock_path) -srv.listen(1) -conn, _ = srv.accept() -served = 0 -while True: - msg = recv(conn) - if msg is None: - break - tool, args = msg - if tool == "__shutdown__": - # The real worker closes its database here; a clean exit is the signal - # the client waits for rather than killing us. - open(sock_path + ".clean", "w").write("shutdown") - break - served += 1 - if tool == "drop": - conn.close() # vanish mid-conversation - break - if tool == "boom": - send(conn, (False, "the tool exploded")) - continue - if tool == "slow": - time.sleep(float(args.get("secs", 0.2))) - send(conn, (True, {"tool": tool, "args": args, "n": served})) - continue - send(conn, (True, {"tool": tool, "args": args, "n": served, - "binary": os.path.basename(binary)})) -sys.exit(0) -''' - - -def _install_fake(tmpdir: str) -> str: - repo = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) - path = os.path.join(tmpdir, "fake_worker.py") - with open(path, "w", encoding="utf-8") as fh: - fh.write(FAKE_WORKER % {"repo": repo}) - wc._WORKER_PY = path - return path - - -def client(tmpdir, **kw): - """A client wired to the fake worker, running under THIS interpreter. - - ``python=`` matters: the real constructor probes three interpreters for - ``import ida_pro_mcp`` and that is both slow and beside the point here. - """ - binary = os.path.join(tmpdir, "target.bin") - if not os.path.exists(binary): - with open(binary, "wb") as fh: - fh.write(b"\x7fELF" + b"\0" * 60) - return wc.WorkerClient(binary, python=sys.executable, **kw) - - -# --------------------------------------------------------------------------- # -def t_roundtrip(tmp): - c = client(tmp) - try: - c.connect(timeout=30) - r = c.call("survey_binary", depth=2) - check("a call round-trips through the socket", - r["tool"] == "survey_binary" and r["args"] == {"depth": 2}, str(r)) - check("the worker got the binary path we asked for", - r["binary"] == "target.bin", str(r)) - check("pid is exposed for memory accounting", isinstance(c.pid, int)) - r2 = c.call("second") - check("the connection is reused, not respawned per call", - r2["n"] == 2, f"n={r2['n']}") - finally: - c.close(grace=5) - - -def t_envelope(tmp): - """domain.decompile reads result.structuredContent -- keep that shape.""" - c = client(tmp) - try: - c.connect(timeout=30) - env = c.call_envelope("decompile", addr="0x1000") - inner = env["result"]["structuredContent"] - check("call_envelope wraps the payload the way domain.py unwraps it", - inner["tool"] == "decompile" and inner["args"] == {"addr": "0x1000"}, - str(env)) - finally: - c.close(grace=5) - - -def t_tool_error(tmp): - c = client(tmp) - try: - c.connect(timeout=30) - try: - c.call("boom") - check("a failing tool raises IDAToolError", False, "no exception") - except IDAToolError as e: - check("a failing tool raises IDAToolError", True) - check("the error names the tool", e.tool == "boom", f"tool={e.tool!r}") - check("and carries the worker's message", - "exploded" in e.message, e.message) - # A tool error is not a transport error: the connection must survive it, - # or one bad decompile would tear down the session. - r = c.call("after") - check("the connection survives a tool error", r["tool"] == "after", str(r)) - finally: - c.close(grace=5) - - -def t_dropped_socket(tmp): - """The reconnect trigger. The app catches IDAConnectionError and reconnects; - if a drop raised something else, or left _sock set, it would instead surface - as a crash or as every later call failing.""" - c = client(tmp) - try: - c.connect(timeout=30) - try: - c.call("drop") - check("a dropped socket raises IDAConnectionError", False, - "no exception") - except IDAConnectionError: - check("a dropped socket raises IDAConnectionError", True) - except Exception as e: # noqa: BLE001 - check("a dropped socket raises IDAConnectionError", False, - f"got {type(e).__name__}: {e}") - check("the dead socket is cleared, so a retry can reconnect", - c._sock is None) - finally: - c.close(grace=5) - - -def t_startup_crash(tmp): - """A worker that dies before binding must say WHY. - - 'worker exited (code 3)' on its own is indistinguishable from a bug in the - app; the real cause is the last meaningful line of its log, and the licence - banner must not be mistaken for it. - """ - os.environ["FAKE_MODE"] = "die" - try: - c = client(tmp) - try: - c.connect(timeout=30) - check("a worker that exits during startup raises", False, - "connect() returned") - except IDAConnectionError as e: - msg = str(e) - check("a worker that exits during startup raises", True) - check("the exit code is reported", "code 3" in msg, msg) - check("the WORKER-FATAL line is surfaced", - "wedged" in msg, msg) - check("the licence banner is not mistaken for the error", - "thank you" not in msg.lower(), msg) - check("the full log path is offered", ".log" in msg, msg) - finally: - os.environ.pop("FAKE_MODE", None) - - -def t_connect_timeout(tmp): - """A worker that never binds must give up, not block the UI forever.""" - os.environ["FAKE_MODE"] = "hang" - try: - c = client(tmp) - t0 = time.time() - try: - c.connect(timeout=0.4) - check("connect() gives up on a worker that never binds", False, - "returned") - except IDAConnectionError as e: - took = time.time() - t0 - check("connect() gives up on a worker that never binds", True) - check("it honours the timeout it was given", took < 10, f"{took:.1f}s") - check("and says so", "in time" in str(e), str(e)) - finally: - # grace=0: it is sleeping by design, don't wait it out. - c.close(grace=0) - finally: - os.environ.pop("FAKE_MODE", None) - - -def t_progress(tmp): - """connect() reports progress while analysis runs -- that callback is the - only thing on screen during a long open.""" - os.environ["FAKE_MODE"] = "hang" - seen = [] - try: - c = client(tmp) - try: - c.connect(timeout=0.6, progress=seen.append) - except IDAConnectionError: - pass - finally: - c.close(grace=0) - finally: - os.environ.pop("FAKE_MODE", None) - check("connect() reports progress while waiting", bool(seen), - f"{len(seen)} callbacks") - check("progress names the binary being analysed", - any("target.bin" in s for s in seen), str(seen[:1])) - - -def t_serialized(tmp): - """One socket, many UI threads. - - The app fires calls from several worker threads over one client. The frames - are length-prefixed pickle with no request ids, so if two calls interleaved - on the wire each would read the other's reply -- silently, as wrong data - rather than an error. The lock is the only thing preventing that, so this - checks every thread gets its own answer back. - """ - c = client(tmp) - try: - c.connect(timeout=30) - out, errs = {}, [] - - def go(i): - try: - out[i] = c.call("slow", secs=0.05, tag=i) - except Exception as e: # noqa: BLE001 - errs.append(e) - - threads = [threading.Thread(target=go, args=(i,)) for i in range(8)] - t0 = time.time() - for t in threads: - t.start() - for t in threads: - t.join(30) - took = time.time() - t0 - check("concurrent calls all completed", len(out) == 8 and not errs, - f"{len(out)} results, errors={errs[:1]}") - check("each thread got ITS OWN reply, not another's", - all(out[i]["args"]["tag"] == i for i in out), - str({i: out[i]["args"].get("tag") for i in sorted(out)})) - check("calls were serialized, not interleaved", - took >= 8 * 0.05, f"{took:.2f}s for 8 x 0.05s") - check("the worker saw every call exactly once", - sorted(r["n"] for r in out.values()) == list(range(1, 9)), - str(sorted(r["n"] for r in out.values()))) - finally: - c.close(grace=5) - - -def t_clean_shutdown(tmp): - """close() must let the worker close its database. - - A hard kill leaves the .i64 unpacked into .id0/.id1/... and the database - then fails to reopen. So close() sends __shutdown__ and WAITS; only a truly - stuck worker gets signalled. - """ - c = client(tmp) - c.connect(timeout=30) - sock_path = c._sock_path - proc = c._proc - c.close(grace=15) - check("close() sends __shutdown__ rather than killing", - os.path.exists(sock_path + ".clean")) - check("and waits for the worker to exit on its own", - proc.poll() == 0, f"returncode={proc.poll()}") - - -def t_call_after_close(tmp): - """A closed client must stay closed. - - call() reconnects when _sock is None, which is what makes a dropped socket - recoverable -- but it made an explicitly CLOSED client resurrect too, and - spawn a whole new idalib worker to serve one stray call. Teardown and - binary-switch both close while @work threads are in flight, so quitting - during a decompile left a fresh process re-opening the .i64 we had just - released. (Verified before the fix: pid 1066961 -> 1066962.) - """ - c = client(tmp) - c.connect(timeout=30) - pid = c.pid - c.close(grace=5) - try: - c.call("zombie") - check("a call after close() does not resurrect the worker", False, - f"call succeeded; pid {pid} -> {c.pid}") - except IDAConnectionError as e: - check("a call after close() does not resurrect the worker", True) - check("and says the client was closed", "closed" in str(e), str(e)) - check("no second worker was spawned", c.pid == pid, f"{pid} -> {c.pid}") - # ... but an explicit reconnect still revives it: that is how the app - # recovers from a worker that segfaulted. - c.connect(timeout=30) - r = c.call("revived") - check("connect() revives a closed client", r["tool"] == "revived", str(r)) - c.close(grace=5) - - -def t_worker_python_override(tmp): - """$IDATUI_WORKER_PYTHON wins, and the answer is cached. - - Without the override the constructor probes interpreters with a subprocess - each, which is why the override exists at all. - """ - wc._worker_python_cache = None - os.environ["IDATUI_WORKER_PYTHON"] = sys.executable - try: - got = wc._find_worker_python() - check("$IDATUI_WORKER_PYTHON is honoured", got == sys.executable, got) - finally: - os.environ.pop("IDATUI_WORKER_PYTHON", None) - wc._worker_python_cache = None - missing = "/nonexistent/python-that-is-not-there" - os.environ["IDATUI_WORKER_PYTHON"] = missing - try: - got = wc._find_worker_python() - check("an override that doesn't exist falls back instead of crashing", - got != missing and os.path.exists(got), got) - finally: - os.environ.pop("IDATUI_WORKER_PYTHON", None) - wc._worker_python_cache = None - - -def t_session_shims(tmp): - """The single-DB worker still has to answer the session questions the app - inherited from the old multi-session HTTP client.""" - c = client(tmp) - try: - c.connect(timeout=30) - sess = c.list_sessions() - check("list_sessions describes the one open database", - len(sess) == 1 and sess[0].filename == "target.bin" - and sess[0].is_active, str(sess)) - c.set_db("chosen") - check("set_db/resolve_db round-trip", c.resolve_db() == "chosen") - ka = c.keepalive() - ka.start() - ka.stop() - check("keepalive is a no-op the app can still drive", - ka.beats == 0 and ka.failures == 0) - h = c.health() - check("health answers even though the fake has no server_health tool", - isinstance(h, dict) and h, str(h)) - finally: - c.close(grace=5) - - -def t_log_tail(tmp): - """_log_tail picks the real error out of IDA's noise.""" - c = client(tmp) - with open(c._log_path, "w", encoding="utf-8") as fh: - fh.write("Thank you for using IDA\n" - "Licensed to: somebody\n" - "[MCP] registering tools\n" - "WORKER-FATAL: Failed to open database\n") - check("_log_tail surfaces WORKER-FATAL over the banner", - c._log_tail() == "Failed to open database", repr(c._log_tail())) - with open(c._log_path, "w", encoding="utf-8") as fh: - fh.write("Thank you for using IDA\nsomething odd happened\n") - tail = c._log_tail() - check("without a FATAL line it skips the banner and keeps the rest", - "odd happened" in tail and "Thank you" not in tail, repr(tail)) - os.unlink(c._log_path) - check("a missing log is reported, not raised", - "no worker log" in c._log_tail(), repr(c._log_tail())) - - -def main() -> int: - import tempfile - tests = [t_roundtrip, t_envelope, t_tool_error, t_dropped_socket, - t_startup_crash, t_connect_timeout, t_progress, t_serialized, - t_clean_shutdown, t_call_after_close, t_worker_python_override, - t_session_shims, t_log_tail] - with tempfile.TemporaryDirectory(prefix="idatui-wc-") as tmp: - _install_fake(tmp) - for fn in tests: - print(f"\n{fn.__name__}") - try: - fn(tmp) - except Exception as e: # noqa: BLE001 -- isolate one test's crash - import traceback - check(f"{fn.__name__} did not crash", False, - f"{type(e).__name__}: {e}") - traceback.print_exc() - print(f"\n{PASS} passed, {FAIL} failed") - return 1 if FAIL else 0 - - -if __name__ == "__main__": - raise SystemExit(main()) @@ -12,26 +12,68 @@ wheels = [ ] [[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" }, +] |