diff options
| author | user <user@clank> | 2026-08-07 15:14:30 +0200 |
|---|---|---|
| committer | user <user@clank> | 2026-08-07 15:14:30 +0200 |
| commit | 72fce7da1a1fd6527e389ffeb0f951157523589a (patch) | |
| tree | 3f08a83f0d99f3d3fc7069b7be00d235bab82817 | |
| parent | Stop tracking 157MB of core dumps, and ignore them (diff) | |
| parent | docs: upstream findings for the ida-codemode maintainers (diff) | |
| download | ida-tui-72fce7da1a1fd6527e389ffeb0f951157523589a.tar.gz ida-tui-72fce7da1a1fd6527e389ffeb0f951157523589a.tar.xz ida-tui-72fce7da1a1fd6527e389ffeb0f951157523589a.zip | |
Merge the IDA Code Mode port
Replaces the private idalib worker (idatui/worker.py + worker_client.py, with
server/patch_server.py injecting tools into ida-pro-mcp) with an ordinary
client of ida_codemode.client.DatabaseHandle. A database open by an IDA GUI is
reused; otherwise Code Mode starts or shares a managed idalib worker. The TUI
no longer owns an IDA process, and closing it releases only its lease.
Based on Duncan Ogilvie's port, rebased onto ~150 commits of local work it
predated. The rebase itself was mechanical; landing it was not. Nine defects
had to be fixed before the feature set was whole again, none of which the
patch's own tests could catch:
- DatabaseHandle.open() takes image_base, not loading_address: every
connect() would have raised TypeError on the first call
- five operations our tree had grown were simply missing (flowchart, so the
graph view was dead; op_format/pc_nums/pc_num_format, so 'o'/'O' were;
survey_binary)
- set_comments wrote only the disassembly comment, so comments never
appeared in pseudocode
- xref_query returned rows in raw IDA order, and 'follow the call' silently
followed the fall-through instead
- rename accepted one edit per category, so bulk symbol import was dead
- decompile ran decomp_map's full per-column ctree sweep to fill in a
per-line address anchor
- heads shipped without operand extents or the digest protocol
- the package became unimportable without ida_codemode installed, which
killed the offline test suites
Verified against the pre-codemode tag rather than against assumptions: the
full suite is 788 passed / 0 failed, and the pilot's 301 checks match the old
backend exactly. Performance is within 2x on the listing hot path and faster
on decompile, disasm and connect, after fixing two runtime costs that are
documented for upstream in docs/CODEMODE_UPSTREAM.md.
Test runtime came down from ~9m20s to 115s along the way -- not by removing
checks, but by removing four kinds of waiting-on-a-guess that were also
hiding real failures.
40 files changed, 9215 insertions, 2812 deletions
diff --git a/.fastfeedback/SPEED.md b/.fastfeedback/SPEED.md new file mode 100644 index 0000000..e402ea1 --- /dev/null +++ b/.fastfeedback/SPEED.md @@ -0,0 +1,241 @@ +# SPEED — ida-tui-maybe + +_Read this BEFORE running anything slow._ This repo does **not** use pytest for its +suites; every suite is a standalone script with its own runner and its own tally line +(`N passed, M failed` / `N checks, M failed`). `speedscan`'s generic pytest template is +wrong here — these are the real commands. + +--- + +## The two interpreters + +| python | has | use for | +|---|---|---| +| `python3` (system) | stdlib only | every **pure** suite. No IDA, no Code Mode. | +| `~/ida-venv/bin/python` | textual + idapro + ida_codemode | anything marked `ida`, and the TUI itself | + +`tests/run.py --list` prints `pure` / `ida` per file. A `pure` file must keep running +under system `python3` — that is a house rule, and it is why `idatui/codemode_client.py` +defers its `ida_codemode` import instead of doing it at module top. + +--- + +## Escalation ladder for THIS repo + +**Cheap gates (seconds, run these first)** +```bash +python3 -c "import ast;ast.parse(open('idatui/codemode_client.py').read())" # ~0.05s syntax +python3 tests/run.py --fast # every pure suite, 302 checks, 0.6s +python3 tests/run.py graph -x # one suite by substring, fail-fast +``` + +**Narrow test — DEFAULT RUNG.** The pilot suite takes `--only <substr,...>`: +```bash +~/ida-venv/bin/python tests/test_scenarios.py /tmp/scratch_bin --only rename # ~10s +~/ida-venv/bin/python tests/test_scenarios.py --list # scenario names +``` +`--only` matches substrings, so `--only listing` runs listing_view + listing_name_addr + +listing_make_string + listing_struct_expand. Name the scenario exactly to stay narrow. + +**Full gate — BACKGROUND ONLY, once, at the end** (see next section for the launch form): +```bash +~/ida-venv/bin/python tests/run.py # every suite, serial on purpose +``` + +--- + +## Coverage map — which suite catches what + +The pilot suite is big and slow and it is **not** the whole story. Two backend bugs +this port shipped were invisible to `test_scenarios.py` and only fell out of the +smaller RPC/UI suites: + +| suite | ~time | catches | +|---|---|---| +| `test_scenarios.py` | 166s | the TUI end to end: nav, views, listing, graph, opfmt, rename-one | +| `test_rawimage_rpc.py` | ~60s | **`rename_many` (batch rename), define/blob workflow, opfmt over RPC** | +| `test_blob_ui.py` | ~40s | raw-image/blob UI | +| `test_project_ui.py` | ~8s | multi-binary project UI | +| `test_trace_rpc.py` / `test_trace_ui.py` | ~30s | Tenet trace integration | + +**If you touch a backend operation, grep for its callers and run the suite that owns +them — not just the pilot.** `rename` is exercised one-edit-at-a-time by the pilot and +as a *batch* only by `test_rawimage_rpc.py`, which is exactly where the list-vs-dict +bug hid. + +--- + +## Backgrounding in THIS harness — the gotcha that cost a turn + +`bgrun run ... ` **blocks the agent tool call** even though bgrun itself returns in ms. +`bgrun` line 64 is `( "$@" >"$LOG" 2>&1; echo $? >"$EXIT" ) &` — the command's own stdio +goes to the log, but the *backgrounded subshell* still holds the inherited stdout fd, and +the harness's bash waits for EOF on that pipe, not for the parent to exit. + +**Always launch detached with all three fds redirected:** +```bash +SK="$HOME/.the assistant/agent/skills/fast-feedback/scripts" +timeout 20 setsid "$SK/bgrun" run fullgate -- ~/ida-venv/bin/python tests/run.py \ + </dev/null >/tmp/bgrun.out 2>&1 ; cat /tmp/bgrun.out # returns in ~7ms +timeout 20 "$SK/bgrun" check fullgate # non-blocking, safe +``` +`bgrun check` / `bgrun list` do **not** need this treatment; only `run` does. + +**Never `pkill -f <pattern>` here.** The agent's own shell command line contains the +pattern you are matching, so `pkill -f "tests/run.py"` kills the very shell issuing it +and the rest of the command silently never runs. Kill by **pid** (`bgrun list`, +`pgrep -af` first), then gate on it with `waitfor pid-gone <pid>`. + +**Killing a pilot leaves a Code Mode worker behind** for its `--lease-grace` (20s); it +self-exits, and it is holding only that run's temp `.i64`, so it does not block a new +run. Do not `kill -9` it — a hard-killed IDA wedges its database. + +--- + +## The four ways a test here wastes minutes + +Every slow suite in this repo was slow for one of these, not for doing real work. +Check them before optimising anything else. + +1. **A wait on a signal that can no longer happen.** `wait(lambda: lst.model is not + old, ..., 60)` was the idiom for "the edit landed". The perf work made edits KEEP + the listing's walk and re-render in place, so the model object is never replaced: + every one of those waits sat out its full timeout and the check afterwards passed + **vacuously**. This alone was 30s in blob_ui and 4x60s in thumb_ui. + → Gate on what the check is about (the row is code, the status says Thumb, the + function is in the index), via `settle(app, pred)` from `idatui._sync`. +2. **A wait on a signal that is set too early.** `app._t` is assigned when the key is + handled; the navigation it starts runs in a worker. Waiting on it and then reading + the cursor is a race master won (single-digit ms backend) and Code Mode loses. + → `settle()` gates on quiescence *and* the predicate, which is what you want. +3. **Regenerated fixtures.** `os.urandom` into a fresh `TemporaryDirectory` means new + bytes at a new path every run, so the pristine-database cache can never apply and + full auto-analysis is paid forever. → `_fixtures.synthetic(name, build)` writes + deterministic bytes to a stable path; `staged()` then caches the analysis. + Determinism is also correctness: "this blob has no functions" must not depend on luck. +4. **Deleting a database and reopening the same path.** Safe when the TUI owned a + private worker; under Code Mode the previous phase's worker still holds the lease + for its grace period, so the delete races a live owner and the reopen yields no + listing. → Give each phase its own temp copy (`fresh_copy`). Never `os.remove` an + `.i64` a suite is about to reopen. + +**`settle(app, pred, timeout=...)`** (`idatui/_sync.py`) is the one true gate: it drains +the message pump, waits for workers, and returns the moment `pred` holds. It is what the +live RPC layer uses, so tests and driver agree on what "done" means. A bare +`settle(app)` (no pred) means "the app finished reacting" — the right gate when the edit +may legitimately do nothing (e.g. carving random bytes). + +## Measured timings (2026-08-07, this box, warm) + +| command | time | notes | +|---|---|---| +| `python3 tests/run.py --fast` | **0.6s** | 302 checks, all pure suites | +| `python3 tests/test_graph.py` | ~1s | 86 checks, pure layout engine | +| `python3 tests/test_codemode_client.py` | ~0.1s | 14 checks, fakes the DatabaseHandle | +| `~/ida-venv/bin/python tests/test_rawimage_rpc.py` | ~60s | 21 checks, owns batch rename | +| pilot boot (first scenario) | ~10-25s | opens the DB; seeded from `.pristine.i64` | +| pilot `--only graph` | ~10s | 50 checks | +| **pilot full (`test_scenarios.py`)** | **62.2s** | 56 scenarios, 301 checks (was 166s) | +| `test_trace_ui.py` | 19.4s | 39 checks | +| `test_rawimage_rpc.py` | 7.5s | 21 checks, owns batch rename | +| `test_project_ui.py` | 8.1s | 30 checks | +| `test_thumb_ui.py` | **8.5s** | 20 checks — was 313s AND crashing | +| `test_blob_ui.py` | **4.1s** | 30 checks — was 39.8s | +| `tests/run.py` (everything) | **115s** | 788 checks. was ~9m20s. BACKGROUND ONLY | + +IDA-suite total went from ~9m21s to ~2m16s (4.1x) by fixing the four wastes above — +not by removing a single check. + +### Backend performance (worker vs Code Mode), after the two transport fixes + +| op | worker | codemode | | +|---|---|---|---| +| `heads` 200 rows | 2.65ms | 5.85ms | 2.2x | +| `heads` expect-hit | 2.16ms | 4.61ms | 2.1x | +| `disasm` 200 | 9.03ms | 5.25ms | **0.6x** | +| `decompile` cold | 162.6ms | 30.7ms | **0.2x** | +| `decomp_map` | 45.8ms | 47.9ms | 1.0x | +| empty round trip | ~0.07ms | **2.0ms** | the floor | + +Two things dominated and are fixed (see `idatui/codemode_client.py:_script`): +`to_jsonable` walking every returned object (snippets now return one +pre-serialised JSON string), and `sys.settrace` — the runtime installs a trace +that returns itself, i.e. LINE tracing in every frame, which made +`ida_bytes.get_flags` 52x slower than native. The snippet detaches it and +restores it in a finally; `IDATUI_CODEMODE_TRACE=1` keeps the stock behaviour. + +**What is left is the 2ms round-trip floor, and it is NOT ours.** Measured against +the same worker, same connection: + +| | cost | whose | +|---|---|---| +| `GET /health` (no execute_sync) | 0.165ms | HTTP transport | +| `execute_python("result = 1")` | 2.025ms | + `ida_kernwin.execute_sync` | + +So HTTP is 7% of the floor and marshalling an operation onto IDA's main thread is +92%. The worker runs IDA's own `kernwin.serve()` dispatch loop, so that latency is +inside IDA, not something Code Mode exposes a knob for. + +**Do not re-chase this by batching operations.** The call volume is already +minimal, measured on `targets/bash`: + +| flow | wall | calls | +|---|---|---| +| open a listing | 3.2ms | 1 | +| scroll 2000 rows | 87.9ms | 8 | +| rename + re-render those rows | 36.2ms | 4 | +| graph of a **1060-block** function | 171.0ms | 4 | +| decompile a 17785-byte function | 10806ms | 1 | + +Block-coalescing, the page digest and the prefetch caches already collapse the +bursts, so a batch endpoint would save single-digit milliseconds on flows that +cost hundreds. And the big number is pure Hex-Rays: that same decompile is +10723ms on the worker backend (0.8% apart) — there is no transport in it at all. + + +_Superseded note:_ **What is left is the 2ms round-trip floor.** A trivial op (`data_type`, +`force_recompile`, one xref query) is ~2.5ms wall clock and looks like 40x +against an in-process worker. That is fixed by making FEWER calls, not faster +ones — which is what the digest/`expect` path does for the listing. + +Benchmark harness: `/tmp/cmport/bench.py` + `compare.py` (backend-agnostic; it +picks whichever client the checked-out tree has, so it runs on master too). + +--- + +## Known-flaky + +- `test_scenarios.py --only listing_view` (**codemode backend only**). Passed 2 runs, + failed the 3rd, on an unchanged tree and a fresh scratch DB. Symptom: `listing shows + data heads (not just code)` reports `{'label','sep','funchdr','code'}` and + `total=1500`, i.e. the model is read before its pages have materialised. Adding *any* + delay (even a single `stderr.write`) makes it pass, which is the signature of a settle + race, not a logic bug. The scenario waits on `c.lst.total > 0`, which becomes true + before the rows exist. **Do not chase this as a regression without reproducing it + 3+ times.** Real fix would be a wait on materialised rows, not on `total`. + +--- + +## Slow traps + +- **Never run the pilot suite against `targets/echo` directly.** It edits the database. + The suite already stages a scratch copy (`tests/_fixtures.py: staged()`), seeded from + `<binary>.pristine.i64` which nothing writes back to. Pass a **copy** in `/tmp` when + invoking by hand, or you are testing your own history. +- **`tests/run.py` runs suites serially on purpose.** Running them 4-up took the suite + from 153s to 296s and got three killed mid-analysis (idalib contends hard). Do not + "optimise" it with parallelism. +- **Cold analysis dominates a first run.** `.pristine.i64` turns ~30s of auto-analysis + into a file copy; if it is missing or older than the binary it is rebuilt. A suite that + suddenly takes minutes longer is usually rebuilding that cache. +- **Under load, idalib gets SIGKILLed mid-analysis** and it looks like a hang or empty + output. Check `uptime` before believing a failure. +- **`experiments/*.py` need `PYTHONPATH=$PWD`** under `~/ida-venv/bin/python` (they are + scripts, not a package entry point). + +## Comparing against master (backend A/B) + +The port replaces the whole backend, so "is this a regression?" means running the same +scenario on both. `git stash -u; git checkout master; <run>; git checkout -; git stash pop` +works, but **use a different scratch binary per branch** (`/tmp/x_master`, `/tmp/x_port`) +so the two runs never share a `.pristine.i64` or a live Code Mode instance. @@ -19,3 +19,5 @@ bin/ # core dumps (idalib/SWIG can segfault under differential probes) core core.* +.fastfeedback/ +tests/.synthetic/ @@ -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/CODEMODE_UPSTREAM.md b/docs/CODEMODE_UPSTREAM.md new file mode 100644 index 0000000..f903599 --- /dev/null +++ b/docs/CODEMODE_UPSTREAM.md @@ -0,0 +1,269 @@ +# Findings from porting a real client to IDA Code Mode + +Notes for the `ida-codemode` maintainers, gathered while porting **ida-tui** (a +Textual TUI frontend for IDA) from a private idalib worker to +`ida_codemode.client.DatabaseHandle`. + +Everything below is measured, not inferred. Where we worked around something, the +workaround is named so you can judge whether the library should make it +unnecessary. + +**Environment:** ida-codemode 0.3.1, IDA 9.4 (idalib), Linux, single managed +worker backend, quiet box. Target for timings: `targets/echo` unless stated. + +**What the client does**, for scale: it renders a continuous disassembly listing, +pseudocode, a CFG graph view and a hex view, paging over the database as the user +scrolls. It is latency-sensitive in a way an agent-driven MCP client is not — a +keypress must repaint. It issues ~1–8 operations per user action. + +--- + +## 1. `timeout_trace` enables line tracing in every frame — 52x on IDA calls + +**Highest-impact item by a wide margin.** + +`runtime.py` wraps every `execute_python` in `sys.settrace(timeout_trace)` to +enforce the deadline. `timeout_trace` ends with `return timeout_trace`, and +returning a trace function from a `'call'` event asks CPython to trace **every +line of that frame**. So every line of every function the snippet touches pays a +Python-level callback, and the specialising interpreter is disabled throughout. + +Measured inside the worker, same process, same database: + +| | traced (stock) | untraced | native idalib | +|---|---|---|---| +| `ida_bytes.get_flags(ea)` | 5.49 µs | 0.106 µs | 0.119 µs | +| our 200-row listing page | 20.2 ms | 2.0 ms | — | + +Untraced matches a plain idalib process, so the trace hook accounts for +essentially all of it. For us this was the single largest cost in the port — +larger than HTTP, serialisation and IDA itself combined. + +Reproduce inside any `execute_python`: + +```python +import sys, time, ida_bytes +def bench(): + t = time.perf_counter() + for _ in range(20000): ida_bytes.get_flags(0x1000) + return (time.perf_counter() - t) / 20000 * 1e6 +traced = bench() +old = sys.gettrace(); sys.settrace(None) +try: untraced = bench() +finally: sys.settrace(old) +result = {"traced_us": traced, "untraced_us": untraced} +``` + +**Suggested fixes, cheapest first** + +1. `return None` from `timeout_trace` instead of itself. You keep `'call'`-event + deadline checks — which is enough to interrupt anything that calls a function + — and drop per-line tracing entirely. +2. On 3.12+, use `sys.monitoring` with only the events you need; it is designed + for exactly this and is far cheaper than `settrace`. +3. Or drop the trace and rely on the `threading.Timer` → + `ida_kernwin.set_cancelled()` path you already have, accepting that a + pure-Python loop with no calls in it cannot be interrupted. + +**Our workaround** (we would rather not ship it): the snippet detaches the trace +and restores it in a `finally`. That gives up deadline enforcement for +pure-Python loops inside our own code; your native cancel timer is unaffected and +still fires. Every client that does real work per call will eventually find this +and do the same, which is an argument for fixing it in the runtime. + +--- + +## 2. `to_jsonable` dominates any large result + +`execute_python` runs `to_jsonable()` over whatever the snippet returns. Our +answers are already JSON-safe and they are big — a 200-row listing page is +roughly 10k small objects. + +| | cost | +|---|---| +| `to_jsonable(page)` | 66.2 ms | +| `json.dumps(page, separators=(",",":"))` — same data | 0.58 ms | +| serialised size | 34.9 KB | + +That is 114x, and it was 72% of the page's total cost before we changed it. + +**Suggested fixes** + +- Fast-path values that are already JSON-safe (a cheap recursive type check that + bails to the original object beats rebuilding it), or +- let a snippet opt out by returning an already-serialised payload — a documented + envelope such as `{"__json__": "<...>"}`, or simply passing `str`/`bytes` + through untouched. + +**Our workaround:** snippets `json.dumps` inside the database process and return +one string, which the client parses. `to_jsonable` then walks a single scalar. +Cost went 66.2 ms → ~0.6 ms. It works, but every client with a large result set +has to discover and re-implement it. + +--- + +## 3. The per-operation floor is `execute_sync`, not HTTP + +Same worker, same connection, 200 iterations: + +| | cost | +|---|---| +| `GET /health` (no `execute_sync`) | **0.165 ms** | +| `execute_python("result = 1")` | **2.025 ms** | + +HTTP framing is ~7% of the floor; marshalling the operation onto IDA's main +thread is the other ~93%. The worker runs IDA's own `kernwin.serve()`, so this is +plausibly IDA's dispatch latency rather than anything you control — but it is +worth **documenting**, because it sets a hard 2 ms per-operation budget that +shapes how a client must be designed. + +It did not hurt us (our call volume is 1–8 per user action; 4 calls to build a +1060-block graph), but a client that makes one call per row or per symbol will be +20–100x slower than an in-process one and the authors will not know why. + +**Suggested fixes:** document the floor; and consider a batch endpoint — accept +`[{op, args}, ...]` and dispatch them within a single `execute_sync` — which +would let chatty clients amortise it without redesigning around it. + +--- + +## 4. Loader switches on an existing database are a FATAL, not an error + +Opening a target that already has an `.i64`, while passing spawn-only options, +kills the worker: + +``` +FATAL ERROR: @0:636[] +Switch '-b400' can be used only when loading a new file +``` + +The client sees only: + +``` +IDAConnectionError: idalib worker launcher <pid> exited with status 1 +``` + +This is easy to hit and hard to diagnose: it is the natural second run of +anything that opens a raw blob (`processor=`/`image_base=`/`file_type=` are +recorded in the database the first run produced). Our test suite hit it as a +crash five minutes into a run. + +**Suggested fixes** + +- In `DatabaseHandle.open()`, when the resolved IDB already exists and + `new_database` is not set, either ignore the spawn-only options or raise a + typed error naming them — before handing them to IDA. +- Propagate the worker's fatal text into the client exception. The message + already exists on the worker's stderr; losing it turns a one-line fix into a + bisect. + +**Our workaround:** the client checks whether the expected IDB exists and drops +`processor`/`image_base`/`file_type` when it does. + +--- + +## 5. Deleting or replacing an IDB under a live lease fails silently + +A suite that did "delete the `.i64`, reopen the same path" (safe when it owned a +private worker) now races the previous worker's lease grace. The reopen produced +a handle that never became usable, with no error — just a database with no +listing, and every wait timing out. + +**Suggested fixes** + +- Detect that the IDB backing a registered instance has been removed or replaced + and fail loudly (the registry already holds `idb_key`). +- Expose a **public** "wait until this database is released" primitive. We needed + one and ended up reaching into `registry.REGISTRY_DIR` and `FileLock` to build + it, which is not an API we should be depending on. +- Document the lease-grace window as part of the lifecycle contract. + +--- + +## 6. No close-without-save, and no rollback + +A managed worker saves when its final lease closes. A GUI handle leaves GUI state +as-is. Neither gives a client a way to say "discard what I did". + +ida-tui had a "discard & quit" that we could not port; it is now "leave as-is & +quit", and we cannot honestly promise the user their edits are not persisted. + +**Suggested fixes:** a close policy on a lease the client created +(`close(save=False)`), or a transaction/rollback API, or a documented +disposable-copy pattern that clients can follow. + +--- + +## 7. No change notification for shared databases + +The lease reports liveness, not mutations. If a GUI user or another Code Mode +client renames or retypes while we are attached, our materialised caches (name +generation, decompilation, listing pages) are silently stale. Our own edits +invalidate correctly; someone else's cannot. + +**Suggested fix — cheap and sufficient:** a monotonic database revision counter, +bumped on any mutating operation and exposed on `/health` (and ideally on the +lease event stream). Clients can then invalidate by comparing one integer. A full +change feed would be better but is much more work; the counter alone would make +shared editing safe for every caching client. + +--- + +## 8. Package exports and API surface stability + +`ida_codemode/__init__.py` exports nothing, so a library consumer must import +from submodules: + +```python +from ida_codemode.client import DatabaseHandle, ClientError, RemoteError, InstanceDisconnectedError +from ida_codemode.registry import REGISTRY_DIR, FileLock, RegistryEntry, canonical_path, idb_key, scan_instances +from ida_codemode.resolver import IdbBusy, expected_idb_path +``` + +Some of those are clearly internals (`FileLock`, `REGISTRY_DIR`) that we only +touch because no public equivalent exists (see §5). + +**Suggested fix:** export `DatabaseHandle` and the public exception types from the +package root, and mark the intended-public registry helpers explicitly. It also +makes "what is API and what is internal" answerable, which right now it is not. + +--- + +## 9. A testing note: `DatabaseHandle.open()`'s 30 keyword-only options + +The port we started from called `open(..., loading_address=...)`. The real +parameter is `image_base`. Every `connect()` would have raised `TypeError` on the +first call, and its contract tests passed anyway, because a hand-written fake +handle accepts `**kwargs`. + +Not a library bug — but with 30 keyword-only options it is a very easy mistake, +and it is invisible to exactly the offline tests people write. + +**Suggested fix:** ship `py.typed` and/or a `Protocol` for the handle, so a fake +can be checked against the real signature and a typo is caught statically. (We +added a test asserting our kwargs are a subset of +`inspect.signature(DatabaseHandle.open).parameters`, which is a poor substitute.) + +--- + +## Priority, from a client author's view + +| # | item | impact | fixable by you? | +|---|---|---|---| +| 1 | `timeout_trace` line tracing | 52x on IDA calls, 10x on real operations | yes, one line | +| 2 | `to_jsonable` on large results | 114x on serialisation | yes | +| 7 | no change/revision counter | correctness for shared editing | yes, cheap | +| 4 | loader switches fatal on reopen | crashes, hard to diagnose | yes | +| 5 | replaced/deleted IDB under lease | silent hang | yes | +| 6 | no close-without-save | a feature we had to drop | design question | +| 8 | package exports | forces internal imports | yes, trivial | +| 3 | 2 ms `execute_sync` floor | shapes client design | document; maybe batch | +| 9 | typed handle for fakes | catches a whole bug class | yes | + +Items 1 and 2 together were the difference between "the port is 35x slower than +the private worker it replaced" and "the port is within 2x, and faster on several +operations". Both are in the runtime, not in client code — which is why they are +worth fixing centrally rather than leaving each client to rediscover. + +Happy to supply the benchmark harness (it is backend-agnostic and runs against +both our old worker and Code Mode), or to test a patch. 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..88d7b31 --- /dev/null +++ b/idatui/codemode_client.py @@ -0,0 +1,1362 @@ +"""Client adapter from ida-tui's domain operations to IDA Code Mode. + +``DatabaseHandle`` is the lifecycle boundary: it discovers an already-registered +GUI database, reuses a shared managed idalib worker, or starts one when needed. +The TUI never owns or terminates an IDA process. Closing this client releases +only its lease. + +The Code Mode transport intentionally exposes one broad operation, +``execute_python``. ``CodeModeClient.invoke`` turns the small, address-centric +operations needed by the paging layer into self-contained snippets. The +snippets prefer the public ``ida-domain`` ``db`` object. A handful of features +that ida-domain does not currently expose (IDA-coloured listing rows, creating +instructions, ARM T-state, and detailed Hex-Rays line maps/failures) use the +IDAPython modules that Code Mode deliberately makes importable. +""" +from __future__ import annotations + +import hashlib +import json +import os +import shlex +import threading +import time +from pathlib import Path +from textwrap import dedent, indent +from typing import Any + +from .errors import IDAConnectionError, IDATimeoutError, IDAToolError, Session + +# ida_codemode is imported EAGERLY-IF-PRESENT but never at hard import cost. +# +# The paging/graph/trace layers and their offline test suites must keep importing +# `idatui` on a machine with no IDA and no Code Mode installed -- that is the +# house rule the stdlib-only worker client used to satisfy for free, and +# `tests/run.py --fast` (257 checks, any python3) depends on it. A hard top-level +# import here makes the whole package unimportable, so the failure is deferred to +# the first operation that genuinely needs the library. +_CODEMODE_ERROR: Exception | None = None +try: + from ida_codemode.client import ( + ClientError, + DatabaseHandle, + InstanceDisconnectedError, + RemoteError, + ) + from ida_codemode.registry import ( + REGISTRY_DIR, + FileLock, + RegistryEntry, + canonical_path, + idb_key, + scan_instances, + ) + from ida_codemode.resolver import IdbBusy, expected_idb_path +except ImportError as _exc: # library absent: usable only for offline layers + _CODEMODE_ERROR = _exc + # Bound to None rather than left undefined so the names stay patchable: the + # offline contract tests inject a fake DatabaseHandle here. + ClientError = InstanceDisconnectedError = RemoteError = None # type: ignore[assignment,misc] + DatabaseHandle = RegistryEntry = FileLock = None # type: ignore[assignment,misc] + REGISTRY_DIR = canonical_path = idb_key = scan_instances = None # type: ignore[assignment] + IdbBusy = expected_idb_path = None # type: ignore[assignment] + + +def _require_codemode() -> None: + """Raise an actionable error when the Code Mode library is missing. + + Gated on the binding, not on the original import result, so a test that + injects a fake ``DatabaseHandle`` exercises the real adapter logic. + """ + if DatabaseHandle is None: + raise IDAConnectionError( + "ida-codemode-mcp is not installed in this environment " + f"({_CODEMODE_ERROR}). Install it (e.g. `uv sync`, or " + "`pip install -e ../ida-codemode-mcp`) so ida-tui can lease a " + "database.") from _CODEMODE_ERROR + + +def database_owner(idb_path: str, staged_path: str | None = None): + """The registry entry that owns ``idb_path``/``staged_path``, else None. + + Returns None when the Code Mode library is absent: with no library there is + no client in this environment that could be holding the database, and the + IDA-free layers (project staging) must keep working. Registry errors that + happen WITH the library installed still propagate -- those mean "we could + not determine ownership", which is not the same as "nobody owns it". + """ + if DatabaseHandle is None: + return None + expected_key = idb_key(idb_path) + staged = canonical_path(staged_path) if staged_path else None + for item in scan_instances(timeout=0.5): + entry = item.entry + if entry.idb_key == expected_key: + return entry + if staged and entry.exe_path and canonical_path(entry.exe_path) == staged: + return entry + return None + + +def registered_database(path: str, output_database: str | None = None) -> bool: + """Whether a live/lock-held Code Mode instance owns this target.""" + _require_codemode() + source = canonical_path(path) + expected = canonical_path(output_database) if output_database else expected_idb_path(source) + expected_key = idb_key(expected) + for instance in scan_instances(timeout=0.5): + entry = instance.entry + if entry.idb_key == expected_key: + return True + if not output_database and entry.backend == "gui" and entry.exe_path: + if canonical_path(entry.exe_path) == source: + return True + return False + + +class _NoopKeepAlive: + """Compatibility shim: the DatabaseHandle's SSE lease is the heartbeat.""" + + def __init__(self) -> None: + self.beats = self.failures = 0 + + def start(self) -> "_NoopKeepAlive": + return self + + def stop(self) -> None: + pass + + +def _parse_load_args(value: str) -> tuple[str | None, int | None, str | None]: + """Translate ida-tui's legacy first-open switches to Code Mode options. + + Code Mode has typed options for processor, natural loading address and file + type. It deliberately has no arbitrary command-line escape hatch; reject + switches we cannot represent instead of silently loading a blob wrongly. + """ + processor: str | None = None + loading_address: int | None = None + file_type: str | None = None + unsupported: list[str] = [] + try: + words = shlex.split(value or "", posix=os.name != "nt") + except ValueError as exc: + raise ValueError(f"invalid IDA load options: {exc}") from exc + for word in words: + if word.startswith("-p") and len(word) > 2: + processor = word[2:] + elif word.startswith("-b") and len(word) > 2: + try: + # IDA's -b is in 16-byte paragraphs. DatabaseHandle expects the + # natural address, which is the safer public API. + loading_address = int(word[2:], 16) << 4 + except ValueError as exc: + raise ValueError(f"invalid IDA loading address: {word!r}") from exc + elif word.startswith("-T") and len(word) > 2: + file_type = word[2:] + else: + unsupported.append(word) + if unsupported: + joined = " ".join(unsupported) + raise ValueError( + "ida-codemode cannot represent arbitrary IDA load options: " + f"{joined!r}; use processor/base/file type options instead" + ) + return processor, loading_address, file_type + + +#: Key of the pre-serialised payload envelope. See _script(). +_PACKED = "__idatui_json__" + +#: Serialise the answer INSIDE the database process and hand back one string. +#: +#: Code Mode runs to_jsonable() over whatever a snippet returns, walking the +#: whole structure to make it JSON-safe. Our answers are already JSON-safe, and +#: they are big: a 200-row listing page is ~10k small objects, which costs 66ms +#: to walk -- 72% of the page's total cost, and 114x what json.dumps of the very +#: same data costs (0.58ms). Returning a STRING makes that walk O(1); the client +#: parses it, which it was going to do at the transport layer anyway. +_PACK_EPILOGUE = ( + '\n{"' + _PACKED + '": json.dumps(result, separators=(",", ":"), default=str)}\n' +) + + +#: Keep Code Mode's per-line trace hook installed while our snippet runs. +#: Set IDATUI_CODEMODE_TRACE=1 to restore the stock behaviour. +_KEEP_TRACE = os.environ.get("IDATUI_CODEMODE_TRACE", "") not in ("", "0") + + +def _script(args: dict[str, Any], body: str) -> str: + """Bind JSON arguments without interpolating user text into Python code. + + Also runs the body with Code Mode's trace hook detached, which is worth an + order of magnitude. The runtime wraps every execute_python in + sys.settrace(timeout_trace), and that trace function RETURNS ITSELF, which + turns on line tracing in every frame it sees -- so every line of every + function we call pays a Python-level callback. Measured on this box: + ida_bytes.get_flags is 0.106us untraced (0.119us in a plain idalib process) + and 5.49us traced, 52x; a 200-row listing page is 2.0ms untraced and 20.2ms + traced. That single hook was the whole residual gap against the old worker. + + What this gives up: the deadline is no longer enforced for a pure-Python + loop inside our snippet. The runtime's OTHER cancellation path -- a + threading.Timer that calls ida_kernwin.set_cancelled() -- is independent of + the trace and still fires, so a long IDA operation is still interruptible; + and every operation here is bounded by its own count/limit argument. The + trace is restored in a finally, so a raising snippet cannot leak the change. + """ + encoded = json.dumps(args, ensure_ascii=False, separators=(",", ":")) + head = f"import json\na = json.loads({encoded!r})\n" + if _KEEP_TRACE: + return f"{head}{dedent(body).strip()}\n{_PACK_EPILOGUE}" + return ( + f"{head}" + "import sys\n" + "_idatui_trace = sys.gettrace()\n" + "sys.settrace(None)\n" + "try:\n" + f"{indent(dedent(body).strip(), ' ')}\n" + ' _idatui_packed = {"' + _PACKED + '": json.dumps(' + 'result, separators=(",", ":"), default=str)}\n' + "finally:\n" + " sys.settrace(_idatui_trace)\n" + "_idatui_packed\n" + ) + + +_OPERATIONS: dict[str, str] = { + "list_funcs": r''' +import fnmatch +queries = a.get("queries") or [{}] +q = queries[0] +offset, count = max(0, int(q.get("offset", 0))), max(1, int(q.get("count", 500))) +pattern = str(q.get("filter") or "").lower() +if pattern and not any(ch in pattern for ch in "*?["): pattern = "*" + pattern + "*" +rows = [] +for fn in db.functions.get_all(): + name = db.functions.get_name(fn) or f"sub_{int(fn.start_ea):X}" + if pattern and not fnmatch.fnmatchcase(name.lower(), pattern): continue + rows.append({"addr": hex(int(fn.start_ea)), "name": name, + "size": int(fn.end_ea) - int(fn.start_ea)}) +page = rows[offset:offset + count] +result = {"result": [{"data": page, "next_offset": offset + len(page), "total": len(rows)}]} +result +''', + "disasm": r''' +ea = int(str(a["addr"]), 16) +fn = db.functions.get_at(ea) +if fn is None: + result = {"instructions": [], "total_instructions": 0, "instruction_count": 0} +else: + instructions = list(db.functions.get_instructions(fn)) + limit = max(1, int(a.get("max_instructions", len(instructions) or 1))) + rows = [{"addr": hex(int(insn.ea)), "instruction": db.instructions.get_disassembly(insn)} + for insn in instructions[:limit]] + result = {"instructions": rows, "total_instructions": len(instructions), + "instruction_count": len(instructions)} +result +''', + "file_regions": r''' +import idaapi +rows = [] +for seg in db.segments.get_all(): + try: file_off = int(idaapi.get_fileregion_offset(seg.start_ea)) + except Exception: file_off = -1 + if file_off < 0 or file_off >= (1 << 48): file_off = -1 + rows.append({"start": hex(int(seg.start_ea)), "end": hex(int(seg.end_ea)), + "file_off": file_off, "name": db.segments.get_name(seg) or ""}) +result = {"regions": rows} +result +''', + "read_raw": r''' +import ida_bytes +ea, size = int(str(a["addr"]), 16), max(0, int(a["size"])) +raw = ida_bytes.get_bytes(ea, size) or b"" +raw = raw[:size] + b"\xff" * max(0, size - len(raw)) +data = bytearray(raw) +for index, value in enumerate(data): + if value == 0xFF and not ida_bytes.is_loaded(ea + index): data[index] = 0 +result = {"addr": a["addr"], "hex": bytes(data).hex(), "n": len(data)} +result +''', + "get_bytes": r''' +rows = [] +for region in a.get("regions", []): + ea, size = int(str(region["addr"]), 16), int(region["size"]) + raw = db.bytes.get_bytes_at(ea, size) or b"" + rows.append({"addr": region["addr"], "data": " ".join(f"{b:02x}" for b in raw)}) +result = {"result": rows} +result +''', + "search_structs": r''' +needle = str(a.get("filter") or "").lower() +rows = [] +for tif in db.types.get_all(): + name = tif.get_type_name() or "" + if not name or needle not in name.lower() or not tif.is_udt(): continue + members = list(db.types.get_udt_members(tif)) + rows.append({"name": name, "size": int(tif.get_size()), "is_union": bool(tif.is_union()), + "cardinality": len(members), "ordinal": int(tif.get_ordinal())}) +result = {"result": rows} +result +''', + "type_inspect": r''' +rows = [] +for query in a.get("queries", []): + name = str(query.get("name") or "") + tif = db.types.get_by_name(name) + if tif is None: + rows.append({"name": name, "error": "type not found"}); continue + members = [{"name": m.name, "type": m.type.dstr() or str(m.type), + "offset": int(m.offset), "size": int(m.size)} + for m in db.types.get_udt_members(tif)] if tif.is_udt() else [] + rows.append({"name": name, "size": int(tif.get_size()), "is_union": bool(tif.is_union()), + "members": members}) +result = {"result": rows} +result +''', + "declare_type": r''' +import ida_typeinf +decls = a.get("decls", "") +if isinstance(decls, str): decls = [decls] +rows = [] +for declaration in decls: + try: + errors = int(db.types.parse_declarations(ida_typeinf.get_idati(), declaration)) + rows.append({"ok": errors == 0, **({} if errors == 0 else {"error": f"{errors} parse error(s)"})}) + except Exception as exc: + rows.append({"ok": False, "error": str(exc)}) +result = {"result": rows} +result +''', + "del_type": r''' +import ida_typeinf +name = str(a["name"]) +ok = bool(ida_typeinf.del_named_type(ida_typeinf.get_idati(), name, ida_typeinf.NTF_TYPE)) +result = {"name": name, "deleted": ok, **({} if ok else {"error": f"Type {name!r} not found or could not be deleted"})} +result +''', + "func_types": r''' +import ida_typeinf +ea = int(str(a["addr"]), 16) +fn = db.functions.get_at(ea) +if fn is None: + result = {"addr": a["addr"], "error": "no function at address"} +else: + pseudo = db.pseudocode.decompile(fn) + name = db.functions.get_name(fn) or "" + tif = pseudo.get_func_type() + try: prototype = ida_typeinf.print_tinfo("", 0, 0, ida_typeinf.PRTYPE_1LINE, tif, name, "") if tif else "" + except Exception: prototype = tif.dstr() if tif else "" + lvars = [{"name": var.name, "type": var.type_info.dstr() if var.type_info else "", + "is_arg": bool(var.is_arg)} for var in pseudo.local_variables] + result = {"addr": hex(int(fn.start_ea)), "name": name, + "prototype": (prototype or "").strip(), "lvars": lvars} +result +''', + "set_lvar_type": r''' +import ida_typeinf +ea, variable, declaration = int(str(a["addr"]), 16), str(a["variable"]), str(a["type"]) +fn = db.functions.get_at(ea) +if fn is None: + result = {"error": "no function at address"} +else: + pseudo = db.pseudocode.decompile(fn) + var = pseudo.find_local_variable(variable) + if var is None: + result = {"error": f"local variable {variable!r} not found"} + else: + try: + tif = db.types.parse_one_declaration(ida_typeinf.get_idati(), declaration) + accepted = bool(var.set_type(tif)) + saved = bool(pseudo.save_local_variable_info(var, save_type=True)) if accepted else False + result = {"addr": hex(int(fn.start_ea)), "variable": variable, + "type": declaration, "ok": accepted and saved} + except Exception as exc: + result = {"error": f"bad type {declaration!r}: {exc}"} +result +''', + "set_type": r''' +from ida_domain.types import TypeApplyFlags +rows = [] +for edit in a.get("edits", []): + ea = int(str(edit["addr"]), 16) + declaration = str(edit.get("signature") or edit.get("type") or "") + try: + ok = bool(db.types.apply_declaration_at(ea, declaration, TypeApplyFlags.DEFINITE)) + rows.append({"addr": hex(ea), "ok": ok, **({} if ok else {"error": "IDA rejected the type"})}) + except Exception as exc: + rows.append({"addr": hex(ea), "ok": False, "error": str(exc)}) +result = {"result": rows} +result +''', + "data_type": r''' +ea = int(str(a["addr"]), 16) +try: + tif = db.types.get_at(ea) + fn = db.functions.get_at(ea) + result = {"addr": hex(ea), "name": db.names.get_at(ea) or "", + "type": tif.dstr() if tif else "", "size": int(db.heads.size(ea)) if db.heads.is_head(ea) else 0, + "is_func": bool(fn)} +except Exception as exc: + result = {"addr": hex(ea), "error": str(exc)} +result +''', + "force_recompile": r''' +import ida_hexrays +rows = [] +for item in a.get("items", []): + ea = int(str(item["addr"]), 16) + ida_hexrays.mark_cfunc_dirty(ea, False) + rows.append({"addr": hex(ea), "ok": True}) +result = {"result": rows} +result +''', + "undefine": r''' +import ida_bytes +rows = [] +for item in a.get("items", []): + ea = int(str(item["addr"]), 16) + size = max(1, int(item.get("size") or ida_bytes.get_item_size(ea) or 1)) + ok = bool(ida_bytes.del_items(ea, ida_bytes.DELIT_SIMPLE, size)) + rows.append({"addr": hex(ea), "ok": ok, **({} if ok else {"error": "delete items failed"})}) +result = {"result": rows} +result +''', + "define_code": r''' +import ida_ua +rows = [] +for item in a.get("items", []): + ea = int(str(item["addr"]), 16); size = int(ida_ua.create_insn(ea)) + rows.append({"addr": hex(ea), "ok": size > 0, "size": size, + **({} if size > 0 else {"error": "instruction did not decode"})}) +result = {"result": rows} +result +''', + "define_func": r''' +rows = [] +for item in a.get("items", []): + ea = int(str(item["addr"]), 16); ok = bool(db.functions.create(ea)) + rows.append({"addr": hex(ea), "ok": ok, **({} if ok else {"error": "IDA refused the function"})}) +result = {"result": rows} +result +''', + "make_data": r''' +import ida_bytes, ida_idaapi, ida_typeinf +from ida_domain.types import TypeApplyFlags +rows = [] +for item in a.get("items", []): + ea, declaration = int(str(item["addr"]), 16), str(item["type"]) + try: + tif = db.types.parse_one_declaration(ida_typeinf.get_idati(), declaration) + size = max(1, int(tif.get_size())) + saved_names = [(addr, name) for addr, name in db.names.get_all() + if ea <= int(addr) < ea + size] + ida_bytes.del_items(ea, ida_bytes.DELIT_EXPAND | ida_bytes.DELIT_DELNAMES, + max(size, int(ida_bytes.get_item_size(ea) or 1))) + created = bool(ida_bytes.create_data(ea, ida_bytes.FF_BYTE, size, ida_idaapi.BADADDR)) + ok = created and bool(db.types.apply_at(tif, ea, TypeApplyFlags.DEFINITE)) + for address, name in saved_names: + db.names.set_name(int(address), name) + if ok and item.get("name"): ok = bool(db.names.set_name(ea, str(item["name"]))) + rows.append({"addr": hex(ea), "ok": ok, "size": size, + **({} if ok else {"error": "IDA rejected the data type"})}) + except Exception as exc: + rows.append({"addr": hex(ea), "ok": False, "error": str(exc)}) +result = {"result": rows} +result +''', + "make_string": r''' +from ida_domain.strings import StringType +ea, length = int(str(a["addr"]), 16), max(0, int(a.get("length", 0))) +kind = {"c": StringType.C, "c16": StringType.C_16, "c32": StringType.C_32, + "pascal": StringType.PASCAL}.get(str(a.get("kind", "c")).lower(), StringType.C) +import ida_bytes +try: + ida_bytes.del_items(ea, ida_bytes.DELIT_SIMPLE, length if length > 0 else 1) +except Exception: + pass +try: + ok = bool(db.bytes.create_string_at(ea, length or None, kind)) + text = db.bytes.get_string_at(ea) or "" if ok else "" + result = {"addr": hex(ea), "ok": ok, "size": int(db.heads.size(ea)) if ok else 0, "text": text} +except Exception as exc: + result = {"addr": hex(ea), "ok": False, "error": str(exc)} +result +''', + "list_strings": r''' +from ida_domain.strings import StringListConfig +offset, count, min_len = max(0, int(a.get("offset", 0))), max(1, int(a.get("count", 2000))), max(1, int(a.get("min_len", 4))) +if offset == 0 or a.get("refresh"): + from ida_domain.strings import StringType + db.strings.rebuild(StringListConfig(string_types=list(StringType), min_len=min_len, + only_ascii_7bit=False)) +items = list(db.strings.get_all()) +page = items[offset:offset + count] +rows = [] +for item in page: + try: text = str(item) + except Exception: text = item.contents.decode("utf-8", "replace") if item.contents else "" + rows.append({"addr": hex(int(item.address)), "text": text, "len": int(item.length), "type": item.type.name}) +result = {"strings": rows, "total": len(items), "next_offset": offset + len(rows)} +result +''', + "list_linkage": r''' +imports = [{"addr": hex(int(item.address)), "name": item.name, "module": item.module_name} + for item in db.imports.get_all_imports() if item.name] +exports = [{"addr": hex(int(item.address)), "name": item.name, "ordinal": int(item.ordinal)} + for item in db.entries.get_all() if item.name] +result = {"imports": imports, "exports": exports, + "n_imports": len(imports), "n_exports": len(exports)} +result +''', + "lookup_funcs": r''' +rows = [] +for query in a.get("queries", []): + raw = str(query) + try: ea = int(raw, 16) + except ValueError: + fn = db.functions.get_by_name(raw); ea = int(fn.start_ea) if fn else None + else: fn = db.functions.get_at(ea) + if fn is None: + rows.append({"query": raw, "fn": None}) + else: + rows.append({"query": raw, "fn": {"addr": hex(int(fn.start_ea)), + "name": db.functions.get_name(fn) or f"sub_{int(fn.start_ea):X}", + "size": int(fn.end_ea) - int(fn.start_ea)}}) +result = {"result": rows} +result +''', + "resolve_names": r''' +import ida_idaapi, ida_name +rows = [] +for query in a.get("queries", []): + name = str(query).strip(); ea = ida_name.get_name_ea(ida_idaapi.BADADDR, name) + rows.append({"query": name, "ea": hex(int(ea)) if ea != ida_idaapi.BADADDR else None}) +result = {"result": rows} +result +''', + # Ours: the coarse code/data type plus a fine `kind` (call/jump/flow, + # read/write/offset/text/info) that the xref dialog draws its badges from. + # Deliberately NOT sorted -- the dialog lists xrefs in IDA's own order. + "xref_types": r''' +import idaapi, idautils, ida_bytes, ida_funcs, ida_xref +code_kind = {ida_xref.fl_CF: "call", ida_xref.fl_CN: "call", ida_xref.fl_JF: "jump", + ida_xref.fl_JN: "jump", ida_xref.fl_F: "flow"} +data_kind = {ida_xref.dr_O: "offset", ida_xref.dr_W: "write", ida_xref.dr_R: "read", + ida_xref.dr_T: "text", ida_xref.dr_I: "info"} +def _kind(xr): + return (code_kind if xr.iscode else data_kind).get(xr.type, "code" if xr.iscode else "data") +def _fn(ea): + f = ida_funcs.get_func(ea) + return {"addr": hex(int(f.start_ea)), "name": ida_funcs.get_func_name(f.start_ea) or ""} if f else None +queries = a.get("queries") or [] +all_results = [] +for query in queries: + query = query if isinstance(query, dict) else {"addr": query} + raw = str(query.get("addr", "")).strip() + direction = str(query.get("direction", "to") or "to").lower() + include_fn = bool(query.get("include_fn", True)) + dedup = bool(query.get("dedup", True)) + try: count = int(query.get("count", 2000) or 2000) + except (TypeError, ValueError): count = 2000 + try: target = int(raw, 16) + except ValueError: target = idaapi.get_name_ea(idaapi.BADADDR, raw) + rows = [] + if target is not None and target != idaapi.BADADDR and ida_bytes.is_mapped(target): + if direction in ("to", "both"): + for xr in idautils.XrefsTo(target, 0): + row = {"direction": "to", "addr": hex(int(xr.frm)), "from": hex(int(xr.frm)), + "to": hex(int(target)), "type": "code" if xr.iscode else "data", "kind": _kind(xr)} + if include_fn: row["fn"] = _fn(xr.frm) + rows.append(row) + if direction in ("from", "both"): + for xr in idautils.XrefsFrom(target, 0): + row = {"direction": "from", "addr": hex(int(xr.to)), "from": hex(int(target)), + "to": hex(int(xr.to)), "type": "code" if xr.iscode else "data", "kind": _kind(xr)} + if include_fn: row["fn"] = _fn(xr.to) + rows.append(row) + if dedup: + seen, deduped = set(), [] + for r in rows: + k = (r["direction"], r["from"], r["to"], r["kind"]) + if k in seen: continue + seen.add(k); deduped.append(r) + rows = deduped + rows = rows[:count] + all_results.append({"query": raw, "data": rows, "next_offset": None}) +result = {"result": all_results} +result +''', + # Mirrors the tool ida-tui was written against, ORDER INCLUDED. The rows are + # sorted by the far-end address and deduped by default, and the pseudocode + # follow's address fallback silently depends on it: at a call site the raw + # IDA order yields the ordinary-flow xref (the next instruction) first, so an + # unsorted result makes "follow the call" land on the following line instead. + "xref_query": r''' +import idaapi, idautils, ida_bytes, ida_funcs +def _fn(ea): + f = ida_funcs.get_func(ea) + return {"addr": hex(int(f.start_ea)), "name": ida_funcs.get_func_name(f.start_ea) or ""} if f else None +queries = a.get("queries") or [] +all_results = [] +for query in queries: + raw = str(query.get("addr", "")).strip() + direction = str(query.get("direction", "both") or "both").lower() + if direction not in ("to", "from", "both"): direction = "both" + xref_type = str(query.get("xref_type", "any") or "any").lower() + if xref_type not in ("any", "code", "data"): xref_type = "any" + include_fn = bool(query.get("include_fn", True)) + dedup = bool(query.get("dedup", True)) + sort_by = str(query.get("sort_by", "addr") or "addr") + descending = bool(query.get("descending", False)) + try: offset = max(0, int(query.get("offset", 0) or 0)) + except (TypeError, ValueError): offset = 0 + try: count = max(0, min(int(query.get("count", 200) or 200), 5000)) + except (TypeError, ValueError): count = 200 + try: + try: target = int(raw, 16) + except ValueError: + target = idaapi.get_name_ea(idaapi.BADADDR, raw) + if target == idaapi.BADADDR: raise ValueError(f"Failed to resolve address/name: {raw}") + if not ida_bytes.is_mapped(target): raise ValueError(f"Address not mapped: {raw}") + rows = [] + if direction in ("to", "both"): + for xr in idautils.XrefsTo(target, 0): + kind = "code" if xr.iscode else "data" + if xref_type != "any" and kind != xref_type: continue + row = {"direction": "to", "addr": hex(int(xr.frm)), "from": hex(int(xr.frm)), + "to": hex(int(target)), "type": kind} + if include_fn: row["fn"] = _fn(xr.frm) + rows.append(row) + if direction in ("from", "both"): + for xr in idautils.XrefsFrom(target, 0): + kind = "code" if xr.iscode else "data" + if xref_type != "any" and kind != xref_type: continue + row = {"direction": "from", "addr": hex(int(xr.to)), "from": hex(int(target)), + "to": hex(int(xr.to)), "type": kind} + if include_fn: row["fn"] = _fn(xr.to) + rows.append(row) + if dedup: + seen, deduped = set(), [] + for row in rows: + key = (row["direction"], row["from"], row["to"], row["type"]) + if key in seen: continue + seen.add(key); deduped.append(row) + rows = deduped + if sort_by == "type": + rows.sort(key=lambda r: (str(r.get("type", "")), int(str(r["addr"]), 16)), reverse=descending) + else: + rows.sort(key=lambda r: int(str(r["addr"]), 16), reverse=descending) + page = rows[offset:offset + count] if count else rows[offset:] + nxt = offset + len(page) + all_results.append({"target": raw, "resolved_addr": hex(int(target)), "direction": direction, + "xref_type": xref_type, "data": page, + "next_offset": nxt if nxt < len(rows) else None, + "total": len(rows), "error": None}) + except Exception as exc: + all_results.append({"target": raw, "resolved_addr": None, "direction": direction, + "xref_type": xref_type, "data": [], "next_offset": None, + "total": 0, "error": str(exc)}) +result = {"result": all_results} +result +''', + # A comment must land in BOTH views, and the pseudocode half is not a + # simple set: db.comments.set_at() alone leaves the pseudocode unchanged. + # Hex-Rays comments are anchored to a ctree location (treeloc_t), and an + # anchor the ctree does not actually own is dropped as an "orphan" -- so the + # itp slot has to be searched until one sticks, exactly as IDA's own UI does. + # Without it a comment silently never appears in the decompilation. + "set_comments": r''' +import idaapi, idc, ida_hexrays +rows = [] +for item in a.get("items", []): + addr_s = str(item.get("addr", "")) + text = str(item.get("comment") or "") + try: + ea = int(addr_s, 16) + if not idaapi.set_cmt(ea, text, False): + rows.append({"addr": addr_s, + "error": f"Failed to set disassembly comment at {hex(ea)}"}) + continue + if not ida_hexrays.init_hexrays_plugin(): + rows.append({"addr": addr_s}); continue + try: + cfunc = ida_hexrays.decompile(ea) + except Exception: + cfunc = None + if cfunc is None: + rows.append({"addr": addr_s}); continue + if ea == cfunc.entry_ea: + # The signature line carries no ctree item: it is a function comment. + idc.set_func_cmt(ea, text, True) + cfunc.refresh_func_ctext() + rows.append({"addr": addr_s}); continue + eamap = cfunc.get_eamap() + if ea not in eamap: + rows.append({"addr": addr_s, + "error": f"Failed to set decompiler comment at {hex(ea)}"}) + continue + nearest_ea = eamap[ea][0].ea + if cfunc.has_orphan_cmts(): + cfunc.del_orphan_cmts(); cfunc.save_user_cmts() + tl = idaapi.treeloc_t(); tl.ea = nearest_ea + placed = False + for itp in range(idaapi.ITP_SEMI, idaapi.ITP_COLON): + tl.itp = itp + cfunc.set_user_cmt(tl, text) + cfunc.save_user_cmts() + cfunc.refresh_func_ctext() + if not cfunc.has_orphan_cmts(): + placed = True; break + cfunc.del_orphan_cmts(); cfunc.save_user_cmts() + rows.append({"addr": addr_s} if placed else + {"addr": addr_s, + "error": f"Failed to set decompiler comment at {hex(ea)}"}) + except Exception as exc: + rows.append({"addr": addr_s, "error": str(exc)}) +result = {"result": rows} +result +''', + # Every category takes EITHER one edit or a LIST of them, and the answer is + # one row per edit. The port accepted only a single dict, so any batch path + # (rpc rename_many applying a whole symbol file, which is the entire point of + # that verb) died with "list indices must be integers or slices, not str" and + # reported the failure against addr=null. Mirrors the real tool: conflict + # detection before the write, dry_run/allow_overwrite/stop_on_error, per-row + # addr/old/name, and a summary counting EDITS rather than categories. + "rename": r''' +import idaapi, ida_hexrays, ida_name +batch = a.get("batch") or {} +dry_run = bool(batch.get("dry_run", False)) +allow_overwrite = bool(batch.get("allow_overwrite", False)) +stop_on_error = bool(batch.get("stop_on_error", False)) + +def _items(value): + if value is None: return [] + if isinstance(value, dict): return [value] + if isinstance(value, list): return [i for i in value if isinstance(i, dict)] + return [] + +def _set_name_checked(ea, new): + conflict = idaapi.get_name_ea(idaapi.BADADDR, new) + if conflict != idaapi.BADADDR and conflict != ea and not allow_overwrite: + return False, f"can't rename at {hex(ea)} as {new!r}: name already used at {hex(conflict)}" + if dry_run: + return True, None + flags = idaapi.SN_CHECK + if allow_overwrite: flags |= int(getattr(idaapi, "SN_FORCE", 0)) + if not idaapi.set_name(ea, new, flags): + return False, (f"Rename failed at {hex(ea)}: IDA rejected name {new!r} " + "(invalid identifier or internal conflict)") + return True, None + +def _refresh_ctext(fn_addr): + # A renamed function must invalidate Hex-Rays' cache, which is per function + # and persisted in the .i64: without this the pseudocode keeps calling the + # old name forever while every other readback reports the new one. + if not ida_hexrays.init_hexrays_plugin(): return + failure = ida_hexrays.hexrays_failure_t() + cfunc = ida_hexrays.decompile_func(fn_addr, failure, ida_hexrays.DECOMP_WARNINGS) + if cfunc: cfunc.refresh_func_ctext() + +out = {}; ok_count = failed = 0; halted = False +for category in ("func", "data", "local", "stack"): + if category not in batch: continue + rows = [] + for edit in _items(batch.get(category)): + try: + if category == "func": + addr_text = edit.get("addr") or edit.get("func_addr") or edit.get("func") + new = edit.get("name") or edit.get("new") or edit.get("new_name") + if not addr_text or not new: + row = {"addr": addr_text, "name": new, + "error": "Function rename requires addr + name"} + else: + ea = int(str(addr_text), 16) + fn = idaapi.get_func(ea) + if fn is None: + row = {"addr": addr_text, "name": new, "error": "Function not found"} + else: + old = idaapi.get_name(fn.start_ea) or None + ok, err = _set_name_checked(fn.start_ea, str(new)) + row = {"addr": addr_text, "old": old, "name": str(new)} + if err: row["error"] = err + if dry_run: row["dry_run"] = True + if ok and not dry_run: _refresh_ctext(fn.start_ea) + elif category == "data": + addr_text = edit.get("addr") + old = edit.get("old") or edit.get("old_name") + new = edit.get("new") or edit.get("new_name") or edit.get("name") + if not new and new != "": + row = {"old": old, "new": None, + "error": "Global rename requires target and new name"} + else: + if addr_text is not None: + ea = int(str(addr_text), 16) + old = old or (idaapi.get_name(ea) or None) + else: + ea = idaapi.get_name_ea(idaapi.BADADDR, str(old or "")) + if ea == idaapi.BADADDR: + row = {"old": old, "new": str(new), "error": f"Global {old!r} not found"} + else: + # An empty new name CLEARS the label; that is a real + # request (tests revert with it), not a missing argument. + if str(new) == "": + ok = bool(ida_name.set_name(ea, "", idaapi.SN_CHECK)) + err = None if ok else f"Failed to clear the name at {hex(ea)}" + else: + ok, err = _set_name_checked(ea, str(new)) + row = {"addr": hex(ea), "old": old, "new": str(new)} + if err: row["error"] = err + if dry_run: row["dry_run"] = True + else: + fa, old, new = edit.get("func_addr"), edit.get("old"), edit.get("new") + if not fa or not old or not new: + row = {"old": old, "new": new, + "error": f"{category} rename requires func_addr + old + new"} + else: + ea = int(str(fa), 16) + pseudo = db.pseudocode.decompile(ea) + var = pseudo.find_local_variable(str(old)) + if var is None: + row = {"func_addr": fa, "old": old, "new": new, + "error": f"no local {old!r} in that function"} + elif dry_run: + row = {"func_addr": fa, "old": old, "new": new, "dry_run": True} + else: + var.set_user_name(str(new)) + ok = bool(pseudo.save_local_variable_info(var, save_name=True)) + row = {"func_addr": fa, "old": old, "new": new} + if not ok: row["error"] = "IDA rejected the local variable name" + except Exception as exc: + row = {"addr": edit.get("addr"), "error": str(exc)} + rows.append(row) + if row.get("error"): failed += 1 + else: ok_count += 1 + if row.get("error") and stop_on_error: + halted = True; break + out[category] = rows + if halted: break +out["summary"] = {"ok": ok_count, "failed": failed} +if dry_run: out["summary"]["dry_run"] = True +if halted: out["summary"]["halted"] = True +result = out +result +''', +} + + +_OPERATIONS["define_code_run"] = r''' +import ida_bytes, ida_idp, ida_segment, ida_ua, idaapi +ea, limit = int(str(a["addr"]), 16), max(1, min(int(a.get("limit", 20000)), 200000)) +seg = ida_segment.getseg(ea) +if seg is None: + result = {"addr": a["addr"], "error": "no segment", "count": 0} +else: + start, count, stopped, hi = ea, 0, "limit", int(seg.end_ea) + while count < limit: + if ea >= hi: stopped = "segment"; break + flags = ida_bytes.get_flags(ea) + if ida_bytes.is_code(flags) or ida_bytes.is_data(flags): stopped = "defined"; break + size = int(ida_ua.create_insn(ea)) + if size <= 0: stopped = "undecodable"; break + count += 1 + insn = ida_ua.insn_t() + if ida_ua.decode_insn(insn, ea) > 0: + try: is_ret = bool(ida_idp.is_ret_insn(insn)) + except Exception: is_ret = False + if is_ret or (insn.get_canon_feature() & idaapi.CF_STOP): + ea += size; stopped = "flow"; break + ea += size + result = {"start": hex(start), "end": hex(ea), "count": count, "stopped": stopped} +result +''' + + +_OPERATIONS["define_func_run"] = r''' +import ida_bytes, ida_funcs, ida_segment +ea = int(str(a["addr"]), 16) +fn = db.functions.get_at(ea) +if fn is not None and int(fn.start_ea) == ea: + result = {"addr": hex(ea), "ok": True, "start": hex(ea), "end": hex(int(fn.end_ea)), "how": "existed"} +else: + automatic = bool(db.functions.create(ea)) + if not automatic: + seg = db.segments.get_at(ea); end = ea; hi = int(seg.end_ea) if seg else ea + while end < hi and ida_bytes.is_code(ida_bytes.get_flags(end)): + nxt = int(ida_bytes.get_item_end(end)) + if nxt <= end: break + end = nxt + ok = bool(end > ea and ida_funcs.add_func(ea, end)) + else: ok = True + fn = db.functions.get_at(ea) + result = ({"addr": hex(ea), "ok": True, "start": hex(int(fn.start_ea)), + "end": hex(int(fn.end_ea)), "how": "auto" if automatic else "explicit-end"} + if ok and fn is not None else + {"addr": hex(ea), "ok": False, "error": f"IDA refused a function at {ea:#x}"}) +result +''' + + +_OPERATIONS["set_thumb"] = r''' +import ida_bytes, ida_ida, ida_idp, ida_segment, ida_segregs +ea = int(str(a["addr"]), 16); treg = ida_idp.str2reg("T") +seg = ida_segment.getseg(ea) +if treg is None or treg < 0: + result = {"addr": hex(ea), "error": "no T register (not an ARM database)"} +elif seg is None: + result = {"addr": hex(ea), "error": "no segment"} +else: + current = ida_segregs.get_sreg(ea, treg) + current = 0 if current in (None, 0xFFFFFFFF, -1) else int(current) + want = {"on": 1, "off": 0}.get(str(a.get("mode", "toggle")).lower(), 0 if current else 1) + changed = False + if want and seg.bitness != 1: + ida_segment.set_segm_addressing(seg, 1); changed = True + size = max(int(ida_bytes.get_item_size(ea)), 2) + ida_bytes.del_items(ea, 0, size) + ok = bool(ida_segregs.split_sreg_range(ea, treg, want, ida_segregs.SR_user)) + now = ida_segregs.get_sreg(ea, treg) + result = {"addr": hex(ea), "thumb": bool(now), "was": bool(current), "ok": ok, + "bitness": ida_segment.getseg(ea).bitness, "forced_32bit": changed, + "db_64bit": bool(ida_ida.inf_get_app_bitness() == 64 and want)} +result +''' + + +_OPERATIONS["thumb_scan"] = r''' +import ida_bytes, ida_funcs, ida_idp, ida_segment, ida_segregs, ida_ua +lo, hi = int(str(a["start"]), 16), int(str(a["end"]), 16) +apply, limit = bool(a.get("apply", True)), int(a.get("limit", 512)) +treg = ida_idp.str2reg("T"); found = []; applied = 0; cursor = lo +while cursor + 4 <= hi and len(found) < limit: + at = cursor; value = int(ida_bytes.get_dword(cursor)); cursor += 4 + if not value & 1: continue + target = value & ~1; seg = ida_segment.getseg(target) + if seg is None or not (seg.perm & ida_segment.SEGPERM_EXEC or seg.perm == 0): continue + flags = ida_bytes.get_flags(target) + if ida_bytes.is_data(flags): continue + item = {"at": hex(at), "value": hex(value), "target": hex(target), + "was_code": bool(ida_bytes.is_code(flags))}; found.append(item) + if not apply: continue + if treg is not None and treg >= 0: ida_segregs.split_sreg_range(target, treg, 1, ida_segregs.SR_user) + if not ida_bytes.is_code(ida_bytes.get_flags(target)): + ida_bytes.del_items(target, 0, 2) + if ida_ua.create_insn(target) <= 0: item["decoded"] = False; continue + item["decoded"] = True; item["function"] = bool(db.functions.get_at(target) or db.functions.create(target)); applied += 1 +result = {"start": hex(lo), "end": hex(hi), "found": found, "applied": applied, "n": len(found)} +result +''' + + +_OPERATIONS["decomp_error"] = r''' +import ida_hexrays, ida_ida +ea = int(str(a["addr"]), 16); fn = db.functions.get_at(ea) +result = {"addr": hex(ea), "bitness": ida_ida.inf_get_app_bitness()} +if fn is None: + result["reason"] = "no function here" +else: + try: + failure = ida_hexrays.hexrays_failure_t(); cfunc = ida_hexrays.decompile_func(fn, failure) + if cfunc is not None: result["reason"] = "" + else: + result.update({"reason": failure.desc() or f"error {failure.code}", + "code": int(failure.code), "errea": hex(int(failure.errea))}) + except Exception as exc: result["reason"] = f"{type(exc).__name__}: {exc}" +result +''' + +# `heads` and the operand-format tools are the port's IDAPython island: the +# continuous listing's presentation model (undefined runs, colour spans, operand +# extents, banners, struct members, the digest protocol) and IDA/Hex-Rays number +# formats have no ida-domain surface. Rather than paraphrase ~1100 lines of +# performance-tuned, behaviour-sensitive code into string literals, they stay +# real, diffable source in idatui/remote_tools.py and are shipped to the database +# process as text. Read once at import; the file ships beside this module. +_REMOTE_LIB = (Path(__file__).with_name("remote_tools.py")).read_text(encoding="utf-8") + +#: Versioned by content, so editing remote_tools.py re-installs it instead of +#: silently running the copy a long-lived worker already has. +_REMOTE_MODULE = "_idatui_remote_" + hashlib.sha1( + _REMOTE_LIB.encode("utf-8")).hexdigest()[:12] + +#: Sent back when the database process has not got the library yet; the client +#: installs it and retries once. Amortised, a worker receives it exactly once. +_NEED_LIB = "__idatui_needs_remote_lib__" + +#: Installs the library as a real module in the database process. Persisting it +#: in sys.modules is what makes the module-level caches (the tag maps, and the +#: line-render lru_cache the listing's throughput depends on) survive between +#: calls -- execute_python builds a fresh namespace every time, so a library +#: exec'd inline is rebuilt, and its caches thrown away, on every single call. +_INSTALL_LIB = f''' +import sys, types +_m = types.ModuleType({_REMOTE_MODULE!r}) +exec(compile(a["source"], {_REMOTE_MODULE!r}, "exec"), _m.__dict__) +sys.modules[{_REMOTE_MODULE!r}] = _m +result = True +result +''' + + +def _remote_op(call: str) -> str: + """A snippet that calls one of the carried-over tools by its real signature. + + Costs one short request: the library is imported from the database process's + own sys.modules, not shipped again. + """ + return (f"import sys\n" + f"_m = sys.modules.get({_REMOTE_MODULE!r})\n" + f"result = {{{_NEED_LIB!r}: True}} if _m is None else _m.{call}\n" + f"result\n") + + +_OPERATIONS["op_format"] = _remote_op( + 'op_format(addr=a["addr"], mode=a.get("mode", "cycle"),' + ' col=int(a.get("col", -1)), n=int(a.get("n", -1)))') +_OPERATIONS["pc_nums"] = _remote_op('pc_nums(addr=a["addr"])') +_OPERATIONS["decompile"] = _remote_op( + 'decompile(addr=a["addr"],' + ' include_addresses=bool(a.get("include_addresses", True)))') +_OPERATIONS["decomp_map"] = _remote_op('decomp_map(addr=a["addr"])') +_OPERATIONS["pc_num_format"] = _remote_op( + 'pc_num_format(addr=a["addr"], mode=a.get("mode", "cycle"),' + ' line=int(a.get("line", -1)), col=int(a.get("col", -1)),' + ' ea=a.get("ea", ""), opnum=int(a.get("opnum", -1)))') + +# The listing walker itself. Replaces the port's re-implementation, which +# rendered no per-operand extents (so no keypress could say which literal it +# would reformat) and had no digest/expect support (so every page was re-sent +# after any edit), and whose span walk was the per-character loop our own +# version had already been rewritten to avoid. +_HEADS = _remote_op( + 'heads(addr=a["addr"], count=int(a.get("count", 200)),' + ' offset=int(a.get("offset", 0)), end=a.get("end", ""),' + ' back=bool(a.get("back", False)), annotate=bool(a.get("annotate", False)),' + ' expect=a.get("expect", ""))') + + +# The graph view's only backend call. Blocks are address RANGES, never text: +# the client re-renders them with `heads`, so boxes reuse the exact listing rows +# (colours, operand marks, trail painting) instead of growing a second renderer. +# +# ida-domain exposes no basic-block/edge-kind surface, so this stays on ida_gdl. +_OPERATIONS["flowchart"] = r''' +import ida_funcs, ida_gdl +ea = int(str(a["addr"]), 16) +fn = ida_funcs.get_func(ea) +if fn is None: + result = {"addr": hex(ea), "error": "no function at that address", "blocks": []} +else: + fc = ida_gdl.FlowChart(fn, flags=ida_gdl.FC_PREDS) + index, order = {}, [] + for bb in fc: + index[bb.start_ea] = len(order) + order.append(bb) + blocks = [] + for bb in order: + sl = [s for s in bb.succs() if s.start_ea in index] + succs = [] + for s in sl: + # Edge kind is what the graph view colours by: an n-way dispatch is + # "switch", a successor that is literally the next address falls + # through, anything else is a taken branch. + if len(sl) > 2: kind = "switch" + elif s.start_ea == bb.end_ea: kind = "fall" + else: kind = "jump" + succs.append([index[s.start_ea], kind]) + blocks.append({"id": index[bb.start_ea], "start": hex(int(bb.start_ea)), + "end": hex(int(bb.end_ea)), "succs": succs}) + result = {"addr": hex(ea), + "func": {"addr": hex(int(fn.start_ea)), "end": hex(int(fn.end_ea)), + "name": ida_funcs.get_func_name(fn.start_ea) or ""}, + "entry": index.get(fn.start_ea, 0), "blocks": blocks} +result +''' + +# Only ever reached as domain.py's fallback when file_regions yields nothing. +_OPERATIONS["survey_binary"] = r''' +segments = [] +for seg in db.segments.get_all(): + segments.append({"start": hex(int(seg.start_ea)), "end": hex(int(seg.end_ea)), + "name": db.segments.get_name(seg) or ""}) +result = {"segments": segments} +result +''' + + +class CodeModeClient: + """A leased GUI/idalib database accessed through ``ida_codemode``.""" + + def __init__( + self, + binary_path: str, + *, + ttl: int = 0, + load_args: str = "", + processor: str | None = None, + loading_address: int | None = None, + file_type: str | None = None, + output_database: str | None = None, + spawn: bool = True, + new_database: bool = False, + ) -> None: + del ttl # managed-worker lifetime is lease-based, not idle-TTL based + self._path = os.path.abspath(os.path.expanduser(binary_path)) + parsed_processor, parsed_address, parsed_file_type = _parse_load_args(load_args) + self._processor = processor or parsed_processor + self._loading_address = loading_address if loading_address is not None else parsed_address + self._file_type = file_type or parsed_file_type + self._output_database = output_database + self._spawn = spawn + self._new_database = new_database + self._handle: DatabaseHandle | None = None + self._last_entry: RegistryEntry | None = None + self._connect_lock = threading.Lock() + + def _database_exists(self) -> bool: + """Whether the IDB this open would target is already on disk. + + Its loader switches are baked in, so they must not be sent again. + """ + try: + target = self._output_database or expected_idb_path(self._path) + except Exception: # noqa: BLE001 -- resolver unavailable: assume fresh + return False + return bool(target) and os.path.exists(target) + + def connect(self, timeout: float = 1800.0, progress=None) -> "CodeModeClient": + _require_codemode() + with self._connect_lock: + if self._handle is not None and self._handle.connected: + return self + if progress: + progress(f"discovering Code Mode database for {os.path.basename(self._path)}…") + try: + # A Ctrl+L reload releases its current managed-worker lease, but + # that worker remains registered during Code Mode's final-lease + # grace period. Retry only that known handoff window. A GUI or + # another long-lived client remains busy and yields a clear + # failure rather than being modified underneath its owner. + deadline = time.monotonic() + min(timeout, 60.0) + while True: + try: + # Loader switches describe how to IMPORT a raw file and + # are recorded in the database it produces. Sending them + # again for a database that already exists is a FATAL + # error in IDA itself ("Switch '-b400' can be used only + # when loading a new file"), which kills the worker + # before it can report anything useful. So: describe the + # import only when there is an import to describe. + fresh = self._new_database or not self._database_exists() + handle = DatabaseHandle.open( + self._path, + spawn=self._spawn, + timeout=max(0.1, timeout), + output_database=self._output_database, + processor=self._processor if fresh else None, + # DatabaseHandle calls this image_base and wants the + # natural (16-byte aligned) address; it does the + # conversion to IDA's paragraph-based -b itself. + image_base=self._loading_address if fresh else None, + file_type=self._file_type if fresh else None, + new_database=self._new_database, + ) + break + except IdbBusy: + if not self._new_database or time.monotonic() >= deadline: + raise + if progress: + progress("waiting for the previous Code Mode lease to close…") + # Remember the record before managed shutdown withdraws + # its JSON. The lifetime lock remains held until IDA has + # actually closed the IDB; waiting on it avoids racing a + # replacement worker into the old process's file lock. + expected = canonical_path( + self._output_database or expected_idb_path(self._path) + ) + owners = [item.entry for item in scan_instances(timeout=0.5) + if item.entry.idb_key == idb_key(expected)] + if owners: + self._wait_for_entry_release( + owners[0], max(0.0, deadline - time.monotonic()) + ) + else: + time.sleep(0.2) + if progress: + backend = handle.entry.backend + progress(f"attached to {backend} database; waiting for auto-analysis…") + handle.wait_autoanalysis(timeout=timeout) + except Exception as exc: # normalize the dependency's transport errors + raise self._connection_error(exc) from exc + self._handle = handle + self._last_entry = handle.entry + return self + + @staticmethod + def _connection_error(exc: BaseException) -> IDAConnectionError: + return IDAConnectionError(str(exc) or type(exc).__name__) + + @property + def connected(self) -> bool: + return self._handle is not None and self._handle.connected + + @property + def pid(self) -> int | None: + return self._handle.entry.pid if self._handle is not None else None + + @property + def backend(self) -> str | None: + return self._handle.entry.backend if self._handle is not None else None + + def execute_python(self, code: str, *, timeout: float | None = None) -> Any: + if not self.connected: + self.connect() + handle = self._handle + if handle is None: + raise IDAConnectionError("Code Mode database is not connected") + try: + response = handle.execute_python(code, timeout=timeout) + except RemoteError as exc: + details = exc.details or {} + message = str(exc) + if details.get("traceback"): + message += f"\n{details['traceback']}" + if exc.code == "operation_timeout": + raise IDATimeoutError(message) from exc + raise IDAToolError("execute_python", message) from exc + except (InstanceDisconnectedError, ClientError) as exc: + raise self._connection_error(exc) from exc + if not isinstance(response, dict) or "result" not in response: + raise IDAToolError("execute_python", "Code Mode returned an invalid execution result") + return response["result"] + + @staticmethod + def _unpack(answer: Any) -> Any: + """Undo _PACK_EPILOGUE. Anything else passes through untouched.""" + if isinstance(answer, dict) and _PACKED in answer: + return json.loads(answer[_PACKED]) + return answer + + def invoke(self, operation: str, *, timeout: float | None = None, **args) -> Any: + """Execute one TUI domain operation through Code Mode.""" + if operation in ("idb_save", "save"): + return self.save_database() + if operation in ("server_health", "ping", "health", "state"): + return self.health() + body = _HEADS if operation == "heads" else _OPERATIONS.get(operation) + if body is None: + raise IDAToolError(operation, f"unknown ida-tui Code Mode operation: {operation}") + try: + answer = self._unpack(self.execute_python(_script(args, body), timeout=timeout)) + if isinstance(answer, dict) and answer.get(_NEED_LIB): + # First call against this database process (or a restarted one). + self.execute_python(_script({"source": _REMOTE_LIB}, _INSTALL_LIB), + timeout=timeout) + answer = self._unpack( + self.execute_python(_script(args, body), timeout=timeout)) + return answer + except IDAToolError as exc: + if exc.tool == "execute_python": + raise IDAToolError(operation, exc.message) from exc + raise + + # Temporary source compatibility for external drivers/tests that used the + # old WorkerClient. Application code uses the accurately named invoke(). + call = invoke + + def save_database(self) -> dict[str, Any]: + if not self.connected: + self.connect() + handle = self._handle + if handle is None: + raise IDAConnectionError("Code Mode database is not connected") + try: + return handle.save_database() + except RemoteError as exc: + raise IDAToolError("save_database", str(exc)) from exc + except (InstanceDisconnectedError, ClientError) as exc: + raise self._connection_error(exc) from exc + + def health(self) -> dict[str, Any]: + if not self.connected: + self.connect() + assert self._handle is not None + entry = self._handle.entry + module = os.path.basename(entry.exe_path or entry.idb_path or self._path) + return { + "ok": self._handle.connected, + "module": module, + "backend": entry.backend, + "record_id": entry.record_id, + "input_path": entry.exe_path, + "idb_path": entry.idb_path, + } + + def keepalive(self, interval: float = 120.0) -> _NoopKeepAlive: + del interval + return _NoopKeepAlive() + + def resolve_db(self) -> str: + if not self.connected: + self.connect() + assert self._handle is not None + return self._handle.entry.record_id + + def set_db(self, db: str | None) -> None: + del db # one handle is permanently bound to one registered database + + def list_sessions(self) -> list[Session]: + if not self.connected: + self.connect() + assert self._handle is not None + entry = self._handle.entry + path = entry.exe_path or entry.idb_path or self._path + return [Session(session_id=entry.record_id, filename=os.path.basename(path), + input_path=path, is_active=True)] + + def close(self, grace: float = 0.0) -> None: + del grace + with self._connect_lock: + handle, self._handle = self._handle, None + if handle is not None: + self._last_entry = handle.entry + handle.close() # release our lease; never close a GUI/other client's DB + + @staticmethod + def _wait_for_entry_release(entry: "RegistryEntry", timeout: float) -> bool: + _require_codemode() + path = REGISTRY_DIR / f"{entry.record_id}.lock" + deadline = time.monotonic() + max(0.0, timeout) + while True: + lock = FileLock(path) + try: + if lock.try_acquire(): + return True + except OSError: + pass + finally: + lock.close() + if time.monotonic() >= deadline: + return False + time.sleep(min(0.1, deadline - time.monotonic())) + + def wait_released(self, timeout: float = 45.0) -> bool: + """Wait until a managed instance releases its lifetime lock. + + Normal application shutdown must not wait: another client may retain the + worker. This is an explicit test/maintenance helper for deleting a + temporary IDB safely after this client closes. GUI instances return + ``False`` immediately because clients never own their lifetime. + """ + entry = self._last_entry + if entry is None or entry.backend != "idalib": + return False + return self._wait_for_entry_release(entry, timeout) + + def __enter__(self) -> "CodeModeClient": + return self.connect() + + def __exit__(self, *exc) -> None: + self.close() diff --git a/idatui/domain.py b/idatui/domain.py index aedade9..6fe73c2 100644 --- a/idatui/domain.py +++ b/idatui/domain.py @@ -1,19 +1,15 @@ -"""Domain / paging layer: address-centric models over the raw MCP client. +"""Domain / paging layer: address-centric models over IDA Code Mode. This is where the "millions of lines" problem is solved, so the TUI widgets only ever see a viewport-sized slice. Every hard-won constraint from ``docs/PAGING_FINDINGS.md`` is encoded here: -* Per-call caps are silent (over the cap the server returns 10, not a clamp), so - we clamp page sizes ourselves: ``LIST_PAGE`` / ``DISASM_BLOCK`` <= the caps. -* ``next_offset`` is unreliable; we paginate by advancing ``len(data)``. -* ``disasm offset=N`` is O(N) with no resumable cursor, so windowed disassembly - is **block-cached** (revisits are free) and **prefetches** the next block on a - background thread (the client is concurrency-safe). -* ``include_total`` scans the whole function (~200ms on monsters); totals are - fetched once and cached. -* ``decompile`` can hard-fail on huge functions as a *soft* error (``code`` is - null); that is surfaced as data, not an exception. +* Page sizes remain bounded so remote execution returns viewport-scale JSON. +* Pagination advances by the number of rows actually returned. +* Deep head walks are block-cached (revisits are free) and neighboring blocks + prefetch through the thread-safe Code Mode client. +* Expensive function totals are fetched once and cached. +* Decompilation failures are surfaced as data, not application crashes. Everything here is synchronous and thread-safe. The TUI runs these calls from Textual worker threads; the internal prefetch pool is separate and small. @@ -22,10 +18,8 @@ Textual worker threads; the internal prefetch pool is separate and small. from __future__ import annotations import bisect -import json import re import threading -import urllib.request from concurrent.futures import ThreadPoolExecutor from collections.abc import Sequence from dataclasses import dataclass, field, replace @@ -36,7 +30,7 @@ from . import diag from .errors import IDAToolError if TYPE_CHECKING: # type hint only - from .worker_client import WorkerClient # noqa: F401 + from .codemode_client import CodeModeClient # Clamps derived from measured caps (list ~700, disasm ~500). Margin included. LIST_PAGE = 500 @@ -66,7 +60,7 @@ class Func: def from_raw(cls, d: dict) -> "Func": addr = _as_int(d["addr"]) name = d.get("name") - # An unnamed function (server returns null/empty) must still have a + # An unnamed function must still have a # usable string name — synthesize IDA's sub_ADDR so every consumer # (palette, sort, rename prefill) can treat name as a str. if not name: @@ -93,7 +87,7 @@ class Line: class Head(NamedTuple): - """One flat-listing item (from the ``heads`` server tool): a code + """One flat-listing item (from the Code Mode ``heads`` operation): a code instruction, a data item, or an undefined byte run. A ``NamedTuple`` rather than a dataclass because this is by far the @@ -114,7 +108,7 @@ class Head(NamedTuple): name: str | None = None raw: bytes | None = None # opcode/item bytes (filled in for code by the model) #: [(kind, text)] from IDA's own colour tags — mnem/reg/num/name/str/punct/… - #: None when the worker didn't provide them (older worker, or the spans + #: None when Code Mode didn't provide them (or the spans #: disagreed with the plain text, in which case the text wins). #: #: Held exactly as it came off the wire, and **read-only**. The worker @@ -301,7 +295,7 @@ class FunctionIndex: """A lazily-paginated, cached view of the function list. Loads pages of ``LIST_PAGE`` on demand, advancing by ``len(data)`` (never by - ``next_offset``). A single index instance corresponds to one server-side + ``next_offset``). A single index instance corresponds to one remote ``filter`` glob (``None`` = all functions). """ @@ -321,7 +315,7 @@ class FunctionIndex: query: dict = {"offset": offset, "count": LIST_PAGE} if self.filter: query["filter"] = self.filter - data = _query_data(self._prog.client.call("list_funcs", queries=[query])) + data = _query_data(self._prog.client.invoke("list_funcs", queries=[query])) added = 0 with self._lock: for d in data: @@ -431,7 +425,7 @@ class DisasmModel: code function this equals the heads row count that backs the lines.""" if self._total is not None: return self._total - payload = self._prog.client.call( + payload = self._prog.client.invoke( "disasm", addr=hex(self.ea), max_instructions=1, include_total=True ) total = payload.get("total_instructions") @@ -492,7 +486,7 @@ class DisasmModel: # The function disasm view is a listing filtered to the function: fetch a # block of heads (one per instruction for code). Over-fetch one row so # the block knows where its last instruction ends (opcode-byte sizing). - payload = self._prog.client.call( + payload = self._prog.client.invoke( "heads", addr=hex(self.ea), offset=b * self.BLOCK, count=self.BLOCK + 1, **self._end_kw(), ) @@ -633,15 +627,15 @@ class ListingModel: """A flat, IDA-style disassembly *listing* over one segment: code, data and undefined heads interleaved, unlike ``DisasmModel`` (one function, code only). - Backed by the injected ``heads`` server tool, which walks item heads and - renders each via ``generate_disasm_line``. The segment is walked lazily in + Backed by the Code Mode adapter's ``heads`` operation, which walks item heads + and renders each via ``generate_disasm_line``. The segment is walked lazily in forward pages (``FunctionIndex`` style); line index == position in the walked head list. Random access to an address is O(distance-from-seg-start) the first time (then cached) — the same tradeoff as ``disasm offset=N``. Grows on demand as the viewport scrolls. Synchronous + thread-safe. """ - PAGE = 500 # heads per server call (well under the tool's 2000 cap) + PAGE = 500 # viewport-scale heads per Code Mode execution def __init__(self, program: "Program", seg_start: int, seg_end: int, name: str | None = None): @@ -754,7 +748,7 @@ class ListingModel: if self._done or self._next is None: return 0 frm = self._next - payload = self._prog.client.call( + payload = self._prog.client.invoke( "heads", addr=hex(frm), count=self.PAGE, annotate=True) rows = payload.get("heads", []) if isinstance(payload, dict) else [] cur = payload.get("cursor", {}) if isinstance(payload, dict) else {} @@ -1019,7 +1013,7 @@ class ListingModel: # the expectation rather than asking first means a page that HAS changed # still costs one round trip. try: - payload = self._prog.client.call( + payload = self._prog.client.invoke( "heads", addr=hex(addr), count=self.PAGE, annotate=True, expect="" if want_digest is None else str(want_digest)) except Exception: # noqa: BLE001 -- keep the old text rather than blank @@ -1233,7 +1227,7 @@ class HexModel: class Program: """The bound analysis session: models, caches, and a small prefetch pool.""" - def __init__(self, client: "WorkerClient", prefetch_workers: int = 2): + def __init__(self, client: "CodeModeClient", prefetch_workers: int = 2): self.client = client self._pool = ThreadPoolExecutor( max_workers=prefetch_workers, thread_name_prefix="idatui-prefetch" @@ -1257,7 +1251,7 @@ class Program: self._sections: list[tuple[int, int, str]] | None = None self._fileregions: list[tuple[int, int, int]] | None = None self._hexmodel: "HexModel | None" = None - self._no_read_raw = False # set if the server lacks the read_raw tool + self._no_read_raw = False # compatibility fallback for alternate clients self._lock = threading.Lock() # -- prefetch plumbing ------------------------------------------------- # @@ -1284,17 +1278,14 @@ class Program: """Sorted raw segment map [(start, end, file_off, name)] — the single source for sections()/file_regions()/image_range. Cached. - Uses the injected ``file_regions`` tool (a plain segment walk, ~ms). - This deliberately AVOIDS ``survey_binary``, which also computes function - counts / strings / stats and takes *seconds* on a large IDB (it was the - cause of the multi-second hex-pane open). Falls back to survey_binary - only if the injected tool is missing. + Uses the Code Mode adapter's ``file_regions`` operation (a plain segment + walk, ~ms), avoiding broad binary surveys on the hex-pane open path. """ if self._segments_cache is not None: return self._segments_cache segs: list[tuple[int, int, int, str]] = [] try: - r = self.client.call("file_regions") + r = self.client.invoke("file_regions") for d in (r.get("regions", []) if isinstance(r, dict) else []): if isinstance(d, dict) and "start" in d: segs.append((_as_int(d["start"]), _as_int(d["end"]), @@ -1303,7 +1294,7 @@ class Program: segs = [] if not segs: # older server without file_regions -> survey_binary (slow) try: - sb = self.client.call("survey_binary") + sb = self.client.invoke("survey_binary") for s in (sb.get("segments", []) if isinstance(sb, dict) else []): try: segs.append((_as_int(s["start"]), _as_int(s["end"]), -1, @@ -1345,8 +1336,7 @@ class Program: def file_regions(self) -> list[tuple[int, int, int]]: """Sorted [(start, end, file_off)] mapping loaded segments to raw file - offsets (file_off == -1 for non-file-backed, e.g. .bss). Cached; needs - the injected ``file_regions`` server tool.""" + offsets (file_off == -1 for non-file-backed, e.g. .bss). Cached.""" if self._fileregions is not None: return self._fileregions regions = [(s, e, fo) for s, e, fo, _nm in self._segments()] @@ -1364,15 +1354,14 @@ class Program: def read_bytes(self, ea: int, n: int) -> bytes: """Raw bytes [ea, ea+n) from IDA (gaps read as zero). - Fast path: the injected ``read_raw`` tool returns one contiguous hex - string (C-speed both ends). Falls back to the stock ``get_bytes`` (a - per-byte '0x..'-with-spaces string) on an older server without it. + The Code Mode adapter returns one contiguous hex string (C-speed in IDA). + A legacy ``get_bytes`` decoding fallback remains for alternate clients. """ if n <= 0: return b"" if not self._no_read_raw: try: - r = self.client.call("read_raw", addr=hex(ea), size=int(n)) + r = self.client.invoke("read_raw", addr=hex(ea), size=int(n)) h = r.get("hex") if isinstance(r, dict) else None if isinstance(h, str): out = bytes.fromhex(h) @@ -1386,7 +1375,7 @@ class Program: except (ValueError, KeyError): pass # malformed hex -> fall through to the legacy decoder try: - r = self.client.call("get_bytes", regions=[{"addr": hex(ea), "size": int(n)}]) + r = self.client.invoke("get_bytes", regions=[{"addr": hex(ea), "size": int(n)}]) except IDAToolError: return b"\x00" * n res = r.get("result", []) if isinstance(r, dict) else [] @@ -1437,7 +1426,7 @@ class Program: def list_structs(self, filter: str = "") -> list[Struct]: """All local structs/unions (optionally name-substring filtered), sorted by name.""" - payload = self.client.call("search_structs", filter=filter) + payload = self.client.invoke("search_structs", filter=filter) res = payload.get("result", []) if isinstance(payload, dict) else [] out = [Struct.from_raw(d) for d in res if isinstance(d, dict) and d.get("name") @@ -1447,9 +1436,9 @@ class Program: def struct_source(self, name: str) -> str: """A C definition for ``name`` reconstructed from its member layout - (the server exposes members, not printable source). Faithful to IDA's + (the remote operation exposes members, not printable source). Faithful to IDA's field names/types; array dims are moved after the field name.""" - payload = self.client.call( + payload = self.client.invoke( "type_inspect", queries=[{"name": name, "include_members": True}]) res = payload.get("result", []) if isinstance(payload, dict) else [] info = res[0] if res and isinstance(res[0], dict) else {} @@ -1472,7 +1461,7 @@ class Program: def declare_type(self, decl: str) -> str | None: """Create or update a C type. Returns None on success, else the parse error. (Re-declaring a name updates it in place.)""" - payload = self.client.call("declare_type", decls=decl) + payload = self.client.invoke("declare_type", decls=decl) res = payload.get("result", []) if isinstance(payload, dict) else [] if res and isinstance(res[0], dict): return res[0].get("error") @@ -1481,10 +1470,9 @@ class Program: # -- function / variable types ---------------------------------------- # def func_types(self, ea: int) -> FuncTypes | None: """Structured decompiler types for the function at ``ea`` (prototype + - local variables). None if ``ea`` isn't a decompilable function. Requires - the injected ``func_types`` server tool.""" + local variables). None if ``ea`` isn't a decompilable function.""" try: - r = self.client.call("func_types", addr=hex(ea)) + r = self.client.invoke("func_types", addr=hex(ea)) except IDAToolError: return None if not isinstance(r, dict) or r.get("error"): @@ -1497,7 +1485,7 @@ class Program: def set_function_type(self, ea: int, signature: str) -> str | None: """Set a function's prototype. None on success, else an error string.""" - r = self.client.call("set_type", edits=[{"addr": hex(ea), "signature": signature}]) + r = self.client.invoke("set_type", edits=[{"addr": hex(ea), "signature": signature}]) res = r.get("result", []) if isinstance(r, dict) else [] row = res[0] if res and isinstance(res[0], dict) else {} if row.get("ok"): @@ -1506,9 +1494,9 @@ class Program: def data_type(self, ea: int) -> dict | None: """Current type info for a data item/global: {addr,name,type,size,is_func}. - None if the tool is unavailable or the address isn't mapped.""" + None if the operation fails or the address isn't mapped.""" try: - r = self.client.call("data_type", addr=hex(ea)) + r = self.client.invoke("data_type", addr=hex(ea)) except IDAToolError: return None if not isinstance(r, dict) or r.get("error"): @@ -1517,7 +1505,7 @@ class Program: def set_data_type(self, ea: int, decl: str) -> str | None: """Set a global/data item's type. None on success, else an error string.""" - r = self.client.call( + r = self.client.invoke( "set_type", edits=[{"kind": "global", "addr": hex(ea), "type": decl}]) res = r.get("result", []) if isinstance(r, dict) else [] row = res[0] if res and isinstance(res[0], dict) else {} @@ -1526,9 +1514,9 @@ class Program: return row.get("error") or "failed to set the type" def set_lvar_type(self, fn_ea: int, var: str, ty: str) -> str | None: - """Set a decompiler local variable's type (via the injected server tool). + """Set a decompiler local variable's type through ida-domain pseudocode. None on success, else an error string.""" - r = self.client.call("set_lvar_type", addr=hex(fn_ea), variable=var, type=ty) + r = self.client.invoke("set_lvar_type", addr=hex(fn_ea), variable=var, type=ty) if isinstance(r, dict) and r.get("error"): return r["error"] if isinstance(r, dict) and not r.get("ok"): @@ -1537,15 +1525,14 @@ class Program: def delete_type(self, name: str) -> str | None: """Delete a named type. Returns None on success, else an error string. - Requires a server-side ``del_type`` tool; if absent, a clear message is - returned instead of raising.""" + Returns a clear error instead of raising when the runtime cannot do it.""" try: - self.client.call("del_type", name=name) + self.client.invoke("del_type", name=name) return None except IDAToolError as e: msg = e.message if "not found" in msg.lower() and "del_type" in msg: - return "delete needs a 'del_type' tool on the ida-pro-mcp server" + return "the connected Code Mode runtime cannot delete local types" return msg # -- disassembly ------------------------------------------------------- # @@ -1559,13 +1546,7 @@ class Program: # -- decompilation ----------------------------------------------------- # def decompile(self, ea: int, refresh: bool = False) -> Decompilation: - """Full pseudocode for a function. - - The server truncates responses over 50KB (strings clipped to 1000 - chars) but caches the full output and exposes it at - ``_meta.ida_mcp.download_url``. We transparently fetch that so the view - always gets the complete body, not a 1KB stub. - """ + """Full pseudocode for a function, returned directly by Code Mode.""" if not refresh: with self._lock: hit = self._decomp.get(ea) @@ -1574,10 +1555,10 @@ class Program: dec, hit_gen = hit if hit_gen == gen: return dec - # Cached before a rename: names may be stale. Drop the server's - # Hex-Rays cache so the refetch reflects the new names. + # Cached before a rename: names may be stale. Drop Hex-Rays' + # cache so the refetch reflects the new names. try: - self.client.call("force_recompile", items=[{"addr": hex(ea)}]) + self.client.invoke("force_recompile", items=[{"addr": hex(ea)}]) except Exception: # noqa: BLE001 pass # Bound the decompile: a function Hex-Rays can't handle tends to stall @@ -1587,7 +1568,10 @@ class Program: # rpcclient socket timeout, and cache the failure below so a re-request # returns instantly instead of re-grinding. try: - envelope = self.client.call_envelope( + # Code Mode returns the complete JSON result directly; unlike the + # old MCP tool transport there is no structured-content envelope or + # out-of-band download URL to unwrap. + payload = self.client.invoke( "decompile", addr=hex(ea), timeout=DECOMPILE_TIMEOUT ) except Exception as e: # noqa: BLE001 -- surface as a failed decompile @@ -1595,15 +1579,6 @@ class Program: with self._lock: self._decomp[ea] = (dec, self._name_gen) return dec - result = envelope.get("result", {}) - payload = result.get("structuredContent") - if payload is None: # fall back to text content - payload = self.client._extract_payload("decompile", result) - meta = (result.get("_meta") or {}).get("ida_mcp") - if isinstance(meta, dict) and meta.get("download_url"): - full = self._fetch_output(meta["download_url"]) - if isinstance(full, dict) and full.get("code"): - payload = full dec = _parse_decompilation(ea, payload) with self._lock: self._decomp[ea] = (dec, self._name_gen) @@ -1681,18 +1656,18 @@ class Program: Undefine first so it works even when the bytes are currently part of a data/align item — ``create_insn`` refuses to carve into a live item.""" try: - self.client.call("undefine", items=[{"addr": hex(ea)}]) + self.client.invoke("undefine", items=[{"addr": hex(ea)}]) except IDAToolError: pass # nothing defined here yet -> just try to create the insn res = self._first_result( - self.client.call("define_code", items=[{"addr": hex(ea)}])) + self.client.invoke("define_code", items=[{"addr": hex(ea)}])) if res.get("error"): raise IDAToolError("define_code", f"@ {ea:#x}: {res['error']}") def decomp_error(self, ea: int) -> str: """Hex-Rays' own reason for refusing ``ea``, or "" if it won't say.""" try: - r = self.client.call("decomp_error", addr=hex(ea)) + r = self.client.invoke("decomp_error", addr=hex(ea)) except IDAToolError: return "" if not isinstance(r, dict): @@ -1711,7 +1686,7 @@ class Program: def thumb_scan(self, start: int, end: int, apply: bool = True) -> dict: """Find Thumb entry points from odd pointers in ``[start, end)``.""" - r = self.client.call("thumb_scan", start=hex(start), end=hex(end), + r = self.client.invoke("thumb_scan", start=hex(start), end=hex(end), apply=bool(apply)) if not isinstance(r, dict) or r.get("error"): raise IDAToolError("thumb_scan", @@ -1720,7 +1695,7 @@ class Program: def set_thumb(self, ea: int, mode: str = "toggle") -> dict: """Switch ARM/Thumb decoding at ``ea``. Returns the resulting state.""" - r = self.client.call("set_thumb", addr=hex(ea), mode=mode) + r = self.client.invoke("set_thumb", addr=hex(ea), mode=mode) if not isinstance(r, dict) or r.get("error"): raise IDAToolError("set_thumb", f"@ {ea:#x}: {(r or {}).get('error', 'failed')}") @@ -1729,11 +1704,11 @@ class Program: def define_code_run(self, ea: int, limit: int = 20000) -> dict: """Disassemble consecutively from ``ea`` until something stops it. - Falls back to a single instruction when the worker predates the tool, so - an old worker degrades to the previous behaviour instead of failing. + Falls back to a single instruction for alternate clients that do not + provide the run operation. """ try: - r = self.client.call("define_code_run", addr=hex(ea), limit=int(limit)) + r = self.client.invoke("define_code_run", addr=hex(ea), limit=int(limit)) except IDAToolError: self.define_code(ea) return {"count": 1, "stopped": "single", "end": hex(ea)} @@ -1745,14 +1720,14 @@ class Program: def define_func(self, ea: int) -> dict: """Create a function starting at ``ea`` (IDA's 'p'). - Prefers the injected tool, which works out the end when IDA can't; - falls back to the plain one for an older worker. + Prefers the Code Mode operation, which works out the end when IDA can't; + falls back to a plain create for alternate clients. """ try: - r = self.client.call("define_func_run", addr=hex(ea)) + r = self.client.invoke("define_func_run", addr=hex(ea)) except IDAToolError: res = self._first_result( - self.client.call("define_func", items=[{"addr": hex(ea)}])) + self.client.invoke("define_func", items=[{"addr": hex(ea)}])) if res.get("error"): raise IDAToolError("define_func", f"@ {ea:#x}: {res['error']}") return {"ok": True, "how": "legacy"} @@ -1766,7 +1741,7 @@ class Program: item: dict = {"addr": hex(ea)} if size: item["size"] = int(size) - res = self._first_result(self.client.call("undefine", items=[item])) + res = self._first_result(self.client.invoke("undefine", items=[item])) if res.get("error"): raise IDAToolError("undefine", f"@ {ea:#x}: {res['error']}") @@ -1776,7 +1751,7 @@ class Program: item: dict = {"addr": hex(ea), "type": type_decl} if name: item["name"] = name - res = self._first_result(self.client.call("make_data", items=[item])) + res = self._first_result(self.client.invoke("make_data", items=[item])) if res.get("ok") is False or res.get("error"): raise IDAToolError( "make_data", f"@ {ea:#x}: {res.get('error') or 'rejected'}") @@ -1784,7 +1759,7 @@ class Program: def make_string(self, ea: int, length: int = 0, kind: str = "c") -> str: """Create a string literal at ``ea`` (IDA's 'A'); auto-length when 0. Returns the decoded contents.""" - r = self.client.call("make_string", addr=hex(ea), length=int(length), kind=kind) + r = self.client.invoke("make_string", addr=hex(ea), length=int(length), kind=kind) res = r if isinstance(r, dict) else {} if not res.get("ok"): raise IDAToolError( @@ -1801,7 +1776,7 @@ class Program: ``cycle``/``back`` (step the stops that make sense for this value) or a format by name. ``show`` reports without changing anything. """ - r = self.client.call("op_format", addr=hex(ea), mode=str(mode), + r = self.client.invoke("op_format", addr=hex(ea), mode=str(mode), col=int(col), n=int(n)) res = r if isinstance(r, dict) else {} if res.get("error"): @@ -1825,7 +1800,7 @@ class Program: if hit is not None and hit[1] == gen: return hit[0] try: - r = self.client.call("pc_nums", addr=hex(fn_ea)) + r = self.client.invoke("pc_nums", addr=hex(fn_ea)) except Exception: # noqa: BLE001 -- an older worker hasn't got the tool r = {} out: dict[int, list[tuple[int, int, str, int, int]]] = {} @@ -1848,7 +1823,7 @@ class Program: listing's format doesn't reach the pseudocode and vice versa, so this is a separate call rather than a flag on ``op_format``. """ - r = self.client.call("pc_num_format", addr=hex(fn_ea), mode=str(mode), + r = self.client.invoke("pc_num_format", addr=hex(fn_ea), mode=str(mode), line=int(line), col=int(col)) res = r if isinstance(r, dict) else {} if res.get("error"): @@ -1865,18 +1840,6 @@ class Program: sec = None return f"{sec} @ {ea:#x}" if sec else f"<no function> @ {ea:#x}" - @staticmethod - def _fetch_output(url: str, timeout: float = 15.0): - """GET the server's cached full-output blob (plain HTTP, not MCP).""" - try: - with urllib.request.urlopen(url, timeout=timeout) as r: - return json.loads(r.read().decode("utf-8", "replace")) - except Exception as e: # noqa: BLE001 -- fall back to the truncated preview - # The user gets CLIPPED pseudocode with no indication that a fetch - # failed rather than the function genuinely being that short. - diag.note(f"decompile: full-body fetch {url}", e) - return None - def strings(self, min_len: int = 4, refresh: bool = False) -> list[StrLit]: """Every string literal in the binary (IDA's Shift+F12 list), paged in full and cached. ``[]`` if the tool is unavailable.""" @@ -1889,7 +1852,7 @@ class Program: offset, page = 0, 2000 while True: try: - payload = self.client.call( + payload = self.client.invoke( "list_strings", offset=offset, count=page, min_len=min_len, refresh=(refresh and offset == 0)) except IDAToolError: @@ -1914,13 +1877,13 @@ class Program: def linkage(self) -> tuple[list[Linkage], list[Linkage]]: """``(imports, exports)`` for this binary, cached. ``([], [])`` if the - tool is unavailable — an old worker must not break the caller.""" + operation is unavailable — an alternate client must not break the caller.""" with self._lock: hit = self._linkage if hit is not None: return hit try: - payload = self.client.call("list_linkage", kind="both") + payload = self.client.invoke("list_linkage", kind="both") except IDAToolError: return ([], []) if not isinstance(payload, dict): @@ -1951,7 +1914,7 @@ class Program: if hit is not None and hit[1] == gen: return hit[0] try: - payload = self.client.call("decomp_map", addr=hex(ea)) + payload = self.client.invoke("decomp_map", addr=hex(ea)) except IDAToolError: return [] lines = payload.get("lines", []) if isinstance(payload, dict) else [] @@ -1981,7 +1944,7 @@ class Program: if hit is not None and hit[1] == gen: return hit[0] try: - payload = self.client.call("flowchart", addr=hex(ea)) + payload = self.client.invoke("flowchart", addr=hex(ea)) except IDAToolError: return None if not isinstance(payload, dict) or payload.get("error"): @@ -2054,7 +2017,7 @@ class Program: for _ in range(64): # bounded: ~128k heads if addr >= hi: break - payload = self.client.call("heads", addr=hex(addr), end=hex(hi), + payload = self.client.invoke("heads", addr=hex(addr), end=hex(hi), count=2000) rows = payload.get("heads", []) if isinstance(payload, dict) else [] if not rows: @@ -2083,13 +2046,13 @@ class Program: # -- cross-references & containing function --------------------------- # def function_of(self, ea: int) -> Func | None: """Return the function containing ``ea`` (resolves mid-function addrs).""" - payload = self.client.call("lookup_funcs", queries=[hex(ea)]) + payload = self.client.invoke("lookup_funcs", queries=[hex(ea)]) res = payload.get("result", []) if isinstance(payload, dict) else [] fn = res[0].get("fn") if res and isinstance(res[0], dict) else None return Func.from_raw(fn) if fn else None def xrefs_from(self, ea: int) -> list[Xref]: - payload = self.client.call( + payload = self.client.invoke( "xref_query", queries=[{"addr": hex(ea), "direction": "from", "include_fn": True}], ) @@ -2101,9 +2064,9 @@ class Program: try: # xref_types adds a fine-grained `kind` (call/read/write/...) for the # xref dialog; fall back to xref_query (code/data only) if absent. - payload = self.client.call("xref_types", queries=q) + payload = self.client.invoke("xref_types", queries=q) except IDAToolError: - payload = self.client.call("xref_query", queries=q) + payload = self.client.invoke("xref_query", queries=q) return _parse_xrefs(payload) # -- address resolution ------------------------------------------------ # @@ -2121,17 +2084,17 @@ class Program: # (loc_/locret_): lookup_funcs would map a label to its *containing* # function's entry, so double-clicking a label jumped to the wrong place. try: - payload = self.client.call("resolve_names", queries=[s]) + payload = self.client.invoke("resolve_names", queries=[s]) res = payload.get("result", []) if isinstance(payload, dict) else [] ea = res[0].get("ea") if res and isinstance(res[0], dict) else None if ea: return _as_int(ea) except IDAToolError: - pass # older server without resolve_names -> fall back below + pass # alternate client without resolve_names -> fall back below # Fall back to function-name resolution (also drives the 'did you mean' # suggestion when the name is unknown). try: - payload = self.client.call("lookup_funcs", queries=[s]) + payload = self.client.invoke("lookup_funcs", queries=[s]) except IDAToolError as e: raise KeyError(f"cannot resolve {target!r}: {e}") from e res = payload.get("result", []) if isinstance(payload, dict) else [] @@ -2169,7 +2132,7 @@ class Program: """Set (empty text clears) the comment at ``ea``; affects both the disasm and decompiler views. Returns the raw payload so the caller can surface a soft per-item error. The caller must invalidate/recompile to see it.""" - return self.client.call("set_comments", items=[{"addr": hex(ea), "comment": text}]) + return self.client.invoke("set_comments", items=[{"addr": hex(ea), "comment": text}]) # -- invalidation (after edits) --------------------------------------- # def invalidate(self, ea: int) -> None: diff --git a/idatui/drive.py b/idatui/drive.py index 1c7d8c2..825111c 100644 --- a/idatui/drive.py +++ b/idatui/drive.py @@ -121,7 +121,8 @@ def cmd_pc(c, args): lines = d["code"].splitlines() if needle: nlow = needle.lower() - lines = [f"{i:4} {l}" for i, l in enumerate(lines) if nlow in l.lower()] + lines = [f"{i:4} {line}" for i, line in enumerate(lines) + if nlow in line.lower()] return "\n".join(lines) or f"(no line matches {needle!r})" return d["code"] diff --git a/idatui/edit_ctl.py b/idatui/edit_ctl.py index 435fab9..80e1109 100644 --- a/idatui/edit_ctl.py +++ b/idatui/edit_ctl.py @@ -242,7 +242,7 @@ class EditController: kind = "stack" batch = {"stack": {"func_addr": hex(cur.ea), "old": old, "new": new}} try: - res = prog.client.call("rename", batch=batch) + res = prog.client.invoke("rename", batch=batch) except IDAToolError as e: app.call_from_thread(app._status, f"rename failed: {e.message}") return @@ -274,7 +274,7 @@ class EditController: app = self.app assert app.program is not None try: - res = app.program.client.call( + res = app.program.client.invoke( "rename", batch={"data": {"addr": hex(addr), "new": name}}) except IDAToolError as e: app.call_from_thread(app._status, f"name failed: {e.message}") diff --git a/idatui/errors.py b/idatui/errors.py index 29b09ae..aaf2dc5 100644 --- a/idatui/errors.py +++ b/idatui/errors.py @@ -1,10 +1,8 @@ -"""Transport-agnostic error hierarchy and the Session model. +"""TUI-facing error hierarchy and lightweight database session model. -These were originally defined in client.py (the ida-pro-mcp HTTP client), but the -idalib worker path (worker_client / domain / app) needs the same exception types -and Session dataclass without dragging in the HTTP transport. They live here so -both backends share one definition; client.py re-exports them for backwards -compatibility with the (deprecated) mcp tooling and the stress tests. +The Code Mode adapter normalizes ``ida_codemode.client`` transport and execution +errors into these types so the domain and Textual layers do not depend on HTTP or +registry implementation details. """ from __future__ import annotations diff --git a/idatui/launch.py b/idatui/launch.py index 35221a1..0c53987 100644 --- a/idatui/launch.py +++ b/idatui/launch.py @@ -1,15 +1,13 @@ -"""One-shot launcher: ``ida-tui foo.elf`` and you're in the TUI. +"""One-shot launcher for the IDA Code Mode-backed TUI. -Spawns a private idalib worker (``idatui.worker``) that opens + auto-analyzes -THIS binary in its own process, talking to the TUI over a unix socket. No shared -supervisor, no HTTP: everything slow (open + analysis) happens behind the TUI's -loading overlay. +A path first resolves to a registered GUI database; when none matches, Code Mode +reuses or starts a managed idalib worker. With no path, a single registered +database is selected automatically. -Usage: +Usage:: - ida-tui /path/to/binary # open a binary and drive it - -Extras: --ttl, --no-keepalive, --rpc (all forwarded to the TUI). + ida-tui /path/to/binary + ida-tui # attach when exactly one database is registered """ from __future__ import annotations @@ -17,12 +15,6 @@ import argparse import os import sys -# The unpacked working-copy files IDA writes next to a `.i64` while a database is -# open. A hard-killed worker leaves them behind and the `.i64` then refuses to -# reopen ("Failed to open database"). Safe to delete when nothing holds the DB. -_LOCK_SUFFIXES = (".id0", ".id1", ".id2", ".nam", ".til") - - def _load_args(load: dict) -> str: """``load`` as IDA switches, for the single-binary path (no project ref). @@ -41,35 +33,19 @@ def _log(msg: str) -> None: print(f"ida-tui: {msg}", file=sys.stderr) -def _sweep_locks(binary: str) -> int: - """Remove stale unpacked DB files next to ``binary``. Returns how many. - - Never touches the ``.i64`` -- that is the real database, and nothing is - saved unless ``idb_save`` was called -- and never the input file itself. - The second guard is not theoretical: ``.til`` is both an unpacked-DB suffix - and the extension of an IDA type library, so ``ida-tui mylib.til`` swept its - own argument out of existence. Same for anything named ``*.id0``/``*.nam``. - """ - keep = os.path.abspath(binary) - stem = os.path.splitext(binary)[0] - n = 0 - for base in (binary, stem): # IDA may key on the full name or the stem - for suf in _LOCK_SUFFIXES: - victim = base + suf - if os.path.abspath(victim) == keep: - continue # that's what the user asked us to open - try: - os.remove(victim) - n += 1 - except OSError: - pass - return n +def _registered_databases() -> tuple[list[dict], list[dict]]: + """Ready and blocked Code Mode registrations, with normalized errors.""" + try: + from ida_codemode.registry import discover_instances + return discover_instances() + except Exception as exc: # discovery diagnostics belong at the CLI boundary + return [], [{"error": str(exc)}] def main(argv: list[str] | None = None) -> int: p = argparse.ArgumentParser( prog="ida-tui", - description="Open a binary in the IDA TUI (private idalib worker).") + description="Open a registered GUI or managed idalib database in the IDA TUI.") p.add_argument("binary", nargs="*", help="binary to open and analyze (several with --project " "creates/extends that project)") @@ -77,9 +53,9 @@ def main(argv: list[str] | None = None) -> int: help="open a multi-binary project (created from the given " "binaries if FILE doesn't exist)") p.add_argument("--ttl", type=int, default=1800, - help="worker idle-TTL seconds (default 1800)") + help="deprecated compatibility option (Code Mode uses leases)") p.add_argument("--no-keepalive", action="store_true", - help="do not run the keepalive heartbeat") + help="deprecated compatibility option (the lease is the heartbeat)") p.add_argument("--rpc", metavar="PATH", help="listen for RPC on this unix socket (puppeteer the TUI)") p.add_argument("--trace", metavar="FILE", @@ -94,7 +70,7 @@ def main(argv: list[str] | None = None) -> int: g.add_argument("--base", metavar="ADDR", help="load address, e.g. 0x8000000 (any base; NOT paragraphs)") g.add_argument("--ida-args", metavar="STR", dest="ida_args", - help="extra IDA command-line switches, passed through as-is") + help="legacy switches; only Code Mode-representable -p/-b/-T are accepted") args = p.parse_args(argv) load: dict = {} @@ -152,27 +128,42 @@ def main(argv: list[str] | None = None) -> int: _log(str(e)) return 2 else: - if len(args.binary) != 1: - _log("give exactly one binary, or use --project for several") + ready, blocked = _registered_databases() + if len(args.binary) > 1: + _log("give at most one binary, or use --project for several") return 2 - binary = os.path.abspath(os.path.expanduser(args.binary[0])) - if not os.path.isfile(binary): - _log(f"no such file: {binary}") + if args.binary: + binary = os.path.abspath(os.path.expanduser(args.binary[0])) + key = os.path.normcase(os.path.realpath(binary)) + registered = any( + key == os.path.normcase(os.path.realpath(str(item.get(field) or ""))) + for item in ready for field in ("exe_path", "idb_path") + if item.get(field) + ) + if not os.path.isfile(binary) and not registered: + _log(f"no such file or registered database: {binary}") + return 2 + elif len(ready) == 1: + item = ready[0] + binary = str(item.get("exe_path") or item.get("idb_path") or "") + _log(f"attaching to registered {item.get('backend')} database: {binary}") + elif not ready: + detail = f" ({blocked[0].get('error')})" if blocked else "" + _log(f"no registered Code Mode database; pass a binary path{detail}") return 2 - if not os.access(os.path.dirname(binary), os.W_OK): - _log(f"directory not writable (IDA writes a .i64 there): " - f"{os.path.dirname(binary)}") + else: + _log("several Code Mode databases are registered; pass one of these paths:") + for item in ready: + _log(f" {item.get('exe_path') or item.get('idb_path')} " + f"[{item.get('backend')}, {item.get('record_id')}]") return 2 - swept = _sweep_locks(binary) # a crashed worker can leave the DB wedged - if swept: - _log(f"cleared {swept} stale lock file(s) from a crashed worker") - # Hand off to the TUI (imported late so --help works without textual). It - # spawns the worker behind its loading overlay while auto-analysis runs. + # Hand off to the TUI (imported late so --help works without Textual). Code + # Mode discovery/opening happens behind its loading overlay. try: from .app import IdaTui except ImportError as e: - _log(f"the TUI needs textual; run with ~/ida-venv/bin/python ({e})") + _log(f"TUI dependencies are missing; run `uv sync` ({e})") return 1 # Ask the terminal about graphics support NOW: the query needs a reply from # stdin, and once Textual starts it reads stdin on its own thread and would diff --git a/idatui/pane.py b/idatui/pane.py index 2592581..c31a93c 100644 --- a/idatui/pane.py +++ b/idatui/pane.py @@ -23,9 +23,9 @@ per pane in the registry, so stop/list/capture/keys keep working across both python -m idatui.pane capture --pane <pane> python -m idatui.pane keys --pane <pane> Escape -Requires: running inside tmux or zellij. Each pane spawns its own private idalib -worker (no shared supervisor). Uses ~/ida-venv/bin/python for the TUI (needs -textual) unless --python / IDATUI_PYTHON says otherwise. +Requires: running inside tmux or zellij. Each pane leases a registered GUI or +shared managed idalib database through Code Mode. Uses ~/ida-venv/bin/python for +the TUI (needs textual) unless --python / IDATUI_PYTHON says otherwise. """ from __future__ import annotations @@ -33,7 +33,6 @@ import argparse import json import os import secrets -import signal import subprocess import sys import time @@ -251,35 +250,9 @@ def _pane_keys(pane: str, keys: list[str], mux: str | None = None) -> None: subprocess.run(["tmux", "send-keys", "-t", pane, *keys], check=True) -# --------------------------------------------------------------------------- # -# idalib worker reaping -# -# ``pane stop`` kills the TUI pane, but a hard-killed pane can leave its private -# idalib worker (idatui/worker.py) running. A worker is only *safe* to reap when -# no idatui pane is live (then every worker is orphaned), which avoids killing an -# in-use analyser. -# --------------------------------------------------------------------------- # -_WORKER_PATTERN = r"idatui/worker\.py" - - -def _worker_pids() -> list[int]: - """PIDs of our private per-pane idalib worker processes (idatui/worker.py), - never our own PID.""" - try: - out = subprocess.run(["pgrep", "-f", _WORKER_PATTERN], - capture_output=True, text=True) - except OSError: - return [] - me = os.getpid() - pids: list[int] = [] - for tok in out.stdout.split(): - try: - pid = int(tok) - except ValueError: - continue - if pid != me: - pids.append(pid) - return pids +# Code Mode owns database process lifetime: a closed pane drops its lease at the +# socket/kernel boundary and Code Mode decides whether a managed worker still +# has clients. There is nothing for the pane layer to reap. def _count_live_panes() -> int: @@ -288,20 +261,9 @@ def _count_live_panes() -> int: def _reap_orphan_workers(force: bool = False) -> int: - """Kill leaked idalib workers when it is safe (no live pane) or ``force``. - - Returns the number of workers signalled. Best-effort; never raises. - """ - if not force and _count_live_panes() > 0: - return 0 - reaped = 0 - for pid in _worker_pids(): - try: - os.kill(pid, signal.SIGKILL) - reaped += 1 - except OSError: - pass - return reaped + """Compatibility no-op: Code Mode workers are shared and lease-managed.""" + del force + return 0 # --------------------------------------------------------------------------- # @@ -332,16 +294,8 @@ def spawn(args) -> int: print(f"error: no such project: {project}", file=sys.stderr) return 2 - # Reap workers leaked by previously-stopped/crashed panes so we don't spawn - # into a full IDA_MCP_MAX_WORKERS (which makes the new TUI hang forever, - # never reaching ready). No-op while any pane is live. - reaped = _reap_orphan_workers() - if reaped: - print(f"reaped {reaped} orphaned idalib worker(s) before spawn", - file=sys.stderr) - - # the command the pane runs: the launcher spawns a private idalib worker for - # this binary and becomes the TUI, so kill-pane tears the whole thing down. + # The pane owns only the TUI. Code Mode's lease cleanup handles crashes; + # kill-pane must never reap a shared GUI/idalib database. if project is not None: # launch takes: --project FILE [binaries...]; extra binaries are added to # the project (and a missing project file is created from them). @@ -393,9 +347,8 @@ def _wait_ready(sock: str, timeout: float, pane: str, stuck_after: float = 45.0, mux: str | None = None) -> dict[str, Any]: """Poll the socket + ping until the TUI reports ready (or timeout). - Emits a one-time hint to stderr if it's still not ready after ``stuck_after`` - seconds, so a wedged idalib worker / full worker pool surfaces a diagnostic - instead of an unexplained silent hang. + Emits a one-time hint if Code Mode discovery/opening is still not ready after + ``stuck_after`` seconds. """ start = time.time() deadline = start + timeout @@ -416,9 +369,8 @@ def _wait_ready(sock: str, timeout: float, pane: str, warned = True why = ("RPC socket not created yet" if not os.path.exists(sock) else "TUI up but analysis not ready") - print(f"still waiting ({int(time.time() - start)}s): {why}. If this " - f"hangs, the idalib worker may be stuck — try " - f"`python -m idatui.pane reap`.", file=sys.stderr) + print(f"still waiting ({int(time.time() - start)}s): {why}. " + f"Check Code Mode registrations and worker logs.", file=sys.stderr) time.sleep(0.4) last = dict(last) last["ready"] = False @@ -515,13 +467,9 @@ def list_panes(args) -> int: def reap(args) -> int: - """Kill leaked idalib workers (safe when no pane is live; --force overrides).""" - live = _count_live_panes() - n = _reap_orphan_workers(force=args.force) - print(json.dumps({"reaped_workers": n, "live_panes": live, "forced": args.force})) - if n == 0 and not args.force and live > 0: - print(f"note: {live} live pane(s) — not reaping in-use workers; pass " - f"--force to reap anyway", file=sys.stderr) + """Deprecated no-op; shared Code Mode workers are managed by leases.""" + print(json.dumps({"reaped_workers": 0, "live_panes": _count_live_panes(), + "forced": args.force, "deprecated": True})) return 0 @@ -624,9 +572,8 @@ def main(argv: list[str]) -> int: ls.add_argument("--prune", action="store_true", help="drop dead panes (and their sockets)") ls.set_defaults(fn=list_panes) - rp = sub.add_parser("reap", help="kill leaked idalib workers (frees worker slots)") - rp.add_argument("--force", action="store_true", - help="reap even while panes are live (may kill an in-use analyser)") + rp = sub.add_parser("reap", help="deprecated no-op (Code Mode uses shared leases)") + rp.add_argument("--force", action="store_true", help=argparse.SUPPRESS) rp.set_defaults(fn=reap) cp = sub.add_parser("capture", help="print a pane's visible screen") diff --git a/idatui/pool.py b/idatui/pool.py index ae37c25..465dff2 100644 --- a/idatui/pool.py +++ b/idatui/pool.py @@ -1,23 +1,16 @@ -"""WorkerPool — keeps a live idalib worker per project binary, within a budget. +"""DatabasePool — LRU leases on Code Mode databases for a project. -One worker process holds exactly one database (idalib is single-DB and -main-thread-only), so a project with N binaries means up to N processes. They are -not cheap and they do not share: a worker on ``bash`` measures ~126 MB RSS / -117 MB PSS, and the database working set dominates for anything larger -(``libcrypto.so.3``'s ``.i64`` alone is 72 MB). +Code Mode may bind a lease to an existing IDA GUI or to a shared managed idalib +worker. The pool therefore owns *client interest*, never an IDA process. Releasing +an LRU entry persists managed IDBs but does not implicitly save a GUI, then closes +only this TUI's lease; other clients and GUI sessions remain alive. Managed workers exit themselves after their final lease. -Residency is therefore bounded by a **memory budget**, not a worker count — a -count is the wrong knob when one project holds both a 50 KB helper and a 6 MB -crypto library. Workers are spawned lazily on first use, kept resident while they -fit, and least-recently-used ones evicted when they don't. Eviction **saves the -database first**, so coming back is a load rather than a re-analysis. - -The pool never evicts the active binary, nor anything pinned. +The historical memory budget remains useful for managed idalib instances, while +GUI process memory is only advisory. The active and pinned databases are never +released to satisfy it. """ from __future__ import annotations -import os - from .project import BinaryRef, Project #: Fallback budget if /proc/meminfo can't be read (MB). @@ -36,11 +29,11 @@ def _total_ram_mb() -> int: def _pss_mb(pid: int | None) -> int: - """Proportional set size of a worker, in MB. + """Proportional set size of the leased instance process, in MB. - PSS (not RSS) is the honest per-worker cost: it splits shared pages between - the processes mapping them. In practice workers share very little, so the two - are close, but PSS is what makes summing across workers meaningful. + PSS is useful for managed idalib workers. For GUI/shared processes it is only + advisory because the TUI neither owns all that memory nor controls process + exit. """ if not pid: return 0 @@ -54,13 +47,19 @@ def _pss_mb(pid: int | None) -> int: return 0 -def _default_spawn(ref: BinaryRef, ttl: int): # pragma: no cover - needs idalib - from .worker_client import WorkerClient - return WorkerClient(ref.staged, ttl=ttl, load_args=ref.load_args) +def _default_spawn(ref: BinaryRef, ttl: int, *, new_database: bool = False): # pragma: no cover - needs IDA + from .codemode_client import CodeModeClient + return CodeModeClient( + ref.staged, + ttl=ttl, + load_args=ref.load_args, + output_database=ref.db, + new_database=new_database, + ) -class WorkerPool: - """Live workers for a project's binaries, keyed by label.""" +class DatabasePool: + """Live Code Mode database leases, keyed by project label.""" def __init__(self, project: Project, *, budget_mb: int | None = None, ttl: int = 1800, spawn=None, mem_fn=None) -> None: @@ -71,6 +70,7 @@ class WorkerPool: self._clients: dict[str, object] = {} self._lru: list[str] = [] # least-recently-used first self._pinned: set[str] = set() + self._recreate: set[str] = set() # Ctrl+L: next attachment creates a fresh IDB self.active: str | None = None # never evicted if budget_mb is None: ram = _total_ram_mb() @@ -98,11 +98,11 @@ class WorkerPool: # -- acquire ----------------------------------------------------------- # def get(self, label: str, progress=None): - """A live client for ``label``, spawning it (and making room) if needed. + """A live client for ``label``, attaching or spawning as needed. - Staging and the scratch sweep happen here: a worker killed hard last time - leaves unpacked ``.id0/.id1/...`` behind, and the database then refuses to - reopen. Nothing else holds this DB (one worker per label), so it is safe. + Do not sweep IDA scratch files here: a registered GUI or another Code + Mode client may own the database. Code Mode's registry locks and health + probes are the authority for safe discovery and stale-record cleanup. """ client = self._clients.get(label) if client is not None: @@ -118,28 +118,29 @@ class WorkerPool: note(f"staging {ref.label}\u2026") self.project.stage(ref) - self.project.sweep_scratch(ref) note(f"opening {ref.label}\u2026") - client = self._spawn(ref, self._ttl) + fresh = label in self._recreate + client = (_default_spawn(ref, self._ttl, new_database=fresh) + if self._spawn is _default_spawn else self._spawn(ref, self._ttl)) connect = getattr(client, "connect", None) if connect is not None: connect(progress=progress) if progress is not None else connect() self._clients[label] = client + self._recreate.discard(label) self._lru.append(label) self._enforce_budget(protect=label) return client def prewarm(self, label: str, progress=None) -> bool: - """Spawn a worker for ``label`` only if it fits the budget AS IT STANDS. + """Attach a database for ``label`` only if it fits the current budget. Pre-warming must never cost residency: evicting a binary the user actually visited to speculatively load one they haven't is a straight downgrade, and the eviction would also throw away that binary's caches. So this refuses rather than making room, and returns False. - The cost of a worker that doesn't exist yet can only be estimated; the - largest resident one is the best evidence available (they are all the - same program with a different database). With nothing resident we have + The cost of a database not attached yet can only be estimated; the + largest resident instance is the best evidence available. With nothing resident we have no evidence at all, so we allow one — that is the case where the budget is certainly free. """ @@ -153,13 +154,19 @@ class WorkerPool: return False self.get(label, progress=progress) # get() enforces the budget protecting the NEW label; if that had to - # evict, our estimate was wrong and the speculative worker is the one + # evict, our estimate was wrong and the speculative lease is the one # that should go — never a binary the user chose. if self.memory_mb() > self.budget_mb and label != self.active: self.evict(label) return False return True + def recreate_on_next_open(self, label: str) -> None: + """Request a fresh IDB after the current lease has been released.""" + if self.project.by_label(label) is None: + raise KeyError(f"no such binary in the project: {label}") + self._recreate.add(label) + def _touch(self, label: str) -> None: if label in self._lru: self._lru.remove(label) @@ -171,16 +178,21 @@ class WorkerPool: self._touch(label) # -- release ----------------------------------------------------------- # - def evict(self, label: str, save: bool = True) -> bool: - """Drop a resident worker, persisting its database first.""" + def evict(self, label: str, save: bool = True, + save_gui: bool = False) -> bool: + """Release a resident lease, persisting a managed database first. + + A budget-driven eviction must not save somebody's GUI implicitly. GUI + saves are reserved for an explicit/defensive ``close_all(save=True)``. + """ client = self._clients.pop(label, None) if client is None: return False if label in self._lru: self._lru.remove(label) - if save: + if save and (save_gui or getattr(client, "backend", None) != "gui"): try: # persist analysis + edits so the next open is a load - client.call("idb_save") + client.save_database() except Exception: # noqa: BLE001 -- evict regardless pass try: @@ -198,7 +210,7 @@ class WorkerPool: return None def _enforce_budget(self, protect: str | None = None) -> int: - """Evict LRU workers until the pool fits its budget. Returns how many.""" + """Release LRU leases until the pool fits its budget. Returns how many.""" n = 0 while self.memory_mb() > self.budget_mb: victim = self._evictable(protect) @@ -210,7 +222,7 @@ class WorkerPool: def close_all(self, save: bool = True) -> None: for label in list(self._clients): - self.evict(label, save=save) + self.evict(label, save=save, save_gui=save) self.active = None # -- introspection ------------------------------------------------------ # @@ -231,5 +243,9 @@ class WorkerPool: return out def __repr__(self) -> str: # pragma: no cover - debug aid - return (f"<WorkerPool {len(self._clients)}/{len(self.project.refs)} resident " + return (f"<DatabasePool {len(self._clients)}/{len(self.project.refs)} resident " f"{self.memory_mb()}/{self.budget_mb}MB active={self.active}>") + + +# Source compatibility for callers that imported the pre-Code-Mode name. +WorkerPool = DatabasePool diff --git a/idatui/project.py b/idatui/project.py index e2fc542..53f3b0a 100644 --- a/idatui/project.py +++ b/idatui/project.py @@ -23,7 +23,8 @@ firmware image, a cleaned build tree). A source whose size/mtime no longer matches the staged copy is re-staged, and its now-stale database is dropped (the DB describes the old bytes). -stdlib-only, like the domain/worker layers — the TUI is the only Textual consumer. +The model has no IDA imports. Staging consults ida_codemode's registry before +replacing files so it never mutates a database owned by a GUI/shared worker. """ from __future__ import annotations @@ -56,7 +57,7 @@ class BinaryRef: #: it is, a raw firmware image doesn't, and IDA defaults to metapc at 0. processor: str = "" # IDA processor name: arm, armb, mipsb, metapc, … base: int = 0 # load address (natural, e.g. 0x8000000) - ida_args: str = "" # escape hatch: extra IDA command-line switches + ida_args: str = "" # legacy -p/-b/-T switches accepted by Code Mode adapter @property def db(self) -> str: @@ -310,13 +311,25 @@ class Project: """Ensure ``ref`` is staged in the sidecar; returns the staged path. Re-staging a changed source drops its database: the DB describes the old - bytes, so keeping it would silently mismatch the disassembly (any renames - in it are lost, which is why callers should say so out loud). + bytes. Refuse while Code Mode reports a GUI/idalib owner; replacing a + staged executable or IDB underneath a shared live instance is corruption. """ if not os.path.isfile(ref.source): raise ProjectError(f"no such binary: {ref.source}") if not self.is_stale(ref): return ref.staged + try: + from .codemode_client import database_owner + owner = database_owner(ref.db, ref.staged) + except Exception as exc: + raise ProjectError( + f"cannot verify Code Mode ownership before staging {ref.label}: {exc}" + ) from exc + if owner is not None: + raise ProjectError( + f"cannot restage {ref.label}: Code Mode instance {owner.record_id} " + f"still owns {owner.idb_path}; close/release it first" + ) os.makedirs(self.bin_dir, exist_ok=True) tmp = ref.staged + ".staging" _unlink(tmp) @@ -338,10 +351,11 @@ class Project: return out def sweep_scratch(self, ref: BinaryRef) -> int: - """Delete IDA's unpacked working files (never the ``.i64``) for ``ref``. + """Delete unpacked working files (never the ``.i64``) for maintenance. - A hard-killed worker leaves them behind and the database then refuses to - reopen. Only safe when no worker holds it. + Runtime paths no longer call this: Code Mode instances are shared, so a + registry owner may still be using these files. Callers must independently + prove that no GUI/idalib instance owns the database. """ return sum(1 for suf in SCRATCH_SUFFIXES if _unlink(ref.staged + suf)) diff --git a/server/patch_server.py b/idatui/remote_tools.py index 6667e12..6fb6436 100644 --- a/server/patch_server.py +++ b/idatui/remote_tools.py @@ -1,285 +1,154 @@ -#!/usr/bin/env python3 -"""Inject idatui's extra ida-pro-mcp tools into the installed server package. +"""The IDAPython ida-tui runs inside the Code Mode sandbox. -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. +Two features have no ida-domain surface at all and are carried over VERBATIM +from the tools ida-tui was developed against (`server/patch_server.py`'s +injected BODY, which the Code Mode port deletes): -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. +* `heads` -- the continuous listing. ida-domain enumerates defined heads and + renders plain disassembly; the listing also needs coalesced undefined runs, + IDA colour-tag spans, PER-OPERAND EXTENTS, function banners, code labels and + expanded struct members, plus the digest/`expect` protocol the paging layer + uses to skip re-sending a page that has not changed. +* `op_format` / `pc_nums` / `pc_num_format` -- `o`/`O`. IDA's operand types and + Hex-Rays' per-(ea, opnum) numforms are separate sets, and neither is exposed. -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). +Keeping the originals rather than paraphrasing them is deliberate: this is the +most performance-tuned and most behaviour-sensitive code in the project (the +span walker is a single regex pass because a per-character loop was the most +expensive thing the listing did, and the cycle only offers stops that change +what you see). A re-implementation drifts from it silently. -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. +This file is SOURCE SHIPPED AS TEXT to the database process; it is never +imported here, because the ida_* modules do not exist in the TUI's interpreter. +`codemode_client` reads it and prepends it to the relevant snippets. Keep it +self-contained: no relative imports, nothing beyond what Code Mode provides. """ -from __future__ import annotations +# ruff: noqa +import re as _re -import importlib.util -import pathlib -import sys +from typing import Annotated # the extracted tool signatures still carry these -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 +class IDAError(Exception): + """The MCP host's error type; the tools raise/catch it by name.""" -@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 +def parse_address(addr): + """ida_pro_mcp.utils.parse_address: hex/decimal string, int, or a symbol.""" + if isinstance(addr, int): + return addr + try: + return int(addr, 0) + except ValueError: + import idaapi + ea = idaapi.get_name_ea(idaapi.BADADDR, str(addr).strip()) + if ea != idaapi.BADADDR: + return ea + raise IDAError(f"Not found: {addr!r}") -@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} +#: Byte-identical to ida_pro_mcp.utils._STRING_OR_SPACES_RE: the pseudocode +#: column coordinates the client holds depend on collapsing exactly the same way. +_IDATUI_STRING_OR_SPACES_RE = _re.compile( + r'"(?:[^"\\]|\\.)*"' # double-quoted string + r"|'(?:[^'\\]|\\.)*'" # single-quoted string / char + r"|[ \t]{2,}" # run of 2+ whitespace (outside strings) +) -@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 compact_whitespace(line: str) -> str: + """ida_pro_mcp.utils.compact_whitespace: collapse runs of 2+ spaces/tabs to + one, preserving string literals.""" + stripped = line.lstrip(" \t") + if not stripped: + return line + lead = line[: len(line) - len(stripped)] - def _tstr(tif): - try: - s = tif.dstr() - if s: - return s - except Exception: - pass - return str(tif) + def _repl(m): + s = m.group() + if s[0] in ('"', "'"): + return s # preserve string content + return " " - 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, - } + return lead + _IDATUI_STRING_OR_SPACES_RE.sub(_repl, stripped) -@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 +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. - 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} + ``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 -@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} +import functools as _idatui_functools -@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 +import os as _idatui_os - 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} +_IDATUI_LINE_CACHE = int(_idatui_os.environ.get("IDATUI_LINE_CACHE") or 65536) -@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. +def _idatui_line_parts(line): + """``(text, spans, ops)`` for one tagged disassembly line -- memoised. - 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 + 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. - 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)} + ``spans`` is None when the tag walk and the plain text disagree about what + the line says (then the text wins and the row renders unhighlighted). + + The returned lists are SHARED between every row that has the same line; + treat them as read-only. Pickle notices the sharing too, so a page of + repetitive disassembly also serialises smaller. + """ + import ida_lines + text = " ".join(ida_lines.tag_remove(line).split()) # collapse the padding + spans, ops = _idatui_spans(line) + # Built from the SAME line as `text`, then whitespace-collapsed identically, + # so the two can never disagree about what the row says. + joined = "".join([t for _k, t in spans]) + if " ".join(joined.split()) != text: + return (text, None, None) + return (text, spans, ops) def _idatui_head_row(ea, flags=None): @@ -323,23 +192,11 @@ def _idatui_head_row(ea, flags=None): 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) @@ -373,11 +230,6 @@ def _idatui_line_parts(line): 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",), @@ -412,11 +264,14 @@ def _idatui_tag_map(): _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 @@ -691,8 +546,6 @@ def _idatui_func_footer_rows(ea, func): ] -@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, @@ -841,634 +694,9 @@ def heads( 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") @@ -1693,8 +921,6 @@ def _idatui_apply_fmt(ea, n, fmt): 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", @@ -1819,27 +1045,31 @@ def op_format( 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) + a column in what the client SHOWS can be mapped back to Hex-Rays' line. + + DEVIATION FROM THE EXTRACTED ORIGINAL, deliberately: this used to be + ``from ida_pro_mcp.ida_mcp.utils import compact_whitespace`` inside a + try/except, with a plain ``[ \\t]{2,}`` regex as the fallback. Under Code + Mode ida_pro_mcp is not installed in the database process, so BOTH halves + of that were wrong: + + * the import failed on every call, and a failed import is never cached, so + each one re-searched the whole of sys.path -- 422 failures per pc_nums + call, which was the majority of its runtime; + * the fallback collapses runs of spaces INSIDE STRING LITERALS, which the + real function preserves. Pseudocode columns are served in these + coordinates, so a line containing a string with two spaces would have put + every literal's mark, and every reformat, on the wrong column. + + The module-level shim above is byte-identical to the original regex, so + call it directly. + """ + return compact_whitespace(line) def _idatui_compact_col(plain, compact, col): @@ -1865,8 +1095,6 @@ def _idatui_uncompact_col(plain, compact, col): 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") @@ -1960,8 +1188,6 @@ def _idatui_pc_nums(cf, sl): return out -@tool -@idasync def pc_nums( addr: Annotated[str, "Function address (or any address inside it)"], ) -> dict: @@ -2009,8 +1235,6 @@ def pc_nums( 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", @@ -2154,105 +1378,171 @@ def pc_num_format( 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``. +def decompile(addr, include_addresses=True): + """Pseudocode for the function at ``addr``, plus the objects it references. - 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. + Faithful to the tool ida-tui was written against, and in particular to its + COST: the per-line address anchor comes from ONE ``get_line_item`` at column + 0 per line. The Code Mode port asked for the full per-column line map (what + ``decomp_map`` is for) purely to fill in that anchor, which is thousands of + ``get_line_item``+``dstr()`` calls per function instead of one per line, and + made every pseudocode open cost the same as opening the split view. - 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 + Text is whitespace-collapsed exactly as the client displays it, because + ``pc_nums`` reports literal columns in those coordinates. """ - import ida_funcs - import ida_gdl + import ida_bytes + import ida_hexrays + import ida_lines + import ida_name + import idaapi try: ea = parse_address(addr) except Exception as e: - return {"addr": str(addr), "error": str(e), "blocks": []} - fn = ida_funcs.get_func(ea) + return {"addr": str(addr), "code": None, "error": str(e)} + fn = idaapi.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, - } -''' + return {"addr": str(addr), "code": None, "error": f"no function at {ea:#x}"} + if not ida_hexrays.init_hexrays_plugin(): + return {"addr": hex(int(fn.start_ea)), "code": None, "error": "no decompiler"} + failure = ida_hexrays.hexrays_failure_t() + try: + cfunc = ida_hexrays.decompile_func(fn, failure) + except Exception as e: + return {"addr": hex(int(fn.start_ea)), "code": None, + "error": f"Decompilation failed at {ea:#x}: {e}"} + if cfunc is None: + return {"addr": hex(int(fn.start_ea)), "code": None, + "error": failure.desc() or f"Decompilation failed at {ea:#x}"} -SNIPPET = f"{BEGIN}\n{BODY.strip()}\n{END}\n" + lines = [] + for sl in cfunc.get_pseudocode(): + head = ida_hexrays.ctree_item_t() + item = ida_hexrays.ctree_item_t() + tail = ida_hexrays.ctree_item_t() + line_ea = None + if include_addresses and cfunc.get_line_item(sl.line, 0, False, head, item, tail): + parts = (item.dstr() or "").split(": ") + if len(parts) == 2: + try: + line_ea = int(parts[0], 16) + except ValueError: + line_ea = None + text = compact_whitespace(ida_lines.tag_remove(sl.line)) + lines.append(f"{text} /*{line_ea:#x}*/" if line_ea is not None else text) + refs, seen = [], set() -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 + class _RefVisitor(ida_hexrays.ctree_visitor_t): + def __init__(self): + ida_hexrays.ctree_visitor_t.__init__(self, ida_hexrays.CV_FAST) + def visit_expr(self, e): + if e.op == ida_hexrays.cot_obj: + target = int(e.obj_ea) + if target != idaapi.BADADDR and target not in seen: + seen.add(target) + try: + raw = ida_bytes.get_strlit_contents(target, -1, 0) + text = raw.decode("utf-8", "replace") if raw else None + except Exception: + text = None + refs.append({"addr": hex(target), + "name": ida_name.get_name(target) or "", + "string": text}) + return 0 -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 + _RefVisitor().apply_to(cfunc.body, None) + except Exception: + pass + return {"addr": hex(int(fn.start_ea)), "code": "\n".join(lines), "refs": refs} -if __name__ == "__main__": - raise SystemExit(main()) +def decomp_map( + addr: Annotated[str, "Function address or name"], +) -> dict: + """Per-pseudocode-line instruction coverage for the split view's region + highlight: for each line, the set of EAs the decompiler attributes to it, + swept across the line's columns via get_line_item. Shape: + {addr, lines:[{ea: primary|None, eas:[hex,...]}, ...]}.""" + import ida_hexrays + import idaapi + try: + ea = int(str(addr), 16) + except ValueError: + ea = idaapi.get_name_ea(idaapi.BADADDR, str(addr).strip()) + func = idaapi.get_func(ea) + if not func: + return {"error": f"no function at {addr}"} + try: + cfunc = ida_hexrays.decompile(func.start_ea) + except Exception as e: # noqa: BLE001 + return {"error": f"decompile failed: {e}"} + if cfunc is None: + return {"error": "decompile failed"} + import ida_lines + # Three things this loop must not do, each measured on real functions (the 25 + # largest of bash went 68.3s -> 6.5s; echo's 60 largest 5.4s -> 0.6s, with + # byte-identical output): + # + # * allocate ctree_item_t's per COLUMN. They are SWIG objects and this is + # the innermost loop; one per call is enough, and head/tail are never + # read, so don't ask for them at all. + # * sweep the TAGGED length. ``x`` is a screen column but ``sl.line`` still + # carries IDA's colour tags, so a 23-column line was swept 124 times. + # * call dstr() per column. It formats a whole 'EA: description' string -- + # 24us a call, which is 79% of this tool. Comparing against the PREVIOUS + # column's item id is not enough: items interleave, so `foo(a, b)` flips + # call -> arg -> call -> arg and every flip re-formats an item already + # seen (106 594 calls for 15 417 lines of bash). Memoise id -> ea for the + # whole function instead: obj_id is unique within a cfunc, so the same id + # always yields the same string, and the result is deduped by ``seen`` + # anyway. Items with no ctree node (it is None) have no id to key on and + # still pay per occurrence. + item = ida_hexrays.ctree_item_t() + tag_remove = ida_lines.tag_remove + get_line_item = cfunc.get_line_item + ea_of_id = {} + lines = [] + for sl in cfunc.get_pseudocode(): + line = sl.line + eas, seen = [], set() + prev_id = None + for x in range(len(tag_remove(line)) + 1): + if not get_line_item(line, x, False, None, item, None): + continue + it = item.it + if it is not None: + oid = it.obj_id + if oid == prev_id: + continue + prev_id = oid + if oid in ea_of_id: + e = ea_of_id[oid] + if e is not None and e not in seen: + seen.add(e) + eas.append(hex(e)) + continue + else: + oid = None + prev_id = None + # Match the /*ea*/ marker's source (decompile_function_safe): the + # item's dstr() is 'EA: description'; get_ea() reports a different ea. + e = None + dstr = item.dstr() + if dstr: + parts = dstr.split(": ", 1) + if len(parts) == 2: + try: + e = int(parts[0], 16) + except ValueError: + e = None + if oid is not None: + ea_of_id[oid] = e + if e is not None and e not in seen: + seen.add(e) + eas.append(hex(e)) + lines.append({"ea": eas[0] if eas else None, "eas": eas}) + return {"addr": hex(func.start_ea), "lines": lines} diff --git a/idatui/rpc.py b/idatui/rpc.py index 698e4cf..225aeb2 100644 --- a/idatui/rpc.py +++ b/idatui/rpc.py @@ -757,7 +757,7 @@ class RpcServer: batch = {"func": ops, "allow_overwrite": bool(overwrite)} # The worker is blocking and single-threaded; off the event loop it goes, # or the TUI freezes for the length of the batch. - res = await asyncio.to_thread(app.program.client.call, "rename", batch=batch) + res = await asyncio.to_thread(app.program.client.invoke, "rename", batch=batch) summary = res.get("summary", {}) if isinstance(res, dict) else {} failed = [r for r in (res.get("func") or []) if isinstance(r, dict) and r.get("error")] if isinstance(res, dict) else [] @@ -768,7 +768,7 @@ class RpcServer: # this a batch import leaves pseudocode calling sub_98C0 forever while # the listing (and every readback) says memset. try: - await asyncio.to_thread(app.program.client.call, "force_recompile") + await asyncio.to_thread(app.program.client.invoke, "force_recompile") except Exception: # noqa: BLE001 -- older worker without the tool pass app.program.bump_names() diff --git a/idatui/worker.py b/idatui/worker.py deleted file mode 100644 index 556e69a..0000000 --- a/idatui/worker.py +++ /dev/null @@ -1,301 +0,0 @@ -"""idatui's own idalib worker — the replacement for the ida-pro-mcp supervisor. - -Opens ONE database in-process (on the main thread, as idalib requires) and -serves ida-pro-mcp's *tool functions* over a unix socket with length-prefixed -pickle. Same tool implementations as the MCP path (we call -``MCP_SERVER.tools.methods[name](**args)`` directly), so return shapes are -byte-identical — but with ~50us/call instead of the HTTP path's ~5ms, and no -supervisor / HTTP / JSON / 50KB-truncation machinery. - - python -m idatui.worker <sock_path> <binary_path> - -The socket only appears once the database is open + analyzed, so a client can -poll ``connect()`` to know when the worker is ready. Requests are served -serially on the main thread (idalib is single-threaded; every tool runs inline -through its own execute_sync, which is a no-op on the main thread). - -Protocol (both directions length-prefixed: 4-byte big-endian len + pickle): - request = (tool_name: str, kwargs: dict) - response = (ok: bool, result_or_error) - tool_name == "__shutdown__" ends the worker. -""" -from __future__ import annotations - -import os -import pickle -import socket -import struct -import sys -import threading -import time -import uuid - -#: Seconds a single tool call may run before it is cancelled. 0 disables the -#: deadline entirely. -TOOL_TIMEOUT_SEC = float(os.environ.get("IDATUI_TOOL_TIMEOUT_SEC") or 60) - -# ida-pro-mcp enforces its own tool deadline by installing a `sys.setprofile` -# hook for the duration of every call, so that a pure-python loop inside a tool -# body can be interrupted. That hook runs a python function on EVERY python call -# and return -- and our tools are exactly the call-heavy kind: `heads` renders -# hundreds of items per request and measured 92us/row with the hook against -# 28us/row without it. A 3.3x tax on the whole backend to bound loops that are -# already bounded by their `count` argument. -# -# So: turn the upstream mechanism off and re-arm the half that does the real -# work ourselves (see _Deadline). ida_kernwin.set_cancelled() is what actually -# frees the IDA main thread -- decompile, auto_wait, find_bytes and friends poll -# user_cancelled() and bail within a poll cycle -- and it costs nothing until it -# fires. -os.environ["IDA_MCP_TOOL_TIMEOUT_SEC"] = "0" - - -class _Deadline: - """A single watchdog thread that cancels a tool call which overruns. - - Arming is two attribute writes, because it is on the path of every call the - TUI makes (a scroll is dozens of them). The watchdog polls instead of being - signalled for the same reason: waking a thread per call costs more than the - 0.25s of granularity it buys on a 60s deadline. - """ - - TICK = 0.25 - - def __init__(self, seconds: float) -> None: - import ida_kernwin - self._kernwin = ida_kernwin - self.seconds = seconds - self._until: float | None = None - t = threading.Thread(target=self._run, name="idatui-deadline", - daemon=True) - t.start() - - def _run(self) -> None: - while True: - time.sleep(self.TICK) - until = self._until - if until is not None and time.monotonic() >= until: - self._until = None - # THREAD_SAFE in the IDA SDK; upstream fires it off a Timer too. - self._kernwin.set_cancelled() - - def arm(self) -> None: - # Clear unconditionally: the flag is sticky, and one left set would make - # every later user_cancelled() true forever. - self._kernwin.clr_cancelled() - self._until = time.monotonic() + self.seconds - - def disarm(self) -> None: - self._until = None - - -# --------------------------------------------------------------------------- # -# framing -# --------------------------------------------------------------------------- # -def _recvn(sock: socket.socket, n: int) -> bytes | None: - buf = bytearray() - while len(buf) < n: - chunk = sock.recv(n - len(buf)) - if not chunk: - return None - buf += chunk - return bytes(buf) - - -def send(sock: socket.socket, obj) -> None: - data = pickle.dumps(obj, protocol=pickle.HIGHEST_PROTOCOL) - sock.sendall(struct.pack(">I", len(data)) + data) - - -def recv(sock: socket.socket): - hdr = _recvn(sock, 4) - if hdr is None: - return None - (n,) = struct.unpack(">I", hdr) - body = _recvn(sock, n) - return None if body is None else pickle.loads(body) - - -# --------------------------------------------------------------------------- # -# worker -# --------------------------------------------------------------------------- # -def _ensure_tools_injected() -> None: - """Inject idatui's custom tools (heads/read_raw/resolve_names/func_types/...) - into the installed ida_pro_mcp, idempotently, so the worker is self-sufficient - (nothing else has to inject these tools first). Must run BEFORE - ida_pro_mcp.ida_mcp is imported (the injected code lives in api_types.py).""" - import importlib.util - repo = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) - patch = os.path.join(repo, "server", "patch_server.py") - if not os.path.exists(patch): - return - try: - spec = importlib.util.spec_from_file_location("_idatui_patch", patch) - mod = importlib.util.module_from_spec(spec) - spec.loader.exec_module(mod) # IDA-free; just defines + patches api_types - mod.main() - except Exception as e: # noqa: BLE001 -- tools may already be present - sys.stderr.write(f"idatui: tool injection skipped: {e}\n") - - -def _has_database(binpath: str) -> bool: - """Whether IDA already has a database for ``binpath``. - - IDA names it ``<file>.i64`` (keeping the extension), but a database made - from ``foo.bin`` can also appear as ``foo.i64`` depending on how it was - created — check both, because guessing wrong here means re-passing load - switches to an existing database, which fails the open. - """ - return (os.path.exists(binpath + ".i64") - or os.path.exists(os.path.splitext(binpath)[0] + ".i64")) - - -def _open_and_register(binpath: str, load_args: str = ""): - """Open the DB (main thread) then import ida-pro-mcp so every @tool registers - against this live database. Returns (tools_dict, module_name, save_fn). - - ``load_args`` is passed to IDA as command-line switches, which is the only - way to tell it how to read a headerless blob: a raw firmware image has no - format to detect, so without ``-p<processor>`` it loads as metapc at 0 and - finds nothing. Ignored once a database exists — the .i64 already records how - it was loaded, and re-passing conflicting switches is how you corrupt one. - """ - _ensure_tools_injected() # before any ida_pro_mcp import - import idapro - idapro.enable_console_messages(False) - args = load_args or None - if args and _has_database(binpath): - # The .i64 already records how this image was loaded. Passing the - # switches again on reopen makes IDA fail outright (rc != 0) — the load - # options belong to the FIRST open only. - args = None - if idapro.open_database(binpath, run_auto_analysis=True, - args=args): # nonzero == failure - if args: - # With load switches in play they are the likeliest culprit by far: - # IDA refuses an unknown -p name with no diagnostic of its own, so - # saying "the database is locked" here sends people hunting a - # problem they don't have. - raise RuntimeError( - f"failed to open {binpath} with load options {args!r}: IDA " - f"rejected them \u2014 an unknown processor name is the usual " - f"cause (see tools/verify_procs.py for the valid ones)") - raise RuntimeError( - f"failed to open {binpath}: the .i64 is likely held by a running " - f"ida-mcp worker (try: pkill -f idalib) or wedged from a crash " - f"(delete its .id0/.id1/.id2/.nam/.til next to the binary)") - import ida_auto - ida_auto.auto_wait() # block until auto-analysis settles (match ida-mcp) - - # importing the package registers all api_*/patched tools against MCP_SERVER - from ida_pro_mcp.ida_mcp import MCP_SERVER # noqa: WPS433 - - import ida_nalt - module = os.path.basename(ida_nalt.get_root_filename() or binpath) - - def save(): - import idc - try: - idc.save_database(idc.get_idb_path(), 0) - except Exception: # noqa: BLE001 - import ida_loader, ida_pro # noqa: WPS433 - ida_loader.save_database(idc.get_idb_path(), 0) - - return MCP_SERVER.tools.methods, module, save - - -def serve(sockpath: str, binpath: str, load_args: str = "") -> None: - tools, module, save = _open_and_register(binpath, load_args) - sid = uuid.uuid4().hex[:8] - deadline = _Deadline(TOOL_TIMEOUT_SEC) if TOOL_TIMEOUT_SEC > 0 else None - - def dispatch(name: str, args: dict): - args = dict(args) - args.pop("database", None) # single-DB worker: no session routing - # session-management shims (were the supervisor's job): - if name in ("idb_open",): - return {"success": True, - "session": {"session_id": sid, "module": module, - "input_path": binpath}} - if name in ("idb_save", "save"): - save() - return {"success": True} - if name in ("server_health", "ping", "health", "state"): - return {"module": module, "ok": True, "session_id": sid} - if name in ("idb_list",): - return {"sessions": [{"session_id": sid, "module": module, - "input_path": binpath}]} - fn = tools.get(name) - if fn is None: - raise KeyError(f"unknown tool: {name!r}") - if deadline is None: - result = fn(**args) - else: - deadline.arm() - try: - result = fn(**args) - finally: - deadline.disarm() - # Match the MCP server's structuredContent: a dict passes through, any - # other return (list/scalar) is wrapped as {"result": ...}. domain.py - # parses that exact shape (e.g. lookup_funcs -> payload["result"]). - return result if isinstance(result, dict) else {"result": result} - - try: - os.unlink(sockpath) - except OSError: - pass - srv = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) - srv.bind(sockpath) - srv.listen(8) - try: - while True: - conn, _ = srv.accept() - try: - while True: - req = recv(conn) - if req is None: - break - name, args = req - if name == "__shutdown__": - return - try: - send(conn, (True, dispatch(name, args))) - except Exception as e: # noqa: BLE001 -- report, keep serving - send(conn, (False, f"{type(e).__name__}: {e}")) - except (ConnectionError, OSError): - pass - finally: - conn.close() - finally: - try: - import idapro - idapro.close_database(save=False) - except Exception: # noqa: BLE001 - pass - try: - os.unlink(sockpath) - except OSError: - pass - - -def main(argv=None) -> None: - argv = argv if argv is not None else sys.argv[1:] - if len(argv) < 2: - sys.stderr.write( - "usage: python -m idatui.worker <sock> <binary> [ida-load-args]\n") - raise SystemExit(2) - try: - serve(argv[0], argv[1], argv[2] if len(argv) > 2 else "") - except SystemExit: - raise - except BaseException as e: # noqa: BLE001 -- surface a clean cause + code 1 - import traceback - sys.stderr.write(f"\nWORKER-FATAL: {type(e).__name__}: {e}\n") - traceback.print_exc() - sys.stderr.flush() - raise SystemExit(1) - - -if __name__ == "__main__": - main() diff --git a/idatui/worker_client.py b/idatui/worker_client.py deleted file mode 100644 index 173cd9d..0000000 --- a/idatui/worker_client.py +++ /dev/null @@ -1,266 +0,0 @@ -"""WorkerClient — a drop-in replacement for ``IDAClient`` backed by our own -idalib worker (``idatui.worker``) over a unix socket instead of ida-pro-mcp's -HTTP/JSON transport. - -It exposes exactly the surface the app/domain use on the client -(``call``/``call_envelope``/``connect``/``set_db``/``resolve_db``/ -``list_sessions``/``health``/``keepalive``/``close``) and returns byte-identical -payloads (the worker calls the same tool functions), so ``domain.py`` and the -app are unchanged — you just construct a WorkerClient instead of an IDAClient. - -Concurrency: the app fires calls from several worker threads over one client; -the worker is single-threaded, so calls are serialized under a lock (the worker -processes one tool at a time anyway — and at ~50us/call that's free). -""" -from __future__ import annotations - -import os -import socket -import subprocess -import sys -import threading -import time -import uuid -from typing import Any - -from .errors import IDAToolError, IDAConnectionError, Session -from .worker import recv as _recv -from .worker import send as _send - -_WORKER_PY = os.path.join(os.path.dirname(os.path.abspath(__file__)), "worker.py") -_worker_python_cache: str | None = None - - -def _find_worker_python() -> str: - """A python that can import ``ida_pro_mcp`` (and thus idalib) — NOT necessarily - the TUI's python. On a typical box the TUI runs under a venv that has textual - + idalib but not ida_pro_mcp, while the system python has idalib + - ida_pro_mcp. Override with IDATUI_WORKER_PYTHON.""" - global _worker_python_cache - if _worker_python_cache: - return _worker_python_cache - override = os.environ.get("IDATUI_WORKER_PYTHON") - candidates = [override] if override else [] - candidates += ["/usr/bin/python", "/usr/bin/python3", sys.executable] - for py in candidates: - if not py or not os.path.exists(py): - continue - try: - r = subprocess.run([py, "-c", "import ida_pro_mcp"], - capture_output=True, timeout=30) - if r.returncode == 0: - _worker_python_cache = py - return py - except Exception: # noqa: BLE001 - continue - return sys.executable # last resort; the worker will report the real error - - -class _NoopKeepAlive: - """The worker is ours and never idles out, so keepalive is a no-op.""" - - def __init__(self) -> None: - self.beats = self.failures = 0 - - def start(self): - return self - - def stop(self) -> None: - pass - - -class WorkerClient: - def __init__(self, binary_path: str, *, ttl: int = 0, - python: str | None = None, load_args: str = "") -> None: - self._bin = os.path.abspath(os.path.expanduser(binary_path)) - self._load_args = load_args or "" # IDA switches for a headerless blob - self._python = python or _find_worker_python() - tag = f"{os.getpid()}-{uuid.uuid4().hex[:8]}" - self._sock_path = f"/tmp/idatui-worker-{tag}.sock" - self._log_path = f"/tmp/idatui-worker-{tag}.log" - self._proc: subprocess.Popen | None = None - self._sock: socket.socket | None = None - self._sid = uuid.uuid4().hex[:8] - self._lock = threading.Lock() # serialize socket use - self._spawn_lock = threading.Lock() - #: Set by close(). A dropped socket is respawned on the next call (the - #: worker segfaulted and we want it back); a CLOSED one must not be. - #: Teardown and binary-switch both close while @work threads are still - #: in flight, so without this, quitting during a decompile spawned a - #: fresh idalib worker that re-opened the database nobody was looking - #: at any more -- a stray process holding the .i64 we just released. - self._closed = False - - # -- lifecycle --------------------------------------------------------- # - def connect(self, timeout: float = 1800.0, progress=None) -> "WorkerClient": - """Spawn the worker (opens + analyzes the DB) and connect once ready.""" - with self._spawn_lock: - self._closed = False # an explicit reconnect revives this client - if self._sock is not None: - return self - if self._proc is None or self._proc.poll() is not None: - # run worker.py as a SCRIPT (not -m idatui.worker) so we don't - # import the textual-dependent idatui package __init__ under the - # IDA python, which usually has no textual. - argv = [self._python, _WORKER_PY, self._sock_path, self._bin] - if self._load_args: - argv.append(self._load_args) - self._proc = subprocess.Popen( - argv, - stdout=open(self._log_path, "wb"), - stderr=subprocess.STDOUT, stdin=subprocess.DEVNULL, - ) - deadline = time.time() + timeout - t0 = time.time() - # Poll fast at first, then back off. A flat 0.2s cost every caller a - # fifth of a second even when the worker was ready in milliseconds - # (a small binary, or a seeded .i64), which is most of the time in - # the tests and noticeable on a re-open. - # - # Backing off all the way to 0.2s was too eager: a seeded database - # is ready at ~250ms, by which point the delay has grown to 134ms, so - # every open waited ~350ms whatever the binary -- the same number for - # a 47KB `echo` and a 1.2MB `bash`, which is what gives a polling - # artefact away. Cap the backoff at 25ms instead: the overshoot on a - # fast open is bounded by that, and 40 probes a second is nothing - # next to an auto-analysis that runs for minutes. - # - # Do NOT be tempted to hold the 5ms rate instead. This poll runs on a - # background thread while the UI thread is drawing, and 200 wakeups a - # second through a cold analysis cost enough GIL time to delay the - # app's own startup -- it left the loading overlay up long enough for - # project mode's first keypress to land on it. - delay = 0.005 - while time.time() < deadline: - try: - s = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) - s.connect(self._sock_path) - self._sock = s - return self - except OSError: - if self._proc.poll() is not None: - raise IDAConnectionError( - f"worker exited (code {self._proc.returncode}): " - f"{self._log_tail()} [full log: {self._log_path}]") - if progress: - progress(f"auto-analyzing {os.path.basename(self._bin)}… " - f"({int(time.time() - t0)}s)") - time.sleep(delay) - delay = min(delay * 1.6, 0.025) - raise IDAConnectionError("worker did not become ready in time") - - @property - def pid(self) -> int | None: - """The worker process id (for memory accounting), or None if not spawned.""" - return self._proc.pid if self._proc is not None else None - - def close(self, grace: float = 20.0) -> None: - """Shut the worker down cleanly. - - After ``__shutdown__`` the worker still has to ``close_database()``, which - re-packs the ``.i64`` and removes the unpacked ``.id0/.id1/...`` scratch. - Signalling it before that finishes is what leaves databases wedged, so - wait out the grace period first and only escalate if it really is stuck. - """ - with self._lock: - s = self._sock - self._sock = None - self._closed = True - if s is not None: - try: - _send(s, ("__shutdown__", {})) - except Exception: # noqa: BLE001 - pass - try: - s.close() - except Exception: # noqa: BLE001 - pass - if self._proc is not None: - try: - self._proc.wait(timeout=grace) # let it close the DB properly - except Exception: # noqa: BLE001 -- TimeoutExpired: it's stuck - try: - self._proc.terminate() - self._proc.wait(timeout=5) - except Exception: # noqa: BLE001 - try: - self._proc.kill() - except Exception: # noqa: BLE001 - pass - - # -- the call surface -------------------------------------------------- # - def call(self, tool: str, *, timeout: float | None = None, **args) -> Any: - if self._closed: - raise IDAConnectionError( - f"{tool}: this worker was closed (call connect() to revive it)") - if self._sock is None: - self.connect() - with self._lock: - s = self._sock - if s is None: - raise IDAConnectionError("worker connection is closed") - try: - _send(s, (tool, args)) - reply = _recv(s) - except (OSError, ConnectionError) as e: - self._sock = None - raise IDAConnectionError(f"worker transport failed: {e}") from e - if reply is None: - self._sock = None - raise IDAConnectionError("worker closed the connection") - ok, payload = reply - if not ok: - raise IDAToolError(tool, str(payload)) - return payload - - def call_envelope(self, tool: str, *, timeout: float | None = None, - **args) -> dict: - # domain.decompile() reads result.structuredContent — mirror that shape. - return {"result": {"structuredContent": self.call(tool, timeout=timeout, - **args)}} - - # -- session shims (single-DB worker) --------------------------------- # - def set_db(self, db: str | None) -> None: - if db: - self._sid = db - - def resolve_db(self) -> str: - return self._sid - - def list_sessions(self) -> list[Session]: - return [Session(session_id=self._sid, - filename=os.path.basename(self._bin), - input_path=self._bin, is_active=True)] - - def health(self) -> dict: - try: - return self.call("server_health") - except IDAToolError: - return {"module": os.path.basename(self._bin), "ok": True} - - def keepalive(self, interval: float = 120.0) -> _NoopKeepAlive: - return _NoopKeepAlive() - - def _log_tail(self, n: int = 400) -> str: - """Last meaningful line(s) of the worker log (skip IDA's licence banner), - so a startup crash surfaces the real cause instead of just 'code 1'.""" - try: - with open(self._log_path, encoding="utf-8", errors="replace") as f: - lines = [ln.strip() for ln in f if ln.strip()] - except OSError: - return "(no worker log)" - # the worker prints a clean 'WORKER-FATAL: ...' line on a startup crash - for ln in reversed(lines): - if ln.startswith("WORKER-FATAL:"): - return ln[len("WORKER-FATAL:"):].strip()[-n:] - skip = ("thank you", "licensed to", "[mcp]", "ida ", "hex-rays") - meaningful = [ln for ln in lines - if not any(s in ln.lower() for s in skip)] - return " | ".join((meaningful or lines)[-3:])[-n:] - - # context manager parity with IDAClient - def __enter__(self) -> "WorkerClient": - return self.connect() - - def __exit__(self, *exc) -> None: - self.close() 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/tests/_fixtures.py b/tests/_fixtures.py index 766ecd5..0b72856 100644 --- a/tests/_fixtures.py +++ b/tests/_fixtures.py @@ -37,6 +37,36 @@ def cache_is_fresh(binary: str) -> bool: return os.path.exists(c) and os.path.getmtime(c) >= os.path.getmtime(binary) +#: Generated targets live here so their pristine caches survive between runs. +#: Gitignored; safe to delete (the next run rebuilds both). +SYNTHETIC_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), ".synthetic") + + +def synthetic(name: str, build) -> str: + """A generated binary at a STABLE path, rebuilt only when its bytes change. + + Generated targets used to be written into a fresh TemporaryDirectory on + every run, which quietly defeated the whole pristine-cache scheme: a new + path with new bytes every time means auto-analysis is paid in full, every + run, forever. `test_blob_ui`'s 64KB blob cost ~40s a run that way. + + ``build()`` must be DETERMINISTIC and return bytes. That is also what makes + the suites reproducible: a blob built from os.urandom can, by luck, contain + something IDA reads as a function, and then a test asserting "no functions" + fails for reasons no one can reproduce. + """ + os.makedirs(SYNTHETIC_DIR, exist_ok=True) + path = os.path.join(SYNTHETIC_DIR, name) + data = build() + if not os.path.exists(path) or open(path, "rb").read() != data: + with open(path, "wb") as fh: # content changed -> cache is stale + fh.write(data) + for stale in (cache_path(path), path + ".i64"): + if os.path.exists(stale): + os.remove(stale) + return path + + async def build_pristine(binary: str, cache: str, app_factory) -> None: """Analyse ``binary`` once and keep the database as a golden copy. @@ -50,7 +80,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_blob_ui.py b/tests/test_blob_ui.py index f60e914..1055e5a 100644 --- a/tests/test_blob_ui.py +++ b/tests/test_blob_ui.py @@ -17,10 +17,13 @@ import sys import tempfile sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) from textual.widgets import Input, Static # noqa: E402 from idatui.app import ConfirmScreen, IdaTui, ListingView # noqa: E402 +from _fixtures import staged, synthetic # noqa: E402 +from idatui._sync import settle # noqa: E402 PASS = FAIL = 0 @@ -46,28 +49,63 @@ async def wait(pred, pilot, t=240.0): return False -async def run() -> int: - with tempfile.TemporaryDirectory() as tmp: - # Random bytes so IDA finds no functions... but with REAL AArch64 - # instructions planted at a known offset. Whether arbitrary random bytes - # happen to decode is chance, and a test that depends on chance tells you - # nothing on the run where it fails. - data = bytearray(os.urandom(64 * 1024)) +#: File offset of the planted instruction run -> ea 0x4000 + PLANTED. +PLANTED = 0x40 + + +def _blob_bytes() -> bytes: + """A DETERMINISTIC pseudo-random blob with real AArch64 instructions planted. + + Seeded, not os.urandom: the bytes must be identical every run or the + pristine-database cache can never apply (this suite used to pay ~40s of + auto-analysis per run because the content, and the path, changed each time). + Determinism also removes a genuine flake -- whether 64KB of chance bytes + contains something IDA reads as a function is luck, and "and really has no + functions" is asserted below. + """ + import random + data = bytearray(random.Random(0xB10BCAFE).randbytes(64 * 1024)) # -parm puts IDA in AArch64 mode, so these are A64 encodings; the ARM32 # spelling of a nop (0xE1A00000) is NOT decodable there and made this # test fail for a reason that had nothing to do with what it checks. - planted = 0x40 # file offset -> ea 0x4040 - for k, insn in enumerate((0xD503201F, # nop - 0xD503201F, # nop - 0xD65F03C0)): # ret <- the run must stop here - data[planted + k * 4:planted + k * 4 + 4] = insn.to_bytes(4, "little") - blob = os.path.join(tmp, "rnd.bin") - with open(blob, "wb") as f: - f.write(bytes(data)) + for k, insn in enumerate((0xD503201F, # nop + 0xD503201F, # nop + 0xD65F03C0)): # ret <- the run must stop here + data[PLANTED + k * 4:PLANTED + k * 4 + 4] = insn.to_bytes(4, "little") + return bytes(data) + + +#: -parm puts IDA in AArch64 mode (the ARM32 spelling of a nop is not decodable +#: there); -b400 sets the image base. The cached database must be built with the +#: SAME switches, so both go through one factory. +BLOB_ARGS = "-parm -b400" + +def _blob_app(path): + return IdaTui(open_path=path, keepalive=False, load_args=BLOB_ARGS) + + +def head_at(lst, ea): + """The listing row for ``ea`` off the LIVE model, or None. + + Always re-reads ``lst.model``: an edit may rebuild the model, and holding + the old object shows pre-edit rows -- which looks exactly like the edit + silently failing. + """ + m = lst.model + if m is None: + return None + i = m.index_of_ea(ea) + return m.get(i) if i is not None and i >= 0 else None + + +async def run() -> int: + blob_src = synthetic("rnd.bin", _blob_bytes) + # staged() analyses once ever and copies the result in on later runs. + async with staged(blob_src, _blob_app) as blob: # Skip the dialog by answering up front; this test is about what # happens AFTER a described blob turns out to contain nothing. - app = IdaTui(open_path=blob, keepalive=False, load_args="-parm -b400") + app = _blob_app(blob) async with app.run_test(size=(140, 44)) as pilot: ok = await wait(lambda: app._func_index is not None and app._func_index.complete, pilot) @@ -126,7 +164,7 @@ async def run() -> int: i >= 0 and m.get(i).ea == 0x4021, f"row={i} ea={m.get(i).ea if i >= 0 else None}") - target = 0x4000 + planted # a NOP we put there ourselves + target = 0x4000 + PLANTED # a NOP we put there ourselves lst.cursor = m.index_of_ea(target) lst._scroll_cursor_into_view() await pilot.pause(0.1) @@ -134,12 +172,18 @@ async def run() -> int: lst._cursor_ea() == target, f"{lst._cursor_ea():#x} want {target:#x}") await pilot.press("c") - await pilot.pause(2.0) - # Defining an item REBUILDS the listing model, so re-read it from the - # view: holding the old object shows the pre-edit rows and looks - # exactly like the edit silently failing. + # settle(), not a fixed sleep AND not a bare predicate: an edit can + # look done for a moment and then be replaced when a queued listing + # rebuild lands, so the gate has to be "the row is code AND the app + # has stopped working". settle() is the same helper the app's own + # RPC layer uses, so tests and driver agree on what "done" means. + await settle(app, lambda: (lambda h: h is not None and h.kind == "code")( + head_at(lst, target)), timeout=30) + # Re-read the model: defining an item rebuilds it, and holding the + # old object shows pre-edit rows -- which looks exactly like the + # edit silently failing. m = lst.model - h = m.get(m.index_of_ea(target)) + h = head_at(lst, target) check("`c` on a chosen byte carves an instruction there", h is not None and h.kind == "code", f"kind={h.kind if h else None} text={h.text if h else None!r}") @@ -184,9 +228,17 @@ async def run() -> int: for ch in "note": await pilot.press(ch) await pilot.press("enter") - await wait(lambda: lst.model is not old and lst.model is not None, - pilot, 30) - await pilot.pause(0.4) + # Wait for the COMMENT ITSELF to show up, not for the model object to + # be replaced: a comment now re-renders the listing in place (the + # walk is kept), so `model is not old` never becomes true and this + # burned its full 30s timeout on every run -- after which the check + # below passed vacuously, because nothing had happened at all. + # The prompt closing plus quiescence is the real end of the edit. + # (The listing re-renders its text lazily, so the comment is not + # necessarily visible in model rows the moment the worker returns -- + # which is why this waits for the app, not for the text.) + await settle(app, lambda: not app.query_one("#comment", Input).display, + timeout=30) check("commenting leaves the view where it was", lst.model.get(round(lst.scroll_offset.y)).ea == ctop and lst._cursor_ea() == ccur, @@ -208,7 +260,12 @@ async def run() -> int: check("scrolled somewhere with rows above us", round(lst.scroll_offset.y) > 0, f"top={lst.scroll_offset.y}") await pilot.press("c") - await pilot.pause(2.5) + # No predicate here on purpose: this spot is random data, so the + # carve may legitimately produce nothing and "the row became code" + # would never hold (it timed out for 30s and then passed anyway). + # What is being checked is that the VIEW did not move, so the gate + # is simply "the app has finished reacting". + await settle(app, timeout=30) m2 = lst.model top_after = m2.get(round(lst.scroll_offset.y)).ea check("carving leaves the scroll position where it was", diff --git a/tests/test_codemode_client.py b/tests/test_codemode_client.py new file mode 100644 index 0000000..2f70ba3 --- /dev/null +++ b/tests/test_codemode_client.py @@ -0,0 +1,162 @@ +"""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 _open_kwargs_are_real(sent: dict): + """(ok, detail) for the kwargs the adapter passes to DatabaseHandle.open. + + Skips (passes) when ida_codemode is not installed, so the file stays pure. + """ + try: + import inspect + from ida_codemode.client import DatabaseHandle as Real + except ImportError: + return True, "ida_codemode not installed - signature not checked" + accepted = set(inspect.signature(Real.open).parameters) + unknown = sorted(set(sent) - accepted) + return not unknown, f"open() rejects {unknown}" + + +def main() -> int: + proc, base, file_type = _parse_load_args("-parm:ARMv7-M -b800000 -TRaw") + check("legacy switches map to typed Code Mode options", + (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["image_base"] == 0x1000, + FakeDatabaseHandle.kwargs) + # A fake that swallows **kwargs cannot catch a keyword the real + # library does not have -- which is exactly how this port shipped + # `loading_address` (the real name is `image_base`) and would have + # raised TypeError on the very first connect. Check the names we + # send against the real signature whenever it is importable. + check("every open() keyword exists in the real library", + *_open_kwargs_are_real(FakeDatabaseHandle.kwargs)) + check("connect waits for Code Mode autoanalysis", + handle.waited == 42, getattr(handle, "waited", None)) + check("progress distinguishes discovery and backend attachment", + 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_thumb_ui.py b/tests/test_thumb_ui.py index 92a19fb..fed08a7 100644 --- a/tests/test_thumb_ui.py +++ b/tests/test_thumb_ui.py @@ -15,12 +15,15 @@ Needs IDA. ~40s. NEEDS_IDA = True import asyncio import os +import shutil import sys +import tempfile sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) from textual.widgets import Static # noqa: E402 +from idatui._sync import settle # noqa: E402 from idatui.app import DecompView, IdaTui, ListingView # noqa: E402 PASS = FAIL = 0 @@ -38,6 +41,36 @@ def check(name, ok, detail=""): print(f" FAIL {name} {detail}") + +#: Every phase gets its OWN copy of the fixture. +#: +#: This suite used to delete <BIN>.i64 and reopen the SAME path for each phase. +#: That was safe when the TUI owned a private worker that died with it; under +#: Code Mode the database is leased and the previous phase's worker can still +#: hold it through its lease grace, so the delete raced a live owner and the +#: next open never produced a listing (the crash this fixed). Separate paths +#: cannot collide, and nothing has to wait for anyone else to let go. +_SCRATCH = [] + + +def fresh_copy(src: str, tag: str) -> str: + d = tempfile.mkdtemp(prefix=f"idatui-thumb-{tag}-") + _SCRATCH.append(d) + dst = os.path.join(d, os.path.basename(src)) + shutil.copy2(src, dst) + return dst + + +def drop_scratch() -> None: + for d in _SCRATCH: + shutil.rmtree(d, ignore_errors=True) + _SCRATCH.clear() + + +def status_of(app) -> str: + return str(app.query_one("#status", Static).render()) + + async def wait(pred, pilot, t=240.0): for _ in range(int(t / 0.05)): await pilot.pause(0.05) @@ -52,13 +85,8 @@ async def wait(pred, pilot, t=240.0): async def run() -> int: # A fresh database every time: the T flag and the segment's addressing mode # are SAVED in the .i64, so a previous run would answer the question for us. - for ext in (".i64", ".id0", ".id1", ".id2", ".nam", ".til"): - try: - os.remove(BIN + ext) - except OSError: - pass - - app = IdaTui(open_path=BIN, keepalive=False, load_args="-parm") + app = IdaTui(open_path=fresh_copy(BIN, "arm"), keepalive=False, + load_args="-parm") async with app.run_test(size=(140, 44)) as pilot: await wait(lambda: app._func_index is not None and app._func_index.complete, pilot) @@ -84,10 +112,14 @@ async def run() -> int: m1 = lst.model await pilot.press("t") - await wait(lambda: lst.model is not m1 and lst.model is not None, pilot, 60) - await pilot.pause(0.5) + # The mode switch announces itself; wait for THAT, plus quiescence. + # `lst.model is not m1` used to be the gate, but an item edit now keeps + # the listing's walk instead of rebuilding it, so the model object is + # never replaced -- every one of these waits sat out its full 60s and + # the suite still "passed", four times over. + await settle(app, lambda: "Thumb" in status_of(app), timeout=60) - status = str(app.query_one("#status", Static).render()) + status = status_of(app) check("the status says it switched to Thumb", "Thumb" in status, status[:90]) # Thumb doesn't exist in AArch64, and -parm on a headerless blob gives a # 64-bit segment, so setting T alone would change nothing and look broken. @@ -114,9 +146,8 @@ async def run() -> int: m2 = lst.model lst.cursor = lst.model.index_of_ea(0) await pilot.press("t") - await wait(lambda: lst.model is not m2 and lst.model is not None, pilot, 60) - await pilot.pause(0.5) - status = str(app.query_one("#status", Static).render()) + await settle(app, lambda: "ARM @" in status_of(app), timeout=60) + status = status_of(app) check("`t` toggles back to ARM", "ARM @" in status, status[:80]) # -- and the reason a carved function wouldn't decompile ---------------- # @@ -126,12 +157,8 @@ async def run() -> int: # disassembly that F5 can never turn into pseudocode. The database's bitness # is fixed at load and cannot be corrected afterwards, so the only honest # thing is to say so. - for ext in (".i64", ".id0", ".id1", ".id2", ".nam", ".til"): - try: - os.remove(BIN + ext) - except OSError: - pass - app = IdaTui(open_path=BIN, keepalive=False, load_args="-parm") # 64-bit + app = IdaTui(open_path=fresh_copy(BIN, "arm64"), keepalive=False, + load_args="-parm") # 64-bit async with app.run_test(size=(140, 44)) as pilot: await wait(lambda: app._func_index is not None and app._func_index.complete, pilot) @@ -143,9 +170,8 @@ async def run() -> int: await pilot.pause(0.3) m = lst.model await pilot.press("t") - await wait(lambda: lst.model is not m and lst.model is not None, pilot, 60) - await pilot.pause(0.5) - status = str(app.query_one("#status", Static).render()) + await settle(app, lambda: "64-bit" in status_of(app), timeout=60) + status = status_of(app) check("a 64-bit database warns that Hex-Rays won't decompile", "64-bit" in status and "decompile" in status, status[:120]) check("and names the fix", "ARMv7-A" in status, status[:120]) @@ -158,9 +184,10 @@ async def run() -> int: await pilot.pause(0.2) mp = lst.model await pilot.press("p") - await wait(lambda: lst.model is not mp and lst.model is not None, pilot, 60) - await wait(lambda: app._func_index is not None - and len(app._func_index) > 0, pilot, 60) + # The function appearing in the index IS the signal; the model identity + # never was one. + await settle(app, lambda: app._func_index is not None + and len(app._func_index) > 0, timeout=60) await pilot.press("tab") await wait(lambda: "cannot decompile" in str(app.query_one("#status", Static).render()), pilot, 90) @@ -177,12 +204,8 @@ async def run() -> int: len(status) < 110, f"{len(status)} chars: {status[:130]}") # -- the whole point: a 32-bit database decompiles ---------------------- # - for ext in (".i64", ".id0", ".id1", ".id2", ".nam", ".til"): - try: - os.remove(BIN + ext) - except OSError: - pass - app = IdaTui(open_path=BIN, keepalive=False, load_args="-parm:ARMv7-A") + app = IdaTui(open_path=fresh_copy(BIN, "armv7a"), keepalive=False, + load_args="-parm:ARMv7-A") async with app.run_test(size=(140, 44)) as pilot: await wait(lambda: app._func_index is not None and app._func_index.complete, pilot) @@ -216,12 +239,8 @@ async def run() -> int: if not os.path.isfile(vec): check("the cortexm fixture exists", False, vec) else: - for ext in (".i64", ".id0", ".id1", ".id2", ".nam", ".til"): - try: - os.remove(vec + ext) - except OSError: - pass - app = IdaTui(open_path=vec, keepalive=False, load_args="-parm:ARMv7-M") + app = IdaTui(open_path=fresh_copy(vec, "cortexm"), keepalive=False, + load_args="-parm:ARMv7-M") async with app.run_test(size=(140, 44)) as pilot: await wait(lambda: app._func_index is not None and app._func_index.complete, pilot) @@ -253,6 +272,8 @@ async def run() -> int: check("the result survives the reload AND the reindex", "3 Thumb entries" in status, status[:90]) + drop_scratch() + print(f"\n{PASS} passed, {FAIL} failed") return 1 if FAIL else 0 diff --git a/tests/test_trace_ui.py b/tests/test_trace_ui.py index bb016eb..5844c90 100644 --- a/tests/test_trace_ui.py +++ b/tests/test_trace_ui.py @@ -20,6 +20,7 @@ sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) from textual.widgets import Input, OptionList, Static # noqa: E402 from _fixtures import staged # noqa: E402 +from idatui._sync import settle # noqa: E402 from idatui.app import (DecompView, IdaTui, ListingView, # noqa: E402 RegWriteScreen, TraceDock) @@ -109,13 +110,20 @@ async def run() -> int: lst.focus() await pilot.pause(0.4) await pilot.press("]") - await wait(lambda: app._t == 1, pilot, 20) + # `app._t` is assigned the moment the key is handled, so it is NOT a + # signal that the VIEW has followed -- the navigation it kicks off + # runs in a worker. Waiting on it and then reading the cursor was a + # race that the (slower) Code Mode backend loses. Gate on the thing + # the check is about. + await settle(app, lambda: app._t == 1 and lst._cursor_ea() == t.ip(1), + timeout=20) check("] steps forward one instruction", app._t == 1, f"t={app._t}") check("the code view follows the trace", lst._cursor_ea() == t.ip(1), f"{lst._cursor_ea()} vs {t.ip(1)}") await pilot.press("[") - await wait(lambda: app._t == 0, pilot, 20) + await settle(app, lambda: app._t == 0 and lst._cursor_ea() == t.ip(0), + timeout=20) check("[ steps backward", app._t == 0, f"t={app._t}") await pilot.press("[") await pilot.pause(0.4) 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" }, +] |