diff options
| author | Duncan Ogilvie <mr.exodia.tpodt@gmail.com> | 2026-08-20 23:42:42 +0200 |
|---|---|---|
| committer | Duncan Ogilvie <mr.exodia.tpodt@gmail.com> | 2026-08-20 23:42:42 +0200 |
| commit | f3715d8de0d255c8b14710acfa120ccb9ea953fd (patch) | |
| tree | 98e79826e4e39e2269b59ccf33fa7c00a1c68c23 | |
| parent | Add Ctrl+R to refresh all views (diff) | |
| download | ida-tui-f3715d8de0d255c8b14710acfa120ccb9ea953fd.tar.gz ida-tui-f3715d8de0d255c8b14710acfa120ccb9ea953fd.tar.xz ida-tui-f3715d8de0d255c8b14710acfa120ccb9ea953fd.zip | |
Adopt idb_events and remote module features from ida-codemode
| -rw-r--r-- | README.md | 13 | ||||
| -rw-r--r-- | docs/CODEMODE_UPSTREAM.md | 139 | ||||
| -rw-r--r-- | experiments/bench_ops.py | 33 | ||||
| -rw-r--r-- | experiments/bench_pack_trace.py | 98 | ||||
| -rw-r--r-- | experiments/profile_client.py | 15 | ||||
| -rw-r--r-- | experiments/profile_remote.py | 111 | ||||
| -rw-r--r-- | idatui/app.py | 238 | ||||
| -rw-r--r-- | idatui/codemode_client.py | 1365 | ||||
| -rw-r--r-- | idatui/domain.py | 649 | ||||
| -rw-r--r-- | idatui/edit_ctl.py | 207 | ||||
| -rw-r--r-- | idatui/pool.py | 22 | ||||
| -rw-r--r-- | idatui/remote_ops.py | 1683 | ||||
| -rw-r--r-- | idatui/remote_tools.py | 610 | ||||
| -rw-r--r-- | tests/test_codemode_client.py | 290 | ||||
| -rw-r--r-- | tests/test_pool.py | 39 | ||||
| -rw-r--r-- | tests/test_scenarios.py | 3303 |
16 files changed, 5833 insertions, 2982 deletions
@@ -51,6 +51,19 @@ uv add --editable ../ida-codemode ida-tui never owns an IDA process — it takes a **lease**. A matching database open in the IDA GUI is reused, otherwise Code Mode starts or shares a managed idalib worker. Quitting drops the lease and leaves everyone else alone. +Changes made in the GUI or another client arrive over Code Mode's IDB event +stream; ida-tui debounces bursts and refreshes its cached views automatically. +On quit with unsaved changes, a final managed-worker lease can discard the +session without saving; GUI-backed or still-shared sessions leave that final +decision with their owner or remaining clients. +If the owning GUI or worker closes, ida-tui never replaces it by spawning a +headless worker implicitly. It keeps the cached view disconnected until a +matching owner is reopened and an attach-only rediscovery succeeds. +Remote operations are typed, source-backed Python functions. ida-codemode +installs content-addressed modules once per handle, so ida-tui keeps normal +refactorable source without paying to resend hot listing/decompiler code. +Operation attribution is also a per-call provider rather than a fixed string, so +it can evolve from `IDA TUI` to labels such as `IDA TUI: alice`. Headerless blobs need a hint, or IDA assumes x86 at address 0 and analyses nothing: diff --git a/docs/CODEMODE_UPSTREAM.md b/docs/CODEMODE_UPSTREAM.md index 9d3e09f..7dff13e 100644 --- a/docs/CODEMODE_UPSTREAM.md +++ b/docs/CODEMODE_UPSTREAM.md @@ -11,17 +11,18 @@ 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. -> **Status against 0.6.1 (upstream `439289f`) — every item re-checked.** +> **Status against the protocol-6 event-stream development tree, based on 0.6.1 +> (upstream `439289f`) — every item re-checked.** > > | item | verdict | > |---|---| > | 1 `timeout_trace` line tracing | ✅ **fixed in 0.3.2** — no `settrace` in the runtime at all | > | 2 `to_jsonable` on large results | ✅ **fixed in 0.3.2** — `dumps_json` C fast path | > | 3 2 ms `execute_sync` floor | ✅ **fixed in 0.3.2, 7.0x** — 2.055 ms → 0.294 ms | -> | 4 loader switches fatal on reopen | ✅ **fixed in 0.5.x** — see the caveat in §4 | -> | 5 IDB replaced under a live lease | ❌ open | -> | 6 no close-without-save | ❌ open | -> | 7 no change notification | ⚠️ partial — `DatabaseEventCallback` exists on `DatabaseManager`, but there is still no revision counter for an *external* caching client | +> | 4 loader switches fatal on reopen | **partial** — normal reopen fixed in 0.5.x; direct `.i64` paths remain [issue #36](https://github.com/HexRaysSA/ida-codemode/issues/36) | +> | 5 IDB replaced under a live lease | **stale** — out-of-band replacement is outside the supported lifecycle, as it is for the IDA GUI | +> | 6 close without save | **fixed in protocol 6** — the final managed-worker lease can choose `shutdown_database(save=False)` | +> | 7 no change notification | ✅ **fixed in protocol 6** — `DatabaseHandle.subscribe_idb_events()` streams revisioned, operation-attributed IDB changes | > | 8 package exports | ✅ **fixed in 0.5.x** — a real `__all__` on the package root | > | 9 no `py.typed` / handle Protocol | ✅ **fixed in 0.5.x** — `ida_codemode/py.typed` ships | > @@ -38,9 +39,10 @@ worker backend, quiet box. Target for timings: `targets/echo` unless stated. > between those two releases. 0.6.1 only collapses the six console scripts into a > single `ida-codemode` command. > -> Both of our client-side workarounds re-measured at **0.99x and 0.97x** on 0.3.2 -> — i.e. nothing. The settrace strip has been deleted; `_PACK_EPILOGUE` is kept -> only for encoder determinism. Harness: `experiments/bench_pack_trace.py`. +> Both client-side workarounds re-measured at **0.99x and 0.97x** on 0.3.2 — +> i.e. nothing — and are deleted. Remote code is now ordinary typed Python, +> installed as content-addressed modules by ida-codemode. Harness: +> `experiments/bench_pack_trace.py`. **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 @@ -108,10 +110,10 @@ and do the same, which is an argument for fixing it in the runtime. ## 2. `to_jsonable` dominates any large result -✅ **FIXED in 0.3.2**, via the first suggested fix below: `serialization.dumps_json` -calls `json.dumps(value, default=to_jsonable)`, so a JSON-safe result never enters -the Python walker. Our packing workaround now measures 0.97x and is retained only -to pin encoder settings, not for speed. +**FIXED in 0.3.2**, via the first suggested fix below: +`serialization.dumps_json` calls `json.dumps(value, default=to_jsonable)`, so a +JSON-safe result never enters the Python walker. Our packing workaround measured +0.97x and has been deleted. `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 @@ -133,10 +135,10 @@ That is 114x, and it was 72% of the page's total cost before we changed it. 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. +**Retired workaround:** snippets used to `json.dumps` inside the database process +and return one string, which the client parsed. The typed remote API now owns +strict argument/result encoding, and ida-tui contains no generated script +strings or packing envelope. --- @@ -181,7 +183,7 @@ would let chatty clients amortise it without redesigning around it. --- -## 4. Loader switches on an existing database are a FATAL, not an error +## 4. Loader switches on an existing database are a FATAL, not an error — one edge remains Opening a target that already has an `.i64`, while passing spawn-only options, kills the worker: @@ -226,57 +228,74 @@ Our workaround is therefore deleted. **One narrow case remains**: the strip need `input_path != source`, so passing an `.i64` path *directly* together with load options (`ida-tui foo.i64 --processor arm`) still forwards the switches and still fatals. Our old guard keyed on "the target IDB exists" and so covered it. It is a -nonsense invocation and no idatui code path generates it — the project layer +nonsense invocation and no ida-tui code path generates it — the project layer always passes `output_database`, and `_needs_load_options` bails when an `.i64` -exists — but if this ever resurfaces as "worker exited with status 1", that is -where it comes from. +exists — but the library boundary should still reject or normalize it rather +than launch a known-fatal IDA command. Tracked upstream as +[issue #36](https://github.com/HexRaysSA/ida-codemode/issues/36). --- -## 5. Deleting or replacing an IDB under a live lease fails silently +## 5. Deleting or replacing an IDB under a live lease — STALE -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. +The original suite deleted an `.i64` while a private worker still had it open, +then immediately reopened the same path. That ownership model no longer applies: +Code Mode databases are shared resources, and the IDA GUI itself does not survive +out-of-band replacement of its open database. Detecting arbitrary filesystem +replacement is therefore not part of the supported lifecycle. -**Suggested fixes** +The actionable lifecycle gaps that originally forced private-registry access are +fixed. `find_database_owner()` and `wait_database_released()` are public exports; +`DatabaseHandle.close(wait_for_database=True)` can wait for a final managed close; +a draining owner remains registered until the IDB is actually closed; and +`new_database=True` refuses to replace a live owner. -- 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. +ida-tui now uses the public owner/release API while recreating a database and no +longer reaches into registry locks. Owner loss is attach-only: ida-tui will +rediscover a replacement GUI or worker, but will never turn a +user-closing-the-GUI action into an implicit headless reopen. There is no +remaining upstream request in this section. --- -## 6. No close-without-save, and no rollback +## 6. Close without save — FIXED in protocol 6 + +`DatabaseHandle.shutdown_database(save=False)` can discard a managed idalib +worker when the requesting handle is its only active lease and no other operation +is running. The server rejects GUI databases and shared workers. -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". +The coherent ownership model is the **final lease**, not necessarily the lease +that spawned the worker. Releasing a non-final lease makes no whole-database save +decision; responsibility transfers to the leases that remain. The final client +can save or discard the shared session. A client that needs its work to survive +regardless of that later decision must call `save_database()` before releasing +its lease. -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. +This does not claim to provide per-client rollback. Discard applies to all +changes since the last database save, and attempting it while another lease is +active is correctly rejected. That is the same ref-counted lifetime model used +by other shared resources and requires no separate starter capability. -**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. +The upstream gap is therefore closed. ida-tui now routes its discard action +through `shutdown_database(save=False)`: a final managed-worker lease discards, +while GUI-backed and still-shared sessions transfer finalization to their owner +or remaining leases. --- -## 7. No change notification for shared databases +## 7. No change notification for shared databases — FIXED in protocol 6 -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. +`DatabaseHandle.subscribe_idb_events()` now returns a closeable iterator over +structured IDB changes. Each event carries a monotonic revision plus +`operation_id`/`operation_label` attribution and an opaque `origin_id`. +`DatabaseHandle.owns_event()` compares that origin with the handle's lease, so a +caching client does not need to generate, retain, or race operation IDs itself. -**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. +ida-tui keeps one subscription for its active database, asks the handle to drop +its own events, and batches peer events behind a 200 ms quiet period. One batch +invalidates the function, listing, decompiler, graph, strings, linkage, segment +and byte caches, then reloads the visible view in place. Closing or switching +databases closes the subscription, so the blocking event reader does not leak. --- @@ -340,10 +359,10 @@ dataclass fields), because the offline contract suite must keep running with no | ~~1~~ | ~~`timeout_trace` line tracing~~ | ~~52x on IDA calls~~ | ✅ fixed in 0.3.2 | | ~~2~~ | ~~`to_jsonable` on large results~~ | ~~114x on serialisation~~ | ✅ fixed in 0.3.2 | | ~~3~~ | ~~2 ms `execute_sync` floor~~ | ~~shapes client design~~ | ✅ fixed in 0.3.2, 7.0x | -| 7 | no change/revision counter | correctness for shared editing | yes, cheap | -| ~~4~~ | ~~loader switches fatal on reopen~~ | ~~crashes, hard to diagnose~~ | ✅ fixed in 0.5.x (one edge, §4) | -| 5 | replaced/deleted IDB under lease | silent hang | yes | -| 6 | no close-without-save | a feature we had to drop | design question | +| ~~7~~ | ~~no change/revision counter~~ | ~~correctness for shared editing~~ | ✅ fixed in protocol 6 | +| 4 | direct `.i64` forwards loader-only options | fatal worker startup | [issue #36](https://github.com/HexRaysSA/ida-codemode/issues/36) | +| ~~5~~ | ~~replaced/deleted IDB under lease~~ | ~~out-of-contract filesystem mutation~~ | **stale** | +| ~~6~~ | ~~no close without save~~ | ~~could not discard a managed session~~ | **fixed in protocol 6: final lease decides** | | ~~8~~ | ~~package exports~~ | ~~forces internal imports~~ | ✅ fixed in 0.5.x | | ~~9~~ | ~~typed handle for fakes~~ | ~~catches a whole bug class~~ | ✅ fixed in 0.5.x (`py.typed` + options dataclass) | @@ -356,13 +375,11 @@ worth fixing centrally rather than leaving each client to rediscover. fixed upstream, and both client-side workarounds could be measured at parity and retired. That is the outcome this document was written for. -**What is left is entirely non-performance**: 0.5.x then fixed the API-surface -items (4, 8, 9) — the package root is a real public API, `py.typed` ships, and -the loader-switch fatal is handled in the resolver. What remains open is -lifecycle: **5** (replacing an IDB under a live lease), **6** (close without -save) and **7** — a monotonic revision counter on `/health`, still the cheapest -large win for any caching client, and still the one thing an *external* client -cannot build for itself. +**What is left is entirely non-performance.** Items 6 through 9 are fixed, and +item 5 is stale because out-of-band replacement is not a supported lifecycle for +either Code Mode or the IDA GUI. One narrow piece remains: **4**, normalize or +reject loader-only options when the source is itself an existing `.i64` +([issue #36](https://github.com/HexRaysSA/ida-codemode/issues/36)). 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/experiments/bench_ops.py b/experiments/bench_ops.py index 178b4bc..9cc51fa 100644 --- a/experiments/bench_ops.py +++ b/experiments/bench_ops.py @@ -27,6 +27,7 @@ ida-codemode is installed **editable** into both venvs, so checking that repo ou swaps the backend under the TUI with no reinstall -- which is what makes this A/B cheap. """ + from __future__ import annotations import argparse @@ -34,6 +35,7 @@ import os import statistics import time +from idatui import remote_ops from idatui.codemode_client import CodeModeClient @@ -59,7 +61,7 @@ def main() -> int: # Work on the biggest function we can find, so the payload-heavy operations # are actually payload-heavy. - index = client.invoke("list_funcs", queries=[{"offset": 0, "count": 60}]) + index = client.call(remote_ops.list_funcs, queries=[{"offset": 0, "count": 60}]) funcs = (index.get("result") or [{}])[0].get("data") or [] if not funcs: print("VERDICT: FAIL - no functions") @@ -71,19 +73,30 @@ def main() -> int: # Synthetic: isolates the per-operation floor (execute_sync marshalling). ("empty round trip", lambda: handle.execute_python("result = 1")), # Payload-dominated: what _PACK_EPILOGUE was written for. - ("list_funcs 500", lambda: client.invoke( - "list_funcs", queries=[{"offset": 0, "count": 500}])), - ("heads 200 (listing page)", lambda: client.invoke( - "heads", addr=ea, count=200, annotate=True)), + ( + "list_funcs 500", + lambda: client.call( + remote_ops.list_funcs, queries=[{"offset": 0, "count": 500}] + ), + ), + ( + "heads 200 (listing page)", + lambda: client.call(remote_ops.heads, addr=ea, count=200, annotate=True), + ), # IDA-work-dominated: Hex-Rays, nothing upstream can move. - ("decompile (warm)", lambda: client.invoke("decompile", addr=ea)), - ("flowchart (graph)", lambda: client.invoke("flowchart", addr=ea)), + ("decompile (warm)", lambda: client.call(remote_ops.decompile, addr=ea)), + ("flowchart (graph)", lambda: client.call(remote_ops.flowchart, addr=ea)), # Round-trip-dominated: small payload, so only the floor matters. - ("xrefs_to", lambda: client.invoke("xref_query", direction="to", addr=ea)), + ( + "xrefs_to", + lambda: client.call(remote_ops.xref_query, direction="to", addr=ea), + ), ] - print(f"# target={os.path.basename(args.target)} func={ea} reps={args.reps} " - f"backend={client.backend}") + print( + f"# target={os.path.basename(args.target)} func={ea} reps={args.reps} " + f"backend={client.backend}" + ) results = {} for name, fn in ops: try: diff --git a/experiments/bench_pack_trace.py b/experiments/bench_pack_trace.py index c6ca980..edbb9f7 100644 --- a/experiments/bench_pack_trace.py +++ b/experiments/bench_pack_trace.py @@ -1,59 +1,49 @@ -"""Measure ``_PACK_EPILOGUE`` against the live ida-codemode runtime. +"""Measure cold installation versus warm calls for typed remote modules. -Our snippets return one pre-serialised JSON STRING instead of a structure, to -dodge to_jsonable()'s Python-level walk of the result (a 200-row listing page -is ~10k small objects). ida-codemode 0.3.2 gave that path a C fast path -- -``serialization.dumps_json`` hands the structure straight to ``json.dumps`` and -only falls back to the walker for values the encoder rejects -- so the packing -now costs a double encode (escaping the whole payload as a string literal) to -avoid a walk that may no longer happen. - -This script answers whether packing still pays. Its sibling question, the -``sys.settrace`` strip, is settled: 0.3.2 deleted the trace hook, the workaround -measured 0.99x, and it has been removed. +Historical note: this file used to benchmark ``_PACK_EPILOGUE``. Application +scripts are no longer strings and packing is gone; the relevant design cost is +now the one-time content-addressed module installation versus steady-state calls. Usage:: - PYTHONPATH=. ~/ida-venv/bin/python experiments/bench_pack_trace.py [FILE] + PYTHONPATH=. python experiments/bench_pack_trace.py [FILE] """ + from __future__ import annotations import argparse -import json import os import statistics import time -from idatui import codemode_client as cc +from idatui import remote_ops from idatui.codemode_client import CodeModeClient -def _time(fn, reps: int) -> tuple[float, float]: - """Best-of and median wall time in ms; best-of resists co-tenant noise.""" +def timed(function, reps: int = 1) -> tuple[object, float]: samples = [] + result = None for _ in range(reps): started = time.perf_counter() - fn() + result = function() samples.append((time.perf_counter() - started) * 1000.0) - return min(samples), statistics.median(samples) + return result, statistics.median(samples) def main() -> int: - ap = argparse.ArgumentParser() - ap.add_argument("target", nargs="?", default="targets/bash") - ap.add_argument("--reps", type=int, default=25) - ap.add_argument("--rows", type=int, default=200) - args = ap.parse_args() + parser = argparse.ArgumentParser() + parser.add_argument("target", nargs="?", default="targets/bash") + parser.add_argument("--reps", type=int, default=25) + parser.add_argument("--rows", type=int, default=200) + args = parser.parse_args() - target = os.path.abspath(args.target) - client = CodeModeClient(target) - client.connect() + client = CodeModeClient(os.path.abspath(args.target)).connect() - # A real listing page: the flow the workaround was tuned for. - # list_funcs answers {"result": [{"data": [...], "total": N}]}. - index = client.invoke("list_funcs", queries=[{"offset": 0, "count": 40}]) + index, operations_cold = timed( + lambda: client.call(remote_ops.list_funcs, queries=[{"offset": 0, "count": 40}]) + ) funcs = (index.get("result") or [{}])[0].get("data") or [] - biggest = max(funcs, key=lambda f: f.get("size") or 0, default=None) + biggest = max(funcs, key=lambda function: function.get("size") or 0, default=None) if not biggest: print("VERDICT: FAIL - no functions") return 1 @@ -61,34 +51,26 @@ def main() -> int: if isinstance(addr, int): addr = hex(addr) - page = lambda: client.invoke( # noqa: E731 - "heads", addr=addr, count=args.rows, annotate=True) - - # _script() reads _PACK_EPILOGUE at CALL time, so both variants share one - # process -- one lease, one warm database, one fair baseline. _unpack() - # passes an unpacked answer through untouched, so plain `result` works. - packed_epilogue = cc._PACK_EPILOGUE - results = {} - for packing in (True, False): - cc._PACK_EPILOGUE = packed_epilogue if packing else "\nresult\n" - payload = page() - rows = len(payload.get("heads", [])) - for _ in range(3): # warm caches; the first sample is always an outlier - page() - results[packing] = (rows, *_time(page, args.reps)) - cc._PACK_EPILOGUE = packed_epilogue + _, operations_warm = timed( + lambda: client.call( + remote_ops.list_funcs, queries=[{"offset": 0, "count": 40}] + ), + args.reps, + ) + page = lambda: client.call( # noqa: E731 + remote_ops.heads, addr=addr, count=args.rows, annotate=True + ) + payload, tools_cold = timed(page) + _, tools_warm = timed(page, args.reps) - size = len(json.dumps(payload, separators=(",", ":"), default=str)) / 1024 - print(f"target {os.path.basename(target)} backend={client.backend}") - print(f"listing page func {addr}, {results[True][0]} rows, " - f"{size:.1f} KiB of JSON, {args.reps} reps") - for packing, label in ((True, "packed string (current)"), - (False, "plain structure")): - rows, best, med = results[packing] - print(f" {label:<26} best {best:7.2f}ms median {med:7.2f}ms" - f" rows={rows}") - print(f" packing buys " - f"{results[False][1] / results[True][1]:.2f}x") + print(f"target {os.path.basename(args.target)} backend={client.backend}") + print(f"function {addr}") + print(f"listing rows {len(payload.get('heads', []))}") + print( + f"operations.py cold {operations_cold:8.3f}ms warm {operations_warm:8.3f}ms" + ) + print(f"remote_tools.py cold {tools_cold:8.3f}ms warm {tools_warm:8.3f}ms") + print(f"tools install overhead {tools_cold / max(tools_warm, 0.001):.2f}x one time") client.close() return 0 diff --git a/experiments/profile_client.py b/experiments/profile_client.py index 77f16e7..bac130e 100644 --- a/experiments/profile_client.py +++ b/experiments/profile_client.py @@ -10,6 +10,7 @@ nothing else here can see it. Time spent in `invoke` is the backend + transport; everything below it in the `tottime` list is ours and is what this file is for. """ + from __future__ import annotations import argparse @@ -19,6 +20,7 @@ import os import pstats import time +from idatui import remote_ops from idatui.codemode_client import CodeModeClient from idatui.domain import Program @@ -28,14 +30,15 @@ def main() -> int: ap.add_argument("binary", nargs="?", default="targets/bash") ap.add_argument("--pages", type=int, default=60) ap.add_argument("--lines", type=int, default=16) - ap.add_argument("--text", action="store_true", - help="load full pages instead of skeletons") + ap.add_argument( + "--text", action="store_true", help="load full pages instead of skeletons" + ) args = ap.parse_args() client = CodeModeClient(os.path.abspath(args.binary)) client.connect() program = Program(client) - regions = client.invoke("file_regions") + regions = client.call(remote_ops.file_regions) rows = regions.get("regions") or regions.get("result") or [] text_seg = next((r for r in rows if ".text" in str(r.get("name", ""))), rows[0]) model = program.listing(int(str(text_seg["start"]), 16)) @@ -57,8 +60,10 @@ def main() -> int: pr.disable() wall = (time.perf_counter() - started) * 1000 - print(f"# {os.path.basename(args.binary)} pages={loaded} " - f"text={want_text} {wall:.0f}ms ({wall/max(loaded,1):.2f}ms/page)") + print( + f"# {os.path.basename(args.binary)} pages={loaded} " + f"text={want_text} {wall:.0f}ms ({wall / max(loaded, 1):.2f}ms/page)" + ) buf = io.StringIO() pstats.Stats(pr, stream=buf).sort_stats("tottime").print_stats(args.lines) print(buf.getvalue()) diff --git a/experiments/profile_remote.py b/experiments/profile_remote.py index ca673db..97f4fb1 100644 --- a/experiments/profile_remote.py +++ b/experiments/profile_remote.py @@ -1,94 +1,71 @@ -"""Profile an operation INSIDE the database process. +"""Profile persistent ida-tui operations inside the IDA process. -`bench_ops.py` says how long an operation takes; this says where that time -goes. The snippet ships cProfile into the Code Mode sandbox, runs the real -remote-library function there in a loop, and returns the stats as text -- so -the split between IDA's own calls and OUR python in `remote_tools.py` is -visible, which no client-side timer can see. +The profiler itself is a typed ``RemoteModule`` function in ``remote_tools.py``; +this file contains no generated Python source or knowledge of remote module names. - PYTHONPATH=. ~/ida-venv/bin/python experiments/profile_remote.py [BINARY] - PYTHONPATH=. ~/ida-venv/bin/python experiments/profile_remote.py --op decompile +Usage:: -Read the `tottime` column: time in that function excluding subcalls. IDA -builtins (generate_disasm_line, get_flags, next_head...) are the floor; a -python frame from ida_tui_remote near the top is ours, and ours is fixable. + PYTHONPATH=. python experiments/profile_remote.py [BINARY] + PYTHONPATH=. python experiments/profile_remote.py --op decompile """ + from __future__ import annotations import argparse import os -import sys - -from idatui.codemode_client import CodeModeClient, _REMOTE_MODULE, _script -# Runs in the database process. `a` is the bound argument dict. -PROFILE = ''' -import cProfile, pstats, io, sys -_m = sys.modules.get(%(mod)r) -if _m is None: - result = {"error": "remote lib not installed yet"} -else: - call = a["call"] - reps = int(a["reps"]) - ns = {"_m": _m, "a": a} - src = "for _ in range(%%d):\\n _m.%%s" %% (reps, call) - code = compile(src, "<profile>", "exec") - pr = cProfile.Profile() - pr.enable() - exec(code, ns) - pr.disable() - buf = io.StringIO() - st = pstats.Stats(pr, stream=buf).sort_stats("tottime") - st.print_stats(int(a["lines"])) - result = {"stats": buf.getvalue(), "total": st.total_tt, "reps": reps} -''' % {"mod": _REMOTE_MODULE} +from idatui import remote_ops +from idatui.codemode_client import CodeModeClient CALLS = { - # One full listing page, exactly as the background grower asks for it. - "heads": 'heads(addr=a["addr"], count=500, annotate=True)', - "heads_plain": 'heads(addr=a["addr"], count=500, annotate=False)', - "heads_skeleton": 'heads(addr=a["addr"], count=500, annotate=True, text=False)', - "decompile": 'decompile(a["addr"])', - "disasm": 'disasm(a["addr"], 500)', + "heads": ("heads", {"count": 500, "annotate": True}), + "heads_plain": ("heads", {"count": 500, "annotate": False}), + "heads_skeleton": ( + "heads", + {"count": 500, "annotate": True, "text": False}, + ), + "decompile": ("decompile", {}), } def main() -> int: - ap = argparse.ArgumentParser() - ap.add_argument("binary", nargs="?", default="targets/bash") - ap.add_argument("--op", default="heads", choices=sorted(CALLS)) - ap.add_argument("--reps", type=int, default=20) - ap.add_argument("--lines", type=int, default=18) - ap.add_argument("--addr", default=None, help="default: the .text start") - args = ap.parse_args() - - client = CodeModeClient(os.path.abspath(args.binary)) - client.connect() + parser = argparse.ArgumentParser() + parser.add_argument("binary", nargs="?", default="targets/bash") + parser.add_argument("--op", default="heads", choices=sorted(CALLS)) + parser.add_argument("--reps", type=int, default=20) + parser.add_argument("--addr", default=None, help="default: the .text start") + args = parser.parse_args() + client = CodeModeClient(os.path.abspath(args.binary)).connect() addr = args.addr if addr is None: - regions = client.invoke("file_regions") + regions = client.call(remote_ops.file_regions) rows = regions.get("regions") or regions.get("result") or [] - text = next((r for r in rows if ".text" in str(r.get("name", ""))), None) + text = next( + (row for row in rows if ".text" in str(row.get("name", ""))), + None, + ) addr = (text or rows[0])["start"] if rows else "0x0" - print(f"# {os.path.basename(args.binary)} op={args.op} addr={addr} reps={args.reps}") - # Prime: the remote lib installs lazily, and its lru_caches must be warm or - # the profile measures cache misses that the real workload never pays. - client.invoke("heads", addr=addr, count=500, annotate=True) + operation, call_args = CALLS[args.op] + call_args = {"addr": addr, **call_args} + print( + f"# {os.path.basename(args.binary)} op={args.op} " + f"addr={addr} reps={args.reps}" + ) - # _script binds the args as JSON and adds the pack epilogue, exactly as a - # real operation is shipped -- so this measures the same path, not a - # special one. - out = client._unpack(client.execute_python(_script( - {"call": CALLS[args.op], "addr": addr, - "reps": args.reps, "lines": args.lines}, PROFILE), timeout=600)) - if "error" in out: - print("FAILED:", out["error"]) - return 1 + # Install the persistent tool module and warm its caches before profiling. + client.call(remote_ops.heads, addr=addr, count=500, annotate=True) + out = client.call( + remote_ops.profile_remote, + operation=operation, + args=call_args, + reps=args.reps, + ) per = out["total"] / out["reps"] * 1000 - print(f"# {out['total']*1000:.0f}ms total, {per:.1f}ms per call\n") + print(f"# {out['total'] * 1000:.0f}ms total, {per:.1f}ms per call\n") print(out["stats"]) + client.close() return 0 diff --git a/idatui/app.py b/idatui/app.py index d7212b1..e4c2936 100644 --- a/idatui/app.py +++ b/idatui/app.py @@ -46,8 +46,7 @@ from textual.widgets import ( ) from textual.widgets.option_list import Option -from . import graph -from . import kittygfx +from . import diag, graph, kittygfx from .edit_ctl import EditController from .prompt import PromptBar from .trace_ctl import TraceController @@ -55,7 +54,7 @@ from . import findings, search from .highlight import CTextArea, highlight_c from .journal import Journal -from .errors import IDAToolError, IDAConnectionError +from .errors import IDAConnectionError from .codemode_client import CodeModeClient, registered_database from .domain import Func, Head, ListingModel, Program, Struct @@ -3749,14 +3748,14 @@ _HELP = ( class QuitScreen(ModalScreen): """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. + A final managed-worker lease can discard the whole session. Shared workers + and GUI databases keep their state: releasing this lease transfers the final + save/discard decision to the remaining client or GUI owner. """ BINDINGS = [ Binding("s", "save", "Save & quit"), - Binding("d", "discard", "Leave & quit"), + Binding("d", "discard", "Discard / leave"), Binding("escape,c", "cancel", "Cancel"), ] @@ -3772,8 +3771,11 @@ class QuitScreen(ModalScreen): body = Text() for label in self._labels: body.append(f" \u2022 {label}\n", _S_LABEL) + body.append( + "\nFinal managed leases discard; shared/GUI sessions stay open.", + _S_DIM) yield Static(body, id="quit-list") - yield Static("s save & quit d leave as-is & quit Esc cancel", + yield Static("s save & quit d discard / leave & quit Esc cancel", id="quit-help") def action_save(self) -> None: @@ -5286,6 +5288,11 @@ class IdaTui(App): self.journal = Journal() self._xref_focus_name: str | None = None self._dirty = False + # One subscription for the active database. CodeModeClient debounces + # bursts off the Textual worker pool; the callback re-enters here on the + # UI thread to invalidate and reload the visible models. + self._idb_event_watch = None + self._idb_refresh_seq = 0 # -- layout ------------------------------------------------------------ # def compose(self) -> ComposeResult: @@ -5447,6 +5454,7 @@ class IdaTui(App): self._ask_load_options(path, label=label) def _release_database(self) -> None: + self._stop_idb_event_watch() if self.program is not None: self.program.close() if self._pool is not None and self._binary is not None: @@ -5625,6 +5633,139 @@ class IdaTui(App): pass return len(text) + # -- live refresh from shared IDB changes ----------------------------- # + def _start_idb_event_watch(self, client: CodeModeClient) -> None: + self._stop_idb_event_watch() + watch = getattr(client, "watch_idb_events", None) + if watch is None: # IDA-free test doubles and pre-event adapters + return + + def changed(events) -> None: # listener thread + try: + self.call_from_thread(self._refresh_idb_events, client, events) + except Exception: # noqa: BLE001 -- app teardown can win this race + pass + + def failed(error: BaseException) -> None: # listener thread + try: + self.call_from_thread(self._idb_event_watch_failed, client, error) + except Exception: # noqa: BLE001 -- app teardown can win this race + pass + + self._idb_event_watch = watch( + changed, on_error=failed, debounce=0.2) + + def _stop_idb_event_watch(self) -> None: + watcher, self._idb_event_watch = self._idb_event_watch, None + if watcher is not None: + watcher.close() + + def _idb_event_watch_failed( + self, client: CodeModeClient, error: BaseException + ) -> None: + if client is not self.client: + return + if isinstance(error, IDAConnectionError): + self._on_connection_lost() + else: + self._status(f"live database refresh stopped: {error}") + + def _listing_event_anchor(self) -> ViewAnchor: + """Capture the listing position even when the split's decompiler has focus.""" + anchor = ViewAnchor(view=self._active) + listing = self.query_one(ListingView) + model = listing.model + if model is None: + return anchor + anchor.cursor_x = listing.cursor_x + anchor.ea = listing._cursor_ea() + top = round(listing.scroll_offset.y) + head = model.cached_line(top) or model.get(top) + anchor.top_ea = getattr(head, "ea", None) + return anchor + + def _refresh_idb_events( + self, client: CodeModeClient, events: tuple[dict, ...] + ) -> None: + """Invalidate once per external edit burst and reload the active surface.""" + program = self.program + if not events or client is not self.client or program is None \ + or program.client is not client: + return + self._idb_refresh_seq += 1 + seq = self._idb_refresh_seq + entry = self._cur + anchor = self._listing_event_anchor() + hex_ea = self.query_one(HexView).cursor_va() if self.is_hex else None + graph_ea = self.query_one(GraphView)._cursor_ea() if self.is_graph else None + decomp = self.query_one(DecompView) + if entry is not None and decomp.loaded_ea == entry.ea: + entry.dec_cursor = decomp.cursor + entry.dec_cursor_x = decomp.cursor_x + entry.dec_scroll_y = round(decomp.scroll_offset.y) + entry.dec_scroll_x = round(decomp.scroll_offset.x) + + program.invalidate_external() + self._status( + f"{len(events)} external database " + f"change{'s' if len(events) != 1 else ''} — refreshing…") + self._reindex_functions() + + if self.is_hex: + self.query_one(HexView).model = None + self._load_hex_model(hex_ea) + return + if self.is_graph and entry is not None: + self._load_graph(entry.ea, graph_ea or entry.ea) + return + if entry is None: + return + + # A split needs both halves rebuilt; a decompiler-only view still keeps + # the hidden listing fresh so Tab does not reveal pre-event rows. + decomp.loaded_ea = None + self._reload_idb_listing(program, seq, entry, anchor) + if self.is_decomp or self._split: + self._show_active() + + @work(thread=True, exclusive=True, group="idb-refresh") + def _reload_idb_listing( + self, program: Program, seq: int, entry: NavEntry, anchor: ViewAnchor + ) -> None: + target = anchor.ea if anchor.ea is not None else entry.ea + model = program.listing(target) + cursor = top = -1 + if model is not None: + model.ensure_ea(target) + cursor, top = self._anchor_rows(anchor, model, target) + fn = program.function_of(entry.ea) + name = fn.name if fn is not None else program.region_label(entry.ea) + self.app.call_from_thread( + self._apply_idb_listing, program, seq, entry, model, + cursor, top, anchor.cursor_x, name, fn is None) + + def _apply_idb_listing( + self, program: Program, seq: int, entry: NavEntry, model, + cursor: int, top: int, cursor_x: int, name: str, is_region: bool, + ) -> None: + if program is not self.program or seq != self._idb_refresh_seq \ + or entry is not self._cur: + return + if model is None: + self._status(f"{entry.ea:#x} is no longer in a loaded segment") + return + entry.name = name + entry.is_region = is_region + entry.cursor = max(cursor, 0) + entry.cursor_x = cursor_x + if top >= 0: + entry.scroll_y = top + self.query_one(ListingView).load( + model, name, cursor=entry.cursor, cursor_x=entry.cursor_x, + scroll_y=top if top >= 0 else None) + if self.is_listing: + self._show_active() + # -- connection loss / recovery --------------------------------------- # def _handle_exception(self, error: BaseException) -> None: """Intercept a lost Code Mode lease so the app can rediscover the DB. @@ -5662,15 +5803,22 @@ class IdaTui(App): @work(thread=True, exclusive=True, group="reconnect") def _reconnect(self) -> None: - # The registered instance disappeared. Rediscover it; Code Mode may find - # a GUI/replacement worker, then we rebuild caches against the new handle. + # Rediscovery is attach-only. If a GUI owner closes its database, a TUI + # must not silently reopen it by spawning a headless worker. try: if self._open_path is None: self.app.call_from_thread(self._reconnect_failed, "no binary to reopen") return - client = CodeModeClient(self._open_path, ttl=self._ttl, - load_args=self._load_args) + if self._project is not None and self._binary is not None: + ref = self._project.by_label(self._binary) + client = CodeModeClient( + ref.staged, ttl=self._ttl, load_args=ref.load_args, + output_database=ref.db, spawn=False) + else: + client = CodeModeClient( + self._open_path, ttl=self._ttl, + load_args=self._load_args, spawn=False) client.connect(progress=lambda m: self.app.call_from_thread( self._conn_note, m)) except Exception as e: # noqa: BLE001 @@ -5679,20 +5827,32 @@ class IdaTui(App): self.app.call_from_thread(self._after_reconnect, client, Program(client)) def _after_reconnect(self, client: "CodeModeClient", program: "Program") -> None: + old_client, old_program = self.client, self.program + self._stop_idb_event_watch() + if old_program is not None: + old_program.close() + if self._pool is not None and self._binary is not None: + self._pool.replace_client(self._binary, old_client, client) + if old_client is not None and old_client is not client: + old_client.close() self.client = client self.program = program + self._start_idb_event_watch(client) self._reconnecting = False + self._dirty = False self._dismiss_conn() - self._status("reconnected \u2014 reloading\u2026") - self._load_functions() # rebuild the function index against the new client + self._status("reattached — reloading persisted state…") + self._load_functions() cur = self._cur - if cur is not None: # refresh the current view with the new program + if cur is not None: self._open_entry(cur, push=False) def _reconnect_failed(self, why: str) -> None: self._reconnecting = False - self._conn_note(f"reconnect failed: {why} \u2014 retry on next action, or 'q'") - self._status(f"reconnect failed: {why}") + note = (f"database owner closed: {why} — reopen it in IDA, then " + "Esc and retry an action; or q to quit") + self._conn_note(note) + self._status(note) # -- connection + initial load ---------------------------------------- # @work(thread=True, exclusive=True, group="connect") @@ -5719,6 +5879,7 @@ class IdaTui(App): return self.client = client self.program = program + self._start_idb_event_watch(client) self._new_database = False self.app.call_from_thread( self._status, f"{module} [{client.backend}] — loading functions…") @@ -6115,11 +6276,13 @@ 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() + # Whole-session discard is legal only for the final managed lease. + # GUI/shared sessions retain state and inherit finalization. + dirty = self._dirty_labels() + self._loading_screen = LoadingScreen( + "discarding", note="finalizing database leases…") + self.push_screen(self._loading_screen) + self._discard_then_exit(dirty) elif choice == "save": # Save with the overlay up: writing a big .i64 takes seconds, and # doing it during teardown would look like a hang with no UI left. @@ -6129,6 +6292,31 @@ class IdaTui(App): # None: cancel, stay put @work(thread=True, exclusive=True, group="save-exit") + def _discard_then_exit(self, dirty: list[str]) -> None: + try: + if self._pool is not None: + transferred = self._pool.discard_changes(dirty) + elif self.client is not None: + transferred = [] if self.client.discard_database() else dirty + else: + transferred = dirty + except Exception as exc: # noqa: BLE001 -- keep the app open on failure + self.app.call_from_thread(self._discard_failed, str(exc)) + return + self.app.call_from_thread(self._finish_discard, transferred) + + def _discard_failed(self, why: str) -> None: + self._dismiss_loading() + self._status(f"discard failed: {why}", priority=True) + + def _finish_discard(self, transferred: list[str]) -> None: + if transferred and self._loading_screen is not None: + labels = ", ".join(transferred) + self._loading_screen.update_note( + f"finalization transferred: {labels}") + self._finish_exit() + + @work(thread=True, exclusive=True, group="save-exit") def _save_then_exit(self) -> None: try: if self._pool is not None: @@ -6140,7 +6328,9 @@ class IdaTui(App): self.app.call_from_thread(self._finish_exit) def _finish_exit(self) -> None: - self._save_on_exit = False # already written above + # Teardown must not save again: save, discard, or ownership transfer was + # already decided by the quit path. + self._save_on_exit = False self._dirty = False self.exit() @@ -6207,6 +6397,7 @@ class IdaTui(App): def _after_switch(self, label, client, program, st, reuse) -> None: # type: ignore[no-untyped-def] self.client = client self.program = program + self._start_idb_event_watch(client) self._binary = label self._pool.set_active(label) self._open_path = self._project.by_label(label).staged @@ -8374,6 +8565,7 @@ class IdaTui(App): # -- teardown ---------------------------------------------------------- # async def on_unmount(self) -> None: + self._stop_idb_event_watch() if self._rpc is not None: await self._rpc.stop() if self._ka is not None: diff --git a/idatui/codemode_client.py b/idatui/codemode_client.py index 8eb66b5..e95238f 100644 --- a/idatui/codemode_client.py +++ b/idatui/codemode_client.py @@ -5,24 +5,20 @@ 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. +Remote operations are ordinary typed Python functions declared in +``idatui.remote_ops``. Code Mode installs their content-addressed modules once +per handle; subsequent calls send only encoded arguments. The optimized +IDAPython listing/decompiler implementation remains real source in +``idatui.remote_tools`` and is installed through the same module interface. """ + from __future__ import annotations -import hashlib -import json import os import shlex import threading import time -from pathlib import Path -from textwrap import dedent +from collections.abc import Callable from typing import Any from .errors import IDAConnectionError, IDATimeoutError, IDAToolError, Session @@ -68,7 +64,8 @@ def _require_codemode() -> None: "ida-codemode is not installed in this environment " f"({_CODEMODE_ERROR}). Install it (e.g. `uv sync`, or " "`pip install ida-codemode`) so ida-tui can lease a " - "database.") from _CODEMODE_ERROR + "database." + ) from _CODEMODE_ERROR def database_owner(idb_path: str, staged_path: str | None = None): @@ -155,1106 +152,119 @@ def _parse_load_args(value: str) -> tuple[str | None, int | None, str | None]: 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. -#: -#: Written when Code Mode ran to_jsonable() over every snippet result, walking -#: the whole structure in Python to make it JSON-safe: a 200-row listing page is -#: ~10k small objects, which cost 66ms to walk -- 72% of the page's total cost, -#: and 114x what json.dumps of the same data cost (0.58ms). -#: -#: ida-codemode 0.3.2 removed that reason: serialization.dumps_json now hands -#: the structure straight to the C encoder and only falls back to the walker for -#: values json.dumps rejects. Re-measured against 0.3.2, packing buys 0.97x on -#: that same page (experiments/bench_pack_trace.py) -- i.e. nothing, because the -#: dodged walk is replaced by a double encode. -#: -#: It is kept anyway, on correctness rather than speed: packing pins OUR encoder -#: settings (compact separators, default=str) inside the database process, so an -#: un-encodable IDA object degrades to repr() at a point we control instead of -#: depending on the runtime's fallback. Delete it if that stops being worth a -#: protocol step -- it is no longer load-bearing for performance. -_PACK_EPILOGUE = ( - '\n{"' + _PACKED + '": json.dumps(result, separators=(",", ":"), default=str)}\n' -) - +class IDBEventListener: + """Debounced, closeable delivery of another client's IDB changes. -def _script(args: dict[str, Any], body: str) -> str: - """Bind JSON arguments without interpolating user text into Python code. - - This used to also run the body with Code Mode's trace hook detached - (sys.settrace(None) + restore), because the runtime wrapped every - execute_python in a trace function that returned ITSELF -- enabling line - tracing in every frame it saw, so every line of every function we called - paid a Python-level callback (ida_bytes.get_flags: 0.106us -> 5.49us, 52x). - - ida-codemode 0.3.2 deleted that hook; cancellation is now a C-level thread - interrupt (runtime._interrupt_thread) that costs nothing while idle. The - workaround measured 0.99x on a 200-row listing page against 0.3.2 -- pure - noise -- so it is gone, and with it the caveat that a pure-Python loop in a - snippet escaped its deadline. See experiments/bench_pack_trace.py. + Code Mode's subscription is a blocking iterator, so one daemon thread reads + it and a second waits for a quiet period before handing a batch to the UI. + Keeping the debounce here avoids a permanent Textual worker (which would + make the app's worker-idle contract impossible) and bounds refresh work to + one pass per edit burst. """ - encoded = json.dumps(args, ensure_ascii=False, separators=(",", ":")) - head = f"import json\na = json.loads({encoded!r})\n" - return f"{head}{dedent(body).strip()}\n{_PACK_EPILOGUE}" - - -_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 -''', - # Everything a person ADDED to the database: comments, non-dummy names, and - # the prototypes they set. - # - # Names come from IDA's name list, which is already an index -- no scan at - # all. Comments have no index, so they need a walk, and the walk is over - # HEADS: `next_that`'s predicate is a *Python* callback (SWIG calls it with - # one argument, so `f_has_cmt` does not even fit), which would be one call - # per BYTE -- 400 million of them on a big image. `max_scan` bounds it and - # reports `truncated` rather than sitting there. - "list_annotations": r''' -import ida_bytes, ida_funcs, ida_lines, ida_nalt, ida_name -import ida_segment, ida_typeinf, idautils -limit = max(1, int(a.get("limit", 4000))) -max_scan = max(1000, int(a.get("max_scan", 2000000))) -comments, names = [], [] -scanned = 0 - -def _line(ea): - try: - txt = ida_lines.generate_disasm_line(ea, ida_lines.GENDSM_REMOVE_TAGS) - except Exception: - txt = "" - return " ".join((txt or "").split()) -for ea, nm in idautils.Names(): - if len(names) >= limit: - break - if not nm or not ida_bytes.has_user_name(ida_bytes.get_flags(ea)): - continue - fn = ida_funcs.get_func(ea) - is_fn = fn is not None and int(fn.start_ea) == int(ea) - proto = None - if is_fn: - try: - ti = ida_typeinf.tinfo_t() - if ida_nalt.get_tinfo(ti, ea): - proto = str(ti) - except Exception: - proto = None - seg = ida_segment.getseg(ea) - names.append({"addr": hex(int(ea)), "name": nm, "func": is_fn, - "size": (int(fn.end_ea - fn.start_ea) if is_fn else 0), - "proto": proto, - "seg": (ida_segment.get_segm_name(seg) if seg else "")}) + def __init__( + self, + client: "CodeModeClient", + callback: Callable[[tuple[dict[str, Any], ...]], None], + *, + on_error: Callable[[BaseException], None] | None = None, + debounce: float = 0.2, + ) -> None: + self._client = client + self._callback = callback + self._on_error = on_error + self._debounce = max(float(debounce), 0.0) + self._condition = threading.Condition() + self._closed = False + self._subscription = None + self._pending: list[dict[str, Any]] = [] + self._deadline = 0.0 + self._reader = threading.Thread( + target=self._read, name="idatui-idb-events", daemon=True + ) + self._deliverer = threading.Thread( + target=self._deliver, name="idatui-idb-refresh", daemon=True + ) + self._deliverer.start() + self._reader.start() -for i in range(ida_segment.get_segm_qty()): - seg = ida_segment.getnseg(i) - if seg is None or len(comments) >= limit or scanned >= max_scan: - continue - for ea in idautils.Heads(seg.start_ea, seg.end_ea): - scanned += 1 - if len(comments) >= limit or scanned >= max_scan: - break - if not ida_bytes.has_cmt(ida_bytes.get_flags(ea)): - continue - for rep in (False, True): - text = ida_bytes.get_cmt(ea, rep) - if text: - fn = ida_funcs.get_func(ea) - comments.append({ - "addr": hex(int(ea)), "text": text, "repeatable": rep, - "line": _line(ea), "seg": ida_segment.get_segm_name(seg), - "func": (ida_funcs.get_func_name(fn.start_ea) if fn else None), - "func_addr": (hex(int(fn.start_ea)) if fn else None)}) + def _report(self, error: BaseException) -> None: + disconnected = DatabaseDisconnectedError + if isinstance(disconnected, type) and isinstance(error, disconnected): + error = self._client._connection_error(error) + with self._condition: + closed = self._closed + if not closed and self._on_error is not None: + self._on_error(error) -# Whole-function comments are not on the byte flags, so the scan cannot see them. -for fn_ea in idautils.Functions(): - fn = ida_funcs.get_func(fn_ea) - if fn is None or len(comments) >= limit: - continue - for rep in (False, True): - text = ida_funcs.get_func_cmt(fn, rep) - if text: - seg = ida_segment.getseg(fn_ea) - comments.append({"addr": hex(int(fn_ea)), "text": text, - "repeatable": rep, "line": "", "whole_func": True, - "seg": (ida_segment.get_segm_name(seg) if seg else ""), - "func": ida_funcs.get_func_name(fn_ea), - "func_addr": hex(int(fn_ea))}) -result = {"comments": comments, "names": names, "scanned": scanned, - "truncated": (len(comments) >= limit or len(names) >= limit - or scanned >= max_scan)} -result -''', - # The findings journal (idatui/journal.py). A netnode blob rides along in - # the .i64, so "what did I work out here" survives closing the database. - "journal_get": r''' -import ida_netnode -n = ida_netnode.netnode(a.get("node", "$ idatui.journal")) -blob = n.getblob(0, "I") if ida_netnode.exist(n) else None -result = {"data": blob.decode("utf-8", "replace") if blob else ""} -result -''', - "journal_put": r''' -import ida_netnode -n = ida_netnode.netnode(a.get("node", "$ idatui.journal"), 0, True) -payload = (a.get("data") or "").encode("utf-8") -n.setblob(payload, 0, "I") -result = {"ok": True, "bytes": len(payload)} -result -''', - # Database-wide search (Ctrl+F), two kinds. - # - # BYTES uses IDA's own `find_bytes`, which already understands the pattern - # language people expect -- "B8 ? ? ? ? 90", nibble wildcards ("48 8? ??") - # and quoted literals -- so we neither parse nor match anything ourselves. - # Iterating is match+1, per its documented contract. - "search_bytes": r''' -import ida_bytes, ida_funcs, ida_idaapi, ida_lines, ida_segment -pat = str(a.get("pattern", "")).strip() -limit = max(1, int(a.get("limit", 500))) -lo = int(a.get("start", 0)) -hi = int(a.get("end", 0)) or ida_idaapi.BADADDR -flags = ida_bytes.BIN_SEARCH_FORWARD | ida_bytes.BIN_SEARCH_NOSHOW -if a.get("case"): - flags |= ida_bytes.BIN_SEARCH_CASE -rows, err, ea = [], None, lo -while len(rows) < limit: - try: - hit = ida_bytes.find_bytes(pat, range_start=ea, range_end=hi, flags=flags) - except Exception as exc: - err = str(exc) or exc.__class__.__name__ - break - if hit is None or hit == ida_idaapi.BADADDR: - break - head = ida_bytes.get_item_head(hit) - fn = ida_funcs.get_func(hit) - seg = ida_segment.getseg(hit) - try: - line = ida_lines.generate_disasm_line(head, ida_lines.GENDSM_REMOVE_TAGS) or "" - except Exception: - line = "" - rows.append({"addr": hex(int(hit)), "head": hex(int(head)), - "line": " ".join(line.split()), - "func": (ida_funcs.get_func_name(fn.start_ea) if fn else None), - "func_addr": (hex(int(fn.start_ea)) if fn else None), - "seg": (ida_segment.get_segm_name(seg) if seg else "")}) - ea = int(hit) + 1 -result = {"hits": rows, "error": err, "truncated": len(rows) >= limit} -result -''', - # TEXT walks the listing the way a person reads it: every head's rendered - # disassembly line, which is why it finds "call cs:__isoc99_scanf" and - # "0deadbeefh" alike. Bounded by max_scan, so a 400MB image reports partial - # results instead of stalling. - "search_text": r''' -import ida_lines, ida_funcs, ida_segment, idautils -import re as _re -q = str(a.get("query", "")) -limit = max(1, int(a.get("limit", 500))) -max_scan = max(1000, int(a.get("max_scan", 3000000))) -ci = (not a.get("case")) and q.islower() # smartcase, like the in-view search -rx, err = None, None -if a.get("regex"): - try: - rx = _re.compile(q, _re.I if ci else 0) - except Exception as exc: - err = "bad regex: " + str(exc) -needle = q.lower() if ci else q -rows, scanned = [], 0 -if err is None and q: - for i in range(ida_segment.get_segm_qty()): - seg = ida_segment.getnseg(i) - if seg is None or len(rows) >= limit or scanned >= max_scan: - continue - for ea in idautils.Heads(seg.start_ea, seg.end_ea): - scanned += 1 - if len(rows) >= limit or scanned >= max_scan: - break - try: - line = ida_lines.generate_disasm_line(ea, ida_lines.GENDSM_REMOVE_TAGS) or "" - except Exception: - continue - # Match what the user SEES, not IDA's column padding: nobody types - # "call" + four spaces + "cs:getenv_ptr". - line = " ".join(line.split()) - hay = line.lower() if ci else line - if (rx.search(line) if rx is not None else (needle in hay)): - fn = ida_funcs.get_func(ea) - rows.append({"addr": hex(int(ea)), "head": hex(int(ea)), - "line": line, - "func": (ida_funcs.get_func_name(fn.start_ea) if fn else None), - "func_addr": (hex(int(fn.start_ea)) if fn else None), - "seg": ida_segment.get_segm_name(seg)}) -result = {"hits": rows, "error": err, "scanned": scanned, - "truncated": len(rows) >= limit or scanned >= max_scan} -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 + def _read(self) -> None: 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)): + subscription = self._client.subscribe_idb_events() + except Exception as exc: # noqa: BLE001 -- surfaced through on_error + self._report(exc) + with self._condition: + self._closed = True + self._pending.clear() + self._condition.notify_all() + return + with self._condition: + if self._closed: + subscription.close() + return + self._subscription = subscription 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. -#: Row count + seek anchors for a whole segment, in ONE call. See -#: remote_tools.segment_index: the alternative is fetching every row. -_OPERATIONS["segment_index"] = _remote_op( - 'segment_index(addr=a["addr"], end=a.get("end", ""),' - ' page_rows=int(a.get("page_rows", 500)), detail=bool(a.get("detail", False)))') - -_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", ""), text=bool(a.get("text", True)))') - + for event in subscription: + with self._condition: + if self._closed: + break + if self._client.owns_event(event): + continue + with self._condition: + if self._closed: + break + self._pending.append(event) + self._deadline = time.monotonic() + self._debounce + self._condition.notify_all() + except Exception as exc: # noqa: BLE001 -- stream failures are recoverable + self._report(exc) + finally: + subscription.close() + with self._condition: + if self._subscription is subscription: + self._subscription = None + self._closed = True + self._pending.clear() + self._condition.notify_all() -# 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 -''' + def _deliver(self) -> None: + while True: + with self._condition: + while not self._closed and not self._pending: + self._condition.wait() + if self._closed: + return + remaining = self._deadline - time.monotonic() + if remaining > 0: + self._condition.wait(remaining) + continue + batch = tuple(self._pending) + self._pending.clear() + try: + self._callback(batch) + except Exception as exc: # noqa: BLE001 -- keep the stream alive + self._report(exc) -# 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 -''' + def close(self) -> None: + """Stop delivery and unblock the subscription reader.""" + with self._condition: + if self._closed: + return + self._closed = True + self._pending.clear() + subscription = self._subscription + self._condition.notify_all() + if subscription is not None: + subscription.close() class CodeModeClient: @@ -1277,7 +287,9 @@ class CodeModeClient: 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._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 @@ -1289,10 +301,17 @@ class CodeModeClient: 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 + handle = self._handle + if handle is not None: + if handle.connected: + return self + raise IDAConnectionError( + "Code Mode database disconnected; explicit rediscovery required" + ) if progress: - progress(f"discovering Code Mode database for {os.path.basename(self._path)}…") + 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 @@ -1321,7 +340,9 @@ class CodeModeClient: if not self._new_database or time.monotonic() >= deadline: raise if progress: - progress("waiting for the previous Code Mode lease to close…") + progress( + "waiting for the previous Code Mode lease to close…" + ) owner = find_database_owner( self._path, output_database=self._output_database, @@ -1336,7 +357,9 @@ class CodeModeClient: time.sleep(0.2) if progress: backend = handle.instance.backend - progress(f"attached to {backend} database; waiting for auto-analysis…") + 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 @@ -1360,61 +383,60 @@ class CodeModeClient: def backend(self) -> str | None: return self._handle.instance.backend if self._handle is not None else None - def execute_python(self, code: str, *, timeout: float | None = None) -> Any: + def owns_event(self, event: dict[str, Any]) -> bool: + """Whether ``event`` was produced through this client's handle.""" + handle = self._handle + return handle is not None and handle.owns_event(event) + + def subscribe_idb_events(self): + """Open Code Mode's closeable IDB-change iterator.""" 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 + return handle.subscribe_idb_events() except (DatabaseDisconnectedError, CodeModeConnectionError) 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 watch_idb_events( + self, + callback: Callable[[tuple[dict[str, Any], ...]], None], + *, + on_error: Callable[[BaseException], None] | None = None, + debounce: float = 0.2, + ) -> IDBEventListener: + """Deliver external IDB changes in debounced batches.""" + return IDBEventListener(self, callback, on_error=on_error, debounce=debounce) - 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}") + def call(self, operation: Callable[..., Any], /, **args) -> Any: + """Execute one source-backed remote declaration through this client.""" + name = getattr(operation, "__name__", "remote 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 + from .remote_ops import bind - # Temporary source compatibility for external drivers/tests that used the - # old WorkerClient. Application code uses the accurately named invoke(). - call = invoke + remote = bind(operation) + except KeyError as exc: + raise IDAToolError( + name, f"remote operation {name!r} is not registered" + ) from exc + if not self.connected: + self.connect() + handle = self._handle + if handle is None: + raise IDAConnectionError("Code Mode database is not connected") + try: + return remote(handle, **args) + except RemoteError as exc: + message = str(exc) + if exc.details.get("traceback"): + message += f"\n{exc.details['traceback']}" + if exc.code == "operation_timeout": + raise IDATimeoutError(message) from exc + raise IDAToolError(name, message) from exc + except (DatabaseDisconnectedError, CodeModeConnectionError) as exc: + raise self._connection_error(exc) from exc def save_database(self) -> dict[str, Any]: if not self.connected: @@ -1429,6 +451,35 @@ class CodeModeClient: except (DatabaseDisconnectedError, CodeModeConnectionError) as exc: raise self._connection_error(exc) from exc + def discard_database(self, timeout: float = 5.0) -> bool: + """Discard a final managed-worker lease; otherwise transfer finalization. + + ``False`` is an expected ownership result: a GUI owns its session, or + another lease still shares the managed worker. A busy final worker is + retried briefly so background reads finishing during quit do not turn a + real discard into an implicit save. + """ + handle = self._handle + if handle is None or not handle.connected: + return False + entry = handle.instance + if entry.backend != "idalib" or not getattr(entry, "managed", False): + return False + deadline = time.monotonic() + max(float(timeout), 0.0) + while True: + try: + handle.shutdown_database(save=False) + return True + except RemoteError as exc: + if exc.code in ("instance_shared", "shutdown_not_supported"): + return False + if exc.code == "instance_busy" and time.monotonic() < deadline: + time.sleep(0.05) + continue + raise IDAToolError("shutdown_database", str(exc)) from exc + except (DatabaseDisconnectedError, CodeModeConnectionError) as exc: + raise self._connection_error(exc) from exc + def health(self) -> dict[str, Any]: if not self.connected: self.connect() @@ -1463,8 +514,14 @@ class CodeModeClient: assert self._handle is not None entry = self._handle.instance 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)] + 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 diff --git a/idatui/domain.py b/idatui/domain.py index 699f33d..b042332 100644 --- a/idatui/domain.py +++ b/idatui/domain.py @@ -22,13 +22,12 @@ import bisect import re import threading from base64 import b64decode +from collections.abc import Callable, Sequence from concurrent.futures import ThreadPoolExecutor -from collections.abc import Sequence from dataclasses import dataclass, field, replace -from typing import NamedTuple -from typing import Callable, TYPE_CHECKING +from typing import TYPE_CHECKING, NamedTuple -from . import diag +from . import remote_ops from .errors import IDAToolError if TYPE_CHECKING: # type hint only @@ -37,8 +36,7 @@ if TYPE_CHECKING: # type hint only # Clamps derived from measured caps (list ~700, disasm ~500). Margin included. LIST_PAGE = 500 DISASM_BLOCK = 256 # instructions per cached/fetched block (<= disasm cap) -HEX_BLOCK = 16384 # bytes per cached/fetched hex block (compact read_raw -> cheap) -DECOMPILE_TIMEOUT = 15.0 # s; cap per decompile so a failing one can't hang the CLI +HEX_BLOCK = 16384 # bytes per cached/fetched hex block (compact read_raw -> cheap) _TRUNC_RE = re.compile(r"\[(\d+) chars total\]\s*$") @@ -107,7 +105,7 @@ class Head(NamedTuple): """ ea: int - kind: str # 'code' | 'data' | 'unknown' | 'member' + kind: str # 'code' | 'data' | 'unknown' | 'member' size: int text: str name: str | None = None @@ -199,10 +197,10 @@ class Ref: @dataclass class Xref: - frm: int # the referencing address - to: int | None # the referenced address - type: str # coarse: "code" | "data" - fn_name: str | None # function containing `frm` + frm: int # the referencing address + to: int | None # the referenced address + type: str # coarse: "code" | "data" + fn_name: str | None # function containing `frm` fn_addr: int | None kind: str | None = None # fine: call/jump/flow/read/write/offset/text/info @@ -218,7 +216,7 @@ class LVar: class FuncTypes: addr: int name: str - prototype: str # e.g. 'int __fastcall foo(int a, char *b)' + prototype: str # e.g. 'int __fastcall foo(int a, char *b)' lvars: list[LVar] @@ -227,7 +225,7 @@ class Struct: name: str size: int is_union: bool - members: int # field count + members: int # field count ordinal: int @classmethod @@ -244,6 +242,7 @@ class Struct: @dataclass(frozen=True) class StrLit: """A string literal IDA found in the binary (the Shift+F12 list).""" + addr: int text: str length: int @@ -271,6 +270,7 @@ class SearchHit: an instruction, so ``head`` is the item to navigate to and ``line`` is what that item renders as. """ + addr: int head: int line: str = "" @@ -287,6 +287,7 @@ class Comment: report can show what was being commented ON without a second round trip. ``whole_func`` marks a function comment rather than an instruction one. """ + addr: int text: str repeatable: bool = False @@ -302,6 +303,7 @@ class NamedItem: """An address carrying a real name -- one you typed, or one the file's own symbols supplied. IDA records both as "user" names and does not remember which was which, so a report must say so rather than claim authorship.""" + addr: int name: str is_func: bool = False @@ -319,6 +321,7 @@ class Linkage: ``name`` is the joinable name; ``raw`` keeps the spelling IDA reported, which is what the user sees in the listing. """ + addr: int name: str module: str = "" @@ -374,7 +377,9 @@ class FunctionIndex: query: dict = {"offset": offset, "count": LIST_PAGE} if self.filter: query["filter"] = self.filter - data = _query_data(self._prog.client.invoke("list_funcs", queries=[query])) + data = _query_data( + self._prog.client.call(remote_ops.list_funcs, queries=[query]) + ) added = 0 with self._lock: for d in data: @@ -468,7 +473,7 @@ class DisasmModel: self._blocks: dict[int, list[Line]] = {} self._total: int | None = None self._ea_list: list[int] | None = None - self._max_raw = 0 # widest opcode length seen (bytes) + self._max_raw = 0 # widest opcode length seen (bytes) self._func_end: int | None = None self._func_end_done = False self._lock = threading.Lock() @@ -484,8 +489,8 @@ 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.invoke( - "disasm", addr=hex(self.ea), max_instructions=1, include_total=True + payload = self._prog.client.call( + remote_ops.disasm, addr=hex(self.ea), max_instructions=1, include_total=True ) total = payload.get("total_instructions") if total is None: @@ -527,7 +532,7 @@ class DisasmModel: nxt = lines[i + 1].ea if i + 1 < len(lines) else last_end length = max(nxt - ln.ea, 0) off = ln.ea - start - b = bytes(data[off:off + length]) + b = bytes(data[off : off + length]) biggest = max(biggest, len(b)) out.append(replace(ln, raw=b)) with self._lock: @@ -538,20 +543,22 @@ class DisasmModel: @staticmethod def _line_from_head(r: dict) -> Line: """Adapt a ``heads`` row to a disasm Line (label = the head's name).""" - return Line(ea=_as_int(r["ea"]), text=r.get("text", ""), - label=r.get("name")) + return Line(ea=_as_int(r["ea"]), text=r.get("text", ""), label=r.get("name")) def _fetch_block(self, b: int) -> list[Line]: # 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.invoke( - "heads", addr=hex(self.ea), offset=b * self.BLOCK, - count=self.BLOCK + 1, **self._end_kw(), + payload = self._prog.client.call( + remote_ops.heads, + addr=hex(self.ea), + offset=b * self.BLOCK, + count=self.BLOCK + 1, + **self._end_kw(), ) rows = payload.get("heads", []) if isinstance(payload, dict) else [] fetched = [self._line_from_head(r) for r in rows] - lines = fetched[:self.BLOCK] + lines = fetched[: self.BLOCK] if len(fetched) > self.BLOCK: end_ea: int | None = fetched[self.BLOCK].ea else: # this block ends the function @@ -607,7 +614,7 @@ class DisasmModel: block = self._get_block(b) lo = start - b * self.BLOCK if b == b0 else 0 hi = end - b * self.BLOCK if b == b1 else self.BLOCK - out.extend(block[max(lo, 0):hi]) + out.extend(block[max(lo, 0) : hi]) if prefetch: self._prefetch_block(b1 + 1) # forward scroll self._prefetch_block(b0 - 1) # backward scroll @@ -699,8 +706,9 @@ class ListingModel: #: _text_gen, which counts up from 0, so such a page always reads as stale. _SKELETON_GEN = -1 - def __init__(self, program: "Program", seg_start: int, seg_end: int, - name: str | None = None): + def __init__( + self, program: "Program", seg_start: int, seg_end: int, name: str | None = None + ): self._prog = program self.seg_start = seg_start self.seg_end = seg_end @@ -714,7 +722,7 @@ class ListingModel: # N bytes PRESENTS as N rows and the text for each is synthesised on # demand. _row_at[i] is the logical row where physical head i starts. self._row_at: list[int] = [] - self._head_eas: list[int] = [] # parallel to _heads, for bisect + self._head_eas: list[int] = [] # parallel to _heads, for bisect #: Which name generation each head's TEXT was rendered at, parallel to #: _heads. A rename bumps :attr:`_text_gen`; the rows themselves stay #: (their addresses and row numbers are unchanged) and are re-rendered a @@ -740,7 +748,7 @@ class ListingModel: #: means something DID move the walk. Program.listing() throws the model #: away when it sees this, so the next read rebuilds from scratch. self.stale_structure = False - self._rows = 0 # total logical rows loaded + self._rows = 0 # total logical rows loaded self._ubytes: dict[int, bytes] = {} # lazily-read bytes for those rows self._next: int | None = seg_start # next address to fetch from self._done = False @@ -788,7 +796,7 @@ class ListingModel: size = int(r.get("size") or 0) if size > 0: off = _as_int(r["ea"]) - lo - raw = bytes(data[off:off + size]) + raw = bytes(data[off : off + size]) if len(raw) > biggest: biggest = len(raw) try: @@ -827,19 +835,26 @@ class ListingModel: """ with self._lock: if self._done and self._heads: - return True # already indexed; re-priming is a no-op + return True # already indexed; re-priming is a no-op try: - idx = self._prog.client.invoke( - "segment_index", addr=hex(self.seg_start), end=hex(self.seg_end), - page_rows=self.PAGE, detail=True) + idx = self._prog.client.call( + remote_ops.segment_index, + addr=hex(self.seg_start), + end=hex(self.seg_end), + page_rows=self.PAGE, + detail=True, + ) except Exception: # noqa: BLE001 -- fall back to streaming return False if not isinstance(idx, dict) or idx.get("error") or "eas" not in idx: return False try: - eas = array.array("Q"); eas.frombytes(b64decode(idx["eas"])) - kinds = array.array("B"); kinds.frombytes(b64decode(idx["kinds"])) - sizes = array.array("I"); sizes.frombytes(b64decode(idx["sizes"])) + eas = array.array("Q") + eas.frombytes(b64decode(idx["eas"])) + kinds = array.array("B") + kinds.frombytes(b64decode(idx["kinds"])) + sizes = array.array("I") + sizes.frombytes(b64decode(idx["sizes"])) except Exception: # noqa: BLE001 return False names = idx.get("kind_names") or [] @@ -881,7 +896,8 @@ class ListingModel: self._page_digest = [None] * len(anchors) self._page_rows = [ (anchors[k + 1][2] if k + 1 < len(anchors) else n) - anchors[k][2] - for k in range(len(anchors))] + for k in range(len(anchors)) + ] self._skeleton = True self._done = True self._next = None @@ -912,8 +928,9 @@ class ListingModel: if self._done or self._next is None: return 0 frm = self._next - payload = self._prog.client.invoke( - "heads", addr=hex(frm), count=self.PAGE, annotate=True, text=text) + payload = self._prog.client.call( + remote_ops.heads, addr=hex(frm), count=self.PAGE, annotate=True, text=text + ) rows = payload.get("heads", []) if isinstance(payload, dict) else [] cur = payload.get("cursor", {}) if isinstance(payload, dict) else {} page = self._build_page(rows, raw=text) @@ -925,8 +942,9 @@ class ListingModel: self._skeleton = True self._page_head.append(len(self._heads)) self._page_addr.append(frm) - self._page_digest.append(payload.get("digest") - if isinstance(payload, dict) else None) + self._page_digest.append( + payload.get("digest") if isinstance(payload, dict) else None + ) self._page_rows.append(len(rows)) for h in page: # Banner/label rows (function headers, separators, code labels) @@ -983,7 +1001,7 @@ class ListingModel: self._ubytes[b0] = blk off = a - b0 take = min(BLK - off, n - len(out)) - chunk = blk[off:off + take] if blk else b"" + chunk = blk[off : off + take] if blk else b"" if not chunk: break out += chunk @@ -1004,8 +1022,9 @@ class ListingModel: ea = h.ea + off b = self._unknown_bytes(ea, 1) text = f"db {b[0]:02X}h" if b else "db ?" - return Head(ea=ea, kind="unknown", size=1, text=text, - name=h.name if off == 0 else None) + return Head( + ea=ea, kind="unknown", size=1, text=text, name=h.name if off == 0 else None + ) def ensure(self, n: int) -> None: """Ensure at least ``n`` logical rows are loaded (or all, if fewer).""" @@ -1022,8 +1041,11 @@ class ListingModel: return idx with self._lock: have = self._rows - last_ea = (self._heads[-1].ea + max(self._heads[-1].size, 1) - 1 - if self._heads else -1) + last_ea = ( + self._heads[-1].ea + max(self._heads[-1].size, 1) - 1 + if self._heads + else -1 + ) done = self._done if done or (have and last_ea >= ea): # Loaded past ea without an exact head hit: return the first head @@ -1086,13 +1108,13 @@ class ListingModel: """ with self._lock: if not (self.seg_start <= ea < self.seg_end): - return True # another segment; nothing moved here + return True # another segment; nothing moved here if len(self._page_head) < 3: - return False # barely walked; a rebuild is cheaper + return False # barely walked; a rebuild is cheaper p = bisect.bisect_right(self._page_addr, ea) - 1 p = max(p - 1, 0) if p <= 0: - return False # the edit is in the first pages + return False # the edit is in the first pages keep = self._page_head[p] if keep <= 0: return False @@ -1110,7 +1132,7 @@ class ListingModel: last = self._heads[-1] self._rows = self._row_at[-1] + self._span(last) self._done = False - self._ubytes.clear() # undefined-run bytes behind the drop point + self._ubytes.clear() # undefined-run bytes behind the drop point return True def invalidate_text(self) -> None: @@ -1156,8 +1178,9 @@ class ListingModel: def _page_bounds(self, p: int) -> tuple[int, int]: """[first, last) head index of page ``p`` (caller holds the lock).""" lo = self._page_head[p] - hi = (self._page_head[p + 1] if p + 1 < len(self._page_head) - else len(self._heads)) + hi = ( + self._page_head[p + 1] if p + 1 < len(self._page_head) else len(self._heads) + ) return lo, hi def _ensure_page(self, p: int) -> int: @@ -1181,13 +1204,20 @@ class ListingModel: # the expectation rather than asking first means a page that HAS changed # still costs one round trip. try: - payload = self._prog.client.invoke( - "heads", addr=hex(addr), count=self.PAGE, annotate=True, - expect="" if want_digest is None else str(want_digest)) + payload = self._prog.client.call( + remote_ops.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 return p + 1 - if (isinstance(payload, dict) and "heads" not in payload - and payload.get("count") == want_rows): + if ( + isinstance(payload, dict) + and "heads" not in payload + and payload.get("count") == want_rows + ): with self._lock: if self._text_gen == gen and len(self._heads) >= hi: for k in range(lo, hi): @@ -1212,8 +1242,9 @@ class ListingModel: # what it once loaded. Leaving it stale is how a literal cycling # hex -> dec -> hex ends up declared "unchanged" while the row still # shows the decimal it was refetched with in between. - self._page_digest[p] = (payload.get("digest") - if isinstance(payload, dict) else None) + self._page_digest[p] = ( + payload.get("digest") if isinstance(payload, dict) else None + ) for k in range(lo, hi): self._head_gen[k] = gen return p + 1 @@ -1225,8 +1256,9 @@ class ListingModel: j, off = self._phys(i) if j < 0: return None - stale = ((self._renamed or self._skeleton) - and self._head_gen[j] != self._text_gen) + stale = (self._renamed or self._skeleton) and self._head_gen[ + j + ] != self._text_gen if not stale: span = self._span(self._heads[j]) h = self._heads[j] @@ -1265,8 +1297,9 @@ class ListingModel: spans = [self._phys(i) for i in range(max(start, 0), max(rows, 0))] heads = self._heads plain = [(j, off, heads[j]) for j, off in spans if j >= 0] - return [self._row_head(j, off) if self._span(h) > 1 else h - for j, off, h in plain] + return [ + self._row_head(j, off) if self._span(h) > 1 else h for j, off, h in plain + ] def index_of_ea(self, ea: int) -> int: with self._lock: @@ -1361,7 +1394,7 @@ class HexModel: if block is None: return (va, None) bo = off - b * self.BLOCK - return (va, block[bo:bo + 16]) + return (va, block[bo : bo + 16]) def ensure(self, r0: int, count: int) -> None: """Blocking: fetch the blocks covering rows [r0, r0+count) if missing.""" @@ -1459,20 +1492,32 @@ class Program: return self._segments_cache segs: list[tuple[int, int, int, str]] = [] try: - r = self.client.invoke("file_regions") - for d in (r.get("regions", []) if isinstance(r, dict) else []): + r = self.client.call(remote_ops.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"]), - int(d.get("file_off", -1)), d.get("name", "") or "")) + segs.append( + ( + _as_int(d["start"]), + _as_int(d["end"]), + int(d.get("file_off", -1)), + d.get("name", "") or "", + ) + ) except IDAToolError: segs = [] if not segs: # older server without file_regions -> survey_binary (slow) try: - sb = self.client.invoke("survey_binary") - for s in (sb.get("segments", []) if isinstance(sb, dict) else []): + sb = self.client.call(remote_ops.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, - s.get("name", "") or "")) + segs.append( + ( + _as_int(s["start"]), + _as_int(s["end"]), + -1, + s.get("name", "") or "", + ) + ) except (KeyError, ValueError, TypeError): continue except Exception: # noqa: BLE001 -- best-effort; callers handle empty @@ -1535,21 +1580,27 @@ class Program: return b"" if not self._no_read_raw: try: - r = self.client.invoke("read_raw", addr=hex(ea), size=int(n)) + r = self.client.call(remote_ops.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) return out[:n] if len(out) >= n else out + b"\x00" * (n - len(out)) except IDAToolError as e: # Tool missing on this server: stop trying it, use get_bytes. - if "read_raw" in str(e) or "Unknown tool" in str(e) or "not found" in str(e): + if ( + "read_raw" in str(e) + or "Unknown tool" in str(e) + or "not found" in str(e) + ): self._no_read_raw = True else: return b"\x00" * n except (ValueError, KeyError): pass # malformed hex -> fall through to the legacy decoder try: - r = self.client.invoke("get_bytes", regions=[{"addr": hex(ea), "size": int(n)}]) + r = self.client.call( + remote_ops.get_bytes, regions=[{"addr": hex(ea), "size": int(n)}] + ) except IDAToolError: return b"\x00" * n res = r.get("result", []) if isinstance(r, dict) else [] @@ -1590,7 +1641,7 @@ class Program: with self._lock: m = self._listings.get(start) if m is not None and m.stale_structure: - m = None # a refresh found the walk had moved; start over + m = None # a refresh found the walk had moved; start over if m is None: m = ListingModel(self, start, end, name) self._listings[start] = m @@ -1600,11 +1651,15 @@ class Program: def list_structs(self, filter: str = "") -> list[Struct]: """All local structs/unions (optionally name-substring filtered), sorted by name.""" - payload = self.client.invoke("search_structs", filter=filter) + payload = self.client.call(remote_ops.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") - and not str(d["name"]).startswith("$")] # skip anonymous UDTs + out = [ + Struct.from_raw(d) + for d in res + if isinstance(d, dict) + and d.get("name") + and not str(d["name"]).startswith("$") + ] # skip anonymous UDTs out.sort(key=lambda s: s.name.lower()) return out @@ -1612,8 +1667,9 @@ class Program: """A C definition for ``name`` reconstructed from its member layout (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.invoke( - "type_inspect", queries=[{"name": name, "include_members": True}]) + payload = self.client.call( + remote_ops.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 {} kw = "union" if info.get("is_union") else "struct" @@ -1635,7 +1691,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.invoke("declare_type", decls=decl) + payload = self.client.call(remote_ops.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") @@ -1646,20 +1702,32 @@ class Program: """Structured decompiler types for the function at ``ea`` (prototype + local variables). None if ``ea`` isn't a decompilable function.""" try: - r = self.client.invoke("func_types", addr=hex(ea)) + r = self.client.call(remote_ops.func_types, addr=hex(ea)) except IDAToolError: return None if not isinstance(r, dict) or r.get("error"): return None - lvars = [LVar(name=lv.get("name", ""), type=lv.get("type", ""), - is_arg=bool(lv.get("is_arg"))) - for lv in r.get("lvars", []) if isinstance(lv, dict)] - return FuncTypes(addr=_as_int(r.get("addr", hex(ea))), name=r.get("name", ""), - prototype=r.get("prototype", ""), lvars=lvars) + lvars = [ + LVar( + name=lv.get("name", ""), + type=lv.get("type", ""), + is_arg=bool(lv.get("is_arg")), + ) + for lv in r.get("lvars", []) + if isinstance(lv, dict) + ] + return FuncTypes( + addr=_as_int(r.get("addr", hex(ea))), + name=r.get("name", ""), + prototype=r.get("prototype", ""), + lvars=lvars, + ) 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.invoke("set_type", edits=[{"addr": hex(ea), "signature": signature}]) + r = self.client.call( + remote_ops.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"): @@ -1670,7 +1738,7 @@ class Program: """Current type info for a data item/global: {addr,name,type,size,is_func}. None if the operation fails or the address isn't mapped.""" try: - r = self.client.invoke("data_type", addr=hex(ea)) + r = self.client.call(remote_ops.data_type, addr=hex(ea)) except IDAToolError: return None if not isinstance(r, dict) or r.get("error"): @@ -1679,8 +1747,10 @@ 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.invoke( - "set_type", edits=[{"kind": "global", "addr": hex(ea), "type": decl}]) + r = self.client.call( + remote_ops.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 {} if row.get("ok"): @@ -1690,7 +1760,9 @@ class Program: def set_lvar_type(self, fn_ea: int, var: str, ty: str) -> str | None: """Set a decompiler local variable's type through ida-domain pseudocode. None on success, else an error string.""" - r = self.client.invoke("set_lvar_type", addr=hex(fn_ea), variable=var, type=ty) + r = self.client.call( + remote_ops.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"): @@ -1701,7 +1773,7 @@ class Program: """Delete a named type. Returns None on success, else an error string. Returns a clear error instead of raising when the runtime cannot do it.""" try: - self.client.invoke("del_type", name=name) + self.client.call(remote_ops.del_type, name=name) return None except IDAToolError as e: msg = e.message @@ -1732,7 +1804,7 @@ class Program: self._pc_nums.pop(ea, None) self._decomp_maps.pop(ea, None) try: - self.client.invoke("force_recompile", items=[{"addr": hex(ea)}]) + self.client.call(remote_ops.force_recompile, items=[{"addr": hex(ea)}]) except Exception: # noqa: BLE001 -- refresh still refetches best-effort pass @@ -1749,22 +1821,15 @@ class Program: # Cached before a rename: names may be stale. Drop Hex-Rays' # cache so the refetch reflects the new names. try: - self.client.invoke("force_recompile", items=[{"addr": hex(ea)}]) + self.client.call( + remote_ops.force_recompile, items=[{"addr": hex(ea)}] + ) except Exception: # noqa: BLE001 pass - # Bound the decompile: a function Hex-Rays can't handle tends to stall - # near the client's default 30s timeout, and the transport retries a - # dropped connection up to max_retries+1 times, re-running the failing - # decompile each time. Cap it so the worst case stays well under the - # rpcclient socket timeout, and cache the failure below so a re-request - # returns instantly instead of re-grinding. + # The typed remote declaration carries a 15-second transport timeout, + # so a function Hex-Rays cannot handle does not stall the UI. try: - # 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 - ) + payload = self.client.call(remote_ops.decompile, addr=hex(ea)) except Exception as e: # noqa: BLE001 -- surface as a failed decompile dec = Decompilation(ea, None, True, f"decompile error: {e}", False, None) with self._lock: @@ -1790,7 +1855,7 @@ class Program: self._name_gen += 1 models = list(self._disasm.values()) listings = list(self._listings.values()) - self._pc_nums.clear() # a reformat moves every literal on its line + self._pc_nums.clear() # a reformat moves every literal on its line for m in models: m.invalidate() for lm in listings: @@ -1832,6 +1897,34 @@ class Program: if self._listings.get(start) is lm: del self._listings[start] + def invalidate_external(self) -> None: + """Drop every cached view of an IDB changed by another client. + + An event may describe a rename, a byte patch, a new function, or a + segment move. Treating an unknown event as text-only risks displaying a + structurally impossible mix of old rows and new metadata, so the + external boundary deliberately invalidates all derived state. The app + debounces event bursts before reaching this method. + """ + with self._lock: + self._name_gen += 1 + models = list(self._disasm.values()) + self._indices.clear() + self._disasm.clear() + self._listings.clear() + self._decomp.clear() + self._pc_nums.clear() + self._decomp_maps.clear() + self._flowcharts.clear() + self._strings = None + self._linkage = None + self._segments_cache = None + self._sections = None + self._fileregions = None + self._hexmodel = None + for model in models: + model.invalidate() + # -- item / function structure edits (IDA c/d/u/p) --------------------- # @staticmethod def _first_result(payload) -> dict: @@ -1847,18 +1940,19 @@ 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.invoke("undefine", items=[{"addr": hex(ea)}]) + self.client.call(remote_ops.undefine, items=[{"addr": hex(ea)}]) except IDAToolError: pass # nothing defined here yet -> just try to create the insn res = self._first_result( - self.client.invoke("define_code", items=[{"addr": hex(ea)}])) + self.client.call(remote_ops.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.invoke("decomp_error", addr=hex(ea)) + r = self.client.call(remote_ops.decomp_error, addr=hex(ea)) except IDAToolError: return "" if not isinstance(r, dict): @@ -1877,19 +1971,22 @@ 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.invoke("thumb_scan", start=hex(start), end=hex(end), - apply=bool(apply)) + r = self.client.call( + remote_ops.thumb_scan, start=hex(start), end=hex(end), apply=bool(apply) + ) if not isinstance(r, dict) or r.get("error"): - raise IDAToolError("thumb_scan", - f"@ {start:#x}: {(r or {}).get('error', 'failed')}") + raise IDAToolError( + "thumb_scan", f"@ {start:#x}: {(r or {}).get('error', 'failed')}" + ) return r def set_thumb(self, ea: int, mode: str = "toggle") -> dict: """Switch ARM/Thumb decoding at ``ea``. Returns the resulting state.""" - r = self.client.invoke("set_thumb", addr=hex(ea), mode=mode) + r = self.client.call(remote_ops.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')}") + raise IDAToolError( + "set_thumb", f"@ {ea:#x}: {(r or {}).get('error', 'failed')}" + ) return r def define_code_run(self, ea: int, limit: int = 20000) -> dict: @@ -1899,13 +1996,16 @@ class Program: provide the run operation. """ try: - r = self.client.invoke("define_code_run", addr=hex(ea), limit=int(limit)) + r = self.client.call( + remote_ops.define_code_run, addr=hex(ea), limit=int(limit) + ) except IDAToolError: self.define_code(ea) return {"count": 1, "stopped": "single", "end": hex(ea)} if not isinstance(r, dict) or r.get("error"): - raise IDAToolError("define_code_run", - f"@ {ea:#x}: {(r or {}).get('error', 'failed')}") + raise IDAToolError( + "define_code_run", f"@ {ea:#x}: {(r or {}).get('error', 'failed')}" + ) return r def define_func(self, ea: int) -> dict: @@ -1915,16 +2015,18 @@ class Program: falls back to a plain create for alternate clients. """ try: - r = self.client.invoke("define_func_run", addr=hex(ea)) + r = self.client.call(remote_ops.define_func_run, addr=hex(ea)) except IDAToolError: res = self._first_result( - self.client.invoke("define_func", items=[{"addr": hex(ea)}])) + self.client.call(remote_ops.define_func, items=[{"addr": hex(ea)}]) + ) if res.get("error"): raise IDAToolError("define_func", f"@ {ea:#x}: {res['error']}") return {"ok": True, "how": "legacy"} if not isinstance(r, dict) or not r.get("ok"): - raise IDAToolError("define_func", - f"@ {ea:#x}: {(r or {}).get('error', 'failed')}") + raise IDAToolError( + "define_func", f"@ {ea:#x}: {(r or {}).get('error', 'failed')}" + ) return r def undefine(self, ea: int, size: int | None = None) -> None: @@ -1932,7 +2034,7 @@ class Program: item: dict = {"addr": hex(ea)} if size: item["size"] = int(size) - res = self._first_result(self.client.invoke("undefine", items=[item])) + res = self._first_result(self.client.call(remote_ops.undefine, items=[item])) if res.get("error"): raise IDAToolError("undefine", f"@ {ea:#x}: {res['error']}") @@ -1942,24 +2044,29 @@ class Program: item: dict = {"addr": hex(ea), "type": type_decl} if name: item["name"] = name - res = self._first_result(self.client.invoke("make_data", items=[item])) + res = self._first_result(self.client.call(remote_ops.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'}") + "make_data", f"@ {ea:#x}: {res.get('error') or 'rejected'}" + ) 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.invoke("make_string", addr=hex(ea), length=int(length), kind=kind) + r = self.client.call( + remote_ops.make_string, addr=hex(ea), length=int(length), kind=kind + ) res = r if isinstance(r, dict) else {} if not res.get("ok"): raise IDAToolError( - "make_string", f"@ {ea:#x}: {res.get('error') or 'rejected'}") + "make_string", f"@ {ea:#x}: {res.get('error') or 'rejected'}" + ) return res.get("text", "") # -- literal display formats (IDA's 'o': hex / dec / char / offset) ---- # - def op_format(self, ea: int, mode: str = "cycle", col: int = -1, - n: int = -1) -> dict: + def op_format( + self, ea: int, mode: str = "cycle", col: int = -1, n: int = -1 + ) -> dict: """Change how the literal at ``ea`` is DISPLAYED in the listing. ``col`` is a column inside the rendered line, which is how the cursor @@ -1967,8 +2074,9 @@ 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.invoke("op_format", addr=hex(ea), mode=str(mode), - col=int(col), n=int(n)) + r = self.client.call( + remote_ops.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"): raise IDAToolError("op_format", f"@ {ea:#x}: {res['error']}") @@ -1991,31 +2099,43 @@ class Program: if hit is not None and hit[1] == gen: return hit[0] try: - r = self.client.invoke("pc_nums", addr=hex(fn_ea)) + r = self.client.call(remote_ops.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]]] = {} for rec in (r or {}).get("nums", []): try: out.setdefault(int(rec["line"]), []).append( - (int(rec["x0"]), int(rec["x1"]), str(rec.get("value", "")), - _as_int(rec["ea"]), int(rec.get("opnum", 0)))) + ( + int(rec["x0"]), + int(rec["x1"]), + str(rec.get("value", "")), + _as_int(rec["ea"]), + int(rec.get("opnum", 0)), + ) + ) except Exception: # noqa: BLE001 -- skip a malformed row continue with self._lock: self._pc_nums[fn_ea] = (out, gen) return out - def pc_num_format(self, fn_ea: int, mode: str = "cycle", line: int = -1, - col: int = -1) -> dict: + def pc_num_format( + self, fn_ea: int, mode: str = "cycle", line: int = -1, col: int = -1 + ) -> dict: """The same, for a number in the DECOMPILATION of ``fn_ea``. Hex-Rays keeps number formats of its own, per (address, operand) — the 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.invoke("pc_num_format", addr=hex(fn_ea), mode=str(mode), - line=int(line), col=int(col)) + r = self.client.call( + remote_ops.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"): raise IDAToolError("pc_num_format", f"@ {fn_ea:#x}: {res['error']}") @@ -2043,22 +2163,30 @@ class Program: offset, page = 0, 2000 while True: try: - payload = self.client.invoke( - "list_strings", offset=offset, count=page, min_len=min_len, - refresh=(refresh and offset == 0)) + payload = self.client.call( + remote_ops.list_strings, + offset=offset, + count=page, + min_len=min_len, + refresh=(refresh and offset == 0), + ) except IDAToolError: return [] rows = payload.get("strings", []) if isinstance(payload, dict) else [] for r in rows: if not isinstance(r, dict): continue - out.append(StrLit( - addr=_as_int(r.get("addr", 0)), - text=r.get("text", ""), - length=int(r.get("len", 0) or 0), - type=r.get("type", "") or "", - )) - total = int(payload.get("total", 0) or 0) if isinstance(payload, dict) else 0 + out.append( + StrLit( + addr=_as_int(r.get("addr", 0)), + text=r.get("text", ""), + length=int(r.get("len", 0) or 0), + type=r.get("type", "") or "", + ) + ) + total = ( + int(payload.get("total", 0) or 0) if isinstance(payload, dict) else 0 + ) if len(rows) < page or len(out) >= total: break offset += len(rows) @@ -2074,27 +2202,39 @@ class Program: if hit is not None: return hit try: - payload = self.client.invoke("list_linkage", kind="both") + payload = self.client.call(remote_ops.list_linkage, kind="both") except IDAToolError: return ([], []) if not isinstance(payload, dict): return ([], []) - imps = [Linkage(addr=_as_int(r.get("addr", 0)), - name=link_name(r.get("name", "")), - module=r.get("module", "") or "", - raw=r.get("name", "") or "") - for r in payload.get("imports", []) if isinstance(r, dict)] - exps = [Linkage(addr=_as_int(r.get("addr", 0)), - name=link_name(r.get("name", "")), - ordinal=int(r.get("ordinal", 0) or 0), - raw=r.get("name", "") or "") - for r in payload.get("exports", []) if isinstance(r, dict)] + imps = [ + Linkage( + addr=_as_int(r.get("addr", 0)), + name=link_name(r.get("name", "")), + module=r.get("module", "") or "", + raw=r.get("name", "") or "", + ) + for r in payload.get("imports", []) + if isinstance(r, dict) + ] + exps = [ + Linkage( + addr=_as_int(r.get("addr", 0)), + name=link_name(r.get("name", "")), + ordinal=int(r.get("ordinal", 0) or 0), + raw=r.get("name", "") or "", + ) + for r in payload.get("exports", []) + if isinstance(r, dict) + ] out = ([i for i in imps if i.name], [e for e in exps if e.name]) with self._lock: self._linkage = out return out - def annotations(self, limit: int = 4000) -> tuple[list["Comment"], list["NamedItem"]]: + def annotations( + self, limit: int = 4000 + ) -> tuple[list["Comment"], list["NamedItem"]]: """``(comments, names)`` -- everything a person added to this database. Not cached: it is the *current* state of your work, and the one caller @@ -2102,44 +2242,64 @@ class Program: no such operation, so an alternate client degrades instead of breaking. """ try: - payload = self.client.invoke("list_annotations", limit=int(limit)) + payload = self.client.call(remote_ops.list_annotations, limit=int(limit)) except IDAToolError: return ([], []) if not isinstance(payload, dict): return ([], []) comments = [ - Comment(addr=_as_int(r.get("addr", 0)), text=str(r.get("text", "")), - repeatable=bool(r.get("repeatable")), - whole_func=bool(r.get("whole_func")), - line=str(r.get("line", "") or ""), - seg=str(r.get("seg", "") or ""), - func=(r.get("func") or None), - func_addr=(_as_int(r["func_addr"]) if r.get("func_addr") else None)) - for r in payload.get("comments", []) if isinstance(r, dict) and r.get("text")] + Comment( + addr=_as_int(r.get("addr", 0)), + text=str(r.get("text", "")), + repeatable=bool(r.get("repeatable")), + whole_func=bool(r.get("whole_func")), + line=str(r.get("line", "") or ""), + seg=str(r.get("seg", "") or ""), + func=(r.get("func") or None), + func_addr=(_as_int(r["func_addr"]) if r.get("func_addr") else None), + ) + for r in payload.get("comments", []) + if isinstance(r, dict) and r.get("text") + ] names = [ - NamedItem(addr=_as_int(r.get("addr", 0)), name=str(r.get("name", "")), - is_func=bool(r.get("func")), size=int(r.get("size", 0) or 0), - proto=(r.get("proto") or None), seg=str(r.get("seg", "") or "")) - for r in payload.get("names", []) if isinstance(r, dict) and r.get("name")] + NamedItem( + addr=_as_int(r.get("addr", 0)), + name=str(r.get("name", "")), + is_func=bool(r.get("func")), + size=int(r.get("size", 0) or 0), + proto=(r.get("proto") or None), + seg=str(r.get("seg", "") or ""), + ) + for r in payload.get("names", []) + if isinstance(r, dict) and r.get("name") + ] return (comments, names) - - def search(self, query: str, mode: str = "text", *, limit: int = 500, - regex: bool = False, case: bool = False, - ) -> tuple[list["SearchHit"], str | None, bool]: + def search( + self, + query: str, + mode: str = "text", + *, + limit: int = 500, + regex: bool = False, + case: bool = False, + ) -> tuple[list["SearchHit"], str | None, bool]: """Search the whole database. Returns ``(hits, error, truncated)``. A failed search is DATA (a message to show), not an exception: a bad regex or an unparsable byte pattern is something the user typed, and the palette wants to say so without unwinding. """ - op = "search_bytes" if mode == "bytes" else "search_text" + operation = ( + remote_ops.search_bytes if mode == "bytes" else remote_ops.search_text + ) args: dict = {"limit": int(limit), "case": bool(case)} if mode == "bytes": # Validate HERE, not just in the UI: IDA's find_bytes answers a # malformed pattern with zero hits and no error, which reads as # "not present" -- the most misleading answer a search can give. from .search import normalise_pattern, pattern_problem + problem = pattern_problem(query) if problem: return ([], problem, False) @@ -2148,29 +2308,32 @@ class Program: args["query"] = query args["regex"] = bool(regex) try: - payload = self.client.invoke(op, **args) + payload = self.client.call(operation, **args) except IDAToolError as e: return ([], str(e), False) if not isinstance(payload, dict): return ([], "the backend returned nothing searchable", False) hits = [ - SearchHit(addr=_as_int(r.get("addr", 0)), - head=_as_int(r.get("head", r.get("addr", 0))), - line=str(r.get("line", "") or ""), - func=(r.get("func") or None), - func_addr=(_as_int(r["func_addr"]) if r.get("func_addr") - else None), - seg=str(r.get("seg", "") or "")) - for r in payload.get("hits", []) if isinstance(r, dict)] + SearchHit( + addr=_as_int(r.get("addr", 0)), + head=_as_int(r.get("head", r.get("addr", 0))), + line=str(r.get("line", "") or ""), + func=(r.get("func") or None), + func_addr=(_as_int(r["func_addr"]) if r.get("func_addr") else None), + seg=str(r.get("seg", "") or ""), + ) + for r in payload.get("hits", []) + if isinstance(r, dict) + ] return (hits, payload.get("error") or None, bool(payload.get("truncated"))) def journal_get(self) -> str: """The findings journal blob stored in this database ('' if none).""" - payload = self.client.invoke("journal_get") + payload = self.client.call(remote_ops.journal_get) return str(payload.get("data", "")) if isinstance(payload, dict) else "" def journal_put(self, data: str) -> None: - self.client.invoke("journal_put", data=str(data)) + self.client.call(remote_ops.journal_put, data=str(data)) def decomp_map(self, ea: int) -> list[list[int]]: """Per-pseudocode-line instruction coverage for the split-view region @@ -2183,12 +2346,15 @@ class Program: if hit is not None and hit[1] == gen: return hit[0] try: - payload = self.client.invoke("decomp_map", addr=hex(ea)) + payload = self.client.call(remote_ops.decomp_map, addr=hex(ea)) except IDAToolError: return [] lines = payload.get("lines", []) if isinstance(payload, dict) else [] - out = [[_as_int(e) for e in (ln.get("eas") or [])] - for ln in lines if isinstance(ln, dict)] + out = [ + [_as_int(e) for e in (ln.get("eas") or [])] + for ln in lines + if isinstance(ln, dict) + ] with self._lock: self._decomp_maps[ea] = (out, gen) return out @@ -2213,7 +2379,7 @@ class Program: if hit is not None and hit[1] == gen: return hit[0] try: - payload = self.client.invoke("flowchart", addr=hex(ea)) + payload = self.client.call(remote_ops.flowchart, addr=hex(ea)) except IDAToolError: return None if not isinstance(payload, dict) or payload.get("error"): @@ -2224,10 +2390,14 @@ class Program: blocks = [] for b in raw: try: - blocks.append(BasicBlock( - id=int(b["id"]), start=_as_int(b["start"]), - end=_as_int(b["end"]), - succs=[(int(d), str(k)) for d, k in (b.get("succs") or [])])) + blocks.append( + BasicBlock( + id=int(b["id"]), + start=_as_int(b["start"]), + end=_as_int(b["end"]), + succs=[(int(d), str(k)) for d, k in (b.get("succs") or [])], + ) + ) except (KeyError, ValueError, TypeError): continue if not blocks: @@ -2239,8 +2409,9 @@ class Program: for b in blocks: # bisect, not a scan per block: a 400-block function against a few # thousand rows is a million comparisons done for nothing. - b.rows = rows[bisect.bisect_left(eas, b.start): - bisect.bisect_left(eas, b.end)] + b.rows = rows[ + bisect.bisect_left(eas, b.start) : bisect.bisect_left(eas, b.end) + ] fcv = Flowchart( func_ea=_as_int(f.get("addr", lo)), name=str(f.get("name") or f"sub_{lo:X}"), @@ -2283,11 +2454,12 @@ class Program: operand marks for free.""" out: list[Head] = [] addr = lo - for _ in range(64): # bounded: ~128k heads + for _ in range(64): # bounded: ~128k heads if addr >= hi: break - payload = self.client.invoke("heads", addr=hex(addr), end=hex(hi), - count=2000) + payload = self.client.call( + remote_ops.heads, addr=hex(addr), end=hex(hi), count=2000 + ) rows = payload.get("heads", []) if isinstance(payload, dict) else [] if not rows: break @@ -2315,27 +2487,34 @@ 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.invoke("lookup_funcs", queries=[hex(ea)]) + payload = self.client.call(remote_ops.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.invoke( - "xref_query", + payload = self.client.call( + remote_ops.xref_query, queries=[{"addr": hex(ea), "direction": "from", "include_fn": True}], ) return _parse_xrefs(payload) def xrefs_to(self, ea: int, limit: int = 2000) -> list[Xref]: - q = [{"addr": hex(ea), "direction": "to", "include_fn": True, - "dedup": True, "count": limit}] + q = [ + { + "addr": hex(ea), + "direction": "to", + "include_fn": True, + "dedup": True, + "count": limit, + } + ] 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.invoke("xref_types", queries=q) + payload = self.client.call(remote_ops.xref_types, queries=q) except IDAToolError: - payload = self.client.invoke("xref_query", queries=q) + payload = self.client.call(remote_ops.xref_query, queries=q) return _parse_xrefs(payload) # -- address resolution ------------------------------------------------ # @@ -2353,7 +2532,7 @@ 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.invoke("resolve_names", queries=[s]) + payload = self.client.call(remote_ops.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: @@ -2363,7 +2542,7 @@ class Program: # Fall back to function-name resolution (also drives the 'did you mean' # suggestion when the name is unknown). try: - payload = self.client.invoke("lookup_funcs", queries=[s]) + payload = self.client.call(remote_ops.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 [] @@ -2389,8 +2568,10 @@ class Program: except Exception: # noqa: BLE001 -- suggestions are strictly optional return "" if not cands: - return (" (no function name contains it; it may be a data symbol or " - "not a function — pass an address like 0x1234)") + return ( + " (no function name contains it; it may be a data symbol or " + "not a function — pass an address like 0x1234)" + ) shown = cands[:5] names = ", ".join(f"{c.name} @ {c.addr:#x}" for c in shown) more = " …" if len(cands) > len(shown) else "" @@ -2401,7 +2582,9 @@ 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.invoke("set_comments", items=[{"addr": hex(ea), "comment": text}]) + return self.client.call( + remote_ops.set_comments, items=[{"addr": hex(ea), "comment": text}] + ) # -- invalidation (after edits) --------------------------------------- # def invalidate(self, ea: int) -> None: @@ -2430,14 +2613,16 @@ def _parse_xrefs(payload) -> list[Xref]: fn = d.get("fn") or {} frm = d.get("from", d.get("addr")) to = d.get("to") - out.append(Xref( - frm=_as_int(frm) if frm is not None else 0, - to=_as_int(to) if to is not None else None, - type=d.get("type", "?"), - fn_name=fn.get("name"), - fn_addr=_as_int(fn["addr"]) if fn.get("addr") else None, - kind=d.get("kind"), - )) + out.append( + Xref( + frm=_as_int(frm) if frm is not None else 0, + to=_as_int(to) if to is not None else None, + type=d.get("type", "?"), + fn_name=fn.get("name"), + fn_addr=_as_int(fn["addr"]) if fn.get("addr") else None, + kind=d.get("kind"), + ) + ) return out @@ -2447,13 +2632,15 @@ def _parse_decompilation(ea: int, payload) -> Decompilation: code = payload.get("code") error = payload.get("error") if not code: - return Decompilation(ea, None, True, error or "decompilation failed", - False, None) + return Decompilation( + ea, None, True, error or "decompilation failed", False, None + ) m = _TRUNC_RE.search(code) truncated = m is not None total_chars = int(m.group(1)) if m else len(code) refs = [ Ref(addr=_as_int(r["addr"]), name=r.get("name", ""), string=r.get("string")) - for r in payload.get("refs", []) if isinstance(r, dict) and "addr" in r + for r in payload.get("refs", []) + if isinstance(r, dict) and "addr" in r ] return Decompilation(ea, code, False, error, truncated, total_chars, refs) diff --git a/idatui/edit_ctl.py b/idatui/edit_ctl.py index 52566ca..89dfc69 100644 --- a/idatui/edit_ctl.py +++ b/idatui/edit_ctl.py @@ -19,6 +19,7 @@ The message handlers and the ``@work`` entry points stay on ``IdaTui``: Textual dispatches ``on_<message>`` by name on the DOMNode, and its worker machinery wants a DOMNode host. They are one-line delegates into here. """ + from __future__ import annotations import re @@ -26,10 +27,10 @@ from typing import TYPE_CHECKING from textual.widgets import DataTable -from . import diag +from . import diag, remote_ops from .errors import IDAToolError -if TYPE_CHECKING: # pragma: no cover +if TYPE_CHECKING: # pragma: no cover from .app import IdaTui _app_mod = None @@ -40,13 +41,18 @@ def _M(): global _app_mod if _app_mod is None: from . import app as _m + _app_mod = _m return _app_mod #: A C type wide enough for N bytes, for prefilling a retype/define prompt. -_BY_SIZE = {1: "unsigned __int8", 2: "unsigned __int16", - 4: "unsigned __int32", 8: "unsigned __int64"} +_BY_SIZE = { + 1: "unsigned __int8", + 2: "unsigned __int16", + 4: "unsigned __int32", + 8: "unsigned __int64", +} class EditController: @@ -169,16 +175,24 @@ class EditController: # If the cursor is on a symbol token (a call/branch target, a data # reference, or this head's own label) rename THAT symbol; otherwise # create/rename a label at the head's address (bare/undefined bytes). - if (word and app._looks_like_symbol(word) and word != mnem - and word.lower() not in M._ASM_KEYWORDS): + if ( + word + and app._looks_like_symbol(word) + and word != mnem + and word.lower() not in M._ASM_KEYWORDS + ): app.prompts.rename.show( f"rename '{word}' — Enter=apply Esc=cancel", - word, ctx=(msg.view, word, None)) + word, + ctx=(msg.view, word, None), + ) else: cur = head.name if (head is not None and head.name) else "" app.prompts.rename.show( f"name @ {ea:#x} — Enter=apply Esc=cancel", - cur, ctx=(msg.view, cur, ea)) + cur, + ctx=(msg.view, cur, ea), + ) return if not msg.name: app._status("nothing to rename under the cursor") @@ -186,11 +200,14 @@ class EditController: if self.is_pseudocode_label(msg.view, msg.name): app._status( f"can't rename pseudocode label '{msg.name}' " - "(Hex-Rays goto labels aren't renamable via the API)") + "(Hex-Rays goto labels aren't renamable via the API)" + ) return app.prompts.rename.show( f"rename '{msg.name}' — Enter=apply Esc=cancel", - msg.name, ctx=(msg.view, msg.name, None)) + msg.name, + ctx=(msg.view, msg.name, None), + ) def submit_rename(self, ctx, value: str) -> None: # type: ignore[no-untyped-def] view, old, addr = ctx @@ -242,7 +259,7 @@ class EditController: kind = "stack" batch = {"stack": {"func_addr": hex(cur.ea), "old": old, "new": new}} try: - res = prog.client.invoke("rename", batch=batch) + res = prog.client.call(remote_ops.rename, batch=batch) except IDAToolError as e: app.call_from_thread(app._status, f"rename failed: {e.message}") return @@ -275,8 +292,9 @@ class EditController: app = self.app assert app.program is not None try: - res = app.program.client.invoke( - "rename", batch={"data": {"addr": hex(addr), "new": name}}) + res = app.program.client.call( + remote_ops.rename, batch={"data": {"addr": hex(addr), "new": name}} + ) except IDAToolError as e: app.call_from_thread(app._status, f"name failed: {e.message}") return @@ -307,11 +325,11 @@ class EditController: lm = app.program.listing(addr) label = name if is_func_start else app.program.region_label(addr) idx = max(lm.ensure_ea(addr), 0) if lm is not None else 0 - app.call_from_thread(self.open_at_named, label, addr, idx, name, - is_func_start) + app.call_from_thread(self.open_at_named, label, addr, idx, name, is_func_start) - def open_at_named(self, label: str, addr: int, idx: int, name: str, - is_func_start: bool = False) -> None: + def open_at_named( + self, label: str, addr: int, idx: int, name: str, is_func_start: bool = False + ) -> None: app = self.app if is_func_start: app.program.bump_names() @@ -328,10 +346,11 @@ class EditController: pseudocode a comment is `// text` before the trailing /*0xEA*/ markers; C has no `//` operator, so the last `//` is unambiguously the comment.""" if isinstance(view, _M().DecompView) and 0 <= view.cursor < len(view._texts): - s = re.sub(r"(?:/\*\s*0x[0-9A-Fa-f]+\s*\*/\s*)+$", "", - view._texts[view.cursor]) + s = re.sub( + r"(?:/\*\s*0x[0-9A-Fa-f]+\s*\*/\s*)+$", "", view._texts[view.cursor] + ) i = s.rfind("//") - return s[i + 2:].strip() if i >= 0 else "" + return s[i + 2 :].strip() if i >= 0 else "" return "" def request_comment(self, msg) -> None: # type: ignore[no-untyped-def] @@ -349,7 +368,9 @@ class EditController: what = "function comment" if func_level else "comment" app.prompts.comment.show( f"{what} @ {ea:#x} — Enter=apply (empty=clear) Esc=cancel", - existing, ctx=(msg.view, ea, existing)) + existing, + ctx=(msg.view, ea, existing), + ) def submit_comment(self, ctx, value: str) -> None: # type: ignore[no-untyped-def] view, ea, existing = ctx @@ -369,10 +390,13 @@ class EditController: app.call_from_thread(app._status, f"comment failed: {e.message}") return data = res.get("result") if isinstance(res, dict) else None - if (isinstance(data, list) and data and isinstance(data[0], dict) - and data[0].get("error")): - app.call_from_thread(app._status, - f"comment failed: {data[0]['error']}") + if ( + isinstance(data, list) + and data + and isinstance(data[0], dict) + and data[0].get("error") + ): + app.call_from_thread(app._status, f"comment failed: {data[0]['error']}") return app.call_from_thread(self.after_comment, ea, text) @@ -429,30 +453,34 @@ class EditController: if dt is not None and not dt.get("is_func"): kind, subject = "data", tgt prefill = dt.get("type") or self.guess_data_type( - dt.get("size") or 0) + dt.get("size") or 0 + ) # 3) fall back to the current function itself if kind is None and ft is not None: kind, subject, prefill = "func", app._cur.ea, ft.prototype if kind is None: - app.call_from_thread(app._status, - "nothing to retype under the cursor") + app.call_from_thread(app._status, "nothing to retype under the cursor") return - app.call_from_thread(self.open_retype, view, kind, subject, - word or "", prefill) + app.call_from_thread(self.open_retype, view, kind, subject, word or "", prefill) - def open_retype(self, view, kind: str, subject: int, word: str, - prefill: str) -> None: # type: ignore[no-untyped-def] + def open_retype( + self, view, kind: str, subject: int, word: str, prefill: str + ) -> None: # type: ignore[no-untyped-def] label = "prototype" if kind == "func" else f"type for '{word}'" - self.app.prompts.retype.show(f"{label} — Enter=apply Esc=cancel", - prefill, ctx=(view, kind, subject, word)) + self.app.prompts.retype.show( + f"{label} — Enter=apply Esc=cancel", + prefill, + ctx=(view, kind, subject, word), + ) def submit_retype(self, ctx, value: str) -> None: # type: ignore[no-untyped-def] view, kind, subject, word = ctx if view is not None and value: self.app._do_retype(kind, subject, word, value) - def do_retype(self, kind: str, subject: int, word: str, - new: str) -> None: # worker context + def do_retype( + self, kind: str, subject: int, word: str, new: str + ) -> None: # worker context app = self.app assert app.program is not None if kind == "func": @@ -473,8 +501,9 @@ class EditController: app.program.bump_names() self.reload_active_code() app._dirty = True - app.journal.record("retype", getattr(app._cur, "ea", None), word, - {"kind": kind}) + app.journal.record( + "retype", getattr(app._cur, "ea", None), word, {"kind": kind} + ) what = "prototype" if kind == "func" else f"'{word}'" app._status(f"retyped {what} (Ctrl+S to save)") @@ -498,15 +527,17 @@ class EditController: f"data type @ {ea:#x} (e.g. int, char[16], my_struct)" " — Enter=apply Esc=cancel", self.default_data_type(head) if head is not None else "int", - ctx=(view, ea)) + ctx=(view, ea), + ) def submit_make_data(self, ctx, value: str) -> None: # type: ignore[no-untyped-def] view, ea = ctx if view is not None and value: self.app._do_make_data(ea, value, self.app._anchor()) - def do_make_data(self, ea: int, type_decl: str, - anchor=None) -> None: # worker context + def do_make_data( + self, ea: int, type_decl: str, anchor=None + ) -> None: # worker context app = self.app assert app.program is not None try: @@ -522,8 +553,7 @@ class EditController: lm = app.program.listing(ea) idx = max(lm.ensure_ea(ea), 0) if lm is not None else 0 _cur, top = app._anchor_rows(anchor, lm, ea) - app.call_from_thread( - app._open_at, ea, name, idx, False, -1, 0, True, None, top) + app.call_from_thread(app._open_at, ea, name, idx, False, -1, 0, True, None, top) app.call_from_thread(self.edit_done, anchor) # -- literal display formats (IDA 'o') --------------------------------- # @@ -554,8 +584,9 @@ class EditController: fn = view.loaded_ea if view.loaded_ea is not None else app._cur.ea app._do_op_format(msg.mode, "decomp", fn, view.cursor_x, view.cursor) - def do_op_format(self, mode: str, where: str, ea: int, col: int, - line: int = -1) -> None: # worker context + def do_op_format( + self, mode: str, where: str, ea: int, col: int, line: int = -1 + ) -> None: # worker context app = self.app assert app.program is not None try: @@ -584,7 +615,9 @@ class EditController: app.call_from_thread( app._status, f"{what}{fmt} {r.get('value') or ''}" - f" [{', '.join(r.get('choices', []))}]", True) + f" [{', '.join(r.get('choices', []))}]", + True, + ) return step = f"{prev} \u2192 {fmt}" if prev and prev != fmt else fmt desc = f"{what}{step}: {text[:96]}" @@ -629,13 +662,17 @@ class EditController: return app._do_edit_item(msg.kind, ea, app._anchor()) - def do_edit_item(self, kind: str, ea: int, - anchor=None) -> None: # worker context + def do_edit_item(self, kind: str, ea: int, anchor=None) -> None: # worker context app = self.app assert app.program is not None - verb = {"code": "defined code", "func": "created function", - "undef": "undefined", "string": "made string", - "thumb": "switched decoding", "thumbscan": "scanned"}[kind] + verb = { + "code": "defined code", + "func": "created function", + "undef": "undefined", + "string": "made string", + "thumb": "switched decoding", + "thumbscan": "scanned", + }[kind] try: if kind == "code": # Keep going until something stops it: one instruction is rarely @@ -646,20 +683,24 @@ class EditController: if n == 0 and why == "defined": # Already code/data here — a no-op, not a failure. Saying # "failed to create instruction" for it would be a lie. - app.call_from_thread( - app._status, f"already defined @ {ea:#x}") + app.call_from_thread(app._status, f"already defined @ {ea:#x}") return if n == 0: - raise IDAToolError("define_code", - f"@ {ea:#x}: Failed to create instruction") + raise IDAToolError( + "define_code", f"@ {ea:#x}: Failed to create instruction" + ) end = int(str(r.get("end", hex(ea))), 0) - reason = {"undecodable": "hit bytes that don't decode", - "flow": "control flow ends here", - "defined": "ran into existing code/data", - "segment": "end of segment", - "limit": "instruction limit"}.get(why, why) - verb = (f"defined {n} instruction{'s' if n != 1 else ''} " - f"({ea:#x}\u2013{end:#x}) \u2014 {reason}") + reason = { + "undecodable": "hit bytes that don't decode", + "flow": "control flow ends here", + "defined": "ran into existing code/data", + "segment": "end of segment", + "limit": "instruction limit", + }.get(why, why) + verb = ( + f"defined {n} instruction{'s' if n != 1 else ''} " + f"({ea:#x}\u2013{end:#x}) \u2014 {reason}" + ) elif kind == "thumbscan": # A vector table is a list of Thumb entry points that IDA won't # follow on a headerless image, because nothing tells it those @@ -668,11 +709,15 @@ class EditController: r = app.program.thumb_scan(ea, ea + 0x400) n, applied = int(r.get("n", 0)), int(r.get("applied", 0)) if not n: - verb = (f"no Thumb entry pointers in {ea:#x}\u2013{ea+0x400:#x}" - " (odd words pointing into the image)") + verb = ( + f"no Thumb entry pointers in {ea:#x}\u2013{ea + 0x400:#x}" + " (odd words pointing into the image)" + ) else: - verb = (f"{n} Thumb entr{'y' if n == 1 else 'ies'} found, " - f"{applied} disassembled") + verb = ( + f"{n} Thumb entr{'y' if n == 1 else 'ies'} found, " + f"{applied} disassembled" + ) elif kind == "thumb": # Switch the mode, then disassemble in it: flipping T and # leaving the bytes undefined shows nothing, and the reason you @@ -686,11 +731,16 @@ class EditController: verb += " (segment set to 32-bit; Thumb needs ARM32)" if r.get("db_64bit"): # Disassembly will look right and F5 will never work. - verb += (" \u26a0 this database is 64-bit, so Hex-Rays " - "won't decompile it \u2014 Ctrl+L and pick " - "arm:ARMv7-A") - verb += (f" \u2014 {n} instruction{'s' if n != 1 else ''}" - if n else " \u2014 still doesn't decode") + verb += ( + " \u26a0 this database is 64-bit, so Hex-Rays " + "won't decompile it \u2014 Ctrl+L and pick " + "arm:ARMv7-A" + ) + verb += ( + f" \u2014 {n} instruction{'s' if n != 1 else ''}" + if n + else " \u2014 still doesn't decode" + ) # falls through to the shared reload: same cache bump, same # anchor restore, same flash. That is the whole point of having # one path. @@ -698,9 +748,11 @@ class EditController: anchor.refresh_functions = True r = app.program.define_func(ea) if r.get("start") and r.get("end"): - verb = (f"created function {r['start']}\u2013{r['end']}" - + (" (end worked out from the code)" - if r.get("how") == "explicit-end" else "")) + verb = f"created function {r['start']}\u2013{r['end']}" + ( + " (end worked out from the code)" + if r.get("how") == "explicit-end" + else "" + ) elif kind == "string": s = app.program.make_string(ea) verb = f"made string ({s[:24]!r})" if s else verb @@ -726,13 +778,14 @@ class EditController: idx = 0 if ea == fn.addr else model.index_of_ea(ea) _cur, top = app._anchor_rows(anchor, model, ea) app.call_from_thread( - app._open_at, fn.addr, fn.name, idx, False, -1, 0, False, - None, top) + app._open_at, fn.addr, fn.name, idx, False, -1, 0, False, None, top + ) else: name = app.program.region_label(ea) lm = app.program.listing(ea) idx = max(lm.ensure_ea(ea), 0) if lm is not None else 0 _cur, top = app._anchor_rows(anchor, lm, ea) app.call_from_thread( - app._open_at, ea, name, idx, False, -1, 0, True, None, top) + app._open_at, ea, name, idx, False, -1, 0, True, None, top + ) app.call_from_thread(self.edit_done, anchor) diff --git a/idatui/pool.py b/idatui/pool.py index 465dff2..573aa75 100644 --- a/idatui/pool.py +++ b/idatui/pool.py @@ -225,6 +225,28 @@ class DatabasePool: self.evict(label, save=save, save_gui=save) self.active = None + def discard_changes(self, labels: list[str]) -> list[str]: + """Discard final managed sessions; return labels whose owner remains. + + A returned label is not an error: its client is attached to a GUI or a + still-shared worker, so releasing our lease transfers finalization to + that session's owner or remaining clients. + """ + transferred: list[str] = [] + for label in labels: + client = self._clients.get(label) + if client is not None and not client.discard_database(): + transferred.append(label) + return transferred + + def replace_client(self, label: str, old, new) -> bool: + """Replace one disconnected lease without changing residency policy.""" + if self._clients.get(label) is not old: + return False + self._clients[label] = new + self._touch(label) + return True + # -- introspection ------------------------------------------------------ # def status(self) -> list[dict]: """Per-binary residency for the switcher UI.""" diff --git a/idatui/remote_ops.py b/idatui/remote_ops.py new file mode 100644 index 0000000..91ffeb7 --- /dev/null +++ b/idatui/remote_ops.py @@ -0,0 +1,1683 @@ +"""Typed remote operations executed through ida-codemode.""" + +from __future__ import annotations +# ruff: noqa + +import threading +from collections.abc import Callable +from pathlib import Path +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from ida_domain import Database + + +def operation_label() -> str: + """Display attribution for the current call; ready for per-user context.""" + return "IDA TUI" + + +def data_type(db: Database, **a: Any) -> Any: + 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)} + return result + + +def declare_type(db: Database, **a: Any) -> Any: + 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} + return result + + +def decomp_error(db: Database, **a: Any) -> Any: + 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}" + return result + + +def define_code(db: Database, **a: Any) -> Any: + 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} + return result + + +def define_code_run(db: Database, **a: Any) -> Any: + 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, + } + return result + + +def define_func(db: Database, **a: Any) -> Any: + 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} + return result + + +def define_func_run(db: Database, **a: Any) -> Any: + 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}", + } + ) + return result + + +def del_type(db: Database, **a: Any) -> Any: + 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"}), + } + return result + + +def disasm(db: Database, **a: Any) -> Any: + 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), + } + return result + + +def file_regions(db: Database, **a: Any) -> Any: + 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} + return result + + +def flowchart(db: Database, **a: Any) -> Any: + 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, + } + return result + + +def force_recompile(db: Database, **a: Any) -> Any: + 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} + return result + + +def func_types(db: Database, **a: Any) -> Any: + 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, + } + return result + + +def get_bytes(db: Database, **a: Any) -> Any: + 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} + return result + + +def journal_get(db: Database, **a: Any) -> Any: + import ida_netnode + + n = ida_netnode.netnode(a.get("node", "$ idatui.journal")) + blob = n.getblob(0, "I") if ida_netnode.exist(n) else None + result = {"data": blob.decode("utf-8", "replace") if blob else ""} + return result + + +def journal_put(db: Database, **a: Any) -> Any: + import ida_netnode + + n = ida_netnode.netnode(a.get("node", "$ idatui.journal"), 0, True) + payload = (a.get("data") or "").encode("utf-8") + n.setblob(payload, 0, "I") + result = {"ok": True, "bytes": len(payload)} + return result + + +def list_annotations(db: Database, **a: Any) -> Any: + import ida_bytes, ida_funcs, ida_lines, ida_nalt, ida_name + import ida_segment, ida_typeinf, idautils + + limit = max(1, int(a.get("limit", 4000))) + max_scan = max(1000, int(a.get("max_scan", 2000000))) + comments, names = [], [] + scanned = 0 + + def _line(ea): + try: + txt = ida_lines.generate_disasm_line(ea, ida_lines.GENDSM_REMOVE_TAGS) + except Exception: + txt = "" + return " ".join((txt or "").split()) + + for ea, nm in idautils.Names(): + if len(names) >= limit: + break + if not nm or not ida_bytes.has_user_name(ida_bytes.get_flags(ea)): + continue + fn = ida_funcs.get_func(ea) + is_fn = fn is not None and int(fn.start_ea) == int(ea) + proto = None + if is_fn: + try: + ti = ida_typeinf.tinfo_t() + if ida_nalt.get_tinfo(ti, ea): + proto = str(ti) + except Exception: + proto = None + seg = ida_segment.getseg(ea) + names.append( + { + "addr": hex(int(ea)), + "name": nm, + "func": is_fn, + "size": (int(fn.end_ea - fn.start_ea) if is_fn else 0), + "proto": proto, + "seg": (ida_segment.get_segm_name(seg) if seg else ""), + } + ) + + for i in range(ida_segment.get_segm_qty()): + seg = ida_segment.getnseg(i) + if seg is None or len(comments) >= limit or scanned >= max_scan: + continue + for ea in idautils.Heads(seg.start_ea, seg.end_ea): + scanned += 1 + if len(comments) >= limit or scanned >= max_scan: + break + if not ida_bytes.has_cmt(ida_bytes.get_flags(ea)): + continue + for rep in (False, True): + text = ida_bytes.get_cmt(ea, rep) + if text: + fn = ida_funcs.get_func(ea) + comments.append( + { + "addr": hex(int(ea)), + "text": text, + "repeatable": rep, + "line": _line(ea), + "seg": ida_segment.get_segm_name(seg), + "func": ( + ida_funcs.get_func_name(fn.start_ea) if fn else None + ), + "func_addr": (hex(int(fn.start_ea)) if fn else None), + } + ) + + # Whole-function comments are not on the byte flags, so the scan cannot see them. + for fn_ea in idautils.Functions(): + fn = ida_funcs.get_func(fn_ea) + if fn is None or len(comments) >= limit: + continue + for rep in (False, True): + text = ida_funcs.get_func_cmt(fn, rep) + if text: + seg = ida_segment.getseg(fn_ea) + comments.append( + { + "addr": hex(int(fn_ea)), + "text": text, + "repeatable": rep, + "line": "", + "whole_func": True, + "seg": (ida_segment.get_segm_name(seg) if seg else ""), + "func": ida_funcs.get_func_name(fn_ea), + "func_addr": hex(int(fn_ea)), + } + ) + result = { + "comments": comments, + "names": names, + "scanned": scanned, + "truncated": ( + len(comments) >= limit or len(names) >= limit or scanned >= max_scan + ), + } + return result + + +def list_funcs(db: Database, **a: Any) -> Any: + 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)} + ] + } + return result + + +def list_linkage(db: Database, **a: Any) -> Any: + 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), + } + return result + + +def list_strings(db: Database, **a: Any) -> Any: + 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)} + return result + + +def lookup_funcs(db: Database, **a: Any) -> Any: + 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} + return result + + +def make_data(db: Database, **a: Any) -> Any: + 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} + return result + + +def make_string(db: Database, **a: Any) -> Any: + 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)} + return result + + +def read_raw(db: Database, **a: Any) -> Any: + 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)} + return result + + +def rename(db: Database, **a: Any) -> Any: + 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 + return result + + +def resolve_names(db: Database, **a: Any) -> Any: + 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} + return result + + +def search_bytes(db: Database, **a: Any) -> Any: + import ida_bytes, ida_funcs, ida_idaapi, ida_lines, ida_segment + + pat = str(a.get("pattern", "")).strip() + limit = max(1, int(a.get("limit", 500))) + lo = int(a.get("start", 0)) + hi = int(a.get("end", 0)) or ida_idaapi.BADADDR + flags = ida_bytes.BIN_SEARCH_FORWARD | ida_bytes.BIN_SEARCH_NOSHOW + if a.get("case"): + flags |= ida_bytes.BIN_SEARCH_CASE + rows, err, ea = [], None, lo + while len(rows) < limit: + try: + hit = ida_bytes.find_bytes(pat, range_start=ea, range_end=hi, flags=flags) + except Exception as exc: + err = str(exc) or exc.__class__.__name__ + break + if hit is None or hit == ida_idaapi.BADADDR: + break + head = ida_bytes.get_item_head(hit) + fn = ida_funcs.get_func(hit) + seg = ida_segment.getseg(hit) + try: + line = ( + ida_lines.generate_disasm_line(head, ida_lines.GENDSM_REMOVE_TAGS) or "" + ) + except Exception: + line = "" + rows.append( + { + "addr": hex(int(hit)), + "head": hex(int(head)), + "line": " ".join(line.split()), + "func": (ida_funcs.get_func_name(fn.start_ea) if fn else None), + "func_addr": (hex(int(fn.start_ea)) if fn else None), + "seg": (ida_segment.get_segm_name(seg) if seg else ""), + } + ) + ea = int(hit) + 1 + result = {"hits": rows, "error": err, "truncated": len(rows) >= limit} + return result + + +def search_structs(db: Database, **a: Any) -> Any: + 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} + return result + + +def search_text(db: Database, **a: Any) -> Any: + import ida_lines, ida_funcs, ida_segment, idautils + import re as _re + + q = str(a.get("query", "")) + limit = max(1, int(a.get("limit", 500))) + max_scan = max(1000, int(a.get("max_scan", 3000000))) + ci = (not a.get("case")) and q.islower() # smartcase, like the in-view search + rx, err = None, None + if a.get("regex"): + try: + rx = _re.compile(q, _re.I if ci else 0) + except Exception as exc: + err = "bad regex: " + str(exc) + needle = q.lower() if ci else q + rows, scanned = [], 0 + if err is None and q: + for i in range(ida_segment.get_segm_qty()): + seg = ida_segment.getnseg(i) + if seg is None or len(rows) >= limit or scanned >= max_scan: + continue + for ea in idautils.Heads(seg.start_ea, seg.end_ea): + scanned += 1 + if len(rows) >= limit or scanned >= max_scan: + break + try: + line = ( + ida_lines.generate_disasm_line(ea, ida_lines.GENDSM_REMOVE_TAGS) + or "" + ) + except Exception: + continue + # Match what the user SEES, not IDA's column padding: nobody types + # "call" + four spaces + "cs:getenv_ptr". + line = " ".join(line.split()) + hay = line.lower() if ci else line + if rx.search(line) if rx is not None else (needle in hay): + fn = ida_funcs.get_func(ea) + rows.append( + { + "addr": hex(int(ea)), + "head": hex(int(ea)), + "line": line, + "func": ( + ida_funcs.get_func_name(fn.start_ea) if fn else None + ), + "func_addr": (hex(int(fn.start_ea)) if fn else None), + "seg": ida_segment.get_segm_name(seg), + } + ) + result = { + "hits": rows, + "error": err, + "scanned": scanned, + "truncated": len(rows) >= limit or scanned >= max_scan, + } + return result + + +def set_comments(db: Database, **a: Any) -> Any: + 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} + return result + + +def set_lvar_type(db: Database, **a: Any) -> Any: + 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}"} + return result + + +def set_thumb(db: Database, **a: Any) -> Any: + 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), + } + return result + + +def set_type(db: Database, **a: Any) -> Any: + 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} + return result + + +def survey_binary(db: Database, **a: Any) -> Any: + 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} + return result + + +def thumb_scan(db: Database, **a: Any) -> Any: + 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), + } + return result + + +def type_inspect(db: Database, **a: Any) -> Any: + 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} + return result + + +def undefine(db: Database, **a: Any) -> Any: + 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} + return result + + +def xref_query(db: Database, **a: Any) -> Any: + 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} + return result + + +def xref_types(db: Database, **a: Any) -> Any: + 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} + return result + + +def op_format(addr: str, mode: str = "cycle", col: int = -1, n: int = -1) -> dict: ... + + +def pc_nums(addr: str) -> dict: ... + + +def decompile(addr, include_addresses=True) -> dict: ... + + +def decomp_map(addr: str) -> dict: ... + + +def pc_num_format( + addr: str, + mode: str = "cycle", + line: int = -1, + col: int = -1, + ea: str = "", + opnum: int = -1, +) -> dict: ... + + +def segment_index( + addr: str, + end: str = "", + page_rows: int = 500, + detail: bool = False, +) -> dict: ... + + +def heads( + addr: str, + count: int = 200, + offset: int = 0, + end: str = "", + back: bool = False, + annotate: bool = False, + expect: str = "", + text: bool = True, +) -> dict: ... + + +def profile_remote(operation: str, args: dict[str, Any], reps: int = 5) -> dict: ... + + +OPERATIONS: dict[str, Callable[..., Any]] = { + "data_type": data_type, + "declare_type": declare_type, + "decomp_error": decomp_error, + "define_code": define_code, + "define_code_run": define_code_run, + "define_func": define_func, + "define_func_run": define_func_run, + "del_type": del_type, + "disasm": disasm, + "file_regions": file_regions, + "flowchart": flowchart, + "force_recompile": force_recompile, + "func_types": func_types, + "get_bytes": get_bytes, + "heads": heads, + "journal_get": journal_get, + "journal_put": journal_put, + "list_annotations": list_annotations, + "list_funcs": list_funcs, + "list_linkage": list_linkage, + "list_strings": list_strings, + "lookup_funcs": lookup_funcs, + "make_data": make_data, + "make_string": make_string, + "op_format": op_format, + "pc_nums": pc_nums, + "decompile": decompile, + "decomp_map": decomp_map, + "pc_num_format": pc_num_format, + "profile_remote": profile_remote, + "read_raw": read_raw, + "rename": rename, + "resolve_names": resolve_names, + "search_bytes": search_bytes, + "search_structs": search_structs, + "search_text": search_text, + "segment_index": segment_index, + "set_comments": set_comments, + "set_lvar_type": set_lvar_type, + "set_thumb": set_thumb, + "set_type": set_type, + "survey_binary": survey_binary, + "thumb_scan": thumb_scan, + "type_inspect": type_inspect, + "undefine": undefine, + "xref_query": xref_query, + "xref_types": xref_types, +} + +_MODULE_DECLARATIONS = frozenset( + ( + heads, + segment_index, + op_format, + pc_nums, + decompile, + decomp_map, + pc_num_format, + profile_remote, + ) +) +_BOUND: dict[Callable[..., Any], Any] | None = None +_BIND_LOCK = threading.Lock() + + +def _bindings() -> dict[Callable[..., Any], Any]: + global _BOUND + with _BIND_LOCK: + if _BOUND is not None: + return _BOUND + from ida_codemode import RemoteModule + + operations_module = RemoteModule( + Path(__file__), operation_label=operation_label, codec="json" + ) + tools_module = RemoteModule( + Path(__file__).with_name("remote_tools.py"), + operation_label=operation_label, + codec="json", + ) + bound: dict[Callable[..., Any], Any] = {} + for declaration in OPERATIONS.values(): + if declaration in _MODULE_DECLARATIONS: + bound[declaration] = tools_module.function( + declaration, + timeout=15.0 if declaration is decompile else None, + ) + else: + bound[declaration] = operations_module.function(declaration) + _BOUND = bound + return bound + + +def bind(function: Callable[..., Any]) -> Any: + """Return the lazily constructed remote callable for one declaration.""" + return _bindings()[function] diff --git a/idatui/remote_tools.py b/idatui/remote_tools.py index 0e26965..db04244 100644 --- a/idatui/remote_tools.py +++ b/idatui/remote_tools.py @@ -23,6 +23,7 @@ imported here, because the ida_* modules do not exist in the TUI's interpreter. `codemode_client` reads it and prepends it to the relevant snippets. Keep it self-contained: no relative imports, nothing beyond what Code Mode provides. """ + # ruff: noqa import re as _re @@ -134,7 +135,7 @@ def _idatui_head_row(ea, flags=None, text=True): """ f = ida_bytes.get_flags(ea) if flags is None else flags - cls = f & _MS_CLS # == is_code(f) / is_data(f), without the calls + cls = f & _MS_CLS # == is_code(f) / is_data(f), without the calls if cls == _FF_CODE: kind = "code" elif cls == _FF_DATA: @@ -211,10 +212,20 @@ _IDATUI_SPAN_KINDS = { # 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"), + "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"), @@ -241,7 +252,7 @@ _IDATUI_TAGS = None _IDATUI_OPND_TAGS = None -_IDATUI_CTL = None # re: a tag = one of three control chars plus its argument +_IDATUI_CTL = None # re: a tag = one of three control chars plus its argument _IDATUI_TAGINFO = None @@ -280,6 +291,7 @@ def _idatui_spans(line): _IDATUI_OPND_TAGS = _idatui_opnd_tag_map() if _IDATUI_CTL is None: import re as _re + # One capturing split gives [text, tag, text, tag, ..., text] in a # single C pass. A per-character python loop over the line used to be # the most expensive thing the `heads` tool did, and a line is ~54 @@ -289,17 +301,18 @@ def _idatui_spans(line): if _IDATUI_TAGINFO is None: _IDATUI_TAGINFO = { tag: (_IDATUI_TAGS.get(tag, "text"), _IDATUI_OPND_TAGS.get(tag)) - for tag in set(_IDATUI_TAGS) | set(_IDATUI_OPND_TAGS)} + for tag in set(_IDATUI_TAGS) | set(_IDATUI_OPND_TAGS) + } taginfo = _IDATUI_TAGINFO plain_tag = ("text", None) on, off, esc = "\x01", "\x02", "\x03" addr_tag = chr(getattr(ida_lines, "COLOR_ADDR", 0x28)) addr_len = int(getattr(ida_lines, "COLOR_ADDR_SIZE", 16)) parts = _IDATUI_CTL.split(line) - spans, stack = [], [] # stack entries: (kind, operand index|None) - kind, opnd = "text", None # state the current run of text belongs to + spans, stack = [], [] # stack entries: (kind, operand index|None) + kind, opnd = "text", None # state the current run of text belongs to pend = "" - skip = 0 # characters of an address payload still due + skip = 0 # characters of an address payload still due i, n = 0, len(parts) while i < n: txt = parts[i] @@ -317,11 +330,11 @@ def _idatui_spans(line): break pair = parts[i] i += 1 - if skip: # a tag INSIDE an address payload: 2 chars + if skip: # a tag INSIDE an address payload: 2 chars skip = skip - 2 if skip > 2 else 0 continue ch = pair[0] - if ch == esc: # escaped literal: keep the char it guards + if ch == esc: # escaped literal: keep the char it guards pend += pair[1] continue tag = pair[1] @@ -337,7 +350,7 @@ def _idatui_spans(line): stack.append((kind, opnd)) kind, o = taginfo.get(tag, plain_tag) if o is not None: - opnd = o # operands nest: an inner colour keeps the operand + opnd = o # operands nest: an inner colour keeps the operand elif stack: kind, opnd = stack.pop() else: @@ -359,7 +372,7 @@ def _idatui_spans(line): prev_space = False out.append([kind, txt, opnd]) continue - if not core: # the span is nothing but padding + if not core: # the span is nothing but padding if not prev_space: prev_space = True out.append([kind, " ", opnd]) @@ -396,7 +409,7 @@ def _idatui_spans(line): ops.append([start, pos, cur]) text = "".join(t for _k, t, _o in out) trimmed = [] - for lo, hi, k in ops: # don't let a range own trailing space + for lo, hi, k in ops: # don't let a range own trailing space while hi > lo and text[hi - 1].isspace(): hi -= 1 while lo < hi and text[lo].isspace(): @@ -435,8 +448,17 @@ def _idatui_rows_digest(rows): sh = seen.get(key) if sh is None: sh = seen[key] = hash(tuple(map(tuple, sp))) - acc = hash((acc, r.get("ea"), r.get("kind"), r.get("size"), - r.get("text"), r.get("name"), sh)) + acc = hash( + ( + acc, + r.get("ea"), + r.get("kind"), + r.get("size"), + r.get("text"), + r.get("name"), + sh, + ) + ) return acc @@ -448,8 +470,12 @@ def _idatui_unknown_row(ea, size): if size <= 1: return _idatui_head_row(ea) - row = {"ea": hex(ea), "kind": "unknown", "size": int(size), - "text": f"db {size} dup(?)"} + 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 @@ -481,8 +507,7 @@ def _idatui_struct_member_rows(ea): 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}) + rows.append({"ea": hex(ea + off), "kind": "member", "size": sz, "text": text}) return rows @@ -494,8 +519,13 @@ def _idatui_func_header_rows(ea): 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}, + { + "ea": hex(ea), + "kind": "funchdr", + "size": 0, + "text": name + " proc", + "name": name, + }, ] @@ -504,8 +534,13 @@ def _idatui_func_footer_rows(ea, func): 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": "funchdr", + "size": 0, + "text": name + " endp", + "name": name, + }, {"ea": hex(ea), "kind": "sep", "size": 0, "text": "; " + "-" * 60}, ] @@ -543,9 +578,12 @@ def _idatui_segment_detail(addr, end, page_rows): except Exception: pass - K_CODE = _IDATUI_KIND_ID["code"]; K_DATA = _IDATUI_KIND_ID["data"] - K_UNK = _IDATUI_KIND_ID["unknown"]; K_SEP = _IDATUI_KIND_ID["sep"] - K_FUNC = _IDATUI_KIND_ID["funchdr"]; K_LABEL = _IDATUI_KIND_ID["label"] + K_CODE = _IDATUI_KIND_ID["code"] + K_DATA = _IDATUI_KIND_ID["data"] + K_UNK = _IDATUI_KIND_ID["unknown"] + K_SEP = _IDATUI_KIND_ID["sep"] + K_FUNC = _IDATUI_KIND_ID["funchdr"] + K_LABEL = _IDATUI_KIND_ID["label"] K_MEMBER = _IDATUI_KIND_ID["member"] eas = array.array("Q") @@ -569,8 +607,8 @@ def _idatui_segment_detail(addr, end, page_rows): # are the same number until a segment contains an undefined run) and then # silently yields pages that do not line up with a refetch. anchors = [] - page_phys = 0 # physical rows emitted into the page being filled - rows = 0 # logical rows so far (what the scrollbar counts) + page_phys = 0 # physical rows emitted into the page being filled + rows = 0 # logical rows so far (what the scrollbar counts) fn = None ea = ida_bytes.get_item_head(lo) while ea != BAD and ea < hi: @@ -584,7 +622,9 @@ def _idatui_segment_detail(addr, end, page_rows): nh = next_head(ea, hi) stop = nh if (nh != BAD and ea < nh <= hi) else hi run = stop - ea - ea_ap(ea); kind_ap(K_UNK); size_ap(run) + ea_ap(ea) + kind_ap(K_UNK) + size_ap(run) rows += run if run > 1 else 1 page_phys += len(eas) - before ea = stop @@ -593,23 +633,32 @@ def _idatui_segment_detail(addr, end, page_rows): fn = get_func(ea) at_start = fn is not None and fn.start_ea == ea if at_start: - for k in (K_SEP, K_SEP, K_FUNC): # blank, banner, `name proc` - ea_ap(ea); kind_ap(k); size_ap(0) + for k in (K_SEP, K_SEP, K_FUNC): # blank, banner, `name proc` + ea_ap(ea) + kind_ap(k) + size_ap(0) rows += 3 elif cls == _FF_CODE and get_ea_name(ea): - ea_ap(ea); kind_ap(K_LABEL); size_ap(0) + ea_ap(ea) + kind_ap(K_LABEL) + size_ap(0) rows += 1 - ea_ap(ea); kind_ap(K_CODE if cls == _FF_CODE else K_DATA) - size_ap(int(get_item_size(ea))); rows += 1 + ea_ap(ea) + kind_ap(K_CODE if cls == _FF_CODE else K_DATA) + size_ap(int(get_item_size(ea))) + rows += 1 if cls == _FF_DATA: for m in _idatui_struct_member_rows(ea): ea_ap(int(m["ea"], 16) if isinstance(m["ea"], str) else m["ea"]) kind_ap(_IDATUI_KIND_ID.get(m.get("kind", "member"), K_MEMBER)) - size_ap(int(m.get("size", 0) or 0)); rows += 1 + size_ap(int(m.get("size", 0) or 0)) + rows += 1 item_end = get_item_end(ea) if fn is not None and item_end >= fn.end_ea: - for k in (K_FUNC, K_SEP): # `name endp`, separator - ea_ap(ea); kind_ap(k); size_ap(0) + for k in (K_FUNC, K_SEP): # `name endp`, separator + ea_ap(ea) + kind_ap(k) + size_ap(0) rows += 2 page_phys += len(eas) - before ea = item_end if item_end > ea else ea + 1 @@ -618,19 +667,28 @@ def _idatui_segment_detail(addr, end, page_rows): # which turns a bytes object into its repr -- 4 characters per byte and # unparseable at the other end. Learned by watching 2.97MB arrive as 11.26MB. import base64 + b64 = base64.b64encode - return {"addr": hex(lo), "end": hex(hi), "rows": rows, "heads": len(eas), - "anchors": anchors, "kind_names": list(_IDATUI_KINDS), - "eas": b64(eas.tobytes()).decode(), - "kinds": b64(kinds.tobytes()).decode(), - "sizes": b64(sizes.tobytes()).decode()} + return { + "addr": hex(lo), + "end": hex(hi), + "rows": rows, + "heads": len(eas), + "anchors": anchors, + "kind_names": list(_IDATUI_KINDS), + "eas": b64(eas.tobytes()).decode(), + "kinds": b64(kinds.tobytes()).decode(), + "sizes": b64(sizes.tobytes()).decode(), + } def segment_index( addr: Annotated[str, "Any address in the segment to index"], end: Annotated[str, "Optional exclusive end address; default = segment end"] = "", page_rows: Annotated[int, "Rows between anchors (default 500)"] = 500, - detail: Annotated[bool, "Also return every row's ea/kind/size as packed arrays"] = False, + detail: Annotated[ + bool, "Also return every row's ea/kind/size as packed arrays" + ] = False, ) -> dict: """How many listing rows a segment has, and where to seek into it. @@ -662,6 +720,7 @@ def segment_index( if detail: return _idatui_segment_detail(addr, end, count) import ida_segment + seg = ida_segment.getseg(start) if not seg: return {"addr": str(addr), "error": "no segment", "rows": 0, "anchors": []} @@ -703,19 +762,24 @@ def segment_index( fn = get_func(ea) n = 1 if fn is not None and fn.start_ea == ea: - n += 3 # blank, banner, `proc` + n += 3 # blank, banner, `proc` elif cls == _FF_CODE and get_ea_name(ea): - n += 1 # loc_XXX label on its own row + n += 1 # loc_XXX label on its own row if cls == _FF_DATA: n += len(_idatui_struct_member_rows(ea)) item_end = get_item_end(ea) if fn is not None and item_end >= fn.end_ea: - n += 2 # `endp` + separator + n += 2 # `endp` + separator rows += n n_heads += 1 ea = item_end if item_end > ea else ea + 1 - return {"addr": hex(lo), "end": hex(hi), "rows": rows, - "heads": n_heads, "anchors": anchors} + return { + "addr": hex(lo), + "end": hex(hi), + "rows": rows, + "heads": n_heads, + "anchors": anchors, + } def heads( @@ -723,10 +787,21 @@ def heads( count: Annotated[int, "Max heads to return (default 200, max 2000)"] = 200, offset: Annotated[int, "Skip first N heads from addr (default 0)"] = 0, end: Annotated[str, "Optional exclusive end address; default = segment end"] = "", - back: Annotated[bool, "Walk backwards: return the count heads ENDING just before addr, in forward order"] = False, - annotate: Annotated[bool, "Emit IDA-style function boundary banner rows (kind sep/funchdr)"] = False, - expect: Annotated[str, "Digest a caller already holds: the rows are omitted when they still hash to it"] = "", - text: Annotated[bool, "Render each row's disassembly text (default true). False = a skeleton page: same rows, same addresses, no text"] = True, + back: Annotated[ + bool, + "Walk backwards: return the count heads ENDING just before addr, in forward order", + ] = False, + annotate: Annotated[ + bool, "Emit IDA-style function boundary banner rows (kind sep/funchdr)" + ] = False, + expect: Annotated[ + str, + "Digest a caller already holds: the rows are omitted when they still hash to it", + ] = "", + text: Annotated[ + bool, + "Render each row's disassembly text (default true). False = a skeleton page: same rows, same addresses, no text", + ] = True, ) -> 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 @@ -740,10 +815,20 @@ def heads( try: start = parse_address(addr) except Exception as e: - return {"addr": str(addr), "error": str(e), "heads": [], "cursor": {"done": True}} + 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}} + return { + "addr": str(addr), + "error": "no segment", + "heads": [], + "cursor": {"done": True}, + } lo, hi = seg.start_ea, seg.end_ea if end: try: @@ -766,7 +851,9 @@ def heads( 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)} + 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 @@ -823,8 +910,9 @@ def heads( # 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}) + out.append( + {"ea": hex(e), "kind": "label", "size": 0, "text": nm + ":", "name": nm} + ) row = dict(row) row["name"] = None out.append(row) @@ -845,7 +933,7 @@ def heads( if len(rows) >= count: more = True break - f = get_flags(ea) # once per head, not once per consumer + f = get_flags(ea) # once per head, not once per consumer rows.extend(_rows_for(ea, f)) # a struct head expands into member rows ea = _advance(ea, f) cursor = {"next": hex(ea)} if more else {"done": True} @@ -871,8 +959,18 @@ def heads( _IDATUI_FMT_CYCLE = ("hex", "dec", "bin", "char", "offset", "default") -_IDATUI_FMT_SETTABLE = ("hex", "dec", "oct", "bin", "char", "offset", "seg", - "float", "stack", "default") +_IDATUI_FMT_SETTABLE = ( + "hex", + "dec", + "oct", + "bin", + "char", + "offset", + "seg", + "float", + "stack", + "default", +) def _idatui_fmt_nibbles(): @@ -880,13 +978,20 @@ def _idatui_fmt_nibbles(): this module is injected into a file that is imported before a database is open.""" return { - "default": ida_bytes.FF_N_VOID, "hex": ida_bytes.FF_N_NUMH, - "dec": ida_bytes.FF_N_NUMD, "char": ida_bytes.FF_N_CHAR, - "seg": ida_bytes.FF_N_SEG, "offset": ida_bytes.FF_N_OFF, - "bin": ida_bytes.FF_N_NUMB, "oct": ida_bytes.FF_N_NUMO, - "enum": ida_bytes.FF_N_ENUM, "forced": ida_bytes.FF_N_FOP, - "stroff": ida_bytes.FF_N_STRO, "stack": ida_bytes.FF_N_STK, - "float": ida_bytes.FF_N_FLT, "custom": ida_bytes.FF_N_CUST, + "default": ida_bytes.FF_N_VOID, + "hex": ida_bytes.FF_N_NUMH, + "dec": ida_bytes.FF_N_NUMD, + "char": ida_bytes.FF_N_CHAR, + "seg": ida_bytes.FF_N_SEG, + "offset": ida_bytes.FF_N_OFF, + "bin": ida_bytes.FF_N_NUMB, + "oct": ida_bytes.FF_N_NUMO, + "enum": ida_bytes.FF_N_ENUM, + "forced": ida_bytes.FF_N_FOP, + "stroff": ida_bytes.FF_N_STRO, + "stack": ida_bytes.FF_N_STK, + "float": ida_bytes.FF_N_FLT, + "custom": ida_bytes.FF_N_CUST, } @@ -932,8 +1037,12 @@ def _idatui_op_value(ea, n): size = 0 return int(v), size size = int(ida_bytes.get_item_size(ea)) - read = {1: ida_bytes.get_byte, 2: ida_bytes.get_word, - 4: ida_bytes.get_dword, 8: ida_bytes.get_qword}.get(size) + read = { + 1: ida_bytes.get_byte, + 2: ida_bytes.get_word, + 4: ida_bytes.get_dword, + 8: ida_bytes.get_qword, + }.get(size) if read is None: return None, size try: @@ -991,9 +1100,9 @@ def _idatui_op_candidates(ea): F = ida_bytes.get_flags(ea) if ida_bytes.is_data(F): - return [0] # a data item's value is operand 0 + return [0] # a data item's value is operand 0 if not ida_bytes.is_code(F): - return [] # undefined bytes: IDA refuses a format outright + return [] # undefined bytes: IDA refuses a format outright insn = ida_ua.insn_t() if ida_ua.decode_insn(insn, ea) <= 0: return [] @@ -1037,7 +1146,7 @@ def _idatui_op_spans(ea, text): if not op: continue i = text.find(op, pos) - if i < 0: # duplicated operand text (mov eax, eax) + if i < 0: # duplicated operand text (mov eax, eax) i = text.find(op) if i < 0: continue @@ -1070,21 +1179,34 @@ def _idatui_apply_fmt(ea, n, fmt): if base in (idaapi.BADADDR, None) or base < 0: base = 0 return bool(ida_offset.op_plain_offset(ea, n, base)), "" - fn = {"hex": ida_bytes.op_hex, "dec": ida_bytes.op_dec, - "oct": ida_bytes.op_oct, "bin": ida_bytes.op_bin, - "char": ida_bytes.op_chr, "seg": ida_bytes.op_seg, - "float": ida_bytes.op_flt, "stack": ida_bytes.op_stkvar}.get(fmt) + fn = { + "hex": ida_bytes.op_hex, + "dec": ida_bytes.op_dec, + "oct": ida_bytes.op_oct, + "bin": ida_bytes.op_bin, + "char": ida_bytes.op_chr, + "seg": ida_bytes.op_seg, + "float": ida_bytes.op_flt, + "stack": ida_bytes.op_stkvar, + }.get(fmt) if fn is None: - return False, (f"can't set {fmt!r} from a name alone" - if fmt in _idatui_fmt_nibbles() else - f"unknown format {fmt!r}") + return False, ( + f"can't set {fmt!r} from a name alone" + if fmt in _idatui_fmt_nibbles() + else f"unknown format {fmt!r}" + ) return bool(fn(ea, n)), "" def op_format( addr: Annotated[str, "Address of the instruction or data item"], - mode: Annotated[str, "cycle | back | show | hex | dec | oct | bin | char | offset | stack | default"] = "cycle", - col: Annotated[int, "Cursor column inside the rendered line (-1: first literal)"] = -1, + mode: Annotated[ + str, + "cycle | back | show | hex | dec | oct | bin | char | offset | stack | default", + ] = "cycle", + col: Annotated[ + int, "Cursor column inside the rendered line (-1: first literal)" + ] = -1, n: Annotated[int, "Operand index; -1 derives it from ``col``"] = -1, ) -> dict: """Change how a literal is DISPLAYED (IDA's 'o' family): hex, decimal, @@ -1126,19 +1248,27 @@ def op_format( # a different operand would make that highlight a lie -- say # which one can be changed instead. where = before[lo:hi].strip() - alt = (f"; the literal on this line is operand {cands[0]} " - f"({_idatui_op_text(ea, before, cands[0])})" - if cands else "") - return {"addr": hex(ea), "n": i, "text": before, - "error": f"operand {i} ({where}) has no format to " - f"change{alt}"} + alt = ( + f"; the literal on this line is operand {cands[0]} " + f"({_idatui_op_text(ea, before, cands[0])})" + if cands + else "" + ) + return { + "addr": hex(ea), + "n": i, + "text": before, + "error": f"operand {i} ({where}) has no format to change{alt}", + } if n < 0: if not cands: F = ida_bytes.get_flags(ea) - why = ("no literal on this line to reformat" - if ida_bytes.is_code(F) or ida_bytes.is_data(F) else - "undefined bytes have no format to change -- define " - "them first ('d' makes data, 'c' makes code)") + why = ( + "no literal on this line to reformat" + if ida_bytes.is_code(F) or ida_bytes.is_data(F) + else "undefined bytes have no format to change -- define " + "them first ('d' makes data, 'c' makes code)" + ) return {"addr": hex(ea), "text": before, "error": why} n = cands[0] @@ -1148,9 +1278,12 @@ def op_format( # The ring is a property of the OPERAND, not of what you last pressed: every # stop is one that changes what you see for this value, and it is the same # ring at every step, so a lap always comes home. - choices = [f for f in _IDATUI_FMT_CYCLE - if (f != "char" or _idatui_printable(value)) - and (f != "offset" or _idatui_offset_worth(value))] + choices = [ + f + for f in _IDATUI_FMT_CYCLE + if (f != "char" or _idatui_printable(value)) + and (f != "offset" or _idatui_offset_worth(value)) + ] # A stack variable is deliberately NOT a stop: ``[rbp+var_40]`` is a frame # member, not a way of writing a number, and IDA's own "is this a stack # variable" test isn't exposed to Python here (calc_stkvar_struc_offset @@ -1160,10 +1293,18 @@ def op_format( mode = str(mode or "cycle").lower() if mode == "show": - return {"addr": hex(ea), "n": n, "format": cur, "prev": cur, - "choices": choices, "text": before, "before": before, - "value": None if value is None else hex(value), - "width": width, "applied": False} + return { + "addr": hex(ea), + "n": n, + "format": cur, + "prev": cur, + "choices": choices, + "text": before, + "before": before, + "value": None if value is None else hex(value), + "width": width, + "applied": False, + } if mode in ("cycle", "back"): step = 1 if mode == "cycle" else -1 if cur in choices: @@ -1176,31 +1317,51 @@ def op_format( else: want = mode if want not in _idatui_fmt_nibbles(): - return {"addr": hex(ea), "n": n, "text": before, - "error": f"unknown format {mode!r}; one of " - + ", ".join(_IDATUI_FMT_SETTABLE)} + return { + "addr": hex(ea), + "n": n, + "text": before, + "error": f"unknown format {mode!r}; one of " + + ", ".join(_IDATUI_FMT_SETTABLE), + } if want == "offset" and not mapped: - return {"addr": hex(ea), "n": n, "text": before, "format": cur, - "error": (f"{'0x%x' % value if value is not None else 'this operand'}" - " isn't a mapped address -- an offset to it would" - " invent a name for nothing")} + return { + "addr": hex(ea), + "n": n, + "text": before, + "format": cur, + "error": ( + f"{'0x%x' % value if value is not None else 'this operand'}" + " isn't a mapped address -- an offset to it would" + " invent a name for nothing" + ), + } ok, err = _idatui_apply_fmt(ea, n, want) if err: - return {"addr": hex(ea), "n": n, "text": before, "format": cur, - "error": err} + return {"addr": hex(ea), "n": n, "text": before, "format": cur, "error": err} got = _idatui_op_fmt(ea, n) - out = {"addr": hex(ea), "n": n, "prev": cur, "format": got, - "requested": want, "applied": bool(ok), "choices": choices, - "before": before, "text": _idatui_line_text(ea), - "value": None if value is None else hex(value), "width": width} + out = { + "addr": hex(ea), + "n": n, + "prev": cur, + "format": got, + "requested": want, + "applied": bool(ok), + "choices": choices, + "before": before, + "text": _idatui_line_text(ea), + "value": None if value is None else hex(value), + "width": width, + } if not ok: out["error"] = f"IDA refused {want} on operand {n}" elif lossy: - out["warn"] = ( - f"operand {n} was {cur} and the ring has no stop there -- " - + (f"'{cur}' sets it again" if cur in _IDATUI_FMT_SETTABLE else - f"{cur} names a type this can't put back, reassign it by hand")) + out["warn"] = f"operand {n} was {cur} and the ring has no stop there -- " + ( + f"'{cur}' sets it again" + if cur in _IDATUI_FMT_SETTABLE + else f"{cur} names a type this can't put back, reassign it by hand" + ) return out @@ -1267,13 +1428,18 @@ def _idatui_lit_extent(plain, x): around the column, which cannot reach a ``)`` or a space.""" if x >= len(plain): return None - if plain[x] == "'": # a character constant: '-' + if plain[x] == "'": # a character constant: '-' end = plain.find("'", x + 1) return (x, end + 1) if end > x else None lo = plain.rfind("'", 0, x) - if lo >= 0 and plain.find("'", x) > x and "'" in plain[lo:x] and \ - plain[lo:x].count("'") == 1 and " " not in plain[lo:x]: - return (lo, plain.find("'", x) + 1) # inside 'c' + if ( + lo >= 0 + and plain.find("'", x) > x + and "'" in plain[lo:x] + and plain[lo:x].count("'") == 1 + and " " not in plain[lo:x] + ): + return (lo, plain.find("'", x) + 1) # inside 'c' if plain[x] not in _IDATUI_LIT_CHARS: return None lo = x @@ -1282,7 +1448,7 @@ def _idatui_lit_extent(plain, x): hi = x while hi < len(plain) and plain[hi] in _IDATUI_LIT_CHARS: hi += 1 - if lo > 0 and plain[lo - 1] == "-": # a unary minus is part of it + if lo > 0 and plain[lo - 1] == "-": # a unary minus is part of it lo -= 1 return (lo, hi) @@ -1321,26 +1487,36 @@ def _idatui_pc_nums(cf, sl): continue nf = e.n.nf opnum = ord(nf.opnum) if isinstance(nf.opnum, str) else int(nf.opnum) - nbytes = (ord(nf.org_nbytes) if isinstance(nf.org_nbytes, str) - else int(nf.org_nbytes)) + nbytes = ( + ord(nf.org_nbytes) if isinstance(nf.org_nbytes, str) else int(nf.org_nbytes) + ) ea = int(e.ea) if ea == idaapi.BADADDR: x = extent[1] - continue # synthesised: nothing to key on + continue # synthesised: nothing to key on nib = (nf.flags >> ida_bytes.get_operand_type_shift(opnum)) & 0xF # Whether this format is the USER's or Hex-Rays' own guess. The nibble # can't say: an untouched number reads back as whatever it happens to # be printed as, and cycling from there would skip that stop forever # (default already looks like it) and never come back to it. loc = ida_hexrays.operand_locator_t(ea, opnum) - user = (ida_hexrays.user_numforms_find(cf.numforms, loc) - != ida_hexrays.user_numforms_end(cf.numforms)) - out.append({"x0": extent[0], "x1": extent[1], "ea": ea, - "opnum": opnum, "value": int(e.n._value), - "nbytes": nbytes, "user": user, - "fmt": _idatui_fmt_name(nib) if user else "default", - "shown": _idatui_fmt_name(nib)}) - x = extent[1] # past this literal, not into it + user = ida_hexrays.user_numforms_find( + cf.numforms, loc + ) != ida_hexrays.user_numforms_end(cf.numforms) + out.append( + { + "x0": extent[0], + "x1": extent[1], + "ea": ea, + "opnum": opnum, + "value": int(e.n._value), + "nbytes": nbytes, + "user": user, + "fmt": _idatui_fmt_name(nib) if user else "default", + "shown": _idatui_fmt_name(nib), + } + ) + x = extent[1] # past this literal, not into it return out @@ -1367,31 +1543,36 @@ def pc_nums( try: cf = ida_hexrays.decompile(f.start_ea) except Exception as e: - return {"addr": hex(f.start_ea), "error": f"decompile failed: {e}", - "nums": []} + return {"addr": hex(f.start_ea), "error": f"decompile failed: {e}", "nums": []} if cf is None: - return {"addr": hex(f.start_ea), "error": "decompilation failed", - "nums": []} + return {"addr": hex(f.start_ea), "error": "decompilation failed", "nums": []} sv = cf.get_pseudocode() out = [] for i in range(len(sv)): plain = ida_lines.tag_remove(sv[i].line) compact = _idatui_compact(plain) for rec in _idatui_pc_nums(cf, sv[i]): - out.append({ - "line": i, - "x0": _idatui_compact_col(plain, compact, rec["x0"]), - "x1": _idatui_compact_col(plain, compact, rec["x1"]), - "ea": hex(rec["ea"]), "opnum": rec["opnum"], - "value": hex(rec["value"]), "fmt": rec["fmt"], - "shown": rec["shown"], "user": bool(rec["user"]), - }) + out.append( + { + "line": i, + "x0": _idatui_compact_col(plain, compact, rec["x0"]), + "x1": _idatui_compact_col(plain, compact, rec["x1"]), + "ea": hex(rec["ea"]), + "opnum": rec["opnum"], + "value": hex(rec["value"]), + "fmt": rec["fmt"], + "shown": rec["shown"], + "user": bool(rec["user"]), + } + ) return {"addr": hex(f.start_ea), "nums": out, "lines": len(sv)} def pc_num_format( addr: Annotated[str, "Function address (or any address inside it)"], - mode: Annotated[str, "cycle | back | show | hex | dec | oct | char | default"] = "cycle", + mode: Annotated[ + str, "cycle | back | show | hex | dec | oct | char | default" + ] = "cycle", line: Annotated[int, "0-based pseudocode line index"] = -1, col: Annotated[int, "Cursor column in the DISPLAYED line (-1: first literal)"] = -1, ea: Annotated[str, "Address of the number instead of line/col"] = "", @@ -1431,8 +1612,9 @@ def pc_num_format( return {"addr": hex(f.start_ea), "error": str(e)} for i in range(len(sv)): for rec in _idatui_pc_nums(cf, sv[i]): - if rec["ea"] == want_ea and (int(opnum) < 0 - or rec["opnum"] == int(opnum)): + if rec["ea"] == want_ea and ( + int(opnum) < 0 or rec["opnum"] == int(opnum) + ): target, line = rec, i break if target: @@ -1446,26 +1628,42 @@ def pc_num_format( target = next((r for r in nums if r["x0"] <= x < r["x1"]), None) target = target or nums[0] else: - return {"addr": hex(f.start_ea), - "error": f"line {line} is outside the {len(sv)}-line decompilation"} + return { + "addr": hex(f.start_ea), + "error": f"line {line} is outside the {len(sv)}-line decompilation", + } if target is None: - return {"addr": hex(f.start_ea), "line": line, - "text": (ida_lines.tag_remove(sv[line].line).strip() - if 0 <= line < len(sv) else ""), - "error": "no number literal on this line"} + return { + "addr": hex(f.start_ea), + "line": line, + "text": ( + ida_lines.tag_remove(sv[line].line).strip() + if 0 <= line < len(sv) + else "" + ), + "error": "no number literal on this line", + } cur, value = target["fmt"], target["value"] - choices = [c for c in _IDATUI_PC_FMT_CYCLE - if c != "char" or _idatui_printable(value)] + choices = [ + c for c in _IDATUI_PC_FMT_CYCLE if c != "char" or _idatui_printable(value) + ] # Same rule as the listing: one ring per literal, every step. A format the # ring can't hold (an enum set in the GUI) is reported on the way out # instead of being kept for one lap and then lost. lossy = cur not in choices and cur != "default" - out = {"addr": hex(f.start_ea), "ea": hex(target["ea"]), - "opnum": target["opnum"], "line": line, "prev": cur, - "format": cur, "shown": target["shown"], "choices": choices, - "value": hex(value), - "before": ida_lines.tag_remove(sv[line].line).strip()} + out = { + "addr": hex(f.start_ea), + "ea": hex(target["ea"]), + "opnum": target["opnum"], + "line": line, + "prev": cur, + "format": cur, + "shown": target["shown"], + "choices": choices, + "value": hex(value), + "before": ida_lines.tag_remove(sv[line].line).strip(), + } mode = str(mode or "cycle").lower() if mode == "show": @@ -1481,13 +1679,16 @@ def pc_num_format( else: want = mode if want in ("bin", "offset", "stack", "seg", "float"): - out["error"] = (f"Hex-Rays has no {want} format for a number " - f"-- set it on the listing instead") + out["error"] = ( + f"Hex-Rays has no {want} format for a number " + f"-- set it on the listing instead" + ) out["text"] = out["before"] return out if want not in ("hex", "dec", "oct", "char", "default"): - out["error"] = (f"unknown format {mode!r}; one of hex, dec, oct, " - f"char, default") + out["error"] = ( + f"unknown format {mode!r}; one of hex, dec, oct, char, default" + ) out["text"] = out["before"] return out @@ -1499,8 +1700,9 @@ def pc_num_format( ida_hexrays.user_numforms_erase(cf.numforms, it) if want != "default": nf = ida_hexrays.number_format_t(target["opnum"]) - nf.flags = ida_bytes.get_operand_flag(_idatui_fmt_nibbles()[want], - target["opnum"]) + nf.flags = ida_bytes.get_operand_flag( + _idatui_fmt_nibbles()[want], target["opnum"] + ) try: nf.org_nbytes = target["nbytes"] except Exception: @@ -1515,14 +1717,18 @@ def pc_num_format( out["format"] = want out["applied"] = True if lossy: - out["warn"] = (f"this number was {cur}, which names a type a radix " - f"can't put back -- reassign it in IDA") + out["warn"] = ( + f"this number was {cur}, which names a type a radix " + f"can't put back -- reassign it in IDA" + ) try: - cf2 = ida_hexrays.decompile(f.start_ea, - flags=ida_hexrays.DECOMP_NO_CACHE) + cf2 = ida_hexrays.decompile(f.start_ea, flags=ida_hexrays.DECOMP_NO_CACHE) sv2 = cf2.get_pseudocode() if cf2 is not None else None - out["text"] = (ida_lines.tag_remove(sv2[line].line).strip() - if sv2 is not None and line < len(sv2) else out["before"]) + out["text"] = ( + ida_lines.tag_remove(sv2[line].line).strip() + if sv2 is not None and line < len(sv2) + else out["before"] + ) except Exception as e: out["text"] = out["before"] out["warn"] = f"re-render failed: {e}" @@ -1557,11 +1763,17 @@ def decompile(addr, include_addresses=True): 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}"} + 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}"} + return { + "addr": hex(int(fn.start_ea)), + "code": None, + "error": failure.desc() or f"Decompilation failed at {ea:#x}", + } lines = [] for sl in cfunc.get_pseudocode(): @@ -1569,7 +1781,9 @@ def decompile(addr, include_addresses=True): 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): + 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: @@ -1595,9 +1809,13 @@ def decompile(addr, include_addresses=True): 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}) + refs.append( + { + "addr": hex(target), + "name": ida_name.get_name(target) or "", + "string": text, + } + ) return 0 try: @@ -1615,6 +1833,7 @@ def decomp_map( swept across the line's columns via get_line_item. Shape: {addr, lines:[{ea: primary|None, eas:[hex,...]}, ...]}.""" import ida_hexrays + try: ea = int(str(addr), 16) except ValueError: @@ -1691,3 +1910,32 @@ def decomp_map( eas.append(hex(e)) lines.append({"ea": eas[0] if eas else None, "eas": eas}) return {"addr": hex(func.start_ea), "lines": lines} + + +def profile_remote(operation, args, reps=5): + """Profile one persistent remote tool entirely inside IDA.""" + import cProfile + import io + import pstats + + functions = { + "heads": heads, + "segment_index": segment_index, + "op_format": op_format, + "pc_nums": pc_nums, + "decompile": decompile, + "decomp_map": decomp_map, + "pc_num_format": pc_num_format, + } + function = functions.get(str(operation)) + if function is None: + raise ValueError(f"unknown profile operation: {operation!r}") + profiler = cProfile.Profile() + for _ in range(max(1, int(reps))): + profiler.runcall(function, **dict(args)) + stats = pstats.Stats(profiler) + output = io.StringIO() + stats.stream = output + stats.sort_stats("tottime") + stats.print_stats(80) + return {"stats": output.getvalue(), "total": stats.total_tt, "reps": reps} diff --git a/tests/test_codemode_client.py b/tests/test_codemode_client.py index 5ab95e2..dd0ae08 100644 --- a/tests/test_codemode_client.py +++ b/tests/test_codemode_client.py @@ -1,16 +1,20 @@ """IDA-free contract tests for the Code Mode client adapter.""" + from __future__ import annotations import os +import queue import sys +import threading +import time import tempfile from dataclasses import dataclass sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) - +from idatui.errors import IDAConnectionError, IDAToolError # noqa: E402 +from idatui import remote_ops # noqa: E402 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 @@ -35,6 +39,35 @@ class FakeEntry: record_id: str = "123-abcdef" exe_path: str = "" idb_path: str = "" + managed: bool = False + + +_CLOSED = object() + + +class FakeSubscription: + def __init__(self) -> None: + self._queue: queue.Queue = queue.Queue() + self.closed = False + + def __iter__(self): + return self + + def __next__(self): + item = self._queue.get() + if isinstance(item, BaseException): + raise item + if item is _CLOSED: + raise StopIteration + return item + + def emit(self, event: dict) -> None: + self._queue.put(event) + + def close(self) -> None: + if not self.closed: + self.closed = True + self._queue.put(_CLOSED) class FakeHandle: @@ -45,32 +78,74 @@ class FakeHandle: self.saved = 0 self.closed = False self.code = "" + self.codes = [] self.code_timeout = None + self.operation_label = None + self.event_origin_id = "fake-handle-origin" + self.owns_checks = 0 + self.subscription = FakeSubscription() + self.shutdown_calls = [] + self.shutdown_error = None def wait_autoanalysis(self, timeout=None): self.waited = timeout return {"complete": True, "status": "complete"} - def execute_python(self, code, timeout=None): + def execute_python( + self, + code, + timeout=None, + *, + operation_id=None, + operation_label=None, + persist_globals=False, + filename=None, + ): self.code = code + self.codes.append(code) self.code_timeout = timeout - return {"result": {"sentinel": 7}, "stdout": "", "stderr": ""} + self.operation_label = operation_label + result = ( + { + "__remote_ida_status__": "ok", + "__remote_ida_value__": {"sentinel": 7}, + } + if ".modules.get(" in code + else True + ) + return {"result": result, "stdout": "", "stderr": ""} + + def subscribe_idb_events(self): + return self.subscription + + def owns_event(self, event): + self.owns_checks += 1 + return event.get("origin_id") == self.event_origin_id def save_database(self): self.saved += 1 return {"saved": True, "idb_path": self.instance.idb_path} + def shutdown_database(self, *, save=True): + self.shutdown_calls.append(save) + if self.shutdown_error is not None: + raise module.RemoteError(self.shutdown_error, self.shutdown_error, 409) + return {"shutting_down": True, "save": save} + def close(self): + self.subscription.close() self.connected = False self.closed = True class FakeDatabaseHandle: opened = None + opens = 0 kwargs = None @classmethod def open(cls, path, **kwargs): + cls.opens += 1 cls.opened = path cls.kwargs = kwargs return FakeHandle(path) @@ -98,6 +173,10 @@ class FakeBusy(Exception): """Stand-in for DatabaseBusyError: `except None` is a TypeError.""" +class FakeDisconnected(Exception): + """Stand-in for DatabaseDisconnectedError in stdlib-only runs.""" + + def _open_kwargs_are_real(sent: dict): """(ok, detail) for the kwargs the adapter passes to DatabaseHandle.open. @@ -132,9 +211,11 @@ def _option_fields_are_real(options): 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)) + 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: @@ -148,8 +229,10 @@ def main() -> int: # (this file must keep running under a stdlib-only python3). original_options = module.DatabaseOpenOptions original_busy = module.DatabaseBusyError + original_disconnected = module.DatabaseDisconnectedError module.DatabaseOpenOptions = original_options or FakeOpenOptions module.DatabaseBusyError = original_busy or FakeBusy + module.DatabaseDisconnectedError = original_disconnected or FakeDisconnected try: with tempfile.TemporaryDirectory() as tmp: path = os.path.join(tmp, "sample.bin") @@ -159,48 +242,193 @@ def main() -> int: 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( + "connect delegates database discovery to DatabaseHandle.open", + FakeDatabaseHandle.opened == path and handle is not None, + ) options = FakeDatabaseHandle.kwargs["options"] - check("typed loader options cross the dependency boundary", - options.processor == "arm:ARMv7-A" - and options.image_base == 0x1000, - options) - check("every open option exists in the real library", - *_option_fields_are_real(options)) + check( + "typed loader options cross the dependency boundary", + options.processor == "arm:ARMv7-A" and options.image_base == 0x1000, + options, + ) + check( + "every open option exists in the real library", + *_option_fields_are_real(options), + ) # 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") + 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.call( + remote_ops.list_funcs, queries=[{"offset": 0, "count": 2}] + ) + check( + "remote operation returns its JSON result", + result == {"sentinel": 7}, + result, + ) + check( + "operation source is real Python installed through ida-domain", + any("db.functions.get_all()" in code for code in handle.codes), + handle.codes[0][:200], + ) + check( + "remote operations attribute IDB events to IDA TUI", + handle.operation_label == "IDA TUI", + handle.operation_label, + ) + batches = [] + delivered = threading.Event() + + def changed(batch): + batches.append(batch) + delivered.set() + + watcher = client.watch_idb_events(changed, debounce=0.05) + handle.subscription.emit({"event_name": "renamed", "origin_id": "peer-1"}) + handle.subscription.emit( + {"event_name": "cmt_changed", "origin_id": "peer-2"} + ) + check( + "event bursts produce one debounced refresh", + delivered.wait(1) and len(batches) == 1 and len(batches[0]) == 2, + batches, + ) + delivered.clear() + handle.subscription.emit( + {"event_name": "renamed", "origin_id": handle.event_origin_id} + ) + time.sleep(0.1) + check( + "the listener uses handle ownership to ignore its own events", + not delivered.is_set() + and len(batches) == 1 + and handle.owns_checks >= 3, + (batches, handle.owns_checks), + ) + handle.subscription.emit( + {"event_name": "byte_patched", "origin_id": "peer-3"} + ) + watcher.close() + time.sleep(0.1) + check( + "closing drops a pending debounced refresh", + not delivered.is_set() and len(batches) == 1, + batches, + ) + 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) + check( + "GUI leases transfer rather than claiming discard", + client.discard_database() is False and handle.shutdown_calls == [], + handle.shutdown_calls, + ) + handle.instance = FakeEntry( + backend="idalib", managed=True, exe_path=path, idb_path=path + ".i64" + ) + check( + "a final managed lease discards without saving", + client.discard_database() is True and handle.shutdown_calls == [False], + handle.shutdown_calls, + ) + handle.shutdown_error = "instance_shared" + check( + "a shared managed lease transfers finalization", + client.discard_database() is False, + handle.shutdown_calls, + ) + handle.shutdown_error = "instance_busy" + try: + client.discard_database(timeout=0) + except IDAToolError as exc: + check( + "a busy final lease never silently saves", + exc.tool == "shutdown_database", + exc, + ) + else: + check("a busy final lease never silently saves", False) + handle.shutdown_error = None + handle.instance = FakeEntry(exe_path=path, idb_path=path + ".i64") + opens = FakeDatabaseHandle.opens + handle.connected = False + try: + client.health() + except IDAConnectionError as exc: + check( + "a disconnected handle requires explicit rediscovery", + "explicit rediscovery" in str(exc) + and FakeDatabaseHandle.opens == opens, + (exc, FakeDatabaseHandle.opens, opens), + ) + else: + check("a disconnected handle requires explicit rediscovery", False) 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) + check( + "GUI lifetime is never claimed by the client", + client.wait_released(0) is False, + ) + disconnected = CodeModeClient(path).connect() + stream_errors = [] + stream_failed = threading.Event() + + def failed(error): + stream_errors.append(error) + stream_failed.set() + + stream_watch = disconnected.watch_idb_events( + lambda _batch: None, on_error=failed, debounce=0 + ) + disconnected._handle.subscription.emit( + module.DatabaseDisconnectedError("GUI database closed") + ) + check( + "stream disconnects become application connection errors", + stream_failed.wait(1) + and isinstance(stream_errors[0], IDAConnectionError), + stream_errors, + ) + stream_watch.close() + disconnected.close() finally: module.DatabaseHandle = original module.DatabaseOpenOptions = original_options module.DatabaseBusyError = original_busy + module.DatabaseDisconnectedError = original_disconnected client = CodeModeClient(__file__) + + def unknown_operation(): + pass + try: - client.invoke("not-an-operation") + client.call(unknown_operation) except IDAToolError as exc: - check("unknown adapter operations are explicit", exc.tool == "not-an-operation") + check( + "unknown adapter operations are explicit", + exc.tool == "unknown_operation", + ) else: check("unknown adapter operations are explicit", False) diff --git a/tests/test_pool.py b/tests/test_pool.py index 9fc1446..6309e13 100644 --- a/tests/test_pool.py +++ b/tests/test_pool.py @@ -33,11 +33,14 @@ def check(name, cond, detail=""): class FakeClient: """Stands in for a CodeModeClient lease and records saves/closes.""" - def __init__(self, ref, mem=100, backend="idalib"): + def __init__(self, ref, mem=100, backend="idalib", + discardable=True): self.ref = ref self.mem = mem self.backend = backend self.saved = 0 + self.discarded = 0 + self.discardable = discardable self.closed = False self.connected = False @@ -49,6 +52,10 @@ class FakeClient: self.saved += 1 return {"saved": True} + def discard_database(self): + self.discarded += 1 + return self.discardable + def close(self, grace=None): self.closed = True @@ -153,6 +160,36 @@ def main() -> int: except KeyError: check("an unknown label raises KeyError", True) + # -- discard delegates shared/GUI finalization ------------------------ # + discard_made = {} + + def spawn_discard(ref, ttl): + client = FakeClient( + ref, discardable=ref.label != "bin1") + discard_made[ref.label] = client + return client + + discard_pool = DatabasePool( + proj, spawn=spawn_discard, mem_fn=lambda c: c.mem) + discard_pool.get("bin0") + discard_pool.get("bin1") + delegated = discard_pool.discard_changes(["bin0", "bin1"]) + check("discard asks every dirty resident database", + discard_made["bin0"].discarded == 1 + and discard_made["bin1"].discarded == 1, + {k: c.discarded for k, c in discard_made.items()}) + check("discard reports leases whose finalization transferred", + delegated == ["bin1"], delegated) + old = discard_made["bin0"] + replacement = FakeClient(proj.by_label("bin0")) + check("replace_client refuses a stale lease generation", + discard_pool.replace_client("bin0", object(), replacement) is False + and discard_pool.get("bin0") is old) + check("replace_client installs the reattached lease", + discard_pool.replace_client("bin0", old, replacement) is True + and discard_pool.get("bin0") is replacement) + discard_pool.close_all(save=False) + # -- default budget comes from the project's memory_pct ------------------- # pool3 = DatabasePool(proj, spawn=spawn, mem_fn=lambda c: c.mem) check("default budget is derived, not a fixed lease count", diff --git a/tests/test_scenarios.py b/tests/test_scenarios.py index 2acb5ba..45ed51b 100644 --- a/tests/test_scenarios.py +++ b/tests/test_scenarios.py @@ -31,16 +31,34 @@ 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 _fixtures import fast_keys, staged # noqa: E402 -fast_keys() # ~85ms -> ~2ms per keypress; see _fixtures.fast_keys +fast_keys() # ~85ms -> ~2ms per keypress; see _fixtures.fast_keys +from idatui import remote_ops # noqa: E402 from idatui.app import ( # noqa: E402 - ConfirmScreen, DecompView, FunctionsPanel, GraphView, HexView, IdaTui, - HelpScreen, ListingView, QuitScreen, SearchPalette, StringsPalette, - StructEditor, SymbolPalette, XrefsScreen, _HELP, _str_display, + ConfirmScreen, + DecompView, + FunctionsPanel, + GraphView, + HexView, + IdaTui, + HelpScreen, + ListingView, + QuitScreen, + SearchPalette, + StringsPalette, + StructEditor, + SymbolPalette, + XrefsScreen, + _HELP, + _str_display, _word_occurrences, ) from idatui.errors import IDAToolError # noqa: E402 from textual.widgets import ( # noqa: E402 - DataTable, Input, OptionList, Static, TextArea, + DataTable, + Input, + OptionList, + Static, + TextArea, ) from rich.text import Text # noqa: E402 from idatui._sync import settle, wait_for # noqa: E402 @@ -86,8 +104,10 @@ class _Profile: return tp, tw = sum(self.paused.values()), sum(self.waited.values()) tk = sum(self.pressed.values()) - print(f"\nprofile: {tp:.1f}s settling, {tw:.1f}s in waits, " - f"{tk:.1f}s in keystrokes") + print( + f"\nprofile: {tp:.1f}s settling, {tw:.1f}s in waits, " + f"{tk:.1f}s in keystrokes" + ) rows = sorted(self.paused.items(), key=lambda kv: -kv[1])[:8] for name, secs in rows: print(f" pause {secs:5.2f}s {name}") @@ -98,8 +118,10 @@ class _Profile: for name, secs in rows: print(f" keys {secs:5.2f}s {name}") for name, secs in self.expired: - print(f" EXPIRED wait {secs:5.2f}s in {name} " - f"(the check after it may have passed vacuously)") + print( + f" EXPIRED wait {secs:5.2f}s in {name} " + f"(the check after it may have passed vacuously)" + ) PROFILE = _Profile() @@ -113,6 +135,7 @@ def scenario(name): def deco(fn): SCENARIOS.append((name, fn)) return fn + return deco @@ -149,8 +172,12 @@ class Ctx: async def wait(self, pred, t=20.0, step=0.02): t0 = asyncio.get_event_loop().time() ok = await wait_for(pred, self.pilot.pause, t, step) - PROFILE.wait(self.scenario, asyncio.get_event_loop().time() - t0, ok, - sys._getframe(1).f_lineno) + PROFILE.wait( + self.scenario, + asyncio.get_event_loop().time() - t0, + ok, + sys._getframe(1).f_lineno, + ) return ok async def press(self, *keys): @@ -293,8 +320,9 @@ class Ctx: # _cur.ea to match: re-opening a function we already navigated to (its # _cur.ea is stale-true) schedules an async re-navigation, and proceeding # before it lands would leave the cursor parked wherever we last were. - await self.wait(lambda: self.lst.total > 0 - and self.lst._cursor_ea() == fn.addr, t) + await self.wait( + lambda: self.lst.total > 0 and self.lst._cursor_ea() == fn.addr, t + ) if view == "decomp": # F5/Tab only decompiles from a focused code pane, and the listing # may still be settling from the open above — a swallowed Tab used to @@ -304,8 +332,12 @@ class Ctx: self.lst.focus() await self.pause(0.05) await self.press("tab") - if await self.wait(lambda: self.app._active == "decomp" - and self.dec.loaded_ea == fn.addr, max(t / 3, 5)): + if await self.wait( + lambda: ( + self.app._active == "decomp" and self.dec.loaded_ea == fn.addr + ), + max(t / 3, 5), + ): break return fn @@ -332,7 +364,8 @@ class Ctx: # Load is done when the index is complete (robust vs the status line, # which startup auto-land immediately overwrites with the landed fn). await self.wait( - lambda: app._func_index is not None and app._func_index.complete, 60) + lambda: app._func_index is not None and app._func_index.complete, 60 + ) if app._func_index is not None and not app._func_index.complete: app._func_index.load_all() @@ -359,7 +392,7 @@ class Ctx: app._pref = "decomp" if app._active in ("hex", "graph"): app._active = "decomp" - app._graph_sticky = False # else every later scenario rebuilds a graph + app._graph_sticky = False # else every later scenario rebuilds a graph app._split = False await self.pause(0.02) @@ -369,15 +402,26 @@ class Ctx: # --------------------------------------------------------------------------- # @scenario("startup") async def s_startup(c: Ctx): - c.check("function list populated", c.table.row_count > 0, f"rows={c.table.row_count}") + c.check( + "function list populated", c.table.row_count > 0, f"rows={c.table.row_count}" + ) left = c.app.query_one("#left") - c.check("names pane starts hidden (overlay-first)", not left.display, - f"display={left.display}") + c.check( + "names pane starts hidden (overlay-first)", + not left.display, + f"display={left.display}", + ) await c.reveal_pane() - c.check("function pane width is capped (doesn't eat the screen)", - left.size.width <= 44, f"width={left.size.width}") - c.check("function load completed (index complete)", - c.app._func_index is not None and c.app._func_index.complete, c.status()) + c.check( + "function pane width is capped (doesn't eat the screen)", + left.size.width <= 44, + f"width={left.size.width}", + ) + c.check( + "function load completed (index complete)", + c.app._func_index is not None and c.app._func_index.complete, + c.status(), + ) print(f" {c.table.row_count} functions loaded") @@ -398,13 +442,18 @@ async def s_auto_land(c: Ctx): await c.pause(0.2) if fn is not None: await c.wait(lambda: app._cur is not None and app._cur.ea == fn.addr, 20) - c.check("auto-land jumps to the entry function (main) when present", - app._cur is not None and app._cur.ea == fn.addr, - f"entry={fn.name}@{fn.addr:#x} cur={app._cur}") + c.check( + "auto-land jumps to the entry function (main) when present", + app._cur is not None and app._cur.ea == fn.addr, + f"entry={fn.name}@{fn.addr:#x} cur={app._cur}", + ) else: await c.wait(lambda: isinstance(app.screen, SymbolPalette), 10) - c.check("auto-land pops the symbol picker when there's no entry fn", - isinstance(app.screen, SymbolPalette), f"screen={app.screen}") + c.check( + "auto-land pops the symbol picker when there's no entry fn", + isinstance(app.screen, SymbolPalette), + f"screen={app.screen}", + ) app.pop_screen() # guard fires once: a second call is a no-op prev = app._cur @@ -437,17 +486,23 @@ async def s_segment_index(c: Ctx): c.check("segment streamed for comparison", False) return - idx = app.program.client.invoke("segment_index", addr=hex(model.seg_start)) - c.check("segment_index counts exactly what streaming produced", - idx.get("rows") == len(model), - f"index={idx.get('rows')} streamed={len(model)}") - c.check("it reports the same segment", - int(str(idx.get("addr")), 16) == model.seg_start, - f"{idx.get('addr')} vs {model.seg_start:#x}") + idx = app.program.client.call(remote_ops.segment_index, addr=hex(model.seg_start)) + c.check( + "segment_index counts exactly what streaming produced", + idx.get("rows") == len(model), + f"index={idx.get('rows')} streamed={len(model)}", + ) + c.check( + "it reports the same segment", + int(str(idx.get("addr")), 16) == model.seg_start, + f"{idx.get('addr')} vs {model.seg_start:#x}", + ) anchors = idx.get("anchors") or [] - c.check("anchors cover the segment", - len(anchors) >= max(1, len(model) // 500), - f"{len(anchors)} anchors for {len(model)} rows") + c.check( + "anchors cover the segment", + len(anchors) >= max(1, len(model) // 500), + f"{len(anchors)} anchors for {len(model)} rows", + ) # Every anchor must name the address of the row it claims, or seeking to it # would land somewhere else entirely. @@ -456,8 +511,11 @@ async def s_segment_index(c: Ctx): h = model.get(row) if h is None or h.ea != int(str(ea), 16): bad.append((row, ea, hex(h.ea) if h else None)) - c.check("every anchor points at the row it claims", not bad, - f"{len(bad)} wrong, first={bad[:2]}") + c.check( + "every anchor points at the row it claims", + not bad, + f"{len(bad)} wrong, first={bad[:2]}", + ) # A model built from the index must be INDISTINGUISHABLE from a streamed # one. That is the invariant the whole optimisation rests on: _prime builds @@ -471,6 +529,7 @@ async def s_segment_index(c: Ctx): # rename and define things, so that model describes the database as it was # at boot, not as it is now. from idatui.domain import ListingModel # noqa: PLC0415 + args = (app.program, model.seg_start, model.seg_end, model.name) idx_model, streamed = ListingModel(*args), ListingModel(*args) if not idx_model.build_from_index(): @@ -480,16 +539,30 @@ async def s_segment_index(c: Ctx): if streamed.load_next_page(text=False) == 0: break c.check("an index-built model is complete immediately", idx_model.complete) - c.check("index-built model has the streamed row count", - len(idx_model) == len(streamed), f"{len(idx_model)} vs {len(streamed)}") - for field in ("_row_at", "_head_eas", "_by_ea", "_page_head", - "_page_addr", "_page_rows"): + c.check( + "index-built model has the streamed row count", + len(idx_model) == len(streamed), + f"{len(idx_model)} vs {len(streamed)}", + ) + for field in ( + "_row_at", + "_head_eas", + "_by_ea", + "_page_head", + "_page_addr", + "_page_rows", + ): a, b = getattr(idx_model, field), getattr(streamed, field) - c.check(f"index-built {field} matches streaming", a == b, - f"len {len(a)} vs {len(b)}") - c.check("index-built rows carry the same ea/kind/size", - [(h.ea, h.kind, h.size) for h in idx_model._heads] - == [(h.ea, h.kind, h.size) for h in streamed._heads]) + c.check( + f"index-built {field} matches streaming", + a == b, + f"len {len(a)} vs {len(b)}", + ) + c.check( + "index-built rows carry the same ea/kind/size", + [(h.ea, h.kind, h.size) for h in idx_model._heads] + == [(h.ea, h.kind, h.size) for h in streamed._heads], + ) @scenario("reprime_is_free") @@ -521,7 +594,7 @@ async def s_reprime_is_free(c: Ctx): type(client).invoke = counting try: - for _ in range(3): # decomp and back, three times + for _ in range(3): # decomp and back, three times await c.press("tab") await c.pause(0.05) await c.press("tab") @@ -530,8 +603,11 @@ async def s_reprime_is_free(c: Ctx): type(client).invoke = original rebuilds = seen.count("segment_index") - c.check("switching views never rebuilds the segment index", rebuilds == 0, - f"segment_index called {rebuilds}x during 3 view switches: {seen}") + c.check( + "switching views never rebuilds the segment index", + rebuilds == 0, + f"segment_index called {rebuilds}x during 3 view switches: {seen}", + ) @scenario("skeleton_pages") @@ -561,8 +637,11 @@ async def s_skeleton_pages(c: Ctx): if not model.complete or len(model) < 1200: # A target smaller than _prime's horizon has no skeleton pages at all, # so there is nothing to check rather than something broken. - c.check("segment is big enough to have skeleton pages", True, - f"skipped: only {len(model)} rows, _prime renders ~1000") + c.check( + "segment is big enough to have skeleton pages", + True, + f"skipped: only {len(model)} rows, _prime renders ~1000", + ) return c.check("pages were loaded as skeletons", model._skeleton is True) @@ -570,24 +649,32 @@ async def s_skeleton_pages(c: Ctx): deep = max(1200, len(model) - 40) for row in (1200, len(model) // 2, deep): h = model.get(row) - c.check(f"row {row} of a skeleton page has real text", - h is not None and bool((h.text or "").strip()), - f"ea={getattr(h, 'ea', None)} text={getattr(h, 'text', None)!r}") + c.check( + f"row {row} of a skeleton page has real text", + h is not None and bool((h.text or "").strip()), + f"ea={getattr(h, 'ea', None)} text={getattr(h, 'text', None)!r}", + ) # And through the render path the user actually sees, not just the model. lv.cursor = deep lv.refresh() await c.pause(0.1) painted = lv._line_plain(deep) - c.check("a deep row RENDERS with text", - bool(painted and painted.strip()), f"painted={painted!r}") + c.check( + "a deep row RENDERS with text", + bool(painted and painted.strip()), + f"painted={painted!r}", + ) # Materialising must not change the row count or move any address: the # skeleton's structure is what the scrollbar was sized from. before = len(model) model.get(deep) - c.check("materialising a page does not change the row count", - len(model) == before, f"{before} -> {len(model)}") + c.check( + "materialising a page does not change the row count", + len(model) == before, + f"{before} -> {len(model)}", + ) c.check("the walk was not disturbed", not model.stale_structure) @@ -615,41 +702,58 @@ async def s_palette_paging(c: Ctx): # below would pass vacuously against a zero-height page. await c.wait(lambda: ol.scrollable_content_region.height >= 1, 10) page = ol.scrollable_content_region.height - c.check("the palette list has a real viewport to page by", page >= 1, - f"height={page}") + c.check( + "the palette list has a real viewport to page by", page >= 1, f"height={page}" + ) if ol.option_count <= 2: c.check("enough symbols to page through", False, f"n={ol.option_count}") return - c.check("the filter Input holds focus (so the list never sees the key)", - pal.focused is inp, f"focused={type(pal.focused).__name__}") + c.check( + "the filter Input holds focus (so the list never sees the key)", + pal.focused is inp, + f"focused={type(pal.focused).__name__}", + ) ol.highlighted = 0 await c.press("pagedown") down = ol.highlighted or 0 # A page, not a line: the bug this guards against is PgDn falling through to # the Input and moving nothing, or degrading to a single-step cursor move. - c.check("PgDn moves the symbol list by more than one row", down > 1, - f"highlighted={down} page={page} n={ol.option_count}") - c.check("PgDn moves by about a viewport (or lands on the last row)", - down >= min(page, ol.option_count - 1) - 1, - f"highlighted={down} page={page} n={ol.option_count}") + c.check( + "PgDn moves the symbol list by more than one row", + down > 1, + f"highlighted={down} page={page} n={ol.option_count}", + ) + c.check( + "PgDn moves by about a viewport (or lands on the last row)", + down >= min(page, ol.option_count - 1) - 1, + f"highlighted={down} page={page} n={ol.option_count}", + ) await c.press("pageup") - c.check("PgUp comes back to the top", (ol.highlighted or 0) == 0, - f"highlighted={ol.highlighted}") + c.check( + "PgUp comes back to the top", + (ol.highlighted or 0) == 0, + f"highlighted={ol.highlighted}", + ) # Clamping: hammering past the end must settle on the last row, not wrap or # raise. 12 pages clears any list this palette will show. for _ in range(12): await c.press("pagedown") - c.check("PgDn clamps at the last row", - ol.highlighted == ol.option_count - 1, - f"highlighted={ol.highlighted} n={ol.option_count}") + c.check( + "PgDn clamps at the last row", + ol.highlighted == ol.option_count - 1, + f"highlighted={ol.highlighted} n={ol.option_count}", + ) for _ in range(12): await c.press("pageup") - c.check("PgUp clamps at the first row", ol.highlighted == 0, - f"highlighted={ol.highlighted}") + c.check( + "PgUp clamps at the first row", + ol.highlighted == 0, + f"highlighted={ol.highlighted}", + ) await c.press("escape") await c.wait(lambda: not isinstance(app.screen, SymbolPalette), 10) @@ -666,19 +770,28 @@ async def s_xrefs_paging(c: Ctx): await c.open_biggest("listing") await c.press("x") if not await c.wait(lambda: isinstance(app.screen, XrefsScreen), 10): - c.check("xrefs popup opens for the paging test", True, - "skipped: no xrefs at this cursor") + c.check( + "xrefs popup opens for the paging test", + True, + "skipped: no xrefs at this cursor", + ) return scr = app.screen ol = scr.query_one(OptionList) await c.wait(lambda: ol.scrollable_content_region.height >= 1, 10) - c.check("the xrefs list itself has focus", scr.focused is ol, - f"focused={type(scr.focused).__name__}") + c.check( + "the xrefs list itself has focus", + scr.focused is ol, + f"focused={type(scr.focused).__name__}", + ) if ol.option_count > 2: ol.highlighted = 0 await c.press("pagedown") - c.check("PgDn pages the xrefs list natively", (ol.highlighted or 0) > 1, - f"highlighted={ol.highlighted} n={ol.option_count}") + c.check( + "PgDn pages the xrefs list natively", + (ol.highlighted or 0) > 1, + f"highlighted={ol.highlighted} n={ol.option_count}", + ) await c.press("escape") await c.wait(lambda: not isinstance(app.screen, XrefsScreen), 10) @@ -688,40 +801,51 @@ async def s_palette(c: Ctx): app, pilot = c.app, c.pilot await c.press("ctrl+n") pal_open = await c.wait(lambda: isinstance(app.screen, SymbolPalette), 10) - c.check("Ctrl+N opens the symbol palette", pal_open, - f"screen={type(app.screen).__name__}") + c.check( + "Ctrl+N opens the symbol palette", + pal_open, + f"screen={type(app.screen).__name__}", + ) if not pal_open: return pal = app.screen pinp = pal.query_one(Input) pinp.value = "main" await c.wait(lambda: pal._results and pal._results[0][2] == "main", 10) - c.check("palette fuzzy-finds (top result matches the query)", - bool(pal._results) and pal._results[0][2] == "main", - f"top={pal._results[0][2] if pal._results else None}") + c.check( + "palette fuzzy-finds (top result matches the query)", + bool(pal._results) and pal._results[0][2] == "main", + f"top={pal._results[0][2] if pal._results else None}", + ) pinp.value = "eror" # scattered subsequence of 'error' await c.wait(lambda: any(n == "error" for _, _, n in pal._results), 10) - c.check("palette matches a fuzzy subsequence", - any(n == "error" for _, _, n in pal._results), - f"results={[n for _, _, n in pal._results[:4]]}") + c.check( + "palette matches a fuzzy subsequence", + any(n == "error" for _, _, n in pal._results), + f"results={[n for _, _, n in pal._results[:4]]}", + ) # Every other query here is lowercase, which is how a case bug hid for so # long: the name was lowered but the query wasn't, so ONE capital matched # nothing. Invisible on lowercase C symbols, fatal on a library that # capitalises (PEM_read_bio found 0 of 10093 functions in libcrypto). pinp.value = "MAIN" await c.wait(lambda: any(n == "main" for _, _, n in pal._results), 10) - c.check("palette matching is case-insensitive in BOTH directions", - any(n == "main" for _, _, n in pal._results), - f"results={[n for _, _, n in pal._results[:4]]}") + c.check( + "palette matching is case-insensitive in BOTH directions", + any(n == "main" for _, _, n in pal._results), + f"results={[n for _, _, n in pal._results[:4]]}", + ) pinp.value = "main" await c.wait(lambda: pal._results and pal._results[0][2] == "main", 10) want = pal._results[0][1] await c.press("enter") await c.wait(lambda: not isinstance(app.screen, SymbolPalette), 10) await c.wait(lambda: app._cur and app._cur.ea == want, 20) - c.check("selecting a palette entry opens that function", - bool(app._cur) and app._cur.ea == want, - f"cur={app._cur.ea if app._cur else None}") + c.check( + "selecting a palette entry opens that function", + bool(app._cur) and app._cur.ea == want, + f"cur={app._cur.ea if app._cur else None}", + ) # Re-open the function we are ALREADY standing on. That used to append an # identical nav entry, and the extra Esc it bought popped the stack without # changing anything on screen — a dead keypress, which is precisely what @@ -735,8 +859,11 @@ async def s_palette(c: Ctx): await c.press("enter") await c.wait(lambda: not isinstance(app.screen, SymbolPalette), 10) await c.pause(0.4) - c.check("re-opening the current function doesn't stack a duplicate", - len(app._nav) == depth, f"nav {depth} -> {len(app._nav)}") + c.check( + "re-opening the current function doesn't stack a duplicate", + len(app._nav) == depth, + f"nav {depth} -> {len(app._nav)}", + ) await c.press("ctrl+n") await c.wait(lambda: isinstance(app.screen, SymbolPalette), 10) await c.press("escape") @@ -749,43 +876,57 @@ async def s_load_options(c: Ctx): """The dialog must never appear for a file IDA can load itself — the whole suite runs on an ELF, so a false positive here would block every run.""" from idatui.app import LoadOptionsScreen + app = c.app - c.check("no load dialog for a recognised binary", - not isinstance(app.screen, LoadOptionsScreen), - f"screen={type(app.screen).__name__}") - c.check("and the app agrees it shouldn't ask", - not app._should_ask_load_options()) + c.check( + "no load dialog for a recognised binary", + not isinstance(app.screen, LoadOptionsScreen), + f"screen={type(app.screen).__name__}", + ) + c.check("and the app agrees it shouldn't ask", not app._should_ask_load_options()) # The dialog itself, driven directly: it has to come back with switches the # worker can use, and -b has to be paragraphs. from idatui.formats import load_args, needs_load_options, sniff - c.check("the running target sniffs as a real format", - sniff(app._open_path) is not None and not needs_load_options(app._open_path), - f"{sniff(app._open_path)}") - c.check("dialog output converts a base to paragraphs", - load_args("arm", 0x8000000) == "-parm -b800000") + + c.check( + "the running target sniffs as a real format", + sniff(app._open_path) is not None and not needs_load_options(app._open_path), + f"{sniff(app._open_path)}", + ) + c.check( + "dialog output converts a base to paragraphs", + load_args("arm", 0x8000000) == "-parm -b800000", + ) # Tab is a PRIORITY app binding (disasm<->pseudocode), so it fired even with # a modal up and nothing in a dialog could be tabbed to. That is why the load # dialog's address field was unreachable — and it was broken in every other # modal too. from idatui.app import LoadOptionsScreen + app.push_screen(LoadOptionsScreen("/tmp/probe.bin", 1234)) await c.wait(lambda: isinstance(app.screen, LoadOptionsScreen), 10) sc = app.screen first = app.focused await c.press("tab") await c.pause(0.2) - c.check("Tab moves focus inside a modal instead of toggling the view", - app.focused is not first and isinstance(app.screen, LoadOptionsScreen), - f"focus={getattr(app.focused, 'id', None)}") - c.check("Tab in the load dialog lands on the address field", - getattr(app.focused, "id", None) == "load-base", - f"focus={getattr(app.focused, 'id', None)}") + c.check( + "Tab moves focus inside a modal instead of toggling the view", + app.focused is not first and isinstance(app.screen, LoadOptionsScreen), + f"focus={getattr(app.focused, 'id', None)}", + ) + c.check( + "Tab in the load dialog lands on the address field", + getattr(app.focused, "id", None) == "load-base", + f"focus={getattr(app.focused, 'id', None)}", + ) await c.press("tab") await c.pause(0.2) - c.check("Tab again returns to the processor filter", - getattr(app.focused, "id", None) == "pal-input", - f"focus={getattr(app.focused, 'id', None)}") + c.check( + "Tab again returns to the processor filter", + getattr(app.focused, "id", None) == "pal-input", + f"focus={getattr(app.focused, 'id', None)}", + ) await c.press("escape") await c.wait(lambda: not isinstance(app.screen, LoadOptionsScreen), 10) @@ -800,6 +941,7 @@ async def s_asm_highlight(c: Ctx): with the plain text, carried on the Head, and mapped to a style. """ from idatui.app import _S_SPAN + app = c.app await c.open_biggest("listing") lst = c.lst @@ -811,30 +953,40 @@ async def s_asm_highlight(c: Ctx): rows = [lst.model.get(i) for i in range(min(lst.model.loaded(), 400))] rows = [h for h in rows if h is not None] code = [h for h in rows if h.kind == "code"] - c.check("code rows carry IDA's token spans", - code and sum(1 for h in code if h.spans) > len(code) * 0.9, - f"{sum(1 for h in code if h.spans)}/{len(code)} have spans") + c.check( + "code rows carry IDA's token spans", + code and sum(1 for h in code if h.spans) > len(code) * 0.9, + f"{sum(1 for h in code if h.spans)}/{len(code)} have spans", + ) kinds = {k for h in rows for k, _ in (h.spans or ())} # If a tag isn't mapped it renders as body text and nothing says why, so the # ones that carry real meaning are worth asserting explicitly. for want in ("insn", "reg", "punct"): c.check(f"the palette sees {want} tokens", want in kinds, f"{sorted(kinds)}") - c.check("every span kind has a style", - all(k in _S_SPAN for k in kinds), f"unstyled: {sorted(kinds - set(_S_SPAN))}") + c.check( + "every span kind has a style", + all(k in _S_SPAN for k in kinds), + f"unstyled: {sorted(kinds - set(_S_SPAN))}", + ) # Spans must describe the SAME text the row shows, or the row renders # different characters than search/width calculations think it has. - bad = [h for h in rows if h.spans - and "".join(t for _k, t in h.spans) != h.text] - c.check("spans reconstruct the row text exactly", not bad, - f"{[(hex(h.ea), h.text) for h in bad[:2]]}") + bad = [h for h in rows if h.spans and "".join(t for _k, t in h.spans) != h.text] + c.check( + "spans reconstruct the row text exactly", + not bad, + f"{[(hex(h.ea), h.text) for h in bad[:2]]}", + ) # And the mnemonic must be the loudest thing on the line (the column you # scan), not just any styled token. mn = next((h for h in code if h.spans and h.spans[0][0] == "insn"), None) - c.check("the mnemonic is the first span", - mn is not None, f"{code[0].spans if code else None}") + c.check( + "the mnemonic is the first span", + mn is not None, + f"{code[0].spans if code else None}", + ) @scenario("status_names_the_file") @@ -846,26 +998,30 @@ async def s_status_names_the_file(c: Ctx): it belongs to. """ from textual.widgets import Static + app = c.app await c.open_biggest("listing") await c.pause(0.3) name = os.path.basename(app._open_path) status = str(app.query_one("#status", Static).render()) - c.check("the status bar names the open file", - status.startswith(f"[{name}]"), f"{status[:60]!r} (want [{name}])") + c.check( + "the status bar names the open file", + status.startswith(f"[{name}]"), + f"{status[:60]!r} (want [{name}])", + ) # It must survive the messages that WRITE the status, not just the idle one. c.lst.focus() await c.press("down") await c.pause(0.3) status = str(app.query_one("#status", Static).render()) - c.check("and keeps naming it as you move", - status.startswith(f"[{name}]"), status[:60]) + c.check( + "and keeps naming it as you move", status.startswith(f"[{name}]"), status[:60] + ) # The function-count message used to include the module name itself, which # would now read "[echo] echo — 128 functions". - c.check("without saying the name twice", - status.count(name) == 1, status[:70]) + c.check("without saying the name twice", status.count(name) == 1, status[:70]) @scenario("command_palette") @@ -873,10 +1029,12 @@ async def s_command_palette(c: Ctx): app = c.app await c.open_biggest("listing") await c.press("ctrl+p") - opened = await c.wait( - lambda: type(app.screen).__name__ == "CommandPalette", 10) - c.check("Ctrl+P opens the command palette", opened, - f"screen={type(app.screen).__name__}") + opened = await c.wait(lambda: type(app.screen).__name__ == "CommandPalette", 10) + c.check( + "Ctrl+P opens the command palette", + opened, + f"screen={type(app.screen).__name__}", + ) if not opened: return inp = app.screen.query_one(Input) @@ -884,8 +1042,9 @@ async def s_command_palette(c: Ctx): await c.pause(0.5) # let the async search + option list settle await c.press("enter") landed = await c.wait(lambda: app._active == "hex", 10) - c.check("a palette command executes (Hex view opens)", landed, - f"active={app._active}") + c.check( + "a palette command executes (Hex view opens)", landed, f"active={app._active}" + ) if landed: await c.press("backslash") # leave hex await c.wait(lambda: app._active != "hex", 5) @@ -894,21 +1053,32 @@ async def s_command_palette(c: Ctx): @scenario("quit_guard") async def s_quit_guard(c: Ctx): app = c.app - c.check("a clean database reports nothing unsaved", app._dirty_labels() == [], - f"{app._dirty_labels()}") + c.check( + "a clean database reports nothing unsaved", + app._dirty_labels() == [], + f"{app._dirty_labels()}", + ) app._dirty = True # as an edit would - c.check("an edited database is reported unsaved", - len(app._dirty_labels()) == 1, f"{app._dirty_labels()}") + c.check( + "an edited database is reported unsaved", + len(app._dirty_labels()) == 1, + f"{app._dirty_labels()}", + ) await c.press("q") asked = await c.wait(lambda: isinstance(app.screen, QuitScreen), 10) - c.check("quitting with unsaved changes asks first", asked and app.is_running, - f"screen={type(app.screen).__name__} running={app.is_running}") + c.check( + "quitting with unsaved changes asks first", + asked and app.is_running, + f"screen={type(app.screen).__name__} running={app.is_running}", + ) if not asked: return await c.press("escape") await c.wait(lambda: not isinstance(app.screen, QuitScreen), 10) - c.check("Esc cancels the quit and stays put", - app.is_running and not isinstance(app.screen, QuitScreen)) + c.check( + "Esc cancels the quit and stays put", + app.is_running and not isinstance(app.screen, QuitScreen), + ) # leave it clean so the rest of the suite (and teardown) isn't affected app._dirty = False app._save_on_exit = False @@ -918,13 +1088,16 @@ async def s_quit_guard(c: Ctx): async def s_help(c: Ctx): app = c.app st = app.query_one("#status", Static) - c.check("the status line owns the bottom row (no footer cheatsheet)", - st.region.y + st.region.height == app.size.height, - f"status={st.region} screen={app.size}") + c.check( + "the status line owns the bottom row (no footer cheatsheet)", + st.region.y + st.region.height == app.size.height, + f"status={st.region} screen={app.size}", + ) await c.press("f1") opened = await c.wait(lambda: isinstance(app.screen, HelpScreen), 10) - c.check("F1 opens the key cheatsheet", opened, - f"screen={type(app.screen).__name__}") + c.check( + "F1 opens the key cheatsheet", opened, f"screen={type(app.screen).__name__}" + ) if not opened: return cards = app.screen.query(".help-card") @@ -932,17 +1105,26 @@ async def s_help(c: Ctx): txt = " ".join(str(w.render()) for w in cards) # Derived from _HELP, not hardcoded: adding a group is a normal change and # shouldn't fail a test that only meant 'every group is rendered'. - c.check("each key group gets its own card", - titles == {t for t, _ in _HELP}, f"{titles}") - c.check("it documents real bindings", - "set type" in txt and "split view" in txt and "cross-references" in txt - and "refresh the current view" in txt) - c.check("the graph keys are documented", - "control-flow graph" in txt and "minimap" in txt) + c.check( + "each key group gets its own card", titles == {t for t, _ in _HELP}, f"{titles}" + ) + c.check( + "it documents real bindings", + "set type" in txt + and "split view" in txt + and "cross-references" in txt + and "refresh the current view" in txt, + ) + c.check( + "the graph keys are documented", + "control-flow graph" in txt and "minimap" in txt, + ) body = app.screen.query_one("#help-body") - c.check("the cards fit without a scrollbar at a normal size", - body.virtual_size.height <= body.size.height, - f"content={body.virtual_size.height} view={body.size.height}") + c.check( + "the cards fit without a scrollbar at a normal size", + body.virtual_size.height <= body.size.height, + f"content={body.virtual_size.height} view={body.size.height}", + ) await c.press("escape") await c.wait(lambda: not isinstance(app.screen, HelpScreen), 10) c.check("Esc closes it", not isinstance(app.screen, HelpScreen)) @@ -950,8 +1132,9 @@ async def s_help(c: Ctx): # cheatsheet must not be reachable ONLY through F1. await c.press("H") opened_h = await c.wait(lambda: isinstance(app.screen, HelpScreen), 10) - c.check("H opens the cheatsheet too", opened_h, - f"screen={type(app.screen).__name__}") + c.check( + "H opens the cheatsheet too", opened_h, f"screen={type(app.screen).__name__}" + ) if opened_h: await c.press("H") await c.wait(lambda: not isinstance(app.screen, HelpScreen), 10) @@ -963,39 +1146,54 @@ async def s_strings(c: Ctx): app = c.app await c.open_biggest("listing") items = app.program.strings() - c.check("program.strings() lists the binary's literals", len(items) > 3, - f"n={len(items)}") + c.check( + "program.strings() lists the binary's literals", + len(items) > 3, + f"n={len(items)}", + ) if not items: return - c.check("strings carry addr/text/length", - all(s.addr > 0 and s.text and s.length > 0 for s in items[:5]), - f"first={items[0]}") + c.check( + "strings carry addr/text/length", + all(s.addr > 0 and s.text and s.length > 0 for s in items[:5]), + f"first={items[0]}", + ) await c.press("quotation_mark") opened = await c.wait(lambda: isinstance(app.screen, StringsPalette), 25) - c.check('\'"\' opens the strings browser', opened, - f"screen={type(app.screen).__name__}") + c.check( + "'\"' opens the strings browser", opened, f"screen={type(app.screen).__name__}" + ) if not opened: return pal = app.screen - c.check("the browser lists strings", len(pal._results) > 0, - f"results={len(pal._results)}") + c.check( + "the browser lists strings", + len(pal._results) > 0, + f"results={len(pal._results)}", + ) # filter on a fragment of a real (unescaped) literal - target = next((s for s in items - if len(s.text) >= 6 and _str_display(s.text) == s.text), None) + target = next( + (s for s in items if len(s.text) >= 6 and _str_display(s.text) == s.text), None + ) if target is not None: frag = target.text[:6] pal.query_one(Input).value = frag await c.pause(0.2) - ok = (pal._results - and all(frag.lower() in t.lower() for _, _, t in pal._results)) - c.check("filtering narrows to matching strings", bool(ok), - f"frag={frag!r} n={len(pal._results)}") + ok = pal._results and all(frag.lower() in t.lower() for _, _, t in pal._results) + c.check( + "filtering narrows to matching strings", + bool(ok), + f"frag={frag!r} n={len(pal._results)}", + ) want = pal._results[0][1] await c.press("enter") await c.wait(lambda: not isinstance(app.screen, StringsPalette), 10) landed = await c.wait(lambda: c.lst._cursor_ea() == want, 20) - c.check("Enter jumps to the string in the unified listing", landed, - f"cursor={c.lst._cursor_ea()} want={want:#x}") + c.check( + "Enter jumps to the string in the unified listing", + landed, + f"cursor={c.lst._cursor_ea()} want={want:#x}", + ) else: await c.press("escape") @@ -1014,31 +1212,42 @@ async def s_view_modes_all_handled(c: Ctx): the next mode that forgets to appear somewhere. """ from idatui.app import ViewMode + app = c.app fn = await c.open_biggest("listing") try: for mode in ViewMode: app._active = mode - app._show_active() # must not raise for any member + app._show_active() # must not raise for any member await c.pause(0.05) view = app._active_code_view() if mode in ViewMode.code_modes(): - c.check(f"{mode.value}: _active_code_view resolves a widget", - view is not None, f"{mode.value} -> None") - c.check(f"{mode.value}: exactly one predicate is true", - sum((app.is_listing, app.is_decomp, - app.is_hex, app.is_graph)) == 1, - f"{mode.value}: listing={app.is_listing} " - f"decomp={app.is_decomp} hex={app.is_hex} graph={app.is_graph}") - c.check(f"{mode.value}: in_code agrees with code_modes()", - app.in_code == (mode in ViewMode.code_modes()), - f"in_code={app.in_code} for {mode.value}") - c.check("every mode is a plain string over the wire", - all(isinstance(m, str) and m == m.value for m in ViewMode), - str([repr(m) for m in ViewMode])) - c.check("'disasm' is not a mode any more", - "disasm" not in {m.value for m in ViewMode}, - str([m.value for m in ViewMode])) + c.check( + f"{mode.value}: _active_code_view resolves a widget", + view is not None, + f"{mode.value} -> None", + ) + c.check( + f"{mode.value}: exactly one predicate is true", + sum((app.is_listing, app.is_decomp, app.is_hex, app.is_graph)) == 1, + f"{mode.value}: listing={app.is_listing} " + f"decomp={app.is_decomp} hex={app.is_hex} graph={app.is_graph}", + ) + c.check( + f"{mode.value}: in_code agrees with code_modes()", + app.in_code == (mode in ViewMode.code_modes()), + f"in_code={app.in_code} for {mode.value}", + ) + c.check( + "every mode is a plain string over the wire", + all(isinstance(m, str) and m == m.value for m in ViewMode), + str([repr(m) for m in ViewMode]), + ) + c.check( + "'disasm' is not a mode any more", + "disasm" not in {m.value for m in ViewMode}, + str([m.value for m in ViewMode]), + ) finally: # Restore through a real navigation, not by poking _active back. # _show_active() tears down split state and re-points the panes as a @@ -1063,15 +1272,21 @@ async def s_refresh_view(c: Ctx): await c.press("ctrl+r") landed = await c.wait( - lambda: lst.model is not old_model and lst._cursor_ea() == old_ea, 25) + lambda: lst.model is not old_model and lst._cursor_ea() == old_ea, 25 + ) c.check("Ctrl+R rebuilds the listing model", lst.model is not old_model) - c.check("Ctrl+R preserves the cursor address", landed, - f"got={lst._cursor_ea()} want={old_ea}") + c.check( + "Ctrl+R preserves the cursor address", + landed, + f"got={lst._cursor_ea()} want={old_ea}", + ) new_top = round(lst.scroll_offset.y) new_head = lst.model.get(new_top) if lst.model is not None else None - c.check("Ctrl+R preserves the viewport by address", - getattr(new_head, "ea", None) == old_top_ea, - f"got={getattr(new_head, 'ea', None)} want={old_top_ea}") + c.check( + "Ctrl+R preserves the viewport by address", + getattr(new_head, "ea", None) == old_top_ea, + f"got={getattr(new_head, 'ea', None)} want={old_top_ea}", + ) await c.open(fn.addr, "decomp") dec = c.dec @@ -1079,13 +1294,18 @@ async def s_refresh_view(c: Ctx): old_dec_cursor = dec.cursor await c.press("ctrl+r") refreshed = await c.wait( - lambda: c.app.is_decomp and dec.loaded_ea == fn.addr and not dec.loading, - 25) - c.check("Ctrl+R reloads pseudocode without changing views", refreshed, - f"active={c.app._active} loaded={dec.loaded_ea} want={fn.addr:#x}") - c.check("Ctrl+R preserves the pseudocode cursor", - dec.cursor == old_dec_cursor, - f"got={dec.cursor} want={old_dec_cursor}") + lambda: c.app.is_decomp and dec.loaded_ea == fn.addr and not dec.loading, 25 + ) + c.check( + "Ctrl+R reloads pseudocode without changing views", + refreshed, + f"active={c.app._active} loaded={dec.loaded_ea} want={fn.addr:#x}", + ) + c.check( + "Ctrl+R preserves the pseudocode cursor", + dec.cursor == old_dec_cursor, + f"got={dec.cursor} want={old_dec_cursor}", + ) @scenario("split_view") @@ -1113,51 +1333,73 @@ async def s_split_view(c: Ctx): await _split_view_body(c, app, lst, dec) finally: 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") + c.check( + "split view doesn't storm the worker with function lookups", + _lookups["n"] < 500, + f"{_lookups['n']} lookup_funcs calls", + ) async def _split_view_body(c: Ctx, app, lst, dec): await c.open_biggest("listing") await c.press("s") shown = await c.wait(lambda: app._split and lst.display and dec.display, 20) - c.check("'s' enters split view (both panes shown)", shown, - f"split={app._split} lst={lst.display} dec={dec.display}") + c.check( + "'s' enters split view (both panes shown)", + shown, + f"split={app._split} lst={lst.display} dec={dec.display}", + ) loaded = await c.wait(lambda: dec.loaded_ea == app._cur.ea, 25) - c.check("split loads the pseudocode alongside the listing", loaded, - f"loaded={dec.loaded_ea} cur={app._cur.ea if app._cur else None}") + c.check( + "split loads the pseudocode alongside the listing", + loaded, + f"loaded={dec.loaded_ea} cur={app._cur.ea if app._cur else None}", + ) await c.wait(lambda: "[split" in c.status(), 5) - c.check("split view shows a split-aware status", "[split" in c.status(), - f"status={c.status()!r}") + c.check( + "split view shows a split-aware status", + "[split" in c.status(), + f"status={c.status()!r}", + ) # phase 3: the rich per-line instruction map (decomp_map tool, run on the # pilot's real worker) — verify it returns, aligns with the markers, and # bands a whole region for a multi-instruction C line. m = app.program.decomp_map(app._cur.ea) c.check("decomp_map returns per-line ea sets", len(m) > 5, f"lines={len(m)}") - aligned = sum(1 for i in range(min(len(m), len(dec._line_eas))) - if m[i] and dec._line_eas[i] is not None - and dec._line_eas[i] in m[i]) - c.check("decomp_map aligns with the pseudocode markers", aligned >= 3, - f"aligned={aligned}/{len(dec._line_eas)}") + aligned = sum( + 1 + for i in range(min(len(m), len(dec._line_eas))) + if m[i] and dec._line_eas[i] is not None and dec._line_eas[i] in m[i] + ) + c.check( + "decomp_map aligns with the pseudocode markers", + aligned >= 3, + f"aligned={aligned}/{len(dec._line_eas)}", + ) multi = next((i for i, eas in enumerate(m) if len(eas) > 1), None) if multi is not None: app._split_eamap = m - dec.focus() # the decomp must BE the driver for a decomp-driven + dec.focus() # the decomp must BE the driver for a decomp-driven app._active = "decomp" # sync (else its align() re-syncs listing-driven) dec.cursor = multi dec._scroll_cursor_into_view() # key-nav always does; the anchor needs it await c.pause(0.1) app._sync_split("decomp") await c.pause(0.1) - c.check("a multi-instruction C line bands a region (>1 listing row)", - len(lst._link_rows) > 1, - f"line={multi} eas={len(m[multi])} rows={sorted(lst._link_rows)[:8]}") + c.check( + "a multi-instruction C line bands a region (>1 listing row)", + len(lst._link_rows) > 1, + f"line={multi} eas={len(m[multi])} rows={sorted(lst._link_rows)[:8]}", + ) lst.focus() app._active = "listing" await c.pause(0.05) else: - c.check("a multi-instruction C line bands a region (>1 listing row)", - True, "no multi-instruction line in this function (skipped)") + c.check( + "a multi-instruction C line bands a region (>1 listing row)", + True, + "no multi-instruction line in this function (skipped)", + ) # listing drives: move it, the decomp band must track the covering C line lst.focus() for _ in range(6): @@ -1165,48 +1407,60 @@ async def _split_view_body(c: Ctx, app, lst, dec): await c.pause(0.2) lea = lst._cursor_ea() dl = dec._link_line - c.check("listing cursor links the covering pseudocode line", - dl is not None and lea is not None and dec._line_eas[dl] is not None - and dec._line_eas[dl] <= lea, - f"link_line={dl} lea={hex(lea) if lea else None}") + c.check( + "listing cursor links the covering pseudocode line", + dl is not None + and lea is not None + and dec._line_eas[dl] is not None + and dec._line_eas[dl] <= lea, + f"link_line={dl} lea={hex(lea) if lea else None}", + ) # the companion pane sits LEVEL with the driver's cursor (visual coherence): # the linked row lands at the same viewport offset, not merely on-screen. deep = [i for i, eas in enumerate(m) if eas][10:] row = lst.model.ensure_ea(m[deep[0]][0]) if (deep and lst.model) else None if row is not None and row > 12: lst.scroll_to(y=row - 10, animate=False) - await c.pause(0.2) # let the deferred scroll land - lst.cursor = row # driver cursor now at viewport offset 10 + await c.pause(0.2) # let the deferred scroll land + lst.cursor = row # driver cursor now at viewport offset 10 app._sync_split("listing") - await c.pause(0.2) # let the companion's scroll land + await c.pause(0.2) # let the companion's scroll land drv = lst.cursor - round(lst.scroll_offset.y) link, top = dec._link_line, round(dec.scroll_offset.y) # exact, modulo the unavoidable clamps (can't scroll above line 0, nor # past the end when the pseudocode is shorter than the viewport) - want = min(max(0, (link or 0) - drv), - max(0, dec.total - dec._visible_height())) - c.check("the companion pane sits level with the driver's cursor", - link is not None and top == want, - f"driver_row={drv} link={link} dec_top={top} want={want}") + want = min(max(0, (link or 0) - drv), max(0, dec.total - dec._visible_height())) + c.check( + "the companion pane sits level with the driver's cursor", + link is not None and top == want, + f"driver_row={drv} link={link} dec_top={top} want={want}", + ) # a PURE scroll (wheel/scrollbar) moves no cursor — it must still drag # the companion along (anchors on the viewport once the cursor is gone) before_cur, before_dec = lst.cursor, round(dec.scroll_offset.y) lst.scroll_to(y=round(lst.scroll_offset.y) + 30, animate=False) await c.pause(0.35) - c.check("a pure scroll in the driver drags the companion along", - lst.cursor == before_cur - and round(dec.scroll_offset.y) != before_dec, - f"cursor {before_cur}->{lst.cursor} " - f"dec_top {before_dec}->{round(dec.scroll_offset.y)}") + c.check( + "a pure scroll in the driver drags the companion along", + lst.cursor == before_cur and round(dec.scroll_offset.y) != before_dec, + f"cursor {before_cur}->{lst.cursor} " + f"dec_top {before_dec}->{round(dec.scroll_offset.y)}", + ) await c.press("tab") await c.pause(0.1) - c.check("Tab in split focuses the pseudocode pane", app._active == "decomp", - f"active={app._active}") + c.check( + "Tab in split focuses the pseudocode pane", + app._active == "decomp", + f"active={app._active}", + ) # decomp drives: put the cursor on an addressed pseudocode line (past the # variable decls); the listing band must track the covering instruction row. target = next((i for i, e in enumerate(dec._line_eas) if e is not None), None) - c.check("pseudocode has addressed lines", target is not None, - "no /*0xEA*/ markers in the pseudocode") + c.check( + "pseudocode has addressed lines", + target is not None, + "no /*0xEA*/ markers in the pseudocode", + ) if target is not None: dec.cursor = target dec._scroll_cursor_into_view() @@ -1214,27 +1468,39 @@ async def _split_view_body(c: Ctx, app, lst, dec): app._sync_split("decomp") await c.pause(0.1) want = lst.model.ensure_ea(dec._line_eas[target]) - c.check("decomp cursor links the instruction row in the listing", - want in lst._link_rows, - f"link_rows={sorted(lst._link_rows)[:6]} want={want}") + c.check( + "decomp cursor links the instruction row in the listing", + want in lst._link_rows, + f"link_rows={sorted(lst._link_rows)[:6]} want={want}", + ) # and that linked row actually paints a background band (base rows have # no bg; `want` is a deep code row, never the listing's own cursor row) lst.reveal(want) await c.pause(0.05) y = want - round(lst.scroll_offset.y) - banded = (0 <= y < lst.size.height and any( - s.style and s.style.bgcolor is not None for s in lst.render_line(y))) - c.check("the linked instruction row renders a highlight band", banded, - f"y={y} cursor_row={lst.cursor}") + banded = 0 <= y < lst.size.height and any( + s.style and s.style.bgcolor is not None for s in lst.render_line(y) + ) + c.check( + "the linked instruction row renders a highlight band", + banded, + f"y={y} cursor_row={lst.cursor}", + ) await c.press("tab") await c.pause(0.1) - c.check("Tab again focuses the listing pane", app._active == "listing", - f"active={app._active}") + c.check( + "Tab again focuses the listing pane", + app._active == "listing", + f"active={app._active}", + ) # a mouse click on the other pane also makes it the driver (not just Tab) await c.pilot.click(DecompView, offset=(10, 5)) await c.pause(0.15) - c.check("clicking the pseudocode pane makes it the driver", - app._active == "decomp", f"active={app._active}") + c.check( + "clicking the pseudocode pane makes it the driver", + app._active == "decomp", + f"active={app._active}", + ) await c.press("tab") # restore listing as the driver await c.pause(0.1) # cross-function follow: the listing cursor leaving the decompiled function @@ -1250,8 +1516,11 @@ async def _split_view_body(c: Ctx, app, lst, dec): await c.pause(0.1) app._sync_split("listing") # cursor now outside the decompiled fn followed = await c.wait(lambda: dec.loaded_ea == other.addr, 25) - c.check("listing cursor crossing into another function re-syncs the decomp", - followed, f"dec={dec.loaded_ea} want={other.addr}") + c.check( + "listing cursor crossing into another function re-syncs the decomp", + followed, + f"dec={dec.loaded_ea} want={other.addr}", + ) # navigation in split keeps BOTH panes on the (new) function nf = c.find_func(lambda f: f.addr != app._cur.ea and f.size > 80) if nf is not None: @@ -1259,27 +1528,42 @@ async def _split_view_body(c: Ctx, app, lst, dec): await c.type(hex(nf.addr)) await c.press("enter") nav = await c.wait(lambda: app._cur and app._cur.ea == nf.addr, 15) - c.check("goto in split navigates", nav, - f"cur={app._cur.ea if app._cur else None} want={nf.addr}") - both = await c.wait(lambda: dec.loaded_ea == nf.addr and app._split - and lst.display and dec.display, 25) - c.check("split reloads both panes on navigation", both, - f"dec={dec.loaded_ea} split={app._split}") + c.check( + "goto in split navigates", + nav, + f"cur={app._cur.ea if app._cur else None} want={nf.addr}", + ) + both = await c.wait( + lambda: ( + dec.loaded_ea == nf.addr and app._split and lst.display and dec.display + ), + 25, + ) + c.check( + "split reloads both panes on navigation", + both, + f"dec={dec.loaded_ea} split={app._split}", + ) await c.press("s") - gone = await c.wait(lambda: not app._split and lst.display - and not dec.display, 10) - c.check("'s' exits split back to a single view", gone, - f"split={app._split} lst={lst.display} dec={dec.display}") - c.check("exiting split clears the link bands", - not lst._link_rows and dec._link_line is None, - f"rows={lst._link_rows} line={dec._link_line}") + gone = await c.wait(lambda: not app._split and lst.display and not dec.display, 10) + c.check( + "'s' exits split back to a single view", + gone, + f"split={app._split} lst={lst.display} dec={dec.display}", + ) + c.check( + "exiting split clears the link bands", + not lst._link_rows and dec._link_line is None, + f"rows={lst._link_rows} line={dec._link_line}", + ) @scenario("decomp_fallback") async def s_fallback(c: Ctx): app = c.app - failing = next((f for f in reversed(c.all_funcs()) - if c.prog.decompile(f.addr).failed), None) + failing = next( + (f for f in reversed(c.all_funcs()) if c.prog.decompile(f.addr).failed), None + ) if failing is None: c.check("found a decompile-failing function", False) return @@ -1294,15 +1578,23 @@ async def s_fallback(c: Ctx): # was pressed, since the function was opened in the listing. It asserted # nothing, slowly. landed = await c.wait(lambda: _CANNOT_DECOMP in c.status().lower(), 25) - c.check("F5/Tab on an undecompilable function says so", landed, - f"active={app._active} status={c.status()!r}") - c.check("F5/Tab on an undecompilable function falls back to a code view", - app.is_listing and c.dis.display, - f"active={app._active} status={c.status()!r}") + c.check( + "F5/Tab on an undecompilable function says so", + landed, + f"active={app._active} status={c.status()!r}", + ) + c.check( + "F5/Tab on an undecompilable function falls back to a code view", + app.is_listing and c.dis.display, + f"active={app._active} status={c.status()!r}", + ) # a decompilable function F5s into pseudocode await c.open("main", "decomp") - c.check("a decompilable function F5s into pseudocode", - app._active == "decomp" and c.dec.display, f"active={app._active}") + c.check( + "a decompilable function F5s into pseudocode", + app._active == "decomp" and c.dec.display, + f"active={app._active}", + ) @scenario("structs") @@ -1310,95 +1602,144 @@ async def s_structs(c: Ctx): app = c.app await c.press("ctrl+t") se_open = await c.wait(lambda: isinstance(app.screen, StructEditor), 10) - c.check("Ctrl+T opens the struct editor", se_open, - f"screen={type(app.screen).__name__}") + c.check( + "Ctrl+T opens the struct editor", se_open, f"screen={type(app.screen).__name__}" + ) if not se_open: return se = app.screen await c.wait(lambda: bool(se._structs), 15) - c.check("struct editor lists existing structs", len(se._structs) > 0, - f"n={len(se._structs)}") + c.check( + "struct editor lists existing structs", + len(se._structs) > 0, + f"n={len(se._structs)}", + ) ta = se.query_one(TextArea) - tname = next((s.name for s in se._structs if s.name == "timespec"), - se._structs[0].name) + tname = next( + (s.name for s in se._structs if s.name == "timespec"), se._structs[0].name + ) idx = next(i for i, s in enumerate(se._structs) if s.name == tname) se.query_one(OptionList).highlighted = idx se.on_option_list_option_selected(type("E", (), {"option_index": idx})()) await c.wait(lambda: tname in ta.text and "{" in ta.text, 15) - c.check("selecting a struct shows its C definition", - tname in ta.text and "{" in ta.text, f"text={ta.text[:40]!r}") + c.check( + "selecting a struct shows its C definition", + tname in ta.text and "{" in ta.text, + f"text={ta.text[:40]!r}", + ) # The definition is C, so it must be coloured as C (no tree-sitter grammar # for it: idatui.highlight fills TextArea's highlight map from Pygments). names = {n for spans in ta._highlights.values() for _, _, n in spans} - c.check("the C definition is syntax-highlighted", - {"keyword", "name"} <= names, f"names={sorted(names)}") - styled = {s.style.color.name for s in ta.render_line(0) - if s.style and s.style.color} - c.check("highlight styles reach the rendered line", len(styled) > 1, - f"colors={sorted(styled)}") + c.check( + "the C definition is syntax-highlighted", + {"keyword", "name"} <= names, + f"names={sorted(names)}", + ) + styled = { + s.style.color.name for s in ta.render_line(0) if s.style and s.style.color + } + c.check( + "highlight styles reach the rendered line", + len(styled) > 1, + f"colors={sorted(styled)}", + ) app._clipboard = "" se.query_one(TextArea).focus() await c.press("ctrl+y") await c.wait(lambda: app._clipboard == ta.text, 10) - c.check("Ctrl+Y copies the struct definition to the clipboard", - bool(app._clipboard) and app._clipboard == ta.text, - f"clip_len={len(app._clipboard)}") + c.check( + "Ctrl+Y copies the struct definition to the clipboard", + bool(app._clipboard) and app._clipboard == ta.text, + f"clip_len={len(app._clipboard)}", + ) sname = "TuiEdTest" await c.press("ctrl+n") await c.pause(0.05) ta.text = f"struct {sname} {{ int a; char b[8]; }};" await c.press("ctrl+s") await c.wait(lambda: any(s.name == sname for s in se._structs), 15) - c.check("Ctrl+S declares a new struct", - any(s.name == sname for s in se._structs), "not created") + c.check( + "Ctrl+S declares a new struct", + any(s.name == sname for s in se._structs), + "not created", + ) await c.wait(lambda: "\n" in ta.text, 10) - c.check("editing re-highlights the definition", - any(n == "keyword" for spans in ta._highlights.values() - for _, _, n in spans), - f"rows={len(ta._highlights)}") - c.check("save auto-formats the definition in the editor", - ta.text.count("\n") >= 3 and f"struct {sname}" in ta.text - and not se._is_dirty(), f"text={ta.text[:50]!r}") + c.check( + "editing re-highlights the definition", + any(n == "keyword" for spans in ta._highlights.values() for _, _, n in spans), + f"rows={len(ta._highlights)}", + ) + c.check( + "save auto-formats the definition in the editor", + ta.text.count("\n") >= 3 + and f"struct {sname}" in ta.text + and not se._is_dirty(), + f"text={ta.text[:50]!r}", + ) idx = next(i for i, s in enumerate(se._structs) if s.name == sname) se.on_option_list_option_selected(type("E", (), {"option_index": idx})()) await c.wait(lambda: sname in ta.text, 10) ta.text = f"struct {sname} {{ int a; char b[8]; long c; }};" await c.press("ctrl+s") - await c.wait(lambda: next((s.members for s in se._structs if s.name == sname), 0) == 3, 15) - c.check("Ctrl+S updates an existing struct in place", - next((s.members for s in se._structs if s.name == sname), 0) == 3, - "member count not 3") + await c.wait( + lambda: next((s.members for s in se._structs if s.name == sname), 0) == 3, 15 + ) + c.check( + "Ctrl+S updates an existing struct in place", + next((s.members for s in se._structs if s.name == sname), 0) == 3, + "member count not 3", + ) ta.text = f"struct {sname} {{ int a; char b[8]; long c; int __unused; }};" await c.press("ctrl+s") await c.wait(lambda: "save failed" in str(se.query_one("#se-status").render()), 15) st = str(se.query_one("#se-status").render()) - c.check("a rejected save fails loudly, naming the reserved field", - "save failed" in st and "__unused" in st, f"status={st!r}") - c.check("a rejected save leaves the struct unchanged", - next((s.members for s in se._structs if s.name == sname), 0) == 3, "changed") + c.check( + "a rejected save fails loudly, naming the reserved field", + "save failed" in st and "__unused" in st, + f"status={st!r}", + ) + c.check( + "a rejected save leaves the struct unchanged", + next((s.members for s in se._structs if s.name == sname), 0) == 3, + "changed", + ) c.check("a rejected save keeps your edited text", "__unused" in ta.text) other = next(i for i, s in enumerate(se._structs) if s.name != sname) se.on_option_list_option_selected(type("E", (), {"option_index": other})()) guard = await c.wait(lambda: isinstance(app.screen, ConfirmScreen), 10) - c.check("unsaved edits prompt before switching structs", guard, - f"screen={type(app.screen).__name__}") + c.check( + "unsaved edits prompt before switching structs", + guard, + f"screen={type(app.screen).__name__}", + ) await c.press("enter") - await c.wait(lambda: isinstance(app.screen, StructEditor) and not se._is_dirty(), 15) + await c.wait( + lambda: isinstance(app.screen, StructEditor) and not se._is_dirty(), 15 + ) idx = next(i for i, s in enumerate(se._structs) if s.name == sname) se.query_one(OptionList).focus() se.query_one(OptionList).highlighted = idx await c.press("d") confirmed = await c.wait(lambda: isinstance(app.screen, ConfirmScreen), 10) - c.check("delete asks for confirmation", confirmed, - f"screen={type(app.screen).__name__}") + c.check( + "delete asks for confirmation", confirmed, f"screen={type(app.screen).__name__}" + ) await c.press("enter") await c.wait(lambda: isinstance(app.screen, StructEditor), 10) - await c.wait(lambda: not any(s.name == sname for s in se._structs) - or "del_type" in str(se.query_one("#se-status").render()), 15) + await c.wait( + lambda: ( + not any(s.name == sname for s in se._structs) + or "del_type" in str(se.query_one("#se-status").render()) + ), + 15, + ) st = str(se.query_one("#se-status").render()) gone = not any(s.name == sname for s in se._structs) - c.check("confirming delete removes the struct (or reports missing tool)", - gone or "del_type" in st, f"gone={gone} status={st!r}") + c.check( + "confirming delete removes the struct (or reports missing tool)", + gone or "del_type" in st, + f"gone={gone} status={st!r}", + ) se.query_one(OptionList).focus() await c.press("escape") await c.wait(lambda: not isinstance(app.screen, StructEditor), 10) @@ -1416,13 +1757,15 @@ async def s_splash_scaling(c: Ctx): it is placed in, so there was never a reason for all-or-nothing. """ from idatui import kittygfx - from idatui.app import (LOGO_CHROME_ROWS, LOGO_MIN_ROWS, LoadingScreen, - logo_cells) + from idatui.app import LOGO_CHROME_ROWS, LOGO_MIN_ROWS, LoadingScreen, logo_cells app = c.app placed: list[tuple] = [] real_supported, real_upload, real_place = ( - kittygfx.supported, kittygfx.upload, kittygfx.place) + kittygfx.supported, + kittygfx.upload, + kittygfx.place, + ) kittygfx.supported = lambda: True kittygfx.upload = lambda *a, **k: True kittygfx.place = lambda *a, **k: (placed.append(a), True)[1] @@ -1437,28 +1780,41 @@ async def s_splash_scaling(c: Ctx): await c.wait(lambda: scr._cells is not None, 5) room = height - LOGO_CHROME_ROWS has_image = bool(scr.query("#loading-image")) - c.check(f"{width}x{height}: the logo is drawn, not dropped", - has_image and room >= LOGO_MIN_ROWS, - f"image={has_image} room={room}") + c.check( + f"{width}x{height}: the logo is drawn, not dropped", + has_image and room >= LOGO_MIN_ROWS, + f"image={has_image} room={room}", + ) if has_image: cols, rows = scr._cells - c.check(f"{width}x{height}: scaled to the room available", - rows <= room and rows == min(room, logo_cells()[1]), - f"cells={scr._cells} room={room} natural={logo_cells()}") + c.check( + f"{width}x{height}: scaled to the room available", + rows <= room and rows == min(room, logo_cells()[1]), + f"cells={scr._cells} room={room} natural={logo_cells()}", + ) await c.wait(lambda: scr.query_one("#loading-box").region.height > 0, 5) box = scr.query_one("#loading-box").region - c.check(f"{width}x{height}: the box is not clipped", - box.y >= 0 and box.y + box.height <= height, - f"box={box} screen={height}") + c.check( + f"{width}x{height}: the box is not clipped", + box.y >= 0 and box.y + box.height <= height, + f"box={box} screen={height}", + ) app.pop_screen() await c.pause(0.05) - c.check("a full-size pane still gets the artwork's natural size", - logo_cells(999) == logo_cells(), f"{logo_cells(999)}") - c.check("and the image was actually placed each time", len(placed) >= 3, - f"{placed}") + c.check( + "a full-size pane still gets the artwork's natural size", + logo_cells(999) == logo_cells(), + f"{logo_cells(999)}", + ) + c.check( + "and the image was actually placed each time", len(placed) >= 3, f"{placed}" + ) finally: kittygfx.supported, kittygfx.upload, kittygfx.place = ( - real_supported, real_upload, real_place) + real_supported, + real_upload, + real_place, + ) # Every later scenario assumes the suite's own geometry. await c.pilot.resize_terminal(140, 44) await c.pause(0.05) @@ -1478,33 +1834,45 @@ async def s_modal_centering(c: Ctx): import idatui.app as A ours = sorted( - (n for n, v in vars(A).items() - if isinstance(v, type) and issubclass(v, ModalScreen) - and v is not ModalScreen and v.__module__ == A.__name__), - key=str) + ( + n + for n, v in vars(A).items() + if isinstance(v, type) + and issubclass(v, ModalScreen) + and v is not ModalScreen + and v.__module__ == A.__name__ + ), + key=str, + ) c.check("found the app's modal screens", len(ours) >= 8, f"{ours}") styles = A.IdaTui.CSS - c.check("centring is a rule about modals, not a list of them", - "ModalScreen { align: center middle; }" in styles, - "the ModalScreen rule is gone") + c.check( + "centring is a rule about modals, not a list of them", + "ModalScreen { align: center middle; }" in styles, + "the ModalScreen rule is gone", + ) # And prove it REACHES a dialog, rather than just being present in the text. await c.press("ctrl+f") opened = await c.wait(lambda: isinstance(c.app.screen, A.SearchPalette), 10) if not opened: - c.check("the search palette opened", False, - f"screen={type(c.app.screen).__name__}") + c.check( + "the search palette opened", False, f"screen={type(c.app.screen).__name__}" + ) return scr = c.app.screen await c.wait(lambda: scr.query_one("#pal-box").region.height > 0, 5) box = scr.query_one("#pal-box").region above, below = box.y, c.app.size.height - (box.y + box.height) - c.check("the search palette is vertically centred", - box.height > 0 and abs(above - below) <= 1, - f"box={box} screen={c.app.size} above={above} below={below}") + c.check( + "the search palette is vertically centred", + box.height > 0 and abs(above - below) <= 1, + f"box={box} screen={c.app.size} above={above} below={below}", + ) left = box.x right = c.app.size.width - (box.x + box.width) - c.check("and horizontally centred", abs(left - right) <= 1, - f"left={left} right={right}") + c.check( + "and horizontally centred", abs(left - right) <= 1, f"left={left} right={right}" + ) await c.press("escape") await c.wait(lambda: not isinstance(c.app.screen, A.SearchPalette), 5) @@ -1516,8 +1884,9 @@ async def s_db_search(c: Ctx): await c.open("main", "listing") await c.press("ctrl+f") opened = await c.wait(lambda: isinstance(app.screen, SearchPalette), 10) - c.check("Ctrl+F opens the search palette", opened, - f"screen={type(app.screen).__name__}") + c.check( + "Ctrl+F opens the search palette", opened, f"screen={type(app.screen).__name__}" + ) if not opened: return pal = app.screen @@ -1527,29 +1896,40 @@ async def s_db_search(c: Ctx): inp.value = "endbr64" await c.press("enter") await c.wait(lambda: bool(pal._hits), 30) - c.check("a text search finds instructions", len(pal._hits) > 1, - f"n={len(pal._hits)}") - c.check("and it was classified as text", - pal._searched and pal._searched[0] == "text", f"{pal._searched}") - c.check("hits carry the line they matched", - all("endbr64" in h.line for h in pal._hits[:5]), - [h.line for h in pal._hits[:3]]) + c.check( + "a text search finds instructions", len(pal._hits) > 1, f"n={len(pal._hits)}" + ) + c.check( + "and it was classified as text", + pal._searched and pal._searched[0] == "text", + f"{pal._searched}", + ) + c.check( + "hits carry the line they matched", + all("endbr64" in h.line for h in pal._hits[:5]), + [h.line for h in pal._hits[:3]], + ) # -- text with padding: match what is SEEN, not IDA's column spacing ----- # inp.value = "call cs:" await c.press("enter") found = await c.wait(lambda: pal._searched == ("text", "call cs:"), 30) - c.check("a query spanning IDA's column padding still matches", - found and len(pal._hits) > 0, f"n={len(pal._hits)}") + c.check( + "a query spanning IDA's column padding still matches", + found and len(pal._hits) > 0, + f"n={len(pal._hits)}", + ) # -- bytes: the same endbr64, as a pattern ------------------------------- # inp.value = "f3 0f 1e fa" await c.press("enter") await c.wait(lambda: pal._searched and pal._searched[0] == "bytes", 30) - c.check("a hex query is classified as bytes", - pal._searched and pal._searched[0] == "bytes", f"{pal._searched}") - c.check("and finds the same instruction", len(pal._hits) > 1, - f"n={len(pal._hits)}") + c.check( + "a hex query is classified as bytes", + pal._searched and pal._searched[0] == "bytes", + f"{pal._searched}", + ) + c.check("and finds the same instruction", len(pal._hits) > 1, f"n={len(pal._hits)}") # -- wildcards ----------------------------------------------------------- # inp.value = "f3 0f ?? fa" @@ -1562,34 +1942,49 @@ async def s_db_search(c: Ctx): await c.press("enter") await c.pause(0.1) title = str(app.screen.query_one("#pal-box").border_title) - c.check("a malformed byte pattern is refused with a reason", - "not a byte" in title, f"title={title!r}") + c.check( + "a malformed byte pattern is refused with a reason", + "not a byte" in title, + f"title={title!r}", + ) # -- F2 pins the mode against the guess ---------------------------------- # inp.value = "dead" await c.pause(0.05) - c.check("a hex-looking WORD still searches text", - pal._mode_query()[0] == "text", f"{pal._mode_query()}") + c.check( + "a hex-looking WORD still searches text", + pal._mode_query()[0] == "text", + f"{pal._mode_query()}", + ) await c.press("f2") - c.check("F2 forces it to bytes", pal._mode_query()[0] == "bytes", - f"{pal._mode_query()}") + c.check( + "F2 forces it to bytes", pal._mode_query()[0] == "bytes", f"{pal._mode_query()}" + ) # -- Enter on a result navigates ----------------------------------------- # inp.value = "endbr64" - await c.press("f2") # back to text + await c.press("f2") # back to text await c.press("enter") - await c.wait(lambda: bool(pal._hits) and pal._searched - and pal._searched[0] == "text", 30) + await c.wait( + lambda: bool(pal._hits) and pal._searched and pal._searched[0] == "text", 30 + ) target = pal._hits[1] if len(pal._hits) > 1 else pal._hits[0] pal.query_one(OptionList).highlighted = 1 if len(pal._hits) > 1 else 0 await c.press("enter") closed = await c.wait(lambda: not isinstance(app.screen, SearchPalette), 10) - c.check("Enter on a hit closes the palette", closed, - f"screen={type(app.screen).__name__}") + c.check( + "Enter on a hit closes the palette", + closed, + f"screen={type(app.screen).__name__}", + ) landed = await c.wait( - lambda: app._cur is not None and c.lst._cursor_ea() == target.head, 30) - c.check("and lands the cursor on it", landed, - f"cursor={c.lst._cursor_ea()} want={target.head:#x}") + lambda: app._cur is not None and c.lst._cursor_ea() == target.head, 30 + ) + c.check( + "and lands the cursor on it", + landed, + f"cursor={c.lst._cursor_ea()} want={target.head:#x}", + ) @scenario("export_findings") @@ -1612,8 +2007,9 @@ async def s_export_findings(c: Ctx): await c.open(fn.addr, "listing") # Make something to find: a rename and a comment, through the real paths. - app.program.client.invoke( - "rename", batch={"func": {"addr": hex(fn.addr), "name": newname}}) + app.program.client.call( + remote_ops.rename, batch={"func": {"addr": hex(fn.addr), "name": newname}} + ) app.program.bump_names() app.program.set_comment(fn.addr, note) app.program.invalidate(fn.addr) @@ -1629,8 +2025,11 @@ async def s_export_findings(c: Ctx): inp = app.query_one("#export", Input) opened = await c.wait(lambda: inp.display, 5) c.check("Ctrl+E opens the export prompt", opened, f"display={inp.display}") - c.check("the prompt is prefilled with a path beside the binary", - inp.value == default_path(app._open_path), f"value={inp.value!r}") + c.check( + "the prompt is prefilled with a path beside the binary", + inp.value == default_path(app._open_path), + f"value={inp.value!r}", + ) inp.value = out await c.press("enter") written = await c.wait(lambda: os.path.exists(out), 30) @@ -1638,21 +2037,26 @@ async def s_export_findings(c: Ctx): if not written: return doc = open(out, encoding="utf-8").read() - c.check("the report is markdown with the expected sections", - doc.startswith("# Findings") and "## Comments" in doc - and "## Named functions" in doc, doc[:60]) - c.check("a comment written this session is in it", note in doc, - doc[:200]) - c.check("and the function it belongs to is named", newname in doc, - doc[:200]) - c.check("the report is sourced from the journal, not a scan", - "idatui's edit journal" in doc, - [l for l in doc.splitlines() if "**source**" in l]) - c.check("the analyzer's own comments stay out of it", - "switch jump" not in doc and "jumptable" not in doc, - [l for l in doc.splitlines() if "switch" in l][:2]) - c.check("the status line says where it went", - out in c.status(), c.status()) + c.check( + "the report is markdown with the expected sections", + doc.startswith("# Findings") + and "## Comments" in doc + and "## Named functions" in doc, + doc[:60], + ) + c.check("a comment written this session is in it", note in doc, doc[:200]) + c.check("and the function it belongs to is named", newname in doc, doc[:200]) + c.check( + "the report is sourced from the journal, not a scan", + "idatui's edit journal" in doc, + [l for l in doc.splitlines() if "**source**" in l], + ) + c.check( + "the analyzer's own comments stay out of it", + "switch jump" not in doc and "jumptable" not in doc, + [l for l in doc.splitlines() if "switch" in l][:2], + ) + c.check("the status line says where it went", out in c.status(), c.status()) # The journal has to survive the database, or a report is only ever # about the session that happened to be open. from idatui.journal import Journal @@ -1660,14 +2064,17 @@ async def s_export_findings(c: Ctx): app.journal.flush(app.program) reloaded = Journal() reloaded.load(app.program) - c.check("the journal round-trips through the .i64", - fn.addr in reloaded.addresses(), - f"{len(reloaded)} entries, {sorted(reloaded.addresses())[:3]}") + c.check( + "the journal round-trips through the .i64", + fn.addr in reloaded.addresses(), + f"{len(reloaded)} entries, {sorted(reloaded.addresses())[:3]}", + ) finally: # Idempotent: hand the database back exactly as we found it. app.program.set_comment(fn.addr, "") - app.program.client.invoke( - "rename", batch={"func": {"addr": hex(fn.addr), "name": old}}) + app.program.client.call( + remote_ops.rename, batch={"func": {"addr": hex(fn.addr), "name": old}} + ) app.program.bump_names() app.program.invalidate(fn.addr) if os.path.exists(out): @@ -1679,8 +2086,11 @@ async def s_struct_filter(c: Ctx): app = c.app await c.press("ctrl+t") if not await c.wait(lambda: isinstance(app.screen, StructEditor), 10): - c.check("Ctrl+T opens the struct editor", False, - f"screen={type(app.screen).__name__}") + c.check( + "Ctrl+T opens the struct editor", + False, + f"screen={type(app.screen).__name__}", + ) return se = app.screen await c.wait(lambda: bool(se._structs), 15) @@ -1695,8 +2105,11 @@ async def s_struct_filter(c: Ctx): ol.focus() await c.press("slash") opened = await c.wait(lambda: inp.display and app.focused is inp, 5) - c.check("'/' from the list opens the struct filter", opened, - f"display={inp.display} focus={getattr(app.focused, 'id', None)}") + c.check( + "'/' from the list opens the struct filter", + opened, + f"display={inp.display} focus={getattr(app.focused, 'id', None)}", + ) # PgDn from the FILTER: this screen can't use OptionListNav (ctrl+n is "new # type" here), so it forwards through its own guarded _page(). Checked while @@ -1705,31 +2118,42 @@ async def s_struct_filter(c: Ctx): if ol.option_count > 2: ol.highlighted = 0 await c.press("pagedown") - c.check("PgDn pages the struct list from the filter prompt", - (ol.highlighted or 0) > 1, - f"highlighted={ol.highlighted} n={ol.option_count}") + c.check( + "PgDn pages the struct list from the filter prompt", + (ol.highlighted or 0) > 1, + f"highlighted={ol.highlighted} n={ol.option_count}", + ) await c.press("pageup") - c.check("PgUp returns to the first struct", (ol.highlighted or 0) == 0, - f"highlighted={ol.highlighted}") + c.check( + "PgUp returns to the first struct", + (ol.highlighted or 0) == 0, + f"highlighted={ol.highlighted}", + ) for ch in q: await c.press(ch) await c.wait(lambda: len(se._structs) < total, 5) - c.check("typing fuzzy-filters the struct list", - 0 < len(se._structs) < total and - any(s.name == target for s in se._structs), - f"q={q!r} {len(se._structs)}/{total}") + c.check( + "typing fuzzy-filters the struct list", + 0 < len(se._structs) < total and any(s.name == target for s in se._structs), + f"q={q!r} {len(se._structs)}/{total}", + ) cap = str(se.query_one("#se-title", Static).render()) - c.check("the caption counts what the filter kept", - f"{len(se._structs)}/{total}" in cap, f"caption={cap!r}") + c.check( + "the caption counts what the filter kept", + f"{len(se._structs)}/{total}" in cap, + f"caption={cap!r}", + ) # 'd' is the delete binding on this screen: in the prompt it must be a # character, not a destructive verb aimed at the highlighted struct. await c.press("d") await c.pause(0.05) - c.check("'d' in the filter types instead of deleting", - isinstance(app.screen, StructEditor) and inp.value == q + "d", - f"screen={type(app.screen).__name__} value={inp.value!r}") + c.check( + "'d' in the filter types instead of deleting", + isinstance(app.screen, StructEditor) and inp.value == q + "d", + f"screen={type(app.screen).__name__} value={inp.value!r}", + ) await c.press("backspace") await c.wait(lambda: inp.value == q, 5) @@ -1737,34 +2161,45 @@ async def s_struct_filter(c: Ctx): before = ol.highlighted await c.press("down") await c.pause(0.05) - c.check("arrows move the list while the filter has focus", - app.focused is inp and (ol.highlighted != before - or ol.option_count == 1), - f"{before} -> {ol.highlighted} of {ol.option_count}") + c.check( + "arrows move the list while the filter has focus", + app.focused is inp and (ol.highlighted != before or ol.option_count == 1), + f"{before} -> {ol.highlighted} of {ol.option_count}", + ) sel = se._structs[ol.highlighted or 0].name ta = se.query_one(TextArea) ta.text = "" await c.press("enter") loaded = await c.wait(lambda: sel in ta.text, 15) - c.check("Enter in the filter loads the highlighted struct", loaded, - f"want {sel!r} in {ta.text[:40]!r}") + c.check( + "Enter in the filter loads the highlighted struct", + loaded, + f"want {sel!r} in {ta.text[:40]!r}", + ) # Esc backs out one level at a time: definition -> filter -> dialog. await c.press("escape") await c.wait(lambda: app.focused is ol, 5) - c.check("Esc leaves the definition for the list", app.focused is ol, - f"focus={getattr(app.focused, 'id', None)}") + c.check( + "Esc leaves the definition for the list", + app.focused is ol, + f"focus={getattr(app.focused, 'id', None)}", + ) await c.press("escape") cleared = await c.wait(lambda: len(se._structs) == total, 5) - c.check("Esc clears the filter instead of closing", - cleared and not inp.display and isinstance(app.screen, StructEditor), - f"n={len(se._structs)}/{total} display={inp.display}") + c.check( + "Esc clears the filter instead of closing", + cleared and not inp.display and isinstance(app.screen, StructEditor), + f"n={len(se._structs)}/{total} display={inp.display}", + ) await c.press("escape") await c.wait(lambda: not isinstance(app.screen, StructEditor), 10) - c.check("a third Esc closes the editor", - not isinstance(app.screen, StructEditor), - f"screen={type(app.screen).__name__}") + c.check( + "a third Esc closes the editor", + not isinstance(app.screen, StructEditor), + f"screen={type(app.screen).__name__}", + ) @scenario("open_default_view") @@ -1774,9 +2209,11 @@ async def s_open(c: Ctx): app._open_function(fn.addr, fn.name) await c.wait(lambda: app._cur and app._cur.ea == fn.addr, 20) await c.wait(lambda: c.lst.total > 0, 30) - c.check("opening a function shows the linear listing by default", - app._active == "listing" and c.lst.display and c.lst.total > 0, - f"active={app._active} total={c.lst.total}") + c.check( + "opening a function shows the linear listing by default", + app._active == "listing" and c.lst.display and c.lst.total > 0, + f"active={app._active} total={c.lst.total}", + ) print(f" biggest = {fn.name} ({c.lst.total} listing rows)") @@ -1785,15 +2222,19 @@ async def s_disasm_nav(c: Ctx): app, view = c.app, c.dis await c.open_biggest("listing") view.focus() - c.check("first instruction cached", - view.model is not None and view.model.cached_line(0) is not None) + c.check( + "first instruction cached", + view.model is not None and view.model.cached_line(0) is not None, + ) for _ in range(5): await c.press("pagedown") await c.pause(0.025) c.check("pagedown moved the cursor", view.cursor > 0, f"cursor={view.cursor}") await c.wait(lambda: view.model.cached_line(view.cursor) is not None, 15) - c.check("cursor line eventually cached (bg fetch)", - view.model.cached_line(view.cursor) is not None) + c.check( + "cursor line eventually cached (bg fetch)", + view.model.cached_line(view.cursor) is not None, + ) c.check("status shows an address", "@ 0x" in c.status(), c.status()) # Ctrl+Y copies the current code line. view.focus() @@ -1801,14 +2242,20 @@ async def s_disasm_nav(c: Ctx): app._clipboard = "" await c.press("ctrl+y") await c.wait(lambda: app._clipboard == cur_line, 10) - c.check("Ctrl+Y copies the current code line to the clipboard", - bool(cur_line) and app._clipboard == cur_line, f"clip={app._clipboard!r}") + c.check( + "Ctrl+Y copies the current code line to the clipboard", + bool(cur_line) and app._clipboard == cur_line, + f"clip={app._clipboard!r}", + ) # goto-bottom must not hang on a huge function (ctrl+end; plain 'end' now # moves the cursor to end-of-line). await c.press("ctrl+end") await c.pause(0.05) - c.check("goto-bottom lands near end", view.cursor >= view.total - 1, - f"cursor={view.cursor}/{view.total}") + c.check( + "goto-bottom lands near end", + view.cursor >= view.total - 1, + f"cursor={view.cursor}/{view.total}", + ) @scenario("hex") @@ -1822,33 +2269,52 @@ async def s_hex(c: Ctx): await c.press("backslash") await c.wait(lambda: app._active == "hex", 10) hx = c.hex - await c.wait(lambda: hx.model is not None - and hx.model.row(hx.cursor // 16)[1] is not None, 20) - c.check("backslash opens the hex view synced to the code cursor", - app._active == "hex" and code_ea is not None and hx.cursor_va() == code_ea, - f"active={app._active} hexva={hx.cursor_va():#x} ea={code_ea}") + await c.wait( + lambda: hx.model is not None and hx.model.row(hx.cursor // 16)[1] is not None, + 20, + ) + c.check( + "backslash opens the hex view synced to the code cursor", + app._active == "hex" and code_ea is not None and hx.cursor_va() == code_ea, + f"active={app._active} hexva={hx.cursor_va():#x} ea={code_ea}", + ) want = app.program.read_bytes(code_ea, 1) _, rb = hx.model.row(hx.cursor // 16) - c.check("hex shows the actual byte at that address", - rb is not None and rb[hx.cursor % 16] == want[0], - f"got={rb[hx.cursor % 16] if rb else None} want={want[0]}") + c.check( + "hex shows the actual byte at that address", + rb is not None and rb[hx.cursor % 16] == want[0], + f"got={rb[hx.cursor % 16] if rb else None} want={want[0]}", + ) old_va = hx.cursor_va() block = (old_va - hx.model.start) // hx.model.BLOCK old_bytes = hx.model._blocks.get(block) await c.press("ctrl+r") reloaded = await c.wait( - lambda: hx.model._blocks.get(block) is not None - and hx.model._blocks.get(block) is not old_bytes, 20) + lambda: ( + hx.model._blocks.get(block) is not None + and hx.model._blocks.get(block) is not old_bytes + ), + 20, + ) c.check("Ctrl+R refetches the visible hex block", reloaded) - c.check("Ctrl+R preserves the hex cursor", hx.cursor_va() == old_va, - f"got={hx.cursor_va():#x} want={old_va:#x}") + c.check( + "Ctrl+R preserves the hex cursor", + hx.cursor_va() == old_va, + f"got={hx.cursor_va():#x} want={old_va:#x}", + ) await c.press("l") await c.pause(0.1) - c.check("hex cursor steps one byte", hx.cursor_va() == code_ea + 1, - f"va={hx.cursor_va():#x}") + c.check( + "hex cursor steps one byte", + hx.cursor_va() == code_ea + 1, + f"va={hx.cursor_va():#x}", + ) fo = app.program.file_offset(hx.cursor_va()) - c.check("hex carries a file offset for a mapped (.text) address", - fo is not None and hx.model.file_offset(hx.cursor_va()) == fo, f"fo={fo}") + c.check( + "hex carries a file offset for a mapped (.text) address", + fo is not None and hx.model.file_offset(hx.cursor_va()) == fo, + f"fo={fo}", + ) rng = app.program.image_range() target_va = rng[0] + (rng[1] - rng[0]) // 2 await c.press("g") @@ -1856,14 +2322,17 @@ async def s_hex(c: Ctx): await c.type(hex(target_va)) await c.press("enter") await c.wait(lambda: hx.cursor_va() == target_va, 15) - c.check("'g' in the hex view jumps the cursor to an address", - hx.cursor_va() == target_va, f"va={hx.cursor_va():#x} want={target_va:#x}") + c.check( + "'g' in the hex view jumps the cursor to an address", + hx.cursor_va() == target_va, + f"va={hx.cursor_va():#x} want={target_va:#x}", + ) # -- a user scroll freezes the cursor's screen row (points at a new byte) -- await c.press("g") await c.type(hex(rng[0])) await c.press("enter") await c.wait(lambda: hx.cursor_va() == rng[0], 10) - for _ in range(8): # cursor to viewport row 8 (top still 0) + for _ in range(8): # cursor to viewport row 8 (top still 0) await c.press("j") await c.pause(0.1) top0 = round(hx.scroll_offset.y) @@ -1872,24 +2341,33 @@ async def s_hex(c: Ctx): await c.pause(0.15) top1 = round(hx.scroll_offset.y) c.check("hex viewport scrolled", top1 >= top0 + 20, f"top0={top0} top1={top1}") - c.check("hex cursor's screen row stays frozen on scroll", - hx.cursor // 16 - top1 == screen_row, - f"screen_row={screen_row} now={hx.cursor // 16 - top1} top1={top1}") + c.check( + "hex cursor's screen row stays frozen on scroll", + hx.cursor // 16 - top1 == screen_row, + f"screen_row={screen_row} now={hx.cursor // 16 - top1} top1={top1}", + ) PAD = 1 # HexView { padding: 0 1 } -> content is inset one col await c.pilot.click(HexView, offset=(PAD + 19 + 3 * 3, 5)) # hex byte 3, row 5 await c.pause(0.1) - c.check("clicking the hex pane moves the cursor to the clicked byte", - hx.cursor == (top1 + 5) * 16 + 3, - f"cursor={hx.cursor} want={(top1 + 5) * 16 + 3} top1={top1}") + c.check( + "clicking the hex pane moves the cursor to the clicked byte", + hx.cursor == (top1 + 5) * 16 + 3, + f"cursor={hx.cursor} want={(top1 + 5) * 16 + 3} top1={top1}", + ) await c.pilot.click(HexView, offset=(PAD + 70 + 10, 7)) # ascii byte 10, row 7 await c.pause(0.1) - c.check("clicking the ascii pane maps to the right byte", - hx.cursor == (top1 + 7) * 16 + 10, - f"cursor={hx.cursor} want={(top1 + 7) * 16 + 10}") + c.check( + "clicking the ascii pane maps to the right byte", + hx.cursor == (top1 + 7) * 16 + 10, + f"cursor={hx.cursor} want={(top1 + 7) * 16 + 10}", + ) await c.press("backslash") await c.wait(lambda: app._active != "hex", 10) - c.check("backslash returns from hex to the code view", - app._active == "listing", f"active={app._active}") + c.check( + "backslash returns from hex to the code view", + app._active == "listing", + f"active={app._active}", + ) @scenario("filter") @@ -1899,24 +2377,34 @@ async def s_filter(c: Ctx): nfuncs = table.row_count # row selection opens a function fn = c.biggest() - ridx = next((i for i in range(nfuncs) - if int(str(table.get_row_at(i)[0]), 16) == fn.addr), 0) + ridx = next( + (i for i in range(nfuncs) if int(str(table.get_row_at(i)[0]), 16) == fn.addr), 0 + ) table.move_cursor(row=ridx) table.focus() await c.press("enter") await c.wait(lambda: app._cur and app._cur.ea == fn.addr, 20) - c.check("selecting a table row opens that function", - bool(app._cur) and app._cur.ea == fn.addr, f"cur={app._cur.ea if app._cur else None}") + c.check( + "selecting a table row opens that function", + bool(app._cur) and app._cur.ea == fn.addr, + f"cur={app._cur.ea if app._cur else None}", + ) # filter round-trip. Derive the glob from real names: this used to hardcode # 'sub_1*', which matches NOTHING in a binary whose code never reaches # 0x1xxx (echo's functions are sub_2xxx..sub_7xxx) — a deterministic failure # that looked like a flake, and left the table empty for the next scenario. subs = sorted(f.name for f in c.all_funcs() if f.name.startswith("sub_")) term = (subs[0][:5] + "*") if subs else "" - want = sum(1 for f in c.all_funcs() - if fnmatch.fnmatch(f.name.lower(), term.lower())) if term else 0 - c.check("picked a glob that actually matches (test self-check)", - 0 < want < nfuncs, f"term={term!r} want={want} of {nfuncs}") + want = ( + sum(1 for f in c.all_funcs() if fnmatch.fnmatch(f.name.lower(), term.lower())) + if term + else 0 + ) + c.check( + "picked a glob that actually matches (test self-check)", + 0 < want < nfuncs, + f"term={term!r} want={want} of {nfuncs}", + ) table.focus() await c.press("slash") await c.pause(0.05) @@ -1924,32 +2412,44 @@ async def s_filter(c: Ctx): await c.press(ch if ch != "*" else "asterisk") await c.press("enter") filtered = await c.wait(lambda: table.row_count == want, 15) - c.check("filter narrowed the list to exactly the matches", filtered, - f"term={term!r} rows={table.row_count} want={want} of {nfuncs}") + c.check( + "filter narrowed the list to exactly the matches", + filtered, + f"term={term!r} rows={table.row_count} want={want} of {nfuncs}", + ) # pane toggle left = app.query_one("#left", FunctionsPanel) await c.press("ctrl+b") await c.pause(0.05) - c.check("ctrl+b hides functions pane + focuses the code view", - not left.display and isinstance(app.focused, (ListingView, DecompView)), - f"display={left.display} focus={type(app.focused).__name__}") + c.check( + "ctrl+b hides functions pane + focuses the code view", + not left.display and isinstance(app.focused, (ListingView, DecompView)), + f"display={left.display} focus={type(app.focused).__name__}", + ) await c.press("ctrl+b") await c.pause(0.05) - c.check("ctrl+b again restores pane + focuses table", - left.display and isinstance(app.focused, DataTable), - f"display={left.display} focus={type(app.focused).__name__}") + c.check( + "ctrl+b again restores pane + focuses table", + left.display and isinstance(app.focused, DataTable), + f"display={left.display} focus={type(app.focused).__name__}", + ) @scenario("view_toggle") async def s_view_toggle(c: Ctx): app, dis, dec = c.app, c.dis, c.dec await c.open_biggest("decomp") - pc = await c.wait(lambda: dec.display and dec.loaded_ea is not None - and app._active == "decomp", 25) + pc = await c.wait( + lambda: dec.display and dec.loaded_ea is not None and app._active == "decomp", + 25, + ) c.check("pseudocode view shows", pc, f"active={app._active}") c.check("pseudocode has many lines", dec.total > 20, f"lines={dec.total}") - styled = any(seg.style is not None and seg.style.color is not None - for strip in dec._strips[:min(dec.total, 200)] for seg in strip) + styled = any( + seg.style is not None and seg.style.color is not None + for strip in dec._strips[: min(dec.total, 200)] + for seg in strip + ) c.check("pseudocode is syntax-highlighted", styled) # Cancelling a prompt must hand focus back to the pane you were READING. # _code_view() used to choose on _pref, which was only ever "listing", so it @@ -1965,21 +2465,27 @@ async def s_view_toggle(c: Ctx): await c.press("down") await c.press("down") await c.pause(0.2) - c.check("cancelling goto leaves focus in the pseudocode (arrows still work)", - dec.cursor > line0, - f"cursor {line0} -> {dec.cursor} focus={type(app.focused).__name__}") + c.check( + "cancelling goto leaves focus in the pseudocode (arrows still work)", + dec.cursor > line0, + f"cursor {line0} -> {dec.cursor} focus={type(app.focused).__name__}", + ) dec.focus() # Tab only toggles the view from a code pane; elsewhere it's await c.pause(0.05) # focus-next, which would silently leave us in decomp await c.press("tab") # decomp -> listing runs through _toggle_to_listing, a background worker, so # _active only flips once the listing model has loaded. A fixed pause held in # a short run and lost the race in a full one. - switched = await c.wait(lambda: app._active == "listing" and dis.display - and not dec.display, 20) - c.check("tab switches to disassembly", switched, - f"active={app._active} split={app._split} " - f"focus={type(app.focused).__name__} " - f"lst={dis.display} dec={dec.display}") + switched = await c.wait( + lambda: app._active == "listing" and dis.display and not dec.display, 20 + ) + c.check( + "tab switches to disassembly", + switched, + f"active={app._active} split={app._split} " + f"focus={type(app.focused).__name__} " + f"lst={dis.display} dec={dec.display}", + ) # _toggle_to_listing repositions asynchronously; F5 below reads the listing # cursor's ea and no-ops if it isn't on an addressed row yet. await c.wait(lambda: c.lst._cursor_ea() is not None, 10) @@ -1989,34 +2495,54 @@ async def s_view_toggle(c: Ctx): dec.loaded_ea = None app.action_toggle_view() cover = dec._cover_widget - c.check("F5 from the listing raises the 'decompiling…' overlay", - dec.loading and cover is not None and "decomp-loading" in cover.classes - and "decompiling" in str(cover.render()), - f"loading={dec.loading} cover={cover!r}") - await c.wait(lambda: app._active == "decomp" and not dec.loading - and dec._cover_widget is None, 25) - c.check("overlay clears when the decompile finishes", - not dec.loading and dec._cover_widget is None) + c.check( + "F5 from the listing raises the 'decompiling…' overlay", + dec.loading + and cover is not None + and "decomp-loading" in cover.classes + and "decompiling" in str(cover.render()), + f"loading={dec.loading} cover={cover!r}", + ) + await c.wait( + lambda: ( + app._active == "decomp" and not dec.loading and dec._cover_widget is None + ), + 25, + ) + c.check( + "overlay clears when the decompile finishes", + not dec.loading and dec._cover_widget is None, + ) # F5 on an ALREADY-loaded function must still clear the overlay (regression: # the F5-raised overlay had nothing to clear it in the 'already loaded' branch # -> spinner stuck forever). await c.press("tab") # -> listing await c.wait(lambda: app._active == "listing", 10) app.action_toggle_view() # F5 the same, cached function again - cleared = await c.wait(lambda: app._active == "decomp" and not dec.loading - and dec._cover_widget is None, 15) - c.check("re-decompiling an already-loaded function clears the overlay", cleared, - f"loading={dec.loading} cover={dec._cover_widget!r}") + cleared = await c.wait( + lambda: ( + app._active == "decomp" and not dec.loading and dec._cover_widget is None + ), + 15, + ) + c.check( + "re-decompiling an already-loaded function clears the overlay", + cleared, + f"loading={dec.loading} cover={dec._cover_widget!r}", + ) # line-number gutter dec.scroll_to(0, 0, animate=False) await c.pause(0.025) row0 = "".join(seg.text for seg in dec.render_line(0)) - c.check("pseudocode has a numbered gutter (line 1 first)", - dec._gutter > 0 and row0[:dec._gutter].strip() == "1", - f"gutter={dec._gutter} row0={row0[:10]!r}") + c.check( + "pseudocode has a numbered gutter (line 1 first)", + dec._gutter > 0 and row0[: dec._gutter].strip() == "1", + f"gutter={dec._gutter} row0={row0[:10]!r}", + ) # Home/End move along the line here too (they used to scroll to top/bottom). - line = next((i for i, t in enumerate(dec._texts) - if t.startswith(" ") and t.strip()), None) + line = next( + (i for i, t in enumerate(dec._texts) if t.startswith(" ") and t.strip()), None + ) if line is not None: text = dec._texts[line] dec.focus() @@ -2026,28 +2552,37 @@ async def s_view_toggle(c: Ctx): top = round(dec.scroll_offset.y) await c.press("end") await c.pause(0.05) - c.check("<end> in pseudocode goes to end-of-line, not the bottom", - dec.cursor == line and dec.cursor_x == max(len(text) - 1, 0) - and round(dec.scroll_offset.y) == top, - f"line={dec.cursor} col={dec.cursor_x} len={len(text)}") + c.check( + "<end> in pseudocode goes to end-of-line, not the bottom", + dec.cursor == line + and dec.cursor_x == max(len(text) - 1, 0) + and round(dec.scroll_offset.y) == top, + f"line={dec.cursor} col={dec.cursor_x} len={len(text)}", + ) await c.press("home") await c.pause(0.05) - c.check("<home> in pseudocode goes to start-of-line", - dec.cursor == line and dec.cursor_x == 0, - f"line={dec.cursor} col={dec.cursor_x}") + c.check( + "<home> in pseudocode goes to start-of-line", + dec.cursor == line and dec.cursor_x == 0, + f"line={dec.cursor} col={dec.cursor_x}", + ) await c.press("shift+home") await c.pause(0.05) - c.check("<shift+home> skips the indentation", - dec.cursor_x == len(text) - len(text.lstrip()), - f"col={dec.cursor_x} indent={len(text) - len(text.lstrip())}") + c.check( + "<shift+home> skips the indentation", + dec.cursor_x == len(text) - len(text.lstrip()), + f"col={dec.cursor_x} indent={len(text) - len(text.lstrip())}", + ) await c.press("ctrl+end") await c.pause(0.1) - c.check("<ctrl+end> still goes to the bottom", - dec.cursor >= dec.total - 1, f"{dec.cursor}/{dec.total}") + c.check( + "<ctrl+end> still goes to the bottom", + dec.cursor >= dec.total - 1, + f"{dec.cursor}/{dec.total}", + ) await c.press("ctrl+home") await c.pause(0.1) - c.check("<ctrl+home> still goes to the top", dec.cursor == 0, - f"{dec.cursor}") + c.check("<ctrl+home> still goes to the top", dec.cursor == 0, f"{dec.cursor}") @scenario("search") @@ -2067,26 +2602,39 @@ async def s_search(c: Ctx): # the unified listing searches the whole segment (load_all) -> allow time await c.wait(lambda: bool(dis._matches), 45) c.check("search finds matches", len(dis._matches) > 0, f"term={term!r}") - c.check("cursor sits on a match", dis.cursor in dis._matches, f"cursor={dis.cursor}") - c.check("match substring highlighted", - bool(dis._ranges.get(dis.cursor)), str(dis._ranges.get(dis.cursor))) - c.check("search cursor lands on the match's starting column", - bool(dis._ranges.get(dis.cursor)) - and dis.cursor_x == dis._ranges[dis.cursor][0][0], - f"cursor_x={dis.cursor_x} ranges={dis._ranges.get(dis.cursor)}") + c.check( + "cursor sits on a match", dis.cursor in dis._matches, f"cursor={dis.cursor}" + ) + c.check( + "match substring highlighted", + bool(dis._ranges.get(dis.cursor)), + str(dis._ranges.get(dis.cursor)), + ) + c.check( + "search cursor lands on the match's starting column", + bool(dis._ranges.get(dis.cursor)) + and dis.cursor_x == dis._ranges[dis.cursor][0][0], + f"cursor_x={dis.cursor_x} ranges={dis._ranges.get(dis.cursor)}", + ) prev = dis.cursor await c.press("slash") await c.pause(0.05) await c.press("enter") await c.pause(0.05) - c.check("'/' repeats to next match", - dis.cursor != prev and dis.cursor in dis._matches, f"cursor={dis.cursor}") + c.check( + "'/' repeats to next match", + dis.cursor != prev and dis.cursor in dis._matches, + f"cursor={dis.cursor}", + ) await c.press("question_mark") await c.pause(0.05) await c.press("enter") await c.pause(0.05) - c.check("'?' repeats to previous match", dis.cursor in dis._matches, - f"cursor={dis.cursor}") + c.check( + "'?' repeats to previous match", + dis.cursor in dis._matches, + f"cursor={dis.cursor}", + ) # incremental preview + visible bar + Esc cancel si = app.query_one("#search", Input) status = app.query_one("#status", Static) @@ -2094,22 +2642,31 @@ async def s_search(c: Ctx): # display flips synchronously; the REGION only exists once Textual has laid # the prompt out, which is a frame, not a worker. await c.wait(lambda: si.display and si.region.height >= 1, 5) - c.check("search bar visible, status hidden (no overlap)", - si.display and not status.display, f"si={si.display} status={status.display}") - c.check("search input owns the bottom row (nothing overlaps it)", - si.region.height >= 1 - and si.region.y + si.region.height == app.size.height, - f"search={si.region} screen={app.size}") + c.check( + "search bar visible, status hidden (no overlap)", + si.display and not status.display, + f"si={si.display} status={status.display}", + ) + c.check( + "search input owns the bottom row (nothing overlaps it)", + si.region.height >= 1 and si.region.y + si.region.height == app.size.height, + f"search={si.region} screen={app.size}", + ) for ch in term: await c.press(ch) await c.pause(0.05) - c.check("matches highlight incrementally (before Enter)", - len(dis._matches) > 0 and si.value == term, f"val={si.value!r}") + c.check( + "matches highlight incrementally (before Enter)", + len(dis._matches) > 0 and si.value == term, + f"val={si.value!r}", + ) await c.press("escape") await c.pause(0.1) - c.check("Esc cancels: status restored, matches cleared", - status.display and not si.display and not dis._matches, - f"status={status.display} si={si.display} m={len(dis._matches)}") + c.check( + "Esc cancels: status restored, matches cleared", + status.display and not si.display and not dis._matches, + f"status={status.display} si={si.display} m={len(dis._matches)}", + ) @scenario("incr_filter") @@ -2125,19 +2682,30 @@ async def s_incr_filter(c: Ctx): # settling can't see it. Wait for the effect instead of guessing at the # debounce: it returns the moment the rows are rebuilt. await c.wait(lambda: 0 < table.row_count < full, 5) - c.check("filter narrows incrementally as you type", - 0 < table.row_count < full, f"{table.row_count}/{full}") + c.check( + "filter narrows incrementally as you type", + 0 < table.row_count < full, + f"{table.row_count}/{full}", + ) cell = table.get_row_at(0)[1] - c.check("filter highlights matched substring in name", - isinstance(cell, Text) and any(s.style for s in cell.spans), repr(str(cell))) + c.check( + "filter highlights matched substring in name", + isinstance(cell, Text) and any(s.style for s in cell.spans), + repr(str(cell)), + ) await c.press("enter") await c.wait(lambda: isinstance(app.focused, DataTable), 5) - c.check("Enter keeps filter + focuses table", - isinstance(app.focused, DataTable) and table.row_count < full) + c.check( + "Enter keeps filter + focuses table", + isinstance(app.focused, DataTable) and table.row_count < full, + ) await c.press("escape") await c.wait(lambda: table.row_count == full, 5) - c.check("Esc on the list clears the filter", table.row_count == full, - f"{table.row_count}/{full}") + c.check( + "Esc on the list clears the filter", + table.row_count == full, + f"{table.row_count}/{full}", + ) @scenario("follow_xrefs") @@ -2146,10 +2714,13 @@ async def s_follow_xrefs(c: Ctx): await c.open_biggest("listing") dis.focus() lines = dis.model.lines(0, 400, prefetch=False) - call_idx = next((i for i, ln in enumerate(lines) - if ln.text.startswith("call ")), None) + call_idx = next( + (i for i, ln in enumerate(lines) if ln.text.startswith("call ")), None + ) if call_idx is None: - c.check("found a call line to exercise follow/xrefs", False, "no call in first 400") + c.check( + "found a call line to exercise follow/xrefs", False, "no call in first 400" + ) return dis.cursor = call_idx dis.refresh() @@ -2163,8 +2734,11 @@ async def s_follow_xrefs(c: Ctx): # failed about one run in ten with cur == orig, at full speed, looking like # a code regression. await c.wait(lambda: len(app._nav) > depth and app._cur.ea != orig, 25) - c.check("Enter follows the call into another function", - app._cur.ea != orig and len(app._nav) > depth, f"cur={app._cur.ea:#x}") + c.check( + "Enter follows the call into another function", + app._cur.ea != orig and len(app._nav) > depth, + f"cur={app._cur.ea:#x}", + ) await c.press("escape") await c.wait(lambda: app._cur.ea == orig, 15) c.check("Esc returns from the follow", app._cur.ea == orig, f"cur={app._cur.ea:#x}") @@ -2174,15 +2748,19 @@ async def s_follow_xrefs(c: Ctx): opened = await c.wait(lambda: isinstance(app.screen, XrefsScreen), 25) c.check("'x' opens the xrefs popup", opened, f"screen={type(app.screen).__name__}") if opened: - c.check("xrefs popup has entries", - app.screen.query_one(OptionList).option_count >= 1) + c.check( + "xrefs popup has entries", + app.screen.query_one(OptionList).option_count >= 1, + ) await c.press("escape") await c.pause(0.1) c.check("Esc closes the xrefs popup", not isinstance(app.screen, XrefsScreen)) xf = None for cand in c.all_funcs()[:600]: - codex = [x for x in app.program.xrefs_to(cand.addr) if x.type == "code" and x.frm] + codex = [ + x for x in app.program.xrefs_to(cand.addr) if x.type == "code" and x.frm + ] if codex: xf = (cand, codex[0]) break @@ -2200,19 +2778,27 @@ async def s_follow_xrefs(c: Ctx): await c.wait(lambda: not isinstance(app.screen, XrefsScreen), 25) # xref-select lands the listing cursor on the referencing SITE (frm) await c.wait(lambda: c.lst._cursor_ea() == xref.frm, 25) - c.check("xref-select lands the cursor on the referencing site", - c.lst._cursor_ea() == xref.frm, - f"cur_ea={c.lst._cursor_ea()} want={xref.frm:#x}") + c.check( + "xref-select lands the cursor on the referencing site", + c.lst._cursor_ea() == xref.frm, + f"cur_ea={c.lst._cursor_ea()} want={xref.frm:#x}", + ) # F5 at the site decompiles the referencing function c.lst.focus() await c.press("tab") landed = await c.wait( - lambda: (app._active == "decomp" and dec.loaded_ea == xref.fn_addr) - or (app.is_listing - and _CANNOT_DECOMP in c.status().lower()), 25) + lambda: ( + (app._active == "decomp" and dec.loaded_ea == xref.fn_addr) + or (app.is_listing and _CANNOT_DECOMP in c.status().lower()) + ), + 25, + ) if app._active == "decomp": - c.check("F5 at the xref site decompiles the referencing function", - dec.loaded_ea == xref.fn_addr, f"loaded={dec.loaded_ea}") + c.check( + "F5 at the xref site decompiles the referencing function", + dec.loaded_ea == xref.fn_addr, + f"loaded={dec.loaded_ea}", + ) await c.press("tab") await c.wait(lambda: app._active == "listing", 20) @@ -2225,9 +2811,13 @@ async def s_follow_xrefs(c: Ctx): dis.refresh() await c.pause(0.025) await c.press("h") - c.check("h moves the column cursor left", dis.cursor_x == 4, f"x={dis.cursor_x}") + c.check( + "h moves the column cursor left", dis.cursor_x == 4, f"x={dis.cursor_x}" + ) await c.press("l", "l") - c.check("l moves the column cursor right", dis.cursor_x == 6, f"x={dis.cursor_x}") + c.check( + "l moves the column cursor right", dis.cursor_x == 6, f"x={dis.cursor_x}" + ) plain = dis._line_plain(call_idx) or "" m = re.search(r"\b(sub_[0-9A-Fa-f]+)", plain) if m: @@ -2235,24 +2825,32 @@ async def s_follow_xrefs(c: Ctx): dis.cursor_x = m.start(1) + 1 dis.refresh() await c.pause(0.05) - c.check("word-under-cursor is the operand symbol", - dis.word_under_cursor() == m.group(1), - f"{dis.word_under_cursor()!r} vs {m.group(1)!r}") + c.check( + "word-under-cursor is the operand symbol", + dis.word_under_cursor() == m.group(1), + f"{dis.word_under_cursor()!r} vs {m.group(1)!r}", + ) depth = len(app._nav) want = app.program.resolve(m.group(1)) await c.press("enter") await c.wait(lambda: len(app._nav) > depth, 25) - c.check("follows the symbol under the cursor", - app._cur.ea == want, f"cur={app._cur.ea:#x} want={want:#x}") + c.check( + "follows the symbol under the cursor", + app._cur.ea == want, + f"cur={app._cur.ea:#x} want={want:#x}", + ) @scenario("xref_labels") async def s_xref_labels(c: Ctx): from collections import Counter, defaultdict + app, dis, dec = c.app, c.dis, c.dec multi = None for cand in c.all_funcs()[:200]: - callers = Counter(x.fn_addr for x in app.program.xrefs_to(cand.addr) if x.fn_name) + callers = Counter( + x.fn_addr for x in app.program.xrefs_to(cand.addr) if x.fn_name + ) if any(n >= 2 for n in callers.values()): multi = cand break @@ -2273,9 +2871,14 @@ async def s_xref_labels(c: Ctx): if "+0x" in x: nm, off = x.split("+0x", 1) byfn[nm].add(off) - c.check("xref labels distinguish multiple sites in a function by offset", - any(len(offs) >= 2 for offs in byfn.values()), f"locs={locs[:8]}") - c.check("no xref label is a bare '?'", all(x != "?" for x in locs), f"locs={locs[:8]}") + c.check( + "xref labels distinguish multiple sites in a function by offset", + any(len(offs) >= 2 for offs in byfn.values()), + f"locs={locs[:8]}", + ) + c.check( + "no xref label is a bare '?'", all(x != "?" for x in locs), f"locs={locs[:8]}" + ) await c.press("escape") await c.wait(lambda: not isinstance(app.screen, XrefsScreen), 25) # pre-selection: 'x' at a call site highlights that site in the dialog @@ -2283,8 +2886,9 @@ async def s_xref_labels(c: Ctx): for x in app.program.xrefs_to(multi.addr): if x.fn_name and x.type == "code": bycaller[x.fn_addr].append(x) - csites = next((sorted(v, key=lambda x: x.frm) - for v in bycaller.values() if len(v) >= 2), None) + csites = next( + (sorted(v, key=lambda x: x.frm) for v in bycaller.values() if len(v) >= 2), None + ) if not csites: c.check("found a caller with multiple sites for preselect", False) return @@ -2303,9 +2907,11 @@ async def s_xref_labels(c: Ctx): await c.wait(lambda: isinstance(app.screen, XrefsScreen), 25) hl = app.screen.query_one(OptionList).highlighted it = app.screen._items - c.check("xref dialog pre-selects the site it was invoked from", - hl is not None and it[hl][0] == site.frm, - f"hl={hl} frm={hex(it[hl][0]) if hl is not None else None} want={hex(site.frm)}") + c.check( + "xref dialog pre-selects the site it was invoked from", + hl is not None and it[hl][0] == site.frm, + f"hl={hl} frm={hex(it[hl][0]) if hl is not None else None} want={hex(site.frm)}", + ) await c.press("escape") await c.wait(lambda: not isinstance(app.screen, XrefsScreen), 25) @@ -2336,21 +2942,28 @@ async def s_mouse(c: Ctx): return await c.pilot.click(dis, offset=(mcol + 1, mrow)) await c.pause(0.05) - c.check("single click places the cursor on the clicked token", - dis.cursor == mline and dis.word_under_cursor() == msym, - f"cursor={dis.cursor} (want {mline}) word={dis.word_under_cursor()!r}") + c.check( + "single click places the cursor on the clicked token", + dis.cursor == mline and dis.word_under_cursor() == msym, + f"cursor={dis.cursor} (want {mline}) word={dis.word_under_cursor()!r}", + ) depth = len(app._nav) want = app.program.resolve(msym) await c.pilot.click(dis, offset=(mcol + 1, mrow), times=2) await c.wait(lambda: len(app._nav) > depth, 25) - c.check("double-click follows the symbol", app._cur.ea == want, - f"cur={app._cur.ea:#x} want={want:#x}") + c.check( + "double-click follows the symbol", + app._cur.ea == want, + f"cur={app._cur.ea:#x} want={want:#x}", + ) await c.press("escape") await c.wait(lambda: app._cur.ea != want, 20) await c.wait(lambda: dis.total > 0 and dis.cursor == mline, 20) - c.check("back restores the exact line + column", - dis.cursor == mline and dis.word_under_cursor() == msym, - f"cursor={dis.cursor} (want {mline}) word={dis.word_under_cursor()!r}") + c.check( + "back restores the exact line + column", + dis.cursor == mline and dis.word_under_cursor() == msym, + f"cursor={dis.cursor} (want {mline}) word={dis.word_under_cursor()!r}", + ) @scenario("decomp_nav") @@ -2372,23 +2985,32 @@ async def s_decomp_nav(c: Ctx): await c.press("escape") await c.wait(lambda: dec.loaded_ea == fn.addr, 25) await c.pause(0.1) - c.check("pseudocode-view position restored after jump+back", - dec.cursor == drow and dec.word_under_cursor() == dsym, - f"cursor={dec.cursor} (want {drow}) word={dec.word_under_cursor()!r}") + c.check( + "pseudocode-view position restored after jump+back", + dec.cursor == drow and dec.word_under_cursor() == dsym, + f"cursor={dec.cursor} (want {drow}) word={dec.word_under_cursor()!r}", + ) # follow works with a STALE name (post-rename): ea-marker fallback dstale = app.program.resolve(dsym) old_line = dec._texts[drow] if drow < len(dec._texts) else "" 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.invoke("rename", batch={"func": {"addr": hex(dstale), "name": tmp}}) + app.program.client.call( + remote_ops.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.invoke("rename", batch={"func": {"addr": hex(dstale), "name": dsym}}) + 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( + remote_ops.rename, batch={"func": {"addr": hex(dstale), "name": dsym}} + ) app.program.bump_names() @@ -2415,9 +3037,11 @@ async def s_decomp_follow_self(c: Ctx): depth = len(app._nav) await c.press("enter") moved = await c.wait(lambda: len(app._nav) > depth, 25) - c.check("decompiler follows a name not in refs (resolve fallback)", - moved and app._cur.ea == fn.addr, - f"moved={moved} cur={hex(app._cur.ea)} want={hex(fn.addr)}") + c.check( + "decompiler follows a name not in refs (resolve fallback)", + moved and app._cur.ea == fn.addr, + f"moved={moved} cur={hex(app._cur.ea)} want={hex(fn.addr)}", + ) @scenario("sort") @@ -2427,19 +3051,28 @@ async def s_sort(c: Ctx): await c.pilot.click(table, offset=(15, 0)) # Function header await c.pause(0.15) snames = [str(table.get_row_at(i)[1]) for i in range(min(20, table.row_count))] - c.check("click Function header sorts by name", - app._sort_col == 1 and snames == sorted(snames, key=str.lower), - f"sort_col={app._sort_col}") + c.check( + "click Function header sorts by name", + app._sort_col == 1 and snames == sorted(snames, key=str.lower), + f"sort_col={app._sort_col}", + ) first_asc = str(table.get_row_at(0)[1]) await c.pilot.click(table, offset=(15, 0)) # reverse await c.pause(0.15) - c.check("click again reverses the sort", - app._sort_reverse and str(table.get_row_at(0)[1]) != first_asc) + c.check( + "click again reverses the sort", + app._sort_reverse and str(table.get_row_at(0)[1]) != first_asc, + ) await c.pilot.click(table, offset=(3, 0)) # Address header await c.pause(0.15) - saddrs = [int(str(table.get_row_at(i)[0]), 16) for i in range(min(20, table.row_count))] - c.check("click Address header sorts by address", - app._sort_col == 0 and saddrs == sorted(saddrs), f"sort_col={app._sort_col}") + saddrs = [ + int(str(table.get_row_at(i)[0]), 16) for i in range(min(20, table.row_count)) + ] + c.check( + "click Address header sorts by address", + app._sort_col == 0 and saddrs == sorted(saddrs), + f"sort_col={app._sort_col}", + ) @scenario("rename") @@ -2460,18 +3093,34 @@ async def s_rename(c: Ctx): await c.press("n") await c.pause(0.1) ri = app.query_one("#rename", Input) - c.check("'n' opens the rename prompt prefilled with the symbol", - ri.display and ri.value == dsym, f"val={ri.value!r}") + c.check( + "'n' opens the rename prompt prefilled with the symbol", + ri.display and ri.value == dsym, + f"val={ri.value!r}", + ) ri.value = newname await c.press("enter") - await c.wait(lambda: app._func_index.by_addr(dtarget) - and app._func_index.by_addr(dtarget).name == newname, 25) - 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.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"))) + await c.wait( + lambda: ( + app._func_index.by_addr(dtarget) + and app._func_index.by_addr(dtarget).name == newname + ), + 25, + ) + 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( + remote_ops.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 def _find_label(): for i, t in enumerate(dec._texts): @@ -2479,6 +3128,7 @@ async def s_rename(c: Ctx): if mm: return i, mm.start(), mm.group(0) return None + lab = _find_label() if lab is None: await c.open("main", "decomp") @@ -2492,15 +3142,25 @@ async def s_rename(c: Ctx): await c.press("n") await c.pause(0.1) ri2 = app.query_one("#rename", Input) - c.check("renaming a pseudocode label is refused with a clear message", - (not ri2.display) and "label" in c.status().lower(), - f"display={ri2.display} status={c.status()!r}") + c.check( + "renaming a pseudocode label is refused with a clear message", + (not ri2.display) and "label" in c.status().lower(), + f"display={ri2.display} status={c.status()!r}", + ) else: c.check("found a pseudocode label to test", False, "no LABEL_ found") # comment via ';' - cline = next((i for i in range(len(dec._texts)) - if i > 5 and dec._line_ea(i) is not None - and dec._texts[i].strip() and "//" not in dec._texts[i]), None) + cline = next( + ( + i + for i in range(len(dec._texts)) + if i > 5 + and dec._line_ea(i) is not None + and dec._texts[i].strip() + and "//" not in dec._texts[i] + ), + None, + ) if cline is not None: cea = dec._line_ea(cline) dec.focus() @@ -2511,8 +3171,11 @@ async def s_rename(c: Ctx): await c.pause(0.1) ci = app.query_one("#comment", Input) cnote = f"note_{os.getpid()}" - c.check("';' opens the comment prompt on the current line", ci.display, - f"display={ci.display}") + c.check( + "';' opens the comment prompt on the current line", + ci.display, + f"display={ci.display}", + ) ci.value = cnote await c.press("enter") # Gate on the comment showing up, and ONLY that: the extra @@ -2521,9 +3184,14 @@ async def s_rename(c: Ctx): # out its full 25s (9s of wall clock) and then the check below passed # vacuously anyway. await c.wait(lambda: 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.invoke("set_comments", items=[{"addr": hex(cea), "comment": ""}]) + c.check( + "comment appears in the pseudocode after ';'", + any(cnote in t for t in dec._texts), + "comment not shown", + ) + app.program.client.call( + remote_ops.set_comments, items=[{"addr": hex(cea), "comment": ""}] + ) else: c.check("found a pseudocode line to comment", False, "no marker line") @@ -2543,29 +3211,42 @@ async def s_comment_func(c: Ctx): dec.cursor, dec.cursor_x = 0, 2 # the signature line dec.refresh() await c.pause(0.05) - c.check("signature line has no address of its own", dec._line_ea(0) is None, - f"ea={dec._line_ea(0)}") + c.check( + "signature line has no address of its own", + dec._line_ea(0) is None, + f"ea={dec._line_ea(0)}", + ) await c.press("semicolon") await c.pause(0.1) ci = app.query_one("#comment", Input) note = f"fn_note_{os.getpid()}" - c.check("';' on the signature line opens a function-comment prompt", - ci.display and "function comment" in str(ci.placeholder).lower(), - f"display={ci.display} ph={ci.placeholder!r}") + c.check( + "';' on the signature line opens a function-comment prompt", + ci.display and "function comment" in str(ci.placeholder).lower(), + f"display={ci.display} ph={ci.placeholder!r}", + ) # literal '\n' in the comment becomes a real newline -> multi-line render a, b = f"{note}_A", f"{note}_B" ci.value = f"{a}\\n{b}" await c.press("enter") - await c.wait(lambda: dec.loaded_ea == fn.addr - and any(a in t for t in dec._texts) - and any(b in t for t in dec._texts), 25) + await c.wait( + lambda: ( + dec.loaded_ea == fn.addr + and any(a in t for t in dec._texts) + and any(b in t for t in dec._texts) + ), + 25, + ) la = next((i for i, t in enumerate(dec._texts) if a in t), None) lb = next((i for i, t in enumerate(dec._texts) if b in t), None) - c.check("multi-line function comment renders on separate lines", - 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.invoke("set_comments", items=[{"addr": hex(fn.addr), "comment": ""}]) + c.check( + "multi-line function comment renders on separate lines", + 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( + remote_ops.set_comments, items=[{"addr": hex(fn.addr), "comment": ""}] + ) @scenario("retype") @@ -2585,30 +3266,41 @@ async def s_retype(c: Ctx): await c.press("y") await c.wait(lambda: app.query_one("#retype", Input).display, 10) ri = app.query_one("#retype", Input) - c.check("'y' on a function prefills its prototype", - ri.display and ri.value == old_proto, f"val={ri.value!r} want={old_proto!r}") + c.check( + "'y' on a function prefills its prototype", + ri.display and ri.value == old_proto, + f"val={ri.value!r} want={old_proto!r}", + ) ri.value = f"void __fastcall {cf.name}(int zz_retype_arg)" await c.press("enter") - await c.wait(lambda: (lambda f: bool(f) and "zz_retype_arg" in f.prototype)( - app.program.func_types(cf.addr)), 20) + await c.wait( + lambda: (lambda f: bool(f) and "zz_retype_arg" in f.prototype)( + app.program.func_types(cf.addr) + ), + 20, + ) after = app.program.func_types(cf.addr) - c.check("applying a retype changes the function prototype", - after is not None and "zz_retype_arg" in after.prototype, - f"proto={after.prototype if after else None!r}") + c.check( + "applying a retype changes the function prototype", + after is not None and "zz_retype_arg" in after.prototype, + f"proto={after.prototype if after else None!r}", + ) app.program.set_function_type(cf.addr, old_proto) # restore # the retype above kicked off a recompile+reload; let it land before we start # placing the cursor, or the reload resets it under us. app.program.bump_names() - await c.wait(lambda: not dec.loading and dec.loaded_ea == cf.addr - and bool(dec._texts), 25) + await c.wait( + lambda: not dec.loading and dec.loaded_ea == cf.addr and bool(dec._texts), 25 + ) await c.pause(0.3) # -- 'y' on a LOCAL variable retypes that variable, not the prototype --- # fts = app.program.func_types(cf.addr) lv = next((v for v in (fts.lvars if fts else []) if not v.is_arg), None) if lv is not None: - line = next((i for i, t in enumerate(dec._texts) - if _word_occurrences(t, lv.name)), None) + line = next( + (i for i, t in enumerate(dec._texts) if _word_occurrences(t, lv.name)), None + ) if line is not None: col = _word_occurrences(dec._texts[line], lv.name)[0][0] dec.focus() @@ -2618,26 +3310,43 @@ async def s_retype(c: Ctx): await c.press("y") await c.wait(lambda: app.query_one("#retype", Input).display, 10) ri = app.query_one("#retype", Input) - c.check("'y' on a local variable prefills that variable's type", - ri.value == lv.type and lv.name in str(ri.placeholder), - f"val={ri.value!r} want={lv.type!r} ph={ri.placeholder!r}") + c.check( + "'y' on a local variable prefills that variable's type", + ri.value == lv.type and lv.name in str(ri.placeholder), + f"val={ri.value!r} want={lv.type!r} ph={ri.placeholder!r}", + ) ri.value = "unsigned __int64" await c.press("enter") - changed = await c.wait(lambda: (lambda f: bool(f) and any( - v.name == lv.name and v.type == "unsigned __int64" - for v in f.lvars))(app.program.func_types(cf.addr)), 25) - c.check("applying it retypes the local variable", changed, - f"{lv.name}: wanted unsigned __int64") + changed = await c.wait( + lambda: ( + lambda f: ( + bool(f) + and any( + v.name == lv.name and v.type == "unsigned __int64" + for v in f.lvars + ) + ) + )(app.program.func_types(cf.addr)), + 25, + ) + c.check( + "applying it retypes the local variable", + changed, + f"{lv.name}: wanted unsigned __int64", + ) after = app.program.func_types(cf.addr) - c.check("retyping a local leaves the prototype alone", - after is not None and after.prototype == old_proto, - f"proto={after.prototype if after else None!r}") + c.check( + "retyping a local leaves the prototype alone", + after is not None and after.prototype == old_proto, + f"proto={after.prototype if after else None!r}", + ) # -- 'y' on a GLOBAL retypes the global, not the enclosing function ----- # # The lvar retype above recompiled too — settle again, or the scan below # indexes into pseudocode that's about to be replaced. - await c.wait(lambda: not dec.loading and dec.loaded_ea == cf.addr - and bool(dec._texts), 25) + await c.wait( + lambda: not dec.loading and dec.loaded_ea == cf.addr and bool(dec._texts), 25 + ) await c.pause(0.3) # Pick a global that actually appears as a word in the pseudocode — a symbol @@ -2655,8 +3364,11 @@ async def s_retype(c: Ctx): if app.program.func_types(a) is not None: continue d = app.program.data_type(a) or {} - if (d.get("name") and not d.get("is_func") - and "(" not in (d.get("type") or "")): + if ( + d.get("name") + and not d.get("is_func") + and "(" not in (d.get("type") or "") + ): glob = (w, a, d, i) break if glob: @@ -2670,27 +3382,40 @@ async def s_retype(c: Ctx): dec.cursor, dec.cursor_x = line, col + 1 # inside the word dec.refresh() await c.pause(0.05) - c.check("the cursor sits on the global", - dec.word_under_cursor() == gname, - f"word={dec.word_under_cursor()!r} want={gname!r} " - f"line={line} col={col} text={dec._texts[line][:60]!r}") + c.check( + "the cursor sits on the global", + dec.word_under_cursor() == gname, + f"word={dec.word_under_cursor()!r} want={gname!r} " + f"line={line} col={col} text={dec._texts[line][:60]!r}", + ) await c.press("y") await c.wait(lambda: app.query_one("#retype", Input).display, 10) ri = app.query_one("#retype", Input) - c.check("'y' on a global prefills the global's type (not the proto)", - ri.value != old_proto and gname in str(ri.placeholder), - f"val={ri.value!r} ph={ri.placeholder!r}") + c.check( + "'y' on a global prefills the global's type (not the proto)", + ri.value != old_proto and gname in str(ri.placeholder), + f"val={ri.value!r} ph={ri.placeholder!r}", + ) ri.value = "unsigned __int64" await c.press("enter") retyped = await c.wait( - lambda: (app.program.data_type(glob.addr) or {}).get("type") - == "unsigned __int64", 25) - c.check("applying it retypes the global", retyped, - f"type={(app.program.data_type(glob.addr) or {}).get('type')!r}") + lambda: ( + (app.program.data_type(glob.addr) or {}).get("type") + == "unsigned __int64" + ), + 25, + ) + c.check( + "applying it retypes the global", + retyped, + f"type={(app.program.data_type(glob.addr) or {}).get('type')!r}", + ) after = app.program.func_types(cf.addr) - c.check("retyping a global leaves the prototype alone", - after is not None and after.prototype == old_proto, - f"proto={after.prototype if after else None!r}") + c.check( + "retyping a global leaves the prototype alone", + after is not None and after.prototype == old_proto, + f"proto={after.prototype if after else None!r}", + ) if dt.get("type"): # restore app.program.set_data_type(glob.addr, dt["type"]) @@ -2721,10 +3446,12 @@ async def s_scroll_restore(c: Ctx): await c.wait(lambda: app._cur.ea == fb.addr, 20) renders: list[int] = [] _orig_rl = dis.render_line + def _traced(y, _o=_orig_rl): if y == 0: renders.append(round(dis.scroll_offset.y)) return _o(y) + dis.render_line = _traced await c.press("escape") await c.wait(lambda: app._cur.ea == fa.addr, 20) @@ -2733,13 +3460,19 @@ async def s_scroll_restore(c: Ctx): # one happened. Wait for the paint we are actually asserting about (each # poll ticks the screen, so this is ~one frame, not a quarter second). await c.wait(lambda: bool(renders) and renders[-1] == want_sy, 5) - c.check("disasm scroll + cursor restored on back (mid-viewport)", - round(dis.scroll_offset.y) == want_sy and dis.cursor == want_cur and want_rel > 0, - f"scroll={round(dis.scroll_offset.y)} (want {want_sy}) " - f"cursor={dis.cursor} (want {want_cur}) rel_before={want_rel}") - c.check("pane is repainted at the restored scroll (no stale top frame)", - bool(renders) and renders[-1] == want_sy, - f"last repaint scroll={renders[-1] if renders else None} (want {want_sy})") + c.check( + "disasm scroll + cursor restored on back (mid-viewport)", + round(dis.scroll_offset.y) == want_sy + and dis.cursor == want_cur + and want_rel > 0, + f"scroll={round(dis.scroll_offset.y)} (want {want_sy}) " + f"cursor={dis.cursor} (want {want_cur}) rel_before={want_rel}", + ) + c.check( + "pane is repainted at the restored scroll (no stale top frame)", + bool(renders) and renders[-1] == want_sy, + f"last repaint scroll={renders[-1] if renders else None} (want {want_sy})", + ) dis.render_line = _orig_rl @@ -2759,14 +3492,18 @@ async def s_paging(c: Ctx): rel = dis.cursor - round(dis.scroll_offset.y) await c.press("pagedown") await c.pause(0.05) - c.check("PageDown preserves the viewport-relative row", - dis.cursor - round(dis.scroll_offset.y) == rel, - f"rel={dis.cursor - round(dis.scroll_offset.y)} want={rel}") + c.check( + "PageDown preserves the viewport-relative row", + dis.cursor - round(dis.scroll_offset.y) == rel, + f"rel={dis.cursor - round(dis.scroll_offset.y)} want={rel}", + ) await c.press("pageup") await c.pause(0.05) - c.check("PageUp preserves the viewport-relative row", - dis.cursor - round(dis.scroll_offset.y) == rel, - f"rel={dis.cursor - round(dis.scroll_offset.y)} want={rel}") + c.check( + "PageUp preserves the viewport-relative row", + dis.cursor - round(dis.scroll_offset.y) == rel, + f"rel={dis.cursor - round(dis.scroll_offset.y)} want={rel}", + ) @scenario("rename_history") @@ -2816,8 +3553,13 @@ async def s_rename_history(c: Ctx): await c.pause(0.1) app.query_one("#rename", Input).value = hnew await c.press("enter") - await c.wait(lambda: app._func_index.by_addr(htarget) - and app._func_index.by_addr(htarget).name == hnew, 25) + await c.wait( + lambda: ( + app._func_index.by_addr(htarget) + and app._func_index.by_addr(htarget).name == hnew + ), + 25, + ) if app._active != "listing": await c.press("tab") await c.press("escape") @@ -2826,14 +3568,22 @@ async def s_rename_history(c: Ctx): dis.model.lines(hrow, 4, prefetch=False) await c.pause(0.1) hline = dis._line_plain(hrow) - c.check("caller disasm shows renamed callee after 'back'", - hline is not None and hnew in hline, f"line={hline!r}") + c.check( + "caller disasm shows renamed callee after 'back'", + hline is not None and hnew in hline, + f"line={hline!r}", + ) await c.press("tab") await c.wait(lambda: dec.loaded_ea == bea, 25) 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.invoke("rename", batch={"func": {"addr": hex(htarget), "name": hsym}}) + 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( + remote_ops.rename, batch={"func": {"addr": hex(htarget), "name": hsym}} + ) @scenario("region_define") @@ -2850,37 +3600,59 @@ async def s_region_define(c: Ctx): # setup: undefine the whole function so [addr, addr+size) is a bare region c.prog.undefine(addr, size=size) c.prog.bump_items() - c.check("function removed by undefine", - c.prog.function_of(addr) is None, "still a function") + c.check( + "function removed by undefine", + c.prog.function_of(addr) is None, + "still a function", + ) # navigate there via the real 'g' prompt -> opens the flat LISTING view # (a non-function region), not refused await c.goto_ui(hex(addr)) await c.wait(lambda: app._cur is not None and app._cur.ea == addr, 25) - c.check("goto to a non-function address opens the listing view (not refused)", - app._cur is not None and app._cur.is_region - and app._active == "listing" and c.lst.display, - f"cur={app._cur} active={app._active} status={c.status()!r}") + c.check( + "goto to a non-function address opens the listing view (not refused)", + app._cur is not None + and app._cur.is_region + and app._active == "listing" + and c.lst.display, + f"cur={app._cur} active={app._active} status={c.status()!r}", + ) await c.wait(lambda: c.lst.total > 0 and c.lst._cursor_ea() is not None, 25) - c.check("listing renders heads and the cursor sits on the target address", - c.lst.total > 0 and c.lst._cursor_ea() == addr, - f"total={c.lst.total} cur_ea={c.lst._cursor_ea()}") + c.check( + "listing renders heads and the cursor sits on the target address", + c.lst.total > 0 and c.lst._cursor_ea() == addr, + f"total={c.lst.total} cur_ea={c.lst._cursor_ea()}", + ) # the flat listing spans the whole segment, not just this function seg = c.prog.segment_bounds(addr) - c.check("listing spans the whole segment (more heads than one function)", - seg is not None and c.lst.total > 1, f"total={c.lst.total} seg={seg}") + c.check( + "listing spans the whole segment (more heads than one function)", + seg is not None and c.lst.total > 1, + f"total={c.lst.total} seg={seg}", + ) # 'p' on the entry head (re)creates the function c.lst.focus() c.lst.cursor, c.lst.cursor_x = c.lst.model.index_of_ea(addr), 0 await c.pause(0.05) await c.press("p") - await c.wait(lambda: c.prog.function_of(addr) is not None - and app._cur is not None and not app._cur.is_region, 25) - c.check("'p' creates a function and upgrades the listing to a function view", - c.prog.function_of(addr) is not None and not app._cur.is_region - and app._cur.ea == addr and app._active in ("listing", "decomp"), - f"fn={c.prog.function_of(addr)} cur={app._cur} active={app._active}") + await c.wait( + lambda: ( + c.prog.function_of(addr) is not None + and app._cur is not None + and not app._cur.is_region + ), + 25, + ) + c.check( + "'p' creates a function and upgrades the listing to a function view", + c.prog.function_of(addr) is not None + and not app._cur.is_region + and app._cur.ea == addr + and app._active in ("listing", "decomp"), + f"fn={c.prog.function_of(addr)} cur={app._cur} active={app._active}", + ) finally: # idempotency: guarantee the function is back even if a check failed if c.prog.function_of(addr) is None: @@ -2916,46 +3688,83 @@ async def s_listing_view(c: Ctx): # `total > 0` is set from the segment's size before a single page has # materialised, so waiting on it and then reading rows was the suite's # one known flake (it failed roughly one run in three). Wait for a ROW. - await c.wait(lambda: app._cur is not None and app._active == "listing" - and c.lst.total > 0 and c.lst.model is not None - and any(h.kind == "data" for h in c.lst.model.window(0, 40)), 25) - c.check("navigating to a data segment opens the listing view", - app._active == "listing" and c.lst.display and c.lst.total > 0, - f"active={app._active} total={c.lst.total}") + await c.wait( + lambda: ( + app._cur is not None + and app._active == "listing" + and c.lst.total > 0 + and c.lst.model is not None + and any(h.kind == "data" for h in c.lst.model.window(0, 40)) + ), + 25, + ) + c.check( + "navigating to a data segment opens the listing view", + app._active == "listing" and c.lst.display and c.lst.total > 0, + f"active={app._active} total={c.lst.total}", + ) kinds = {h.kind for h in c.lst.model.window(0, 40)} c.check("listing shows data heads (not just code)", "data" in kinds, str(kinds)) # a rendered data line carries the item text (e.g. db/dd/string) c.lst.focus() - first_data = next((i for i in range(min(c.lst.total, 60)) - if c.lst.model.get(i) and c.lst.model.get(i).kind == "data"), None) - c.check("a data head exists in the first screenful", first_data is not None, - f"total={c.lst.total}") + first_data = next( + ( + i + for i in range(min(c.lst.total, 60)) + if c.lst.model.get(i) and c.lst.model.get(i).kind == "data" + ), + None, + ) + c.check( + "a data head exists in the first screenful", + first_data is not None, + f"total={c.lst.total}", + ) if first_data is not None: c.lst.cursor = first_data await c.pause(0.05) plain = c.lst._line_plain(first_data) - c.check("data line renders its item text", bool(plain and plain.strip()), - f"plain={plain!r}") - c.check("listing cursor reports the head address", - c.lst._cursor_ea() == c.lst.model.get(first_data).ea, str(c.lst._cursor_ea())) + c.check( + "data line renders its item text", + bool(plain and plain.strip()), + f"plain={plain!r}", + ) + c.check( + "listing cursor reports the head address", + c.lst._cursor_ea() == c.lst.model.get(first_data).ea, + str(c.lst._cursor_ea()), + ) # backslash from the listing opens hex at the cursor address; and back cur_ea = c.lst._cursor_ea() await c.press("backslash") await c.wait(lambda: app._active == "hex" and c.hex.display, 15) - c.check("backslash from the listing opens the hex view", app._active == "hex", - f"active={app._active}") + c.check( + "backslash from the listing opens the hex view", + app._active == "hex", + f"active={app._active}", + ) await c.press("backslash") await c.wait(lambda: app._active == "listing", 15) - c.check("returning from hex lands back on the listing (not a func view)", - app._active == "listing" and c.lst.display, f"active={app._active}") + c.check( + "returning from hex lands back on the listing (not a func view)", + app._active == "listing" and c.lst.display, + f"active={app._active}", + ) # 'd' defines typed data over an undefined run. Synthesize the run # deterministically: undefine a data head, then re-type it via the prompt. - dhead = next((c.lst.model.get(i) for i in range(min(c.lst.total, 200)) - if c.lst.model.get(i) and c.lst.model.get(i).kind == "data" - and (c.lst.model.get(i).size or 0) >= 4), None) + dhead = next( + ( + c.lst.model.get(i) + for i in range(min(c.lst.total, 200)) + if c.lst.model.get(i) + and c.lst.model.get(i).kind == "data" + and (c.lst.model.get(i).size or 0) >= 4 + ), + None, + ) if dhead is None: c.check("found a data head to re-type", False) return @@ -2971,31 +3780,49 @@ async def s_listing_view(c: Ctx): # win the race, and it started failing the moment page loads got bigger. stale = c.lst.model await c.goto_ui(hex(dea)) - await c.wait(lambda: app._active == "listing" and c.lst.total > 0 - and c.lst.model is not stale - and c.lst.model.index_of_ea(dea) >= 0, 25) + await c.wait( + lambda: ( + app._active == "listing" + and c.lst.total > 0 + and c.lst.model is not stale + and c.lst.model.index_of_ea(dea) >= 0 + ), + 25, + ) ui = c.lst.model.index_of_ea(dea) - c.check("undefining a data head yields an unknown run in the listing", - ui >= 0 and c.lst.model.get(ui).kind == "unknown", - f"kind={c.lst.model.get(ui).kind if ui>=0 else None}") + c.check( + "undefining a data head yields an unknown run in the listing", + ui >= 0 and c.lst.model.get(ui).kind == "unknown", + f"kind={c.lst.model.get(ui).kind if ui >= 0 else None}", + ) c.lst.focus() c.lst.cursor = ui await c.pause(0.05) await c.press("d") await c.pause(0.1) mdi = app.query_one("#makedata", Input) - c.check("'d' opens the make-data prompt prefilled with a type", - mdi.display and bool(mdi.value), f"display={mdi.display} val={mdi.value!r}") + c.check( + "'d' opens the make-data prompt prefilled with a type", + mdi.display and bool(mdi.value), + f"display={mdi.display} val={mdi.value!r}", + ) mdi.value = "char[4]" await c.press("enter") - await c.wait(lambda: app._active == "listing" - and c.lst.model.index_of_ea(dea) >= 0 - and c.lst.model.get(c.lst.model.index_of_ea(dea)) is not None - and c.lst.model.get(c.lst.model.index_of_ea(dea)).kind == "data", 25) + await c.wait( + lambda: ( + app._active == "listing" + and c.lst.model.index_of_ea(dea) >= 0 + and c.lst.model.get(c.lst.model.index_of_ea(dea)) is not None + and c.lst.model.get(c.lst.model.index_of_ea(dea)).kind == "data" + ), + 25, + ) di = c.lst.model.index_of_ea(dea) - c.check("'d' turns the undefined run into a typed data item", - di >= 0 and c.lst.model.get(di).kind == "data", - f"kind={c.lst.model.get(di).kind if di>=0 else None}") + c.check( + "'d' turns the undefined run into a typed data item", + di >= 0 and c.lst.model.get(di).kind == "data", + f"kind={c.lst.model.get(di).kind if di >= 0 else None}", + ) finally: c.prog.bump_items() @@ -3021,8 +3848,11 @@ async def s_listing_name_addr(c: Ctx): if data_ea is None: c.check("found a data segment with a >=2-byte item", False) return - dh = next(h for h in c.prog.listing(data_ea).window(0, 60) - if h.kind == "data" and (h.size or 0) >= 2) + dh = next( + h + for h in c.prog.listing(data_ea).window(0, 60) + if h.kind == "data" and (h.size or 0) >= 2 + ) A = dh.ea newname = f"after_{os.getpid()}" try: @@ -3030,33 +3860,53 @@ async def s_listing_name_addr(c: Ctx): c.prog.make_data(A, "unsigned __int8") c.prog.bump_items() await c.goto_ui(hex(A + 1)) - await c.wait(lambda: app._active == "listing" and app._cur is not None - and app._cur.ea == A + 1, 25) + await c.wait( + lambda: ( + app._active == "listing" + and app._cur is not None + and app._cur.ea == A + 1 + ), + 25, + ) head = c.lst.cur_head() - c.check("cursor lands on the now-undefined byte at addr+1", - head is not None and head.ea == A + 1 and head.kind == "unknown", - f"head={head}") + c.check( + "cursor lands on the now-undefined byte at addr+1", + head is not None and head.ea == A + 1 and head.kind == "unknown", + f"head={head}", + ) # 'n' opens the address-name prompt (even though there's no symbol) c.lst.focus() await c.press("n") await c.pause(0.1) ri = app.query_one("#rename", Input) - c.check("'n' opens the name prompt on an unnamed byte", - ri.display, f"display={ri.display}") + c.check( + "'n' opens the name prompt on an unnamed byte", + ri.display, + f"display={ri.display}", + ) ri.value = newname await c.press("enter") - await c.wait(lambda: app._active == "listing" - and c.lst.model.index_of_ea(A + 1) >= 0 - and c.lst.model.get(c.lst.model.index_of_ea(A + 1)) is not None - and c.lst.model.get(c.lst.model.index_of_ea(A + 1)).name == newname, 25) + await c.wait( + lambda: ( + app._active == "listing" + and c.lst.model.index_of_ea(A + 1) >= 0 + and c.lst.model.get(c.lst.model.index_of_ea(A + 1)) is not None + and c.lst.model.get(c.lst.model.index_of_ea(A + 1)).name == newname + ), + 25, + ) hi = c.lst.model.index_of_ea(A + 1) - c.check("naming a bare byte at addr+1 sticks", - hi >= 0 and c.lst.model.get(hi).name == newname, - f"name={c.lst.model.get(hi).name if hi>=0 else None}") + c.check( + "naming a bare byte at addr+1 sticks", + hi >= 0 and c.lst.model.get(hi).name == newname, + f"name={c.lst.model.get(hi).name if hi >= 0 else None}", + ) finally: # revert: drop the label and restore raw bytes at A try: - c.prog.client.invoke("rename", batch={"data": {"addr": hex(A + 1), "new": ""}}) + c.prog.client.call( + remote_ops.rename, batch={"data": {"addr": hex(A + 1), "new": ""}} + ) except Exception: # noqa: BLE001 pass c.prog.undefine(A, size=8) @@ -3075,8 +3925,9 @@ async def s_listing_make_string(c: Ctx): if lm is None: continue lm.ensure(80) - h = next((h for h in lm.window(0, 80) - if h.kind == "data" and "'" in h.text), None) + h = next( + (h for h in lm.window(0, 80) if h.kind == "data" and "'" in h.text), None + ) if h is not None: target = h.ea break @@ -3088,22 +3939,36 @@ async def s_listing_make_string(c: Ctx): c.prog.undefine(A, size=8) c.prog.bump_items() await c.goto_ui(hex(A)) - await c.wait(lambda: app._active == "listing" and app._cur is not None - and app._cur.ea == A, 25) - c.check("target is undefined before 'a'", - c.lst.cur_head() is not None and c.lst.cur_head().kind == "unknown", - f"head={c.lst.cur_head()}") + await c.wait( + lambda: ( + app._active == "listing" and app._cur is not None and app._cur.ea == A + ), + 25, + ) + c.check( + "target is undefined before 'a'", + c.lst.cur_head() is not None and c.lst.cur_head().kind == "unknown", + f"head={c.lst.cur_head()}", + ) c.lst.focus() await c.press("a") - await c.wait(lambda: c.lst.model.index_of_ea(A) >= 0 - and c.lst.model.get(c.lst.model.index_of_ea(A)) is not None - and c.lst.model.get(c.lst.model.index_of_ea(A)).kind == "data" - and "'" in c.lst.model.get(c.lst.model.index_of_ea(A)).text, 25) + await c.wait( + lambda: ( + c.lst.model.index_of_ea(A) >= 0 + and c.lst.model.get(c.lst.model.index_of_ea(A)) is not None + and c.lst.model.get(c.lst.model.index_of_ea(A)).kind == "data" + and "'" in c.lst.model.get(c.lst.model.index_of_ea(A)).text + ), + 25, + ) hi = c.lst.model.index_of_ea(A) - c.check("'a' creates a string literal at the cursor", - hi >= 0 and c.lst.model.get(hi).kind == "data" - and "'" in c.lst.model.get(hi).text, - f"head={c.lst.model.get(hi) if hi >= 0 else None}") + c.check( + "'a' creates a string literal at the cursor", + hi >= 0 + and c.lst.model.get(hi).kind == "data" + and "'" in c.lst.model.get(hi).text, + f"head={c.lst.model.get(hi) if hi >= 0 else None}", + ) finally: try: c.prog.make_string(A) # restore the original string @@ -3133,25 +3998,38 @@ async def s_listing_struct_expand(c: Ctx): c.check("found a data address for the struct test", False) return try: - c.prog.client.invoke( - "declare_type", - decls=["struct TuiExpandS { int a; char b[4]; short c; };"]) + c.prog.client.call( + remote_ops.declare_type, + decls=["struct TuiExpandS { int a; char b[4]; short c; };"], + ) c.prog.make_data(A, "TuiExpandS") c.prog.bump_items() await c.goto_ui(hex(A)) - await c.wait(lambda: app._active == "listing" and app._cur is not None - and app._cur.ea == A and c.lst.total > 0, 25) + await c.wait( + lambda: ( + app._active == "listing" + and app._cur is not None + and app._cur.ea == A + and c.lst.total > 0 + ), + 25, + ) # the summary head, then member rows for a/b/c si = c.lst.model.index_of_ea(A) members = [c.lst.model.get(si + 1 + k) for k in range(3)] names = [m.text for m in members if m is not None] - c.check("struct global expands into member rows", - all(m is not None and m.kind == "member" for m in members) - and any("a" in t for t in names) and any("b" in t for t in names), - f"members={names}") - c.check("member rows carry field addresses", - members[1] is not None and members[1].ea == A + 4, - f"ea={members[1].ea if members[1] else None:#x} want={A+4:#x}") + c.check( + "struct global expands into member rows", + all(m is not None and m.kind == "member" for m in members) + and any("a" in t for t in names) + and any("b" in t for t in names), + f"members={names}", + ) + c.check( + "member rows carry field addresses", + members[1] is not None and members[1].ea == A + 4, + f"ea={members[1].ea if members[1] else None:#x} want={A + 4:#x}", + ) finally: try: c.prog.undefine(A, size=16) @@ -3169,45 +4047,66 @@ async def s_continuous_view(c: Ctx): fn = await c.open_biggest("listing") fn_ea = fn.addr await c.wait(lambda: app._active == "listing" and c.lst.display, 10) - c.check("a function opens in the continuous listing by default", - app._active == "listing" and c.lst.display - and c.lst._cursor_ea() == fn_ea, - f"active={app._active} disp={c.lst.display} cur_ea={c.lst._cursor_ea()}") + c.check( + "a function opens in the continuous listing by default", + app._active == "listing" and c.lst.display and c.lst._cursor_ea() == fn_ea, + f"active={app._active} disp={c.lst.display} cur_ea={c.lst._cursor_ea()}", + ) # the listing spans the whole segment, not just the function c.lst.model.load_all() seg = c.prog.segment_bounds(fn_ea) seg_rows = len(c.lst.model) # a function's own instruction count is far smaller than the segment fdis = c.prog.disasm(fn_ea, fn.name) - c.check("the continuous listing extends past the function's bounds", - seg_rows > fdis.total(), f"listing={seg_rows} func={fdis.total()}") + c.check( + "the continuous listing extends past the function's bounds", + seg_rows > fdis.total(), + f"listing={seg_rows} func={fdis.total()}", + ) kinds = {c.lst.model.get(i).kind for i in range(seg_rows)} - c.check("continuous listing interleaves code with data/undefined", - "code" in kinds and ("data" in kinds or "unknown" in kinds), str(kinds)) + c.check( + "continuous listing interleaves code with data/undefined", + "code" in kinds and ("data" in kinds or "unknown" in kinds), + str(kinds), + ) # rendering parity with disasm: code lines carry opcode bytes - cidx = next((i for i in range(seg_rows) - if c.lst.model.get(i).kind == "code"), None) - c.check("continuous listing renders opcode bytes (parity with disasm)", - cidx is not None and c.lst.model.get(cidx).raw - and c.lst._op_field(c.lst.model.get(cidx)).strip() != "", - f"raw={c.lst.model.get(cidx).raw if cidx is not None else None!r}") + cidx = next((i for i in range(seg_rows) if c.lst.model.get(i).kind == "code"), None) + c.check( + "continuous listing renders opcode bytes (parity with disasm)", + cidx is not None + and c.lst.model.get(cidx).raw + and c.lst._op_field(c.lst.model.get(cidx)).strip() != "", + f"raw={c.lst.model.get(cidx).raw if cidx is not None else None!r}", + ) # F5/Tab at the function -> decompiler, and back to the same spot c.lst.focus() await c.press("tab") - await c.wait(lambda: (app._active == "decomp" and c.dec.loaded_ea == fn_ea) - or (app.is_listing - and _CANNOT_DECOMP in c.status().lower()), 25) + await c.wait( + lambda: ( + (app._active == "decomp" and c.dec.loaded_ea == fn_ea) + or (app.is_listing and _CANNOT_DECOMP in c.status().lower()) + ), + 25, + ) if app._active == "decomp": - c.check("F5/Tab decompiles the function under the cursor", - c.dec.loaded_ea == fn_ea, f"loaded={c.dec.loaded_ea}") + c.check( + "F5/Tab decompiles the function under the cursor", + c.dec.loaded_ea == fn_ea, + f"loaded={c.dec.loaded_ea}", + ) await c.press("tab") await c.wait(lambda: app._active == "listing", 25) - c.check("F5/Tab in the decompiler returns to the listing at the same ea", - app._active == "listing" and c.lst._cursor_ea() == fn_ea, - f"active={app._active} cur_ea={c.lst._cursor_ea()}") + c.check( + "F5/Tab in the decompiler returns to the listing at the same ea", + app._active == "listing" and c.lst._cursor_ea() == fn_ea, + f"active={app._active} cur_ea={c.lst._cursor_ea()}", + ) else: - c.check("undecompilable function falls back to the listing", - app._active == "listing", f"active={app._active}") + c.check( + "undecompilable function falls back to the listing", + app._active == "listing", + f"active={app._active}", + ) @scenario("func_banners") @@ -3220,18 +4119,27 @@ async def s_func_banners(c: Ctx): c.lst.model.load_all() heads = [c.lst.model.get(i) for i in range(len(c.lst.model))] ci = c.lst.model.index_of_ea(fn.addr) - c.check("navigation to a function lands on its code head, not a banner", - ci >= 0 and c.lst.model.get(ci).kind == "code", - f"kind={c.lst.model.get(ci).kind if ci >= 0 else None}") - c.check("a SUBROUTINE separator banner is present", - any(h.kind == "sep" and "S U B R O U T I N E" in h.text for h in heads)) - c.check("a 'name proc' header is present", - any(h.kind == "funchdr" and h.text.endswith(" proc") for h in heads)) - c.check("a 'name endp' footer is present", - any(h.kind == "funchdr" and h.text.endswith("endp") for h in heads)) + c.check( + "navigation to a function lands on its code head, not a banner", + ci >= 0 and c.lst.model.get(ci).kind == "code", + f"kind={c.lst.model.get(ci).kind if ci >= 0 else None}", + ) + c.check( + "a SUBROUTINE separator banner is present", + any(h.kind == "sep" and "S U B R O U T I N E" in h.text for h in heads), + ) + c.check( + "a 'name proc' header is present", + any(h.kind == "funchdr" and h.text.endswith(" proc") for h in heads), + ) + c.check( + "a 'name endp' footer is present", + any(h.kind == "funchdr" and h.text.endswith("endp") for h in heads), + ) # the proc header for the origin function carries its name - hdr = next((h for h in heads if h.kind == "funchdr" - and h.text == f"{fn.name} proc"), None) + hdr = next( + (h for h in heads if h.kind == "funchdr" and h.text == f"{fn.name} proc"), None + ) c.check("the proc header names the function", hdr is not None, f"fn={fn.name}") @@ -3296,35 +4204,58 @@ async def s_opfmt_listing(c: Ctx): ea = head.ea c.lst.focus() parked = await _park_on(c, ea) - c.check("the cursor is on the literal's line", parked, - f"want {ea:#x}, cursor at {c.lst._cursor_ea():#x}") + c.check( + "the cursor is on the literal's line", + parked, + f"want {ea:#x}, cursor at {c.lst._cursor_ea():#x}", + ) before = head.text try: await c.press("o") - await c.wait(lambda: c.lst.model.index_of_ea(ea) >= 0 - and (c.lst.model.get(c.lst.model.index_of_ea(ea)) or head).text - != before, 25) + await c.wait( + lambda: ( + c.lst.model.index_of_ea(ea) >= 0 + and (c.lst.model.get(c.lst.model.index_of_ea(ea)) or head).text + != before + ), + 25, + ) i = c.lst.model.index_of_ea(ea) after = c.lst.model.get(i).text if i >= 0 else before - c.check("'o' re-renders the literal", after != before, - f"{before!r} -> {after!r} ea={ea:#x} show={show}") - c.check("the status names the format it moved to", - any(f in c.status() for f in show["choices"]), - f"status={c.status()!r} choices={show['choices']}") + c.check( + "'o' re-renders the literal", + after != before, + f"{before!r} -> {after!r} ea={ea:#x} show={show}", + ) + c.check( + "the status names the format it moved to", + any(f in c.status() for f in show["choices"]), + f"status={c.status()!r} choices={show['choices']}", + ) c.check("the change is marked unsaved", app._dirty) # 'O' walks the ring the other way: back to where we started. await c.press("O") - await c.wait(lambda: c.lst.model.index_of_ea(ea) >= 0 - and (c.lst.model.get(c.lst.model.index_of_ea(ea)) or head).text - == before, 25) + await c.wait( + lambda: ( + c.lst.model.index_of_ea(ea) >= 0 + and (c.lst.model.get(c.lst.model.index_of_ea(ea)) or head).text + == before + ), + 25, + ) j = c.lst.model.index_of_ea(ea) - c.check("'O' cycles back", j >= 0 and c.lst.model.get(j).text == before, - f"{c.lst.model.get(j).text if j >= 0 else None!r} want {before!r}") + c.check( + "'O' cycles back", + j >= 0 and c.lst.model.get(j).text == before, + f"{c.lst.model.get(j).text if j >= 0 else None!r} want {before!r}", + ) # An explicit format by name (what the palette/RPC use). r = c.prog.op_format(ea, mode="dec") - c.check("an explicit format renders decimal", - r["format"] == "dec" and str(int(show["value"], 16)) in r["text"], - str(r)) + c.check( + "an explicit format renders decimal", + r["format"] == "dec" and str(int(show["value"], 16)) in r["text"], + str(r), + ) finally: try: c.prog.op_format(ea, mode="default", n=show.get("n", -1)) @@ -3343,24 +4274,33 @@ async def s_opfmt_no_literal(c: Ctx): h = c.lst.model.get(i) return h is not None and h.kind in ("sep", "funchdr", "label") - row = next((i for i in range(c.lst.cursor, min(c.lst.cursor + 400, - len(c.lst.model))) - if _banner(i)), None) + row = next( + ( + i + for i in range(c.lst.cursor, min(c.lst.cursor + 400, len(c.lst.model))) + if _banner(i) + ), + None, + ) if row is None: c.check("found a banner row", False) return c.lst.focus() - for _ in range(20): # hold it against a late-landing open + for _ in range(20): # hold it against a late-landing open c.lst.cursor = row await c.pause(0.05) if c.lst.cursor == row: break - c.check("the cursor is on a banner row", _banner(c.lst.cursor), - f"row={c.lst.cursor}") + c.check( + "the cursor is on a banner row", _banner(c.lst.cursor), f"row={c.lst.cursor}" + ) await c.press("o") await c.wait(lambda: "reformat" in c.status() or "format" in c.status(), 15) - c.check("'o' on a line with no literal explains itself", - "reformat" in c.status(), f"status={c.status()!r}") + c.check( + "'o' on a line with no literal explains itself", + "reformat" in c.status(), + f"status={c.status()!r}", + ) @scenario("opfmt_refusal_is_not_swallowed") @@ -3388,7 +4328,7 @@ async def s_opfmt_refusal_visible(c: Ctx): c.check("the cursor is on the literal's line", False) return try: - await c.press("o") # a success: sets the flash + await c.press("o") # a success: sets the flash await c.wait(lambda: "\u2192" in c.status(), 25) good = c.status() @@ -3401,9 +4341,14 @@ async def s_opfmt_refusal_visible(c: Ctx): h = c.lst.model.get(i) return h is not None and h.kind in ("sep", "funchdr", "label") - row = next((i for i in range(c.lst.cursor, - min(c.lst.cursor + 400, len(c.lst.model))) - if _banner(i)), None) + row = next( + ( + i + for i in range(c.lst.cursor, min(c.lst.cursor + 400, len(c.lst.model))) + if _banner(i) + ), + None, + ) if row is None: c.check("found a banner row to refuse on", False) return @@ -3412,11 +4357,13 @@ async def s_opfmt_refusal_visible(c: Ctx): await c.pause(0.05) if c.lst.cursor == row: break - c.lst.action_op_format("cycle") # no keypress: as the RPC does it + c.lst.action_op_format("cycle") # no keypress: as the RPC does it await c.wait(lambda: c.status() != good, 15) - c.check("a refusal replaces the previous success on the status bar", - c.status() != good and "reformat" in c.status(), - f"still showing {c.status()!r}") + c.check( + "a refusal replaces the previous success on the status bar", + c.status() != good and "reformat" in c.status(), + f"still showing {c.status()!r}", + ) finally: try: c.prog.op_format(head.ea, mode="default", n=show.get("n", -1)) @@ -3430,8 +4377,11 @@ def _styled_cols(strip, style_attr, want): cols, x = [], 0 for seg in strip: st = seg.style - if st is not None and getattr(st, style_attr, None) is not None \ - and str(getattr(st, style_attr)) == want: + if ( + st is not None + and getattr(st, style_attr, None) is not None + and str(getattr(st, style_attr)) == want + ): cols.extend(range(x, x + len(seg.text))) x += len(seg.text) return cols @@ -3447,15 +4397,22 @@ async def s_opfmt_highlight(c: Ctx): (usually the same characters) and used to win. """ from idatui.app import _S_OPERAND + await c.open_biggest("listing") c.lst.model.load_all() lst = c.lst lst.focus() # A row with two operands, so "which one" is a real question. - row = next((i for i in range(lst.cursor, min(lst.cursor + 400, len(lst.model))) - if lst.model.get(i) is not None - and (lst.model.get(i).ops or ()) and len(lst.model.get(i).ops) >= 2), - None) + row = next( + ( + i + for i in range(lst.cursor, min(lst.cursor + 400, len(lst.model))) + if lst.model.get(i) is not None + and (lst.model.get(i).ops or ()) + and len(lst.model.get(i).ops) >= 2 + ), + None, + ) if row is None: c.check("found a row with two operands", False) return @@ -3474,12 +4431,17 @@ async def s_opfmt_highlight(c: Ctx): strip = lst.render_line(row - round(lst.scroll_offset.y)) cols = _styled_cols(strip, "bgcolor", want_bg) seen.append((n, min(cols) if cols else None, max(cols) + 1 if cols else None)) - c.check(f"operand {n} ({h.text[lo:hi]!r}) is marked when the cursor is on it", - cols and min(cols) == base + lo and max(cols) + 1 == base + hi, - f"marked={min(cols) if cols else None}.." - f"{max(cols)+1 if cols else None} want={base+lo}..{base+hi}") - c.check("the mark MOVES between the operands (it isn't the whole line)", - len({s[1] for s in seen}) == len(seen), str(seen)) + c.check( + f"operand {n} ({h.text[lo:hi]!r}) is marked when the cursor is on it", + cols and min(cols) == base + lo and max(cols) + 1 == base + hi, + f"marked={min(cols) if cols else None}.." + f"{max(cols) + 1 if cols else None} want={base + lo}..{base + hi}", + ) + c.check( + "the mark MOVES between the operands (it isn't the whole line)", + len({s[1] for s in seen}) == len(seen), + str(seen), + ) # ... and the marked operand is the one the edit acts on — either it gets # reformatted, or the refusal names that same operand. What must never # happen is a different operand quietly changing. @@ -3488,11 +4450,14 @@ async def s_opfmt_highlight(c: Ctx): try: r = c.prog.op_format(h.ea, mode="show", col=lst.op_col()) got, why = r.get("n"), "" - except IDAToolError as e: # "operand N (rsp) has no format" + except IDAToolError as e: # "operand N (rsp) has no format" m = re.search(r"operand (\d+)", e.message) got, why = (int(m.group(1)) if m else None), e.message - c.check(f"marked operand {n} is the one acted on (or refused)", - got == n, f"marked op{n}, worker said op{got} {why}") + c.check( + f"marked operand {n} is the one acted on (or refused)", + got == n, + f"marked op{n}, worker said op{got} {why}", + ) @scenario("opfmt_sticks_to_its_literal") @@ -3515,8 +4480,7 @@ async def s_opfmt_sticks(c: Ctx): # The second literal must be one whose printed WIDTH changes as it # cycles (0x36u vs 54), or the cursor never falls off it and the # test proves nothing. - if len(recs) >= 2 and recs[0][1] < recs[1][0] \ - and int(recs[1][2], 16) >= 16: + if len(recs) >= 2 and recs[0][1] < recs[1][0] and int(recs[1][2], 16) >= 16: pick = (fn, line, recs) break if pick: @@ -3526,7 +4490,7 @@ async def s_opfmt_sticks(c: Ctx): return fn, line, recs = pick first, second = recs[0], recs[1] - target = (second[3], second[4]) # (ea, opnum) of the literal we mean + target = (second[3], second[4]) # (ea, opnum) of the literal we mean other = (first[3], first[4]) try: # Make it WIDE first (0x30, four characters). The cursor then sits on a @@ -3541,33 +4505,47 @@ async def s_opfmt_sticks(c: Ctx): return dec = c.dec dec.focus() - wide = next((r for r in dec._nums.get(line, ()) - if (r[3], r[4]) == target), None) + wide = next( + (r for r in dec._nums.get(line, ()) if (r[3], r[4]) == target), None + ) if wide is None or wide[1] - wide[0] < 3: - c.check("the literal is now printed wide", False, - f"nums={dec._nums.get(line)}") + c.check( + "the literal is now printed wide", False, f"nums={dec._nums.get(line)}" + ) return - dec.cursor, dec.cursor_x = line, wide[1] - 1 # its LAST character + dec.cursor, dec.cursor_x = line, wide[1] - 1 # its LAST character dec.refresh() await c.pause(0.05) seen = [] for _ in range(3): before = dec._texts[line] await c.press("o") - await c.wait(lambda: dec.loaded_ea == fn.addr - and line < len(dec._texts) - and dec._texts[line] != before, 30) - cur = next(((r[3], r[4]) for r in dec._nums.get(line, ()) - if r[0] <= dec.cursor_x < r[1]), None) + await c.wait( + lambda: ( + dec.loaded_ea == fn.addr + and line < len(dec._texts) + and dec._texts[line] != before + ), + 30, + ) + cur = next( + ( + (r[3], r[4]) + for r in dec._nums.get(line, ()) + if r[0] <= dec.cursor_x < r[1] + ), + None, + ) seen.append(cur) - c.check("every press stays on the literal we started on", - all(s == target for s in seen), - f"target={target} other={other} landed={seen} " - f"line={dec._texts[line].strip()!r}") + c.check( + "every press stays on the literal we started on", + all(s == target for s in seen), + f"target={target} other={other} landed={seen} " + f"line={dec._texts[line].strip()!r}", + ) finally: try: - c.prog.pc_num_format(fn.addr, mode="default", line=line, - col=second[0]) + c.prog.pc_num_format(fn.addr, mode="default", line=line, col=second[0]) except Exception: # noqa: BLE001 pass c.prog.bump_names() @@ -3583,6 +4561,7 @@ async def s_cursor_on_visible(c: Ctx): that hadn't changed. """ from idatui.rpc import cursor_on + app = c.app fn = await c.open_biggest("listing") c.lst.model.load_all() @@ -3591,44 +4570,66 @@ async def s_cursor_on_visible(c: Ctx): await c.pause(0.1) top = round(lst.scroll_offset.y) # A token that occurs both before the viewport and inside it. - here = next((t for t in ("rax", "rsp", "eax", "rbp", "rdi") - if any(t in (lst._line_plain(i) or "") - for i in range(top, min(top + 20, lst.total))) - and any(t in (lst._line_plain(i) or "") for i in range(0, top))), - None) + here = next( + ( + t + for t in ("rax", "rsp", "eax", "rbp", "rdi") + if any( + t in (lst._line_plain(i) or "") + for i in range(top, min(top + 20, lst.total)) + ) + and any(t in (lst._line_plain(i) or "") for i in range(0, top)) + ), + None, + ) if here is None: - c.check("found a token both above and inside the viewport", False, - f"top={top}") + c.check("found a token both above and inside the viewport", False, f"top={top}") return found = cursor_on(app, here) await c.pause(0.1) c.check(f"cursor_on({here!r}) found it", found) vis = round(lst.scroll_offset.y) - c.check("it lands inside the viewport, not thousands of rows above", - vis <= lst.cursor < vis + lst._visible_height(), - f"cursor={lst.cursor} viewport={vis}..{vis + lst._visible_height()}") - c.check("and it searched from the viewport, not from row 0", - lst.cursor >= top, f"cursor={lst.cursor} was top={top}") + c.check( + "it lands inside the viewport, not thousands of rows above", + vis <= lst.cursor < vis + lst._visible_height(), + f"cursor={lst.cursor} viewport={vis}..{vis + lst._visible_height()}", + ) + c.check( + "and it searched from the viewport, not from row 0", + lst.cursor >= top, + f"cursor={lst.cursor} was top={top}", + ) # An explicit line still wins, and lands visibly. - far = next((i for i in range(0, min(top, lst.total)) - if here in (lst._line_plain(i) or "")), None) + far = next( + ( + i + for i in range(0, min(top, lst.total)) + if here in (lst._line_plain(i) or "") + ), + None, + ) if far is not None: cursor_on(app, here, line=far) await c.pause(0.1) v2 = round(lst.scroll_offset.y) - c.check("an explicit line is honoured AND scrolled into view", - lst.cursor == far and v2 <= far < v2 + lst._visible_height(), - f"cursor={lst.cursor} want={far} viewport={v2}") + c.check( + "an explicit line is honoured AND scrolled into view", + lst.cursor == far and v2 <= far < v2 + lst._visible_height(), + f"cursor={lst.cursor} want={far} viewport={v2}", + ) # The `cursor` verb has the same duty: a driver that parks the cursor for # an edit must leave it where the edit can be watched. from idatui.rpc import place_cursor + deep = min(lst.total - 1, 900) place_cursor(lst, deep, 0) await c.pause(0.1) v3 = round(lst.scroll_offset.y) - c.check("`cursor line=` scrolls to what it selected", - v3 <= deep < v3 + lst._visible_height(), - f"cursor={lst.cursor} viewport={v3}..{v3 + lst._visible_height()}") + c.check( + "`cursor line=` scrolls to what it selected", + v3 <= deep < v3 + lst._visible_height(), + f"cursor={lst.cursor} viewport={v3}..{v3 + lst._visible_height()}", + ) @scenario("opfmt_decomp") @@ -3644,7 +4645,7 @@ async def s_opfmt_decomp(c: Ctx): continue for i, txt in enumerate(d.code.split("\n")): m = re.search(r"[=<>+\-*/(,]\s(\d{2,}|0x[0-9A-Fa-f]{2,})\b", txt) - if m and "//" not in txt[:m.start()]: + if m and "//" not in txt[: m.start()]: pick = (fn, i, m.start(1)) break if pick: @@ -3665,14 +4666,24 @@ async def s_opfmt_decomp(c: Ctx): before = dec._texts[line] try: await c.press("o") - await c.wait(lambda: dec.loaded_ea == fn.addr and line < len(dec._texts) - and dec._texts[line] != before, 30) - c.check("'o' re-renders the pseudocode literal", - line < len(dec._texts) and dec._texts[line] != before, - f"{before!r} -> {dec._texts[line] if line < len(dec._texts) else None!r}") - c.check("the status says which format", - any(f in c.status() for f in ("hex", "dec", "oct", "char", "default")), - f"status={c.status()!r}") + await c.wait( + lambda: ( + dec.loaded_ea == fn.addr + and line < len(dec._texts) + and dec._texts[line] != before + ), + 30, + ) + c.check( + "'o' re-renders the pseudocode literal", + line < len(dec._texts) and dec._texts[line] != before, + f"{before!r} -> {dec._texts[line] if line < len(dec._texts) else None!r}", + ) + c.check( + "the status says which format", + any(f in c.status() for f in ("hex", "dec", "oct", "char", "default")), + f"status={c.status()!r}", + ) finally: try: c.prog.pc_num_format(fn.addr, mode="default", line=line, col=col) @@ -3714,10 +4725,13 @@ async def _open_graph(c: Ctx, fn=None, t=60): await c.press("space") gv = app.query_one(GraphView) ok = await c.wait(lambda: app._active == "graph" and gv.lay is not None, t) - c.check("the graph opened", ok, - f"active={app._active} sticky={app._graph_sticky} " - f"focus={type(app.focused).__name__} prompt={app._prompt_active()} " - f"status={c.status()!r}") + c.check( + "the graph opened", + ok, + f"active={app._active} sticky={app._graph_sticky} " + f"focus={type(app.focused).__name__} prompt={app._prompt_active()} " + f"status={c.status()!r}", + ) return fn, gv @@ -3725,35 +4739,58 @@ async def _open_graph(c: Ctx, fn=None, t=60): async def s_graph_open(c: Ctx): app = c.app fn, gv = await _open_graph(c) - c.check("space opens the graph view", app._active == "graph", - f"active={app._active} status={c.status()}") + c.check( + "space opens the graph view", + app._active == "graph", + f"active={app._active} status={c.status()}", + ) if gv.lay is None: return - c.check("the graph has the function's blocks", - len(gv.lay.nodes) == len(gv.fc.blocks) and len(gv.lay.nodes) > 1, - f"nodes={len(gv.lay.nodes)} blocks={len(gv.fc.blocks) if gv.fc else 0}") - c.check("it is the right function", gv.fc is not None and gv.fc.func_ea == fn.addr, - f"{gv.fc.func_ea if gv.fc else None:#x} want {fn.addr:#x}") + c.check( + "the graph has the function's blocks", + len(gv.lay.nodes) == len(gv.fc.blocks) and len(gv.lay.nodes) > 1, + f"nodes={len(gv.lay.nodes)} blocks={len(gv.fc.blocks) if gv.fc else 0}", + ) + c.check( + "it is the right function", + gv.fc is not None and gv.fc.func_ea == fn.addr, + f"{gv.fc.func_ea if gv.fc else None:#x} want {fn.addr:#x}", + ) c.check("the cursor starts on a real address", gv._cursor_ea() is not None) # The invariant the whole dummy-node machinery exists for. boxes = [(n.x, n.y, n.right, n.y + n.h - 1) for n in gv.lay.nodes] - overlap = any(a[0] <= b[2] and b[0] <= a[2] and a[1] <= b[3] and b[1] <= a[3] - for i, a in enumerate(boxes) for b in boxes[i + 1:]) + overlap = any( + a[0] <= b[2] and b[0] <= a[2] and a[1] <= b[3] and b[1] <= a[3] + for i, a in enumerate(boxes) + for b in boxes[i + 1 :] + ) c.check("no two blocks overlap", not overlap) c.check("the status names the graph", "graph" in c.status(), c.status()) old_fc = gv.fc old_ea = gv._cursor_ea() await c.press("ctrl+r") rebuilt = await c.wait( - lambda: app.is_graph and gv.fc is not None and gv.fc is not old_fc - and gv.lay is not None, 30) + lambda: ( + app.is_graph + and gv.fc is not None + and gv.fc is not old_fc + and gv.lay is not None + ), + 30, + ) c.check("Ctrl+R rebuilds the graph", rebuilt) - c.check("Ctrl+R preserves the graph cursor", gv._cursor_ea() == old_ea, - f"got={gv._cursor_ea()} want={old_ea}") + c.check( + "Ctrl+R preserves the graph cursor", + gv._cursor_ea() == old_ea, + f"got={gv._cursor_ea()} want={old_ea}", + ) await c.press("space") await c.wait(lambda: app._active != "graph", 15) - c.check("space returns to the listing", app._active == "listing", - f"active={app._active}") + c.check( + "space returns to the listing", + app._active == "listing", + f"active={app._active}", + ) @scenario("graph_nav") @@ -3766,8 +4803,11 @@ async def s_graph_nav(c: Ctx): start_ea = gv._cursor_ea() await c.press("j") await c.pause(0.05) - c.check("j moves the cursor within the block", gv._cursor_ea() != start_ea, - f"{start_ea:#x} -> {gv._cursor_ea():#x}") + c.check( + "j moves the cursor within the block", + gv._cursor_ea() != start_ea, + f"{start_ea:#x} -> {gv._cursor_ea():#x}", + ) await c.press("k") await c.pause(0.05) c.check("k comes back", gv._cursor_ea() == start_ea) @@ -3777,24 +4817,35 @@ async def s_graph_nav(c: Ctx): if succs: await c.press("J") await c.pause(0.1) - c.check("J follows an edge to a successor block", - gv.cursor_node == succs[0][0], - f"node={gv.cursor_node} want={succs[0][0]}") + c.check( + "J follows an edge to a successor block", + gv.cursor_node == succs[0][0], + f"node={gv.cursor_node} want={succs[0][0]}", + ) await c.press("K") await c.pause(0.1) - c.check("K goes back up an edge", gv.cursor_node == b0, - f"node={gv.cursor_node} want={b0}") + c.check( + "K goes back up an edge", + gv.cursor_node == b0, + f"node={gv.cursor_node} want={b0}", + ) await c.press("0") await c.pause(0.1) - c.check("0 returns to the entry block", gv.cursor_node == gv.fc.entry, - f"node={gv.cursor_node} entry={gv.fc.entry}") + c.check( + "0 returns to the entry block", + gv.cursor_node == gv.fc.entry, + f"node={gv.cursor_node} entry={gv.fc.entry}", + ) # the cursor is always scrolled into view cell = gv._cursor_cell() top, left = int(gv.scroll_offset.y), int(gv.scroll_offset.x) - c.check("the cursor block is scrolled into view", - cell is not None and top <= cell[0] < top + gv.size.height - and left <= cell[1] < left + gv.size.width, - f"cell={cell} scroll=({top},{left}) size={gv.size}") + c.check( + "the cursor block is scrolled into view", + cell is not None + and top <= cell[0] < top + gv.size.height + and left <= cell[1] < left + gv.size.width, + f"cell={cell} scroll=({top},{left}) size={gv.size}", + ) @scenario("graph_zoom") @@ -3810,18 +4861,24 @@ async def s_graph_zoom(c: Ctx): await c.press("z") await c.pause(0.15) seen.append(gv.ZOOMS[gv._zoom]) - c.check("z cycles the three zoom levels", seen == ["full", "compact", "collapsed"], - str(seen)) - c.check("collapsed is much smaller than full", gv.lay.height < full_h, - f"{gv.lay.height} vs {full_h}") + c.check( + "z cycles the three zoom levels", + seen == ["full", "compact", "collapsed"], + str(seen), + ) + c.check( + "collapsed is much smaller than full", + gv.lay.height < full_h, + f"{gv.lay.height} vs {full_h}", + ) c.check("the cursor survives a zoom", gv._cursor_ea() is not None) - c.check("the status still names the function", - gv.fc.name in c.status(), c.status()) + c.check("the status still names the function", gv.fc.name in c.status(), c.status()) await c.press("z") await c.pause(0.15) c.check("and wraps back to full", gv.ZOOMS[gv._zoom] == "full") - c.check("canvas is restored", gv.lay.height == full_h, - f"{gv.lay.height} vs {full_h}") + c.check( + "canvas is restored", gv.lay.height == full_h, f"{gv.lay.height} vs {full_h}" + ) @scenario("graph_engine") @@ -3835,6 +4892,7 @@ async def s_graph_engine(c: Ctx): them), and a missing pytriskel degrades to native instead of raising. """ from idatui import graph_triskel + app = c.app fn, gv = await _open_graph(c) if gv.lay is None: @@ -3842,39 +4900,56 @@ async def s_graph_engine(c: Ctx): return ea = gv._cursor_ea() first = gv.lay.stats["engine"] - c.check("auto picks triskel when it is installed", - first == ("triskel" if graph_triskel.available() else "native"), - f"engine={first} available={graph_triskel.available()}") + c.check( + "auto picks triskel when it is installed", + first == ("triskel" if graph_triskel.available() else "native"), + f"engine={first} available={graph_triskel.available()}", + ) seen = [first] for _ in range(3): await c.press("e") await c.pause(0.2) seen.append(gv.lay.stats["engine"]) - c.check(f"the view survives engine={gv._engine}", - gv.lay is not None and gv.lay.width > 0 and gv.lay.height > 0, - f"{gv.lay.width}x{gv.lay.height}") - c.check(f"the cursor keeps an address on engine={gv._engine}", - gv._cursor_ea() is not None) - c.check(f"the canvas covers every edge on engine={gv._engine}", - all(0 <= col < gv.lay.width and 0 <= row < gv.lay.height - for rt in _routes_of(gv.lay) for row, col in rt), - f"canvas {gv.lay.width}x{gv.lay.height}") + c.check( + f"the view survives engine={gv._engine}", + gv.lay is not None and gv.lay.width > 0 and gv.lay.height > 0, + f"{gv.lay.width}x{gv.lay.height}", + ) + c.check( + f"the cursor keeps an address on engine={gv._engine}", + gv._cursor_ea() is not None, + ) + c.check( + f"the canvas covers every edge on engine={gv._engine}", + all( + 0 <= col < gv.lay.width and 0 <= row < gv.lay.height + for rt in _routes_of(gv.lay) + for row, col in rt + ), + f"canvas {gv.lay.width}x{gv.lay.height}", + ) c.check("e cycles back round", seen[0] == seen[-1], str(seen)) c.check("native was one of them", "native" in seen, str(seen)) - c.check("the status names the engine", "graph:" in c.status() or - gv.fc.name in c.status(), c.status()) + c.check( + "the status names the engine", + "graph:" in c.status() or gv.fc.name in c.status(), + c.status(), + ) if ea is not None: - c.check("the cursor address is unchanged by relayout", - gv._cursor_ea() is not None) + c.check( + "the cursor address is unchanged by relayout", gv._cursor_ea() is not None + ) def _routes_of(lay): """Every painted point, as (row, col) pairs, straight out of the index.""" out = [] for row, runs in lay.painting.hruns.items(): - out.append([(row, lo) for lo, _hi, _s, _e in runs] - + [(row, hi) for _lo, hi, _s, _e in runs]) + out.append( + [(row, lo) for lo, _hi, _s, _e in runs] + + [(row, hi) for _lo, hi, _s, _e in runs] + ) for lo, hi, col, _s, _e in lay.painting.vruns: out.append([(lo, col), (hi, col)]) return out @@ -3890,6 +4965,7 @@ async def s_graph_render(c: Ctx): if gv.lay is None: c.check("graph loaded", False) return + # The layout being ready (`gv.lay`) is not the same as the view having a # SIZE to render into -- that needs a laid-out frame, and reading glyphs # before one lands scrapes an empty canvas. Gate on the paint itself. @@ -3898,17 +4974,28 @@ async def s_graph_render(c: Ctx): await c.wait(lambda: gv.size.height > 0 and "\u250c" in _blob(), 10) blob = _blob() - c.check("boxes are drawn", blob.count("\u250c") >= 1 and blob.count("\u2502") > 4, - f"corners={blob.count(chr(0x250c))} verts={blob.count(chr(0x2502))}") - c.check("edges are drawn", any(ch in blob for ch in "\u25bc\u2570\u256d\u256e\u256f"), - "no edge glyphs on screen") + c.check( + "boxes are drawn", + blob.count("\u250c") >= 1 and blob.count("\u2502") > 4, + f"corners={blob.count(chr(0x250C))} verts={blob.count(chr(0x2502))}", + ) + c.check( + "edges are drawn", + any(ch in blob for ch in "\u25bc\u2570\u256d\u256e\u256f"), + "no edge glyphs on screen", + ) ea = gv._cursor_ea() head = gv.cur_head() - c.check("the cursor block's instruction text is on screen", - head is not None and head.text.split(" ")[0] in blob, - f"mnem={head.text.split(' ')[0] if head else None}") - c.check("the address gutter renders at full zoom", - ea is not None and f"{ea:08X}" in blob, f"ea={ea:#x}") + c.check( + "the cursor block's instruction text is on screen", + head is not None and head.text.split(" ")[0] in blob, + f"mnem={head.text.split(' ')[0] if head else None}", + ) + c.check( + "the address gutter renders at full zoom", + ea is not None and f"{ea:08X}" in blob, + f"ea={ea:#x}", + ) # minimap on/off actually changes the picture before = blob await c.press("m") @@ -3940,20 +5027,26 @@ async def s_graph_click(c: Ctx): top, left = int(gv.scroll_offset.y), int(gv.scroll_offset.x) target = None for n in gv.lay.nodes: - if (n.id != gv.cursor_node and top <= n.y + 1 < top + gv.size.height - 1 - and left <= n.x + 2 < left + gv.size.width - 2): + if ( + n.id != gv.cursor_node + and top <= n.y + 1 < top + gv.size.height - 1 + and left <= n.x + 2 < left + gv.size.width - 2 + ): target = n break if target is None: c.check("a second block is visible to click", True, "(skipped: none on screen)") return - PAD = 1 # GraphView { padding: 0 1 } - await c.pilot.click(GraphView, - offset=(PAD + target.x + 2 - left, target.y + 1 - top)) + PAD = 1 # GraphView { padding: 0 1 } + await c.pilot.click( + GraphView, offset=(PAD + target.x + 2 - left, target.y + 1 - top) + ) await c.pause(0.15) - c.check("clicking a block moves the cursor into it", - gv.cursor_node == target.id, - f"node={gv.cursor_node} want={target.id}") + c.check( + "clicking a block moves the cursor into it", + gv.cursor_node == target.id, + f"node={gv.cursor_node} want={target.id}", + ) @scenario("graph_minimap") @@ -3970,14 +5063,19 @@ async def s_graph_minimap(c: Ctx): # frame -- not just a settled app. await c.wait(lambda: gv._minimap_rect() is not None, 5) rect = gv._minimap_rect() - c.check("the minimap has a hit-box while it's shown", rect is not None, - f"size={gv.size} shown={gv._show_minimap}") + c.check( + "the minimap has a hit-box while it's shown", + rect is not None, + f"size={gv.size} shown={gv._show_minimap}", + ) if rect is None: return left, top, mw, mh = rect - c.check("it sits inside the pane, clear of the scrollbar", - left + mw <= gv.size.width - 1, - f"left={left} w={mw} pane={gv.size.width}") + c.check( + "it sits inside the pane, clear of the scrollbar", + left + mw <= gv.size.width - 1, + f"left={left} w={mw} pane={gv.size.width}", + ) # a big graph, so the overview actually maps to somewhere far away big = c.find_func(lambda f: f.size > 0x300) or fn @@ -4005,33 +5103,45 @@ async def s_graph_minimap(c: Ctx): PAD = 1 await c.pilot.click(GraphView, offset=(PAD + left + mw // 2, top + mh - 2)) await c.pause(0.2) - c.check("clicking low on the minimap scrolls the view down", - gv.scroll_offset.y > 0, f"scroll_y={gv.scroll_offset.y}") + c.check( + "clicking low on the minimap scrolls the view down", + gv.scroll_offset.y > 0, + f"scroll_y={gv.scroll_offset.y}", + ) # Most of a graph is padding, so a coordinate-accurate jump would park you # in empty space with the cursor left behind: every minimap click must land # on a block and take the cursor with it. landed = gv.lay.by_id.get(gv.cursor_node) - c.check("it snaps the cursor onto a real block", - landed is not None and landed.block is not None, - f"node={gv.cursor_node}") - c.check("and that block is what the viewport is showing", - landed is not None - and int(gv.scroll_offset.y) <= landed.y + landed.h - and landed.y <= int(gv.scroll_offset.y) + gv.size.height, - f"node.y={landed.y if landed else None} " - f"scroll={gv.scroll_offset.y} h={gv.size.height}") + c.check( + "it snaps the cursor onto a real block", + landed is not None and landed.block is not None, + f"node={gv.cursor_node}", + ) + c.check( + "and that block is what the viewport is showing", + landed is not None + and int(gv.scroll_offset.y) <= landed.y + landed.h + and landed.y <= int(gv.scroll_offset.y) + gv.size.height, + f"node.y={landed.y if landed else None} " + f"scroll={gv.scroll_offset.y} h={gv.size.height}", + ) low_node = gv.cursor_node # and the top of the minimap brings it back to a block up there await c.pilot.click(GraphView, offset=(PAD + left + mw // 2, top + 1)) await c.pause(0.2) top_node = gv.lay.by_id.get(gv.cursor_node) - c.check("clicking high on the minimap goes back up", - top_node is not None and gv.cursor_node != low_node - and top_node.y < gv.lay.by_id[low_node].y, - f"top={gv.cursor_node} low={low_node}") - c.check("the cursor still has a real address after a minimap jump", - gv._cursor_ea() is not None) + c.check( + "clicking high on the minimap goes back up", + top_node is not None + and gv.cursor_node != low_node + and top_node.y < gv.lay.by_id[low_node].y, + f"top={gv.cursor_node} low={low_node}", + ) + c.check( + "the cursor still has a real address after a minimap jump", + gv._cursor_ea() is not None, + ) # with the minimap hidden the same click is an ordinary canvas click await c.press("m") @@ -4042,16 +5152,19 @@ async def s_graph_minimap(c: Ctx): # Panning into the padding (which is most of the canvas) must not strand # you on a blank screen with nothing to navigate back by. - gv.scroll_to(y=max(gv.lay.height - 1, 0), x=max(gv.lay.width - 1, 0), - animate=False) + gv.scroll_to(y=max(gv.lay.height - 1, 0), x=max(gv.lay.width - 1, 0), animate=False) await c.pause(0.1) - c.check("a pan past the graph leaves the viewport empty", - not gv._viewport_has_block() or True) # setup, not an assertion + c.check( + "a pan past the graph leaves the viewport empty", + not gv._viewport_has_block() or True, + ) # setup, not an assertion gv._snap_into_view() await c.pause(0.1) - c.check("panning into empty padding snaps back to a block", - gv._viewport_has_block(), - f"scroll={gv.scroll_offset} canvas={gv.lay.width}x{gv.lay.height}") + c.check( + "panning into empty padding snaps back to a block", + gv._viewport_has_block(), + f"scroll={gv.scroll_offset} canvas={gv.lay.width}x{gv.lay.height}", + ) @scenario("graph_rename") @@ -4070,8 +5183,10 @@ async def s_graph_rename(c: Ctx): new = f"gtest_{os.getpid()}" await c.press("n") await c.wait(lambda: app.query_one("#rename", Input).display, 10) - c.check("n opens the rename prompt from the graph", - app.query_one("#rename", Input).display) + c.check( + "n opens the rename prompt from the graph", + app.query_one("#rename", Input).display, + ) inp = app.query_one("#rename", Input) inp.value = "" await c.type(new) @@ -4082,12 +5197,16 @@ async def s_graph_rename(c: Ctx): got = app.program.resolve(new) except Exception: # noqa: BLE001 got = None - c.check("the rename reached the database", got == ea, - f"resolve({new}) -> {got if got is None else hex(got)} want {ea:#x}") + c.check( + "the rename reached the database", + got == ea, + 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.invoke( - "rename", batch={"data": {"addr": hex(ea), "new": ""}}) + app.program.client.call( + remote_ops.rename, batch={"data": {"addr": hex(ea), "new": ""}} + ) app.program.bump_names() @@ -4105,17 +5224,30 @@ async def s_graph_sticky(c: Ctx): c.check("a second function exists", True, "(skipped)") return app._goto(other.name) - ok = await c.wait(lambda: app._cur is not None and app._cur.ea == other.addr - and app._active == "graph" - and gv.fc is not None and gv.fc.func_ea == other.addr, 60) - c.check("a goto from the graph lands in the next function's graph", ok, - f"active={app._active} sticky={app._graph_sticky} " - f"cur={app._cur.ea if app._cur else None} " - f"fc={gv.fc.func_ea if gv.fc else None} want={other.addr:#x}") + ok = await c.wait( + lambda: ( + app._cur is not None + and app._cur.ea == other.addr + and app._active == "graph" + and gv.fc is not None + and gv.fc.func_ea == other.addr + ), + 60, + ) + c.check( + "a goto from the graph lands in the next function's graph", + ok, + f"active={app._active} sticky={app._graph_sticky} " + f"cur={app._cur.ea if app._cur else None} " + f"fc={gv.fc.func_ea if gv.fc else None} want={other.addr:#x}", + ) await c.press("space") await c.wait(lambda: app._active != "graph", 15) - c.check("space still leaves graph mode", app._active == "listing", - f"active={app._active}") + c.check( + "space still leaves graph mode", + app._active == "listing", + f"active={app._active}", + ) c.check("and it stops being sticky", app._graph_sticky is False) @@ -4133,8 +5265,9 @@ async def run(binary, only=None): # # So: work on a scratch copy, seeded from a golden database that nothing # ever writes back to. - async with staged(binary, lambda p: IdaTui(open_path=p, keepalive=False), - prefix="idatui-pilot-") as target: + async with staged( + binary, lambda p: IdaTui(open_path=p, keepalive=False), prefix="idatui-pilot-" + ) as target: await _run_on(target, only) @@ -4156,7 +5289,9 @@ async def _run_on(binary, only=None): except _StopSuite: raise except Exception as e: # noqa: BLE001 — isolate: one scenario's crash - print(f"── {name} ({asyncio.get_event_loop().time() - _t0:.1f}s) CRASHED") + 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 @@ -4192,7 +5327,9 @@ def main(argv): if binary is None: # default target for the pilot binary = os.path.join( os.path.dirname(os.path.dirname(os.path.abspath(__file__))), - "targets", "echo") + "targets", + "echo", + ) try: asyncio.run(run(binary, only)) except _StopSuite: |
