From 02d02417800184fb76cd0245cdaa94c437aa4081 Mon Sep 17 00:00:00 2001 From: blasty Date: Fri, 21 Aug 2026 12:14:46 +0200 Subject: reformat: ruff format + import sort, mechanically (see ruff.toml) No behavior. Listed in .git-blame-ignore-revs (next commit). --- experiments/call_census.py | 26 +- experiments/cfg_dump.py | 41 +- experiments/graph_shot.py | 32 +- experiments/graph_smoke.py | 35 +- experiments/graph_spike.py | 68 +- experiments/inproc_spike.py | 112 ++- experiments/modal_shot.py | 72 +- experiments/opfmt_tools.py | 241 +++-- experiments/profile_client.py | 2 +- experiments/splash_place_count.py | 23 +- experiments/worker_smoke.py | 7 +- idatui/__init__.py | 30 +- idatui/__main__.py | 1 + idatui/_sync.py | 1 + idatui/app.py | 1957 ++++++++++++++++++++++++------------- idatui/diag.py | 10 +- idatui/drive.py | 87 +- idatui/errors.py | 1 + idatui/findings.py | 140 ++- idatui/formats.py | 2 +- idatui/graph.py | 196 ++-- idatui/graph_triskel.py | 85 +- idatui/highlight.py | 36 +- idatui/index.py | 26 +- idatui/journal.py | 18 +- idatui/kittygfx.py | 35 +- idatui/launch.py | 118 ++- idatui/pane.py | 314 ++++-- idatui/pool.py | 63 +- idatui/project.py | 92 +- idatui/prompt.py | 3 +- idatui/remote_ops.py | 77 +- idatui/rpc.py | 725 ++++++++++---- idatui/rpcclient.py | 9 +- idatui/search.py | 9 +- idatui/trace.py | 47 +- idatui/trace_ctl.py | 87 +- ruff.toml | 7 + tests/_fixtures.py | 11 +- tests/run.py | 93 +- tests/test_blob_ui.py | 284 ++++-- tests/test_diag.py | 93 +- tests/test_findings.py | 219 +++-- tests/test_formats.py | 170 +++- tests/test_graph.py | 191 +++- tests/test_index.py | 283 ++++-- tests/test_kittygfx.py | 153 ++- tests/test_launch.py | 24 +- tests/test_nexus_client.py | 13 +- tests/test_pool.py | 176 ++-- tests/test_project.py | 229 +++-- tests/test_project_ui.py | 341 ++++--- tests/test_rawimage_rpc.py | 198 ++-- tests/test_scenarios.py | 23 +- tests/test_search.py | 111 ++- tests/test_thumb_ui.py | 216 ++-- tests/test_trace.py | 248 +++-- tests/test_trace_rpc.py | 259 +++-- tests/test_trace_ui.py | 356 ++++--- tests/test_trace_vs_tenet.py | 64 +- tools/demo.py | 75 +- tools/make_logo_ans.py | 28 +- tools/verify_procs.py | 6 +- 63 files changed, 5827 insertions(+), 2872 deletions(-) diff --git a/experiments/call_census.py b/experiments/call_census.py index 17a5f94..4c0f9d3 100644 --- a/experiments/call_census.py +++ b/experiments/call_census.py @@ -11,6 +11,7 @@ Read it as: an action costing 1-8 calls is amortised (the snippet looped inside the database); an action whose call count scales with the number of rows or symbols on screen is a round-trip-per-item bug worth fixing. """ + from __future__ import annotations import asyncio @@ -19,7 +20,9 @@ import os import sys import time -sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "tests")) +sys.path.insert( + 0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "tests") +) from _fixtures import fast_keys, staged # noqa: E402 fast_keys() @@ -64,16 +67,21 @@ class _Span: def __exit__(self, *exc): ms = (time.perf_counter() - self.t0) * 1000 used = self.c.ops - self.ops0 - detail = " ".join(f"{k}x{v}" if v > 1 else k - for k, v in sorted(used.items(), key=lambda kv: -kv[1])) - print(f" {self.label:<34} {self.c.n - self.n0:>3} calls {ms:7.1f}ms {detail}") + detail = " ".join( + f"{k}x{v}" if v > 1 else k + for k, v in sorted(used.items(), key=lambda kv: -kv[1]) + ) + print( + f" {self.label:<34} {self.c.n - self.n0:>3} calls {ms:7.1f}ms {detail}" + ) return False async def main() -> int: binary = os.path.abspath(sys.argv[1] if len(sys.argv) > 1 else "targets/bash") - async with staged(binary, lambda p: IdaTui(open_path=p, keepalive=False), - prefix="idatui-census-") as target: + async with staged( + binary, lambda p: IdaTui(open_path=p, keepalive=False), prefix="idatui-census-" + ) as target: app = IdaTui(open_path=target, keepalive=False) census = Census() try: @@ -107,8 +115,10 @@ async def main() -> int: await pilot.pause(0.05) await pilot.pause(0.4) drained = listing_done() - print(f" {'(grower finished: ' + str(drained) + ')':<34}\n" - f" -- everything below is on a QUIET backend --\n") + print( + f" {'(grower finished: ' + str(drained) + ')':<34}\n" + f" -- everything below is on a QUIET backend --\n" + ) with census.span("scroll one page (pagedown)"): await pilot.press("pagedown") diff --git a/experiments/cfg_dump.py b/experiments/cfg_dump.py index 30b6ba9..bef24e7 100644 --- a/experiments/cfg_dump.py +++ b/experiments/cfg_dump.py @@ -15,6 +15,7 @@ taken branch) or "switch" (n-way). That is exactly the input a graph view needs; everything after this point (layering, ordering, routing, rendering) is pure python and needs no IDA at all. """ + from __future__ import annotations import argparse @@ -34,12 +35,12 @@ def dump(path: str, want: list[str], max_blocks: int) -> list[dict]: if idapro.open_database(local, run_auto_analysis=True) != 0: raise SystemExit(f"failed to open {local}") + import ida_bytes import ida_funcs import ida_gdl import ida_lines - import ida_bytes - import idautils import idaapi + import idautils out = [] try: @@ -60,7 +61,8 @@ def dump(path: str, want: list[str], max_blocks: int) -> list[dict]: ea = bb.start_ea while ea < bb.end_ea and ea != idaapi.BADADDR: txt = ida_lines.tag_remove( - ida_lines.generate_disasm_line(ea, 0) or "") + ida_lines.generate_disasm_line(ea, 0) or "" + ) lines.append(txt.rstrip()) nxt = ida_bytes.next_head(ea, bb.end_ea) if nxt <= ea: @@ -78,13 +80,15 @@ def dump(path: str, want: list[str], max_blocks: int) -> list[dict]: else: kind = "jump" succs.append([index[s.start_ea], kind]) - blocks.append({ - "id": index[bb.start_ea], - "start": bb.start_ea, - "end": bb.end_ea, - "lines": lines, - "succs": succs, - }) + blocks.append( + { + "id": index[bb.start_ea], + "start": bb.start_ea, + "end": bb.end_ea, + "lines": lines, + "succs": succs, + } + ) if max_blocks and len(blocks) > max_blocks: continue out.append({"name": name, "ea": fea, "blocks": blocks}) @@ -98,10 +102,19 @@ def main() -> int: ap = argparse.ArgumentParser() ap.add_argument("binary") ap.add_argument("-o", "--out", default="/tmp/cfg.json") - ap.add_argument("-f", "--func", action="append", default=[], - help="only functions whose name contains this (repeatable)") - ap.add_argument("--max-blocks", type=int, default=0, - help="skip functions with more blocks than this") + ap.add_argument( + "-f", + "--func", + action="append", + default=[], + help="only functions whose name contains this (repeatable)", + ) + ap.add_argument( + "--max-blocks", + type=int, + default=0, + help="skip functions with more blocks than this", + ) args = ap.parse_args() recs = dump(os.path.abspath(args.binary), args.func, args.max_blocks) diff --git a/experiments/graph_shot.py b/experiments/graph_shot.py index e024809..e3eb209 100644 --- a/experiments/graph_shot.py +++ b/experiments/graph_shot.py @@ -6,6 +6,7 @@ The pane a person runs this in is usually too small to judge the layout, and the pilot lays out synchronously at whatever size you ask for -- so this is the way to actually look at the thing. """ + import asyncio import os import shutil @@ -29,21 +30,23 @@ for e in ".i64 .id0 .id1 .id2 .nam .til".split(): except OSError: pass -from idatui.app import IdaTui, GraphView # noqa: E402 -from idatui._sync import wait_for # noqa: E402 -from idatui.rpc import screen_text # noqa: E402 +from idatui._sync import wait_for # noqa: E402 +from idatui.app import GraphView, IdaTui # noqa: E402 +from idatui.rpc import screen_text # noqa: E402 async def main() -> None: app = IdaTui(tmp, keepalive=False) async with app.run_test(size=(cols, rows)) as pilot: - await wait_for(lambda: app.program is not None and app._cur is not None, - pilot.pause, 120) + await wait_for( + lambda: app.program is not None and app._cur is not None, pilot.pause, 120 + ) ea = app.program.resolve(func) fn = app.program.function_of(ea) app._open_function(fn.addr, fn.name) - await wait_for(lambda: app._cur is not None and app._cur.ea == fn.addr, - pilot.pause, 60) + await wait_for( + lambda: app._cur is not None and app._cur.ea == fn.addr, pilot.pause, 60 + ) await pilot.pause(0.2) await pilot.press("space") gv = app.query_one(GraphView) @@ -51,8 +54,9 @@ async def main() -> None: gv._engine = engine gv._relayout() app._graph_status() - ok = await wait_for(lambda: app._active == "graph" and gv.lay is not None, - pilot.pause, 90) + ok = await wait_for( + lambda: app._active == "graph" and gv.lay is not None, pilot.pause, 90 + ) if not ok: print("graph never opened:", app.query_one("#status").render()) return @@ -63,8 +67,14 @@ async def main() -> None: print(screen_text(app)["text"]) print() print("status:", app.query_one("#status").render()) - print("stats :", gv.lay.stats, "canvas", - f"{gv.lay.width}x{gv.lay.height}", "zoom", gv.ZOOMS[gv._zoom]) + print( + "stats :", + gv.lay.stats, + "canvas", + f"{gv.lay.width}x{gv.lay.height}", + "zoom", + gv.ZOOMS[gv._zoom], + ) app._save_on_exit = False diff --git a/experiments/graph_smoke.py b/experiments/graph_smoke.py index 1b6e21e..92fe1f5 100644 --- a/experiments/graph_smoke.py +++ b/experiments/graph_smoke.py @@ -5,6 +5,7 @@ idatui.graph layout, through a real idalib worker. Wants: VERDICT: OK """ + import os import shutil import sys @@ -22,16 +23,17 @@ for e in ".i64 .id0 .id1 .id2 .nam .til".split(): except OSError: pass -from idatui.worker_client import WorkerClient # noqa: E402 -from idatui.domain import Program # noqa: E402 -from idatui import graph as G # noqa: E402 +from idatui.worker_client import WorkerClient # noqa: E402 + +from idatui import graph as G # noqa: E402 +from idatui.domain import Program # noqa: E402 want = sys.argv[1] if len(sys.argv) > 1 else "main" print("spawning worker + opening echo\u2026", flush=True) t = time.time() cl = WorkerClient(tmp) cl.connect(progress=lambda m: None) -print(f" worker ready in {time.time()-t:.2f}s", flush=True) +print(f" worker ready in {time.time() - t:.2f}s", flush=True) prog = Program(cl) ok = True @@ -40,7 +42,7 @@ print(f"resolve({want!r}) = {ea:#x}", flush=True) t = time.time() fc = prog.flowchart(ea) -print(f"flowchart() -> {time.time()-t:.2f}s", flush=True) +print(f"flowchart() -> {time.time() - t:.2f}s", flush=True) if fc is None: print("VERDICT: FAIL (no flowchart)") raise SystemExit(1) @@ -55,8 +57,10 @@ if not nrows: empty = [b for b in fc.blocks if not b.rows] if empty: ok = False - print(f" !! {len(empty)} blocks have NO rows, e.g. " - f"{[hex(b.start) for b in empty[:4]]}") + print( + f" !! {len(empty)} blocks have NO rows, e.g. " + f"{[hex(b.start) for b in empty[:4]]}" + ) b0 = fc.blocks[fc.entry] print(f" entry block {b0.start:#x}-{b0.end:#x}:") @@ -74,20 +78,20 @@ for b in fc.blocks: ok = False print(f" !! block {b.start:#x} has rows outside its range") -blocks = [G.Block(id=b.id, start=b.start, end=b.end, succs=list(b.succs)) - for b in fc.blocks] +blocks = [ + G.Block(id=b.id, start=b.start, end=b.end, succs=list(b.succs)) for b in fc.blocks +] def sizer(b): src_b = fc.blocks[b.id] - w = max([len(f"loc_{b.start:X}")] - + [len(h.text) + 12 for h in src_b.rows]) + 4 + w = max([len(f"loc_{b.start:X}")] + [len(h.text) + 12 for h in src_b.rows]) + 4 return (w, len(src_b.rows) + 3) t = time.time() lay = G.layout(blocks, sizer, entry=fc.entry) -print(f"layout() -> {(time.time()-t)*1000:.0f} ms {lay.stats}", flush=True) +print(f"layout() -> {(time.time() - t) * 1000:.0f} ms {lay.stats}", flush=True) print(f" canvas {lay.width}x{lay.height}") ok &= len(lay.nodes) == len(blocks) @@ -97,9 +101,10 @@ ok &= covered == {b.id for b in blocks} if covered != {b.id for b in blocks}: print(f" !! row index misses {sorted({b.id for b in blocks} - covered)[:5]}") -hits = sum(1 for r in range(min(lay.height, 400)) - if lay.painting.cells_at_row(r, 0, lay.width)) -print(f" {hits} of the first {min(lay.height,400)} rows carry edge cells") +hits = sum( + 1 for r in range(min(lay.height, 400)) if lay.painting.cells_at_row(r, 0, lay.width) +) +print(f" {hits} of the first {min(lay.height, 400)} rows carry edge cells") ok &= hits > 0 cl.close() diff --git a/experiments/graph_spike.py b/experiments/graph_spike.py index e8cf3b8..b3c5dc2 100644 --- a/experiments/graph_spike.py +++ b/experiments/graph_spike.py @@ -11,6 +11,7 @@ milliseconds when you're changing layout heuristics. python3 experiments/graph_spike.py /tmp/cfg-echo.json --func sub_61D0 python3 experiments/graph_spike.py /tmp/cfg-echo.json --stats """ + from __future__ import annotations import argparse @@ -23,8 +24,10 @@ sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) from idatui import graph as G # noqa: E402 COLOR = { - G.E_UNCOND: "\033[38;5;39m", G.E_TRUE: "\033[38;5;40m", - G.E_FALSE: "\033[38;5;203m", G.E_SWITCH: "\033[38;5;178m", + G.E_UNCOND: "\033[38;5;39m", + G.E_TRUE: "\033[38;5;40m", + G.E_FALSE: "\033[38;5;203m", + G.E_SWITCH: "\033[38;5;178m", G.E_BACK: "\033[38;5;135m", } DIM, RESET = "\033[38;5;244m", "\033[0m" @@ -37,11 +40,19 @@ def build(rec: dict, max_lines: int): for b in rec["blocks"]: lines = list(b["lines"]) if max_lines and len(lines) > max_lines: - lines = lines[:max_lines - 1] + [f"... {len(b['lines']) - max_lines + 1} more"] + lines = lines[: max_lines - 1] + [ + f"... {len(b['lines']) - max_lines + 1} more" + ] texts[b["id"]] = lines - blocks = [G.Block(id=b["id"], start=b["start"], end=b["end"], - succs=[(d, k) for d, k in b["succs"]]) - for b in rec["blocks"]] + blocks = [ + G.Block( + id=b["id"], + start=b["start"], + end=b["end"], + succs=[(d, k) for d, k in b["succs"]], + ) + for b in rec["blocks"] + ] def sizer(b: G.Block) -> tuple[int, int]: lines = texts[b.id] @@ -57,7 +68,8 @@ def render(lay: G.Layout, texts: dict[int, list[str]], color: bool) -> str: for row in range(lay.height): cells: dict[int, tuple[str, str]] = {} for col, (ch, kind, _eid) in lay.painting.cells_at_row( - row, 0, lay.width).items(): + row, 0, lay.width + ).items(): cells[col] = (ch, COLOR.get(kind, "")) for n in lay.nodes_at_row(row): label = f"loc_{n.block.start:X}" if n.block else "" @@ -107,34 +119,48 @@ def main() -> int: ap.add_argument("corpus") ap.add_argument("--func", help="function name (default: the smallest)") ap.add_argument("--stats", action="store_true", help="lay out the whole corpus") - ap.add_argument("--max-lines", type=int, default=8, - help="collapse blocks longer than this (0 = never)") + ap.add_argument( + "--max-lines", + type=int, + default=8, + help="collapse blocks longer than this (0 = never)", + ) ap.add_argument("--no-color", action="store_true") - ap.add_argument("--engine", choices=G.ENGINES, default=None, - help="layout backend (default: $IDATUI_GRAPH_ENGINE or auto)") + ap.add_argument( + "--engine", + choices=G.ENGINES, + default=None, + help="layout backend (default: $IDATUI_GRAPH_ENGINE or auto)", + ) args = ap.parse_args() recs = json.load(open(args.corpus)) if args.stats: - print(f"{'blocks':>7} {'nodes':>6} {'dummy':>6} {'layer':>6} " - f"{'canvas':>12} {'ms':>8} name") + print( + f"{'blocks':>7} {'nodes':>6} {'dummy':>6} {'layer':>6} " + f"{'canvas':>12} {'ms':>8} name" + ) tot = 0.0 for r in sorted(recs, key=lambda r: len(r["blocks"])): blocks, sizer, _ = build(r, args.max_lines) lay = G.layout(blocks, sizer, engine=args.engine) s = lay.stats tot += s["ms"] - print(f"{s['blocks']:>7} {s['nodes']:>6} {s['dummies']:>6} " - f"{s['layers']:>6} {lay.width:>5}x{lay.height:<6} " - f"{s['ms']:>8.1f} {r['name']}") + print( + f"{s['blocks']:>7} {s['nodes']:>6} {s['dummies']:>6} " + f"{s['layers']:>6} {lay.width:>5}x{lay.height:<6} " + f"{s['ms']:>8.1f} {r['name']}" + ) print(f"total {tot:.0f} ms over {len(recs)} functions") return 0 if args.func: rec = next((r for r in recs if r["name"] == args.func), None) if rec is None: - print("no such function; have: " - f"{', '.join(r['name'] for r in recs[:20])}", file=sys.stderr) + print( + f"no such function; have: {', '.join(r['name'] for r in recs[:20])}", + file=sys.stderr, + ) return 1 else: rec = min(recs, key=lambda r: len(r["blocks"])) @@ -142,8 +168,10 @@ def main() -> int: blocks, sizer, texts = build(rec, args.max_lines) lay = G.layout(blocks, sizer, engine=args.engine) print(render(lay, texts, color=not args.no_color)) - print(f"\n{rec['name']}: {lay.stats} canvas {lay.width}x{lay.height}", - file=sys.stderr) + print( + f"\n{rec['name']}: {lay.stats} canvas {lay.width}x{lay.height}", + file=sys.stderr, + ) return 0 diff --git a/experiments/inproc_spike.py b/experiments/inproc_spike.py index 469aa21..1b12204 100644 --- a/experiments/inproc_spike.py +++ b/experiments/inproc_spike.py @@ -25,6 +25,7 @@ Key facts this spike encodes (all verified): another process, so the TUI's event loop never blocks. In-process, a slow call (decompile ~150ms, analysis seconds) blocks the UI for its duration. """ + from __future__ import annotations import argparse @@ -60,13 +61,21 @@ class DirectBackend: def __init__(self, path: str) -> None: import idapro + idapro.enable_console_messages(False) t = time.time() rc = idapro.open_database(path, run_auto_analysis=True) self.open_secs = time.time() - t if rc: raise RuntimeError(f"open_database({path!r}) failed rc={rc}") - import ida_bytes, ida_funcs, ida_hexrays, ida_name, idaapi, idautils, idc + import ida_bytes + import ida_funcs + import ida_hexrays + import ida_name + import idaapi + import idautils + import idc + self._idapro = idapro self.idaapi, self.idautils, self.idc = idaapi, idautils, idc self.ida_bytes, self.ida_hexrays = ida_bytes, ida_hexrays @@ -108,6 +117,7 @@ class McpBackend: if REPO not in sys.path: # idalib's init can reset sys.path out from under us sys.path.insert(0, REPO) from idatui.client import IDAClient + self.IDAClient = IDAClient self.c = IDAClient(url, db=db) self.c.connect() @@ -150,8 +160,12 @@ class McpBackend: return r.get("code", "") if isinstance(r, dict) else str(r) def xrefs_to(self, ea): - r = self.c.call("xref_query", queries=[{"addr": hex(ea), "direction": "to", - "include_fn": True, "count": 2000}]) + r = self.c.call( + "xref_query", + queries=[ + {"addr": hex(ea), "direction": "to", "include_fn": True, "count": 2000} + ], + ) res = r.get("result", []) if isinstance(r, dict) else [] refs = res[0].get("refs", []) if res and isinstance(res[0], dict) else [] return [int(x["frm"], 16) for x in refs if x.get("frm")] @@ -200,14 +214,21 @@ def _recv(sock): return _unpack(_recvn(sock, n)) -_WORKER_OPS = ("functions", "resolve", "read_bytes", "disasm_line", - "decompile", "xrefs_to") +_WORKER_OPS = ( + "functions", + "resolve", + "read_bytes", + "disasm_line", + "decompile", + "xrefs_to", +) def _worker_main(sockpath: str, dbpath: str) -> None: """Runs in a child process. Opens idalib on ITS main thread (constraint satisfied), then serves one client serially over a unix socket.""" import socket as sk + direct = DirectBackend(dbpath) ops = {name: getattr(direct, name) for name in _WORKER_OPS} try: @@ -237,10 +258,17 @@ class UnixWorkerBackend: def __init__(self, dbpath: str) -> None: import socket as sk + self.sockpath = f"/tmp/inproc_spike_{os.getpid()}.sock" self.proc = __import__("subprocess").Popen( - [sys.executable, os.path.abspath(__file__), - "--worker", self.sockpath, dbpath]) + [ + sys.executable, + os.path.abspath(__file__), + "--worker", + self.sockpath, + dbpath, + ] + ) deadline = time.time() + 120 self.sock = None while time.time() < deadline: @@ -263,12 +291,23 @@ class UnixWorkerBackend: raise RuntimeError(val) return val - def functions(self): return self._call("functions") - def resolve(self, name): return self._call("resolve", name) - def read_bytes(self, ea, n): return self._call("read_bytes", ea, n) - def disasm_line(self, ea): return self._call("disasm_line", ea) - def decompile(self, ea): return self._call("decompile", ea) - def xrefs_to(self, ea): return self._call("xrefs_to", ea) + def functions(self): + return self._call("functions") + + def resolve(self, name): + return self._call("resolve", name) + + def read_bytes(self, ea, n): + return self._call("read_bytes", ea, n) + + def disasm_line(self, ea): + return self._call("disasm_line", ea) + + def decompile(self, ea): + return self._call("decompile", ea) + + def xrefs_to(self, ea): + return self._call("xrefs_to", ea) def close(self): try: @@ -321,6 +360,7 @@ def bench(target: str, n: int) -> None: try: import socket + socket.create_connection(("127.0.0.1", 8745), 0.3).close() backends["mcp"] = McpBackend() print("benching the running mcp server on :8745 too", flush=True) @@ -338,15 +378,15 @@ def bench(target: str, n: int) -> None: return sample[_nx[0]] ops = { - "resolve(main)": (lambda b: b.resolve("main"), n), - "read_bytes(16)": (lambda b: b.read_bytes(main, 16), n), - "read_bytes(4096)": (lambda b: b.read_bytes(main, 4096), n), - "disasm_line": (lambda b: b.disasm_line(main), n), - "xrefs_to": (lambda b: b.xrefs_to(_next_sample()), min(n, 200)), - "decompile(cached)": (lambda b: b.decompile(main), min(n, 40)), + "resolve(main)": (lambda b: b.resolve("main"), n), + "read_bytes(16)": (lambda b: b.read_bytes(main, 16), n), + "read_bytes(4096)": (lambda b: b.read_bytes(main, 4096), n), + "disasm_line": (lambda b: b.disasm_line(main), n), + "xrefs_to": (lambda b: b.xrefs_to(_next_sample()), min(n, 200)), + "decompile(cached)": (lambda b: b.decompile(main), min(n, 40)), } names = list(backends) - hdr = f"{'op':20}" + "".join(f"{nm+' us':>14}" for nm in names) + hdr = f"{'op':20}" + "".join(f"{nm + ' us':>14}" for nm in names) print(hdr) print("-" * len(hdr)) results: dict[str, dict[str, float]] = {nm: {} for nm in names} @@ -371,7 +411,7 @@ def bench(target: str, n: int) -> None: if nm == "mcp": continue v = results[nm].get(label, float("nan")) - parts.append(f"{nm} {m/v:.0f}x" if v == v and v else f"{nm} -") + parts.append(f"{nm} {m / v:.0f}x" if v == v and v else f"{nm} -") print(f" {label:20} {' '.join(parts)}") for b in backends.values(): @@ -408,10 +448,12 @@ def _build_spike_app(backend: "DirectBackend"): for ea, nm in backend.functions(): ol.add_option(Option(f"{ea:08x} {nm}", id=str(ea))) yield ol - yield Static("select a function, press F5 to decompile inline", - id="code") - yield Static("in-process idalib — every call runs on the UI thread", - id="status") + yield Static( + "select a function, press F5 to decompile inline", id="code" + ) + yield Static( + "in-process idalib — every call runs on the UI thread", id="status" + ) yield Footer() def on_mount(self): @@ -428,13 +470,13 @@ def _build_spike_app(backend: "DirectBackend"): if ea is None: return t = time.perf_counter() - code = backend.decompile(ea) # <-- BLOCKS the event loop + code = backend.decompile(ea) # <-- BLOCKS the event loop dt = (time.perf_counter() - t) * 1e3 - self.query_one("#code", Static).update( - "\n".join(code.splitlines()[:40])) + self.query_one("#code", Static).update("\n".join(code.splitlines()[:40])) self.query_one("#status", Static).update( f"decompiled {backend.idc.get_func_name(ea)} in {dt:.0f} ms " - f"(UI was frozen for those {dt:.0f} ms)") + f"(UI was frozen for those {dt:.0f} ms)" + ) def action_decompile_all(self): funcs = backend.functions() @@ -442,14 +484,15 @@ def _build_spike_app(backend: "DirectBackend"): n = 0 for ea, _ in funcs: try: - backend.decompile(ea) # <-- long, uninterruptible freeze + backend.decompile(ea) # <-- long, uninterruptible freeze n += 1 except Exception: # noqa: BLE001 pass dt = (time.perf_counter() - t) * 1e3 self.query_one("#status", Static).update( f"decompiled {n} funcs in {dt:.0f} ms — the whole UI was frozen " - f"the entire time (no spinner, no input)") + f"the entire time (no spinner, no input)" + ) def action_bytes(self): ea = self._sel_ea() @@ -459,7 +502,8 @@ def _build_spike_app(backend: "DirectBackend"): b = backend.read_bytes(ea, 64) dt = (time.perf_counter() - t) * 1e6 self.query_one("#status", Static).update( - f"read 64 bytes in {dt:.1f} us: {b[:16].hex()}…") + f"read 64 bytes in {dt:.1f} us: {b[:16].hex()}…" + ) return Spike() @@ -472,7 +516,9 @@ def _open_copy(target: str) -> "DirectBackend": os.remove(tmp + e) except OSError: pass - print("opening in-process (blocks the terminal until analysis is done)…", flush=True) + print( + "opening in-process (blocks the terminal until analysis is done)…", flush=True + ) return DirectBackend(tmp) # MAIN THREAD open, before the event loop starts diff --git a/experiments/modal_shot.py b/experiments/modal_shot.py index 2e6cd7a..82f1e2f 100644 --- a/experiments/modal_shot.py +++ b/experiments/modal_shot.py @@ -11,6 +11,7 @@ IDA and no .i64 -- it runs under any python with textual, in about a second. The app it boots is a bare `App` carrying `IdaTui.CSS` and the app theme, which is exactly what the modals resolve their styles against. """ + import asyncio import os import sys @@ -18,16 +19,23 @@ import sys REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) sys.path.insert(0, REPO) -from textual.app import App, ComposeResult # noqa: E402 -from textual.widgets import Static # noqa: E402 - -from idatui.app import ( # noqa: E402 - IDATUI_THEME, IdaTui, BusyScreen, ConfirmScreen, HelpScreen, LoadingScreen, - QuitScreen, StructEditor, XrefsScreen, +from textual.app import App, ComposeResult # noqa: E402 +from textual.widgets import Static # noqa: E402 + +from idatui._sync import wait_for # noqa: E402 +from idatui.app import ( # noqa: E402 + IDATUI_THEME, + BusyScreen, + ConfirmScreen, + HelpScreen, + IdaTui, + LoadingScreen, + QuitScreen, + StructEditor, + XrefsScreen, ) -from idatui.domain import Struct # noqa: E402 -from idatui.rpc import screen_text # noqa: E402 -from idatui._sync import wait_for # noqa: E402 +from idatui.domain import Struct # noqa: E402 +from idatui.rpc import screen_text # noqa: E402 class _StubProgram: @@ -36,14 +44,15 @@ class _StubProgram: _STRUCTS = [ Struct(name="timespec", size=0x10, members=2, is_union=False, ordinal=1), Struct(name="stat", size=0x90, members=15, is_union=False, ordinal=2), - Struct(name="pthread_mutex_t", size=0x28, members=4, is_union=True, - ordinal=3), + Struct(name="pthread_mutex_t", size=0x28, members=4, is_union=True, ordinal=3), ] - _SRC = ("struct timespec\n" - "{\n" - " __time_t tv_sec; /* seconds */\n" - " __syscall_slong_t tv_nsec;\n" - "};\n") + _SRC = ( + "struct timespec\n" + "{\n" + " __time_t tv_sec; /* seconds */\n" + " __syscall_slong_t tv_nsec;\n" + "};\n" + ) def list_structs(self): return list(self._STRUCTS) @@ -57,8 +66,7 @@ def _make(name: str): if name == "structs": return StructEditor(_StubProgram()) if name == "confirm": - return ConfirmScreen("Delete struct 'timespec' ?", - "This cannot be undone.") + return ConfirmScreen("Delete struct 'timespec' ?", "This cannot be undone.") if name == "quit": return QuitScreen(["echo", "libc.so.6"]) if name == "help": @@ -68,10 +76,13 @@ def _make(name: str): if name == "loading": return LoadingScreen("echo", "opening database\u2026") if name == "xrefs": - return XrefsScreen(" xrefs to main", [ - (0x1234, " sub_2297+0x1c call main"), - (0x5678, " _start+0x21 mov rdi, main"), - ]) + return XrefsScreen( + " xrefs to main", + [ + (0x1234, " sub_2297+0x1c call main"), + (0x5678, " _start+0x21 mov rdi, main"), + ], + ) raise SystemExit(f"unknown modal {name!r}; --list to see them") @@ -84,8 +95,9 @@ class _Shot(App): def compose(self) -> ComposeResult: # A little content underneath, so the modal's edge is visible against # something rather than floating on an empty screen. - yield Static("\n".join(" .... the view behind the dialog ...." - for _ in range(60))) + yield Static( + "\n".join(" .... the view behind the dialog ...." for _ in range(60)) + ) def on_mount(self) -> None: self.register_theme(IDATUI_THEME) @@ -108,12 +120,14 @@ async def main() -> None: await pilot.pause() if name == "structs": await wait_for( - lambda: bool(getattr(app.screen, "_structs", None)), - pilot.pause, 5) + lambda: bool(getattr(app.screen, "_structs", None)), pilot.pause, 5 + ) app.screen.on_option_list_option_selected( - type("E", (), {"option_index": 0})()) - await wait_for(lambda: "{" in app.screen.query_one( - "#se-edit").text, pilot.pause, 5) + type("E", (), {"option_index": 0})() + ) + await wait_for( + lambda: "{" in app.screen.query_one("#se-edit").text, pilot.pause, 5 + ) await pilot.pause() print(f"\n=== {name} " + "=" * (cols - len(name) - 5)) print(screen_text(app)["text"]) diff --git a/experiments/opfmt_tools.py b/experiments/opfmt_tools.py index 535688c..7654677 100644 --- a/experiments/opfmt_tools.py +++ b/experiments/opfmt_tools.py @@ -8,6 +8,7 @@ does to the database, what Hex-Rays will and won't render. python3 experiments/opfmt_tools.py # 29 checks, ~40s """ + import importlib.util import os import shutil @@ -30,15 +31,21 @@ if os.path.exists(seed): shutil.copy(seed, tmp + ".i64") import idapro # noqa: E402 + idapro.enable_console_messages(False) assert idapro.open_database(tmp, run_auto_analysis=True) == 0 import ida_auto # noqa: E402 + ida_auto.auto_wait() -import ida_typeinf, idaapi, ida_bytes, ida_lines # noqa: E402,F401 +import ida_bytes +import ida_lines +import ida_typeinf # noqa: E402,F401 +import idaapi spec = importlib.util.spec_from_file_location( - "_patch", os.path.join(REPO, "server", "patch_server.py")) + "_patch", os.path.join(REPO, "server", "patch_server.py") +) patch = importlib.util.module_from_spec(spec) spec.loader.exec_module(patch) @@ -101,14 +108,25 @@ while e < ea + 0x400: if target: break e = ida_bytes.next_head(e, ea + 0x800) -print("target:", hex(target[0]), "n=", target[1], "value", hex(target[2]), - "|", line(target[0])) +print( + "target:", + hex(target[0]), + "n=", + target[1], + "value", + hex(target[2]), + "|", + line(target[0]), +) tea, tn, tv = target r = op_format(addr=hex(tea), mode="show") print(" show:", r) -check("show reports the operand and its choices", - r.get("n") == tn and "hex" in r.get("choices", []), str(r)) +check( + "show reports the operand and its choices", + r.get("n") == tn and "hex" in r.get("choices", []), + str(r), +) check("show doesn't change anything", r.get("applied") is False, str(r)) seen = [] @@ -116,15 +134,25 @@ for i in range(8): r = op_format(addr=hex(tea), mode="cycle") seen.append((r.get("format"), r.get("text"))) print(f" cycle -> {r.get('prev')} -> {r.get('format')}: {r.get('text')}") -check("cycling returns to where it started", - seen[0][0] == seen[len(r.get("choices", []))][0] - if len(seen) > len(r.get("choices", [])) else True, str(seen)) -check("decimal renders differently from hex", - any(s[1] != seen[0][1] for s in seen), str(seen)) +check( + "cycling returns to where it started", + seen[0][0] == seen[len(r.get("choices", []))][0] + if len(seen) > len(r.get("choices", [])) + else True, + str(seen), +) +check( + "decimal renders differently from hex", + any(s[1] != seen[0][1] for s in seen), + str(seen), +) r = op_format(addr=hex(tea), mode="dec") -check("explicit dec sticks", r.get("format") == "dec" and str(tv) in r.get("text", ""), - str(r)) +check( + "explicit dec sticks", + r.get("format") == "dec" and str(tv) in r.get("text", ""), + str(r), +) r = op_format(addr=hex(tea), mode="back") print(" back ->", r.get("format"), r.get("text")) check("back steps the ring the other way", r.get("format") == "hex", str(r)) @@ -142,29 +170,55 @@ while e < ea + 0x800: if off_target: break e = ida_bytes.next_head(e, ea + 0x800) -print("target:", off_target and hex(off_target[0]), "|", - off_target and line(off_target[0])) +print( + "target:", + off_target and hex(off_target[0]), + "|", + off_target and line(off_target[0]), +) if off_target: import ida_name # noqa: E402 + oe, on = off_target v, _w = NS["_idatui_op_value"](oe, on) named = bool(ida_name.get_ea_name(v)) r = op_format(addr=hex(oe), n=on, mode="show") - print(" ", hex(oe), "n=", on, "|", r.get("text"), r.get("choices"), - "target named:", named) - check("an unnamed target is not a cycle stop (it would invent a name)", - ("offset" in r.get("choices", [])) == named, str(r)) + print( + " ", + hex(oe), + "n=", + on, + "|", + r.get("text"), + r.get("choices"), + "target named:", + named, + ) + check( + "an unnamed target is not a cycle stop (it would invent a name)", + ("offset" in r.get("choices", [])) == named, + str(r), + ) r = op_format(addr=hex(oe), n=on, mode="offset") print(" offset ->", r.get("text")) - check("but asking explicitly makes the reference", - r.get("format") == "offset" and "offset" in r.get("text", ""), str(r)) + check( + "but asking explicitly makes the reference", + r.get("format") == "offset" and "offset" in r.get("text", ""), + str(r), + ) r = op_format(addr=hex(oe), n=on, mode="show") - check("and from then on the ring includes it", - "offset" in r.get("choices", []), str(r)) + check( + "and from then on the ring includes it", + "offset" in r.get("choices", []), + str(r), + ) r = op_format(addr=hex(oe), n=on, mode="hex") print(" hex ->", r.get("text")) - check("and back to a number", r.get("format") == "hex" - and "offset" not in r.get("text", ""), str(r)) + check( + "and back to a number", + r.get("format") == "hex" and "offset" not in r.get("text", ""), + str(r), + ) op_format(addr=hex(oe), n=on, mode="default") else: print(" (no literal-that-is-an-address in this function)") @@ -175,22 +229,32 @@ spans = NS["_idatui_op_spans"](tea, txt) print(" text:", repr(txt), "spans:", spans) if len(spans) >= 2: r = op_format(addr=hex(tea), col=spans[0][0], mode="show") - check("a column inside operand 0 picks operand 0 (or the first literal)", - r.get("n") in (spans[0][2], NS["_idatui_op_candidates"](tea)[0]), str(r)) + check( + "a column inside operand 0 picks operand 0 (or the first literal)", + r.get("n") in (spans[0][2], NS["_idatui_op_candidates"](tea)[0]), + str(r), + ) r = op_format(addr=hex(tea), col=spans[-1][0], mode="show") - check("a column inside the last operand picks it", r.get("n") == spans[-1][2], - str(r)) + check( + "a column inside the last operand picks it", r.get("n") == spans[-1][2], str(r) + ) print("\n=== listing: an unmapped value refuses to become an offset ===") r = op_format(addr=hex(tea), n=tn, mode="offset") print(" ", r.get("error") or r) -check("offset on a non-address is refused, not invented", - bool(r.get("error")) or ida_bytes.is_mapped(tv), str(r)) +check( + "offset on a non-address is refused, not invented", + bool(r.get("error")) or ida_bytes.is_mapped(tv), + str(r), +) print("\n=== listing: char is only offered when it renders as one ===") r = op_format(addr=hex(tea), n=tn, mode="show") -check("0x%x isn't offered as a char" % tv, - ("char" in r["choices"]) == NS["_idatui_printable"](tv), str(r)) +check( + "0x%x isn't offered as a char" % tv, + ("char" in r["choices"]) == NS["_idatui_printable"](tv), + str(r), +) print("\n=== listing: a stack variable can be cycled AND put back ===") stk = None @@ -208,23 +272,32 @@ if stk: orig = line(se) r = op_format(addr=hex(se), n=sn, mode="cycle") print(" cycle ->", r.get("format"), r.get("text"), "|", r.get("warn")) - check("leaving a stack variable says so, and how to undo it", - "stack" in (r.get("warn") or ""), str(r)) - ring = [op_format(addr=hex(se), n=sn, mode="cycle") - for _ in range(len(r["choices"]))] + check( + "leaving a stack variable says so, and how to undo it", + "stack" in (r.get("warn") or ""), + str(r), + ) + ring = [ + op_format(addr=hex(se), n=sn, mode="cycle") for _ in range(len(r["choices"])) + ] print(" ring:", [(x["format"], x["text"]) for x in ring]) - check("the ring is the same at every step (a lap comes home)", - [x["format"] for x in ring] == r["choices"][1:] + r["choices"][:1], - f"{[x['format'] for x in ring]} vs {r['choices']}") + check( + "the ring is the same at every step (a lap comes home)", + [x["format"] for x in ring] == r["choices"][1:] + r["choices"][:1], + f"{[x['format'] for x in ring]} vs {r['choices']}", + ) r = op_format(addr=hex(se), n=sn, mode="stack") - check("'stack' puts the frame variable back", - r.get("format") == "stack" and r.get("text") == orig, - f"{r.get('text')!r} want {orig!r}") + check( + "'stack' puts the frame variable back", + r.get("format") == "stack" and r.get("text") == orig, + f"{r.get('text')!r} want {orig!r}", + ) else: print(" (no stack-variable operand found)") print("\n=== data item ===") import ida_segment # noqa: E402 + seg = ida_segment.get_segm_by_name(".data") if seg: de = seg.start_ea @@ -257,8 +330,11 @@ while e < seg.end_ea and ue is None: if ue is not None: r = op_format(addr=hex(ue), mode="cycle") print(" ", hex(ue), "->", r.get("error")) - check("undefined bytes are refused with the fix, not a silent no-op", - "define" in (r.get("error") or ""), str(r)) + check( + "undefined bytes are refused with the fix, not a silent no-op", + "define" in (r.get("error") or ""), + str(r), + ) else: print(" (no undefined bytes)") @@ -268,6 +344,7 @@ print(" show without a line:", r.get("error")) check("no line is an error, not a guess", bool(r.get("error")), str(r)) import ida_hexrays # noqa: E402 + cf = ida_hexrays.decompile(parse_address("main")) sv = cf.get_pseudocode() pcline = None @@ -276,8 +353,13 @@ for i in range(len(sv)): if nums and nums[0]["value"] > 9: pcline = (i, nums[0]) break -print(" line", pcline[0], repr(ida_lines.tag_remove(sv[pcline[0]].line).strip()), - "num:", pcline[1]) +print( + " line", + pcline[0], + repr(ida_lines.tag_remove(sv[pcline[0]].line).strip()), + "num:", + pcline[1], +) i, num = pcline r = pc_num_format(addr="main", line=i, mode="show") check("show finds the literal", r.get("ea") == hex(num["ea"]), str(r)) @@ -293,8 +375,11 @@ print(" bin ->", r.get("error")) check("binary is refused with a reason", bool(r.get("error")), str(r)) r = pc_num_format(addr="main", line=i, mode="default") print(" default ->", r.get("text")) -check("default restores Hex-Rays' own choice", - r.get("text") == ida_lines.tag_remove(sv[i].line).strip(), str(r)) +check( + "default restores Hex-Rays' own choice", + r.get("text") == ida_lines.tag_remove(sv[i].line).strip(), + str(r), +) r = pc_num_format(addr="main", line=i, mode="show") check("and the format reads back as default", r.get("format") == "default", str(r)) print("\n=== pseudocode: the ring visits every stop ===") @@ -303,10 +388,16 @@ for _ in range(len(r.get("choices", [])) * 2): rr = pc_num_format(addr="main", line=i, mode="cycle") ring.append((rr.get("format"), rr.get("text"))) print(" ", [x[0] for x in ring]) -check("every stop in the ring is reached", - set(x[0] for x in ring) == set(r["choices"]), f"{ring} vs {r['choices']}") -check("the ring's renderings are distinct", - len({x[1] for x in ring}) >= len(r["choices"]) - 1, str(ring)) +check( + "every stop in the ring is reached", + set(x[0] for x in ring) == set(r["choices"]), + f"{ring} vs {r['choices']}", +) +check( + "the ring's renderings are distinct", + len({x[1] for x in ring}) >= len(r["choices"]) - 1, + str(ring), +) pc_num_format(addr="main", line=i, mode="default") print("\n=== pseudocode: col picks the literal ===") @@ -327,9 +418,12 @@ if multi: col = len(NS["_idatui_compact"](plain[:x]).rstrip()) if x else 0 r = pc_num_format(addr="main", line=k, col=col, mode="show") print(" col", col, "->", r.get("ea"), r.get("value")) - check("a column selects the number under it", - r.get("value") == hex(nums[1]["value"]) - or r.get("value") == hex(nums[0]["value"]), str(r)) + check( + "a column selects the number under it", + r.get("value") == hex(nums[1]["value"]) + or r.get("value") == hex(nums[0]["value"]), + str(r), + ) else: print(" (no line with two literals)") @@ -342,8 +436,11 @@ if row and row.get("ops"): t = row["text"] for lo, hi, n in row["ops"]: print(f" op{n}: {t[lo:hi]!r}") - check("the extents index the row's own text", - all(t[lo:hi].strip() for lo, hi, n in row["ops"]), str(row["ops"])) + check( + "the extents index the row's own text", + all(t[lo:hi].strip() for lo, hi, n in row["ops"]), + str(row["ops"]), + ) # and they agree with what op_format picks for a column inside them ok = True for lo, hi, n in row["ops"]: @@ -361,8 +458,11 @@ if regop: lo, hi, n = regop r = op_format(addr=hex(tea), col=(lo + hi) // 2, mode="cycle") print(" ", repr(row["text"][lo:hi]), "->", r.get("error")) - check("it names the operand and the one that CAN change", - bool(r.get("error")) and "operand" in r["error"], str(r)) + check( + "it names the operand and the one that CAN change", + bool(r.get("error")) and "operand" in r["error"], + str(r), + ) else: print(" (this instruction has no register-only operand)") @@ -376,15 +476,24 @@ for rec in pn["nums"]: two = next((v for v in multi2.values() if len(v) >= 2), None) if two: import ida_hexrays as _hx + cf2 = _hx.decompile(parse_address("main")) - disp = NS["_idatui_compact"](ida_lines.tag_remove(cf2.get_pseudocode()[two[0]["line"]].line)) + disp = NS["_idatui_compact"]( + ida_lines.tag_remove(cf2.get_pseudocode()[two[0]["line"]].line) + ) print(" line:", repr(disp.strip())) for rec in two: - print(f" x{rec['x0']}..{rec['x1']} = {disp[rec['x0']:rec['x1']]!r} value {rec['value']}") - check("spans land on the literals in the DISPLAYED text", - all(disp[r0["x0"]:r0["x1"]].strip() for r0 in two), str(two)) - check("distinct literals get distinct spans", - two[0]["x0"] != two[1]["x0"], str(two)) + print( + f" x{rec['x0']}..{rec['x1']} = {disp[rec['x0'] : rec['x1']]!r} value {rec['value']}" + ) + check( + "spans land on the literals in the DISPLAYED text", + all(disp[r0["x0"] : r0["x1"]].strip() for r0 in two), + str(two), + ) + check( + "distinct literals get distinct spans", two[0]["x0"] != two[1]["x0"], str(two) + ) print(f"\n{OK} passed, {FAIL} failed") idapro.close_database(save=False) diff --git a/experiments/profile_client.py b/experiments/profile_client.py index 84abd8a..b03ed46 100644 --- a/experiments/profile_client.py +++ b/experiments/profile_client.py @@ -21,8 +21,8 @@ import pstats import time from idatui import remote_ops -from idatui.nexus_client import NexusClient from idatui.domain import Program +from idatui.nexus_client import NexusClient def main() -> int: diff --git a/experiments/splash_place_count.py b/experiments/splash_place_count.py index 6d8322a..4490b9e 100644 --- a/experiments/splash_place_count.py +++ b/experiments/splash_place_count.py @@ -7,6 +7,7 @@ write per progress tick) and tallies the escapes. PYTHONPATH=. ~/ida-venv/bin/python experiments/splash_place_count.py """ + from __future__ import annotations import asyncio @@ -16,10 +17,10 @@ import sys os.environ["IDATUI_KITTY"] = "1" -from textual.app import App, ComposeResult # noqa: E402 -from textual.widgets import Static # noqa: E402 +from textual.app import App, ComposeResult # noqa: E402 +from textual.widgets import Static # noqa: E402 -from idatui import kittygfx # noqa: E402 +from idatui import kittygfx # noqa: E402 SENT: list[str] = [] @@ -29,10 +30,10 @@ def _fake_write(data: str) -> bool: return True -kittygfx._write = _fake_write # type: ignore[assignment] +kittygfx._write = _fake_write # type: ignore[assignment] kittygfx._cell = (9, 22) -from idatui.app import LoadingScreen # noqa: E402 +from idatui.app import LoadingScreen # noqa: E402 def tally() -> dict[str, int]: @@ -79,12 +80,16 @@ async def main() -> None: d = tally() blob = "".join(SENT) cmds = re.findall(r"\x1b_G([^;\x1b]*)", blob) - ids = {dict(kv.split("=", 1) for kv in c.split(",") if "=" in kv).get("p") - for c in cmds if "a=p" in c.split(",")} + ids = { + dict(kv.split("=", 1) for kv in c.split(",") if "=" in kv).get("p") + for c in cmds + if "a=p" in c.split(",") + } onscreen = "unbounded (anonymous)" if None in ids else len(ids) print() - print(f"=> {d['placements (a=p)']} place escapes sent, " - f"{d['deletes (a=d)']} deletes") + print( + f"=> {d['placements (a=p)']} place escapes sent, {d['deletes (a=d)']} deletes" + ) print(f" images actually on screen: {onscreen}") print(" A placement is identified by (image id, placement id). An a=p with") print(" no p= key is ANONYMOUS and stacks a fresh copy every time; with a") diff --git a/experiments/worker_smoke.py b/experiments/worker_smoke.py index 1c2125f..08799a4 100644 --- a/experiments/worker_smoke.py +++ b/experiments/worker_smoke.py @@ -3,18 +3,21 @@ A matching registered GUI is reused; otherwise IDA Nexus starts a managed idalib worker. Usage: ``uv run python experiments/worker_smoke.py FILE``. """ + from __future__ import annotations import os import sys import time -from idatui.nexus_client import NexusClient from idatui.domain import Program +from idatui.nexus_client import NexusClient def main() -> int: - target = os.path.abspath(sys.argv[1] if len(sys.argv) > 1 else "experiments/fibonacci.elf") + target = os.path.abspath( + sys.argv[1] if len(sys.argv) > 1 else "experiments/fibonacci.elf" + ) print(f"attaching IDA Nexus to {target}…", flush=True) started = time.time() client = NexusClient(target) diff --git a/idatui/__init__.py b/idatui/__init__.py index e89d8f6..201fac6 100644 --- a/idatui/__init__.py +++ b/idatui/__init__.py @@ -1,28 +1,28 @@ """idatui — a keyboard-first TUI using shared IDA Nexus databases.""" +from .domain import ( + DISASM_BLOCK, + LIST_PAGE, + Decompilation, + DisasmModel, + Func, + FunctionIndex, + Line, + Program, + Ref, + Struct, +) from .errors import ( - IDAError, IDAConnectionError, - IDATimeoutError, + IDAError, IDAProtocolError, IDARPCError, - IDAToolError, IDASessionError, + IDATimeoutError, + IDAToolError, Session, ) from .nexus_client import NexusClient -from .domain import ( - Program, - FunctionIndex, - DisasmModel, - Func, - Line, - Ref, - Struct, - Decompilation, - LIST_PAGE, - DISASM_BLOCK, -) __all__ = [ "NexusClient", diff --git a/idatui/__main__.py b/idatui/__main__.py index 0090ace..31c1da9 100644 --- a/idatui/__main__.py +++ b/idatui/__main__.py @@ -1,4 +1,5 @@ """``python -m idatui`` -> the one-shot launcher (open a binary in the TUI).""" + import sys from .launch import main diff --git a/idatui/_sync.py b/idatui/_sync.py index b7da338..0b2d9a2 100644 --- a/idatui/_sync.py +++ b/idatui/_sync.py @@ -8,6 +8,7 @@ Two yield strategies feed the same poll loop: under a Pilot (tests) we yield wit ``pilot.pause`` (which also drains the screen); live (RPC) we yield with ``asyncio.sleep`` and drain explicitly via a throwaway ``Pilot(app)``. """ + from __future__ import annotations import asyncio diff --git a/idatui/app.py b/idatui/app.py index 1ee7459..75dc9d1 100644 --- a/idatui/app.py +++ b/idatui/app.py @@ -37,26 +37,28 @@ from textual.containers import Horizontal, Vertical, VerticalScroll from textual.geometry import Region, Size from textual.message import Message from textual.reactive import reactive -from textual.theme import Theme from textual.screen import ModalScreen from textual.scroll_view import ScrollView from textual.strip import Strip +from textual.theme import Theme from textual.widgets import ( - DataTable, Input, OptionList, Static, TextArea, + DataTable, + Input, + OptionList, + Static, + TextArea, ) from textual.widgets.option_list import Option -from . import diag, graph, kittygfx +from . import diag, findings, graph, kittygfx, search +from .domain import Func, Head, ListingModel, Program, Struct from .edit_ctl import EditController -from .prompt import PromptBar -from .trace_ctl import TraceController -from . import findings, search +from .errors import IDAConnectionError from .highlight import CTextArea, highlight_c from .journal import Journal - -from .errors import IDAConnectionError from .nexus_client import NexusClient, registered_database -from .domain import Func, Head, ListingModel, Program, Struct +from .prompt import PromptBar +from .trace_ctl import TraceController # Styles for the disassembly listing. _S_ADDR = Style(color="#6b7684") @@ -68,28 +70,28 @@ _S_INSN = Style(color="#c3cad3") #: text), HUES only where they mean something (numbers, strings, symbols), #: structure recedes so brackets and commas stop competing with operands. _S_SPAN = { - "insn": Style(color="#e8ecf2"), # 15.3:1 mnemonic / directive - "reg": Style(color="#c3cad3"), # 11.0:1 registers = body weight - "num": Style(color="#d8a657"), # 8.2:1 immediates, offsets - "str": Style(color="#9ece6a"), # 9.9:1 string literals - "name": Style(color="#7aa2f7"), # 7.2:1 symbols / xref targets - "seg": Style(color="#93aee0"), # 8.1:1 segment names - "cmt": Style(color="#7c8b9e", italic=True), # 5.2:1 - "punct": Style(color="#626c7a"), # 3.4:1 brackets, commas, +/- - "err": Style(color="#c9762f"), # IDA's own error marker - "text": Style(color="#c3cad3"), # 11.0:1 anything unclassified + "insn": Style(color="#e8ecf2"), # 15.3:1 mnemonic / directive + "reg": Style(color="#c3cad3"), # 11.0:1 registers = body weight + "num": Style(color="#d8a657"), # 8.2:1 immediates, offsets + "str": Style(color="#9ece6a"), # 9.9:1 string literals + "name": Style(color="#7aa2f7"), # 7.2:1 symbols / xref targets + "seg": Style(color="#93aee0"), # 8.1:1 segment names + "cmt": Style(color="#7c8b9e", italic=True), # 5.2:1 + "punct": Style(color="#626c7a"), # 3.4:1 brackets, commas, +/- + "err": Style(color="#c9762f"), # IDA's own error marker + "text": Style(color="#c3cad3"), # 11.0:1 anything unclassified } _S_MNEM = Style(color="#e8ecf2") _S_OPBYTES = Style(color="#5e6875") # raw opcode bytes column _S_DATA = Style(color="#d8a657") _S_UNK = Style(color="#7c8b9e", italic=True) # undefined bytes in the flat listing _S_MEMBER = Style(color="#93aee0") -_S_SEP = Style(color="#5e6875") # function boundary separators / banners +_S_SEP = Style(color="#5e6875") # function boundary separators / banners _S_FUNCHDR = Style(color="#7aa2f7", bold=True) # 'name proc'/'endp' headers -_LST_INDENT = " " # one depth level: function names sit at level 0, code at 1 -_OP_LIMIT = 8 # opcode bytes shown in the 'limited' column mode -_JUMP_CONTEXT = 4 # lines of context kept above a jump target (cursor stays on it) +_LST_INDENT = " " # one depth level: function names sit at level 0, code at 1 +_OP_LIMIT = 8 # opcode bytes shown in the 'limited' column mode +_JUMP_CONTEXT = 4 # lines of context kept above a jump target (cursor stays on it) _SPLIT_MIN_WIDTH = 100 # need room for two usable code panes side by side @@ -117,10 +119,10 @@ class ViewMode(StrEnum): the next one has fewer places to reach. """ - LISTING = "listing" # the unified continuous listing (code + data) - DECOMP = "decomp" # Hex-Rays pseudocode - HEX = "hex" # the hex viewer - GRAPH = "graph" # the CFG graph view + LISTING = "listing" # the unified continuous listing (code + data) + DECOMP = "decomp" # Hex-Rays pseudocode + HEX = "hex" # the hex viewer + GRAPH = "graph" # the CFG graph view #: The two that show a code view over a NavEntry, i.e. where a follow, an #: xref or a rename makes sense. @@ -128,21 +130,47 @@ class ViewMode(StrEnum): def code_modes(cls) -> frozenset["ViewMode"]: return frozenset({cls.LISTING, cls.DECOMP, cls.GRAPH}) + # Tokens that look like identifiers but aren't renamable symbols (so 'n' on them # in the listing names the address instead of trying to rename the token). -_ASM_KEYWORDS = frozenset({ - "db", "dw", "dd", "dq", "dt", "byte", "word", "dword", "qword", "tbyte", - "offset", "short", "near", "far", "ptr", "dup", "cs", "ds", "es", "fs", - "gs", "ss", "align", "public", "assume", "end", -}) +_ASM_KEYWORDS = frozenset( + { + "db", + "dw", + "dd", + "dq", + "dt", + "byte", + "word", + "dword", + "qword", + "tbyte", + "offset", + "short", + "near", + "far", + "ptr", + "dup", + "cs", + "ds", + "es", + "fs", + "gs", + "ss", + "align", + "public", + "assume", + "end", + } +) _S_CURSOR = Style(bgcolor="#2a313c") #: Execution trails. Deliberately faint: they sit UNDER the code palette and #: must not compete with it — the trail says "you came through here", the text #: still has to be readable as code. Now is the loudest because there is exactly #: one of it. _S_TRAIL_NOW = Style(bgcolor="#3f3410") -_S_TRAIL_PAST = Style(bgcolor="#2b1c17") # warm: behind you -_S_TRAIL_FUTURE = Style(bgcolor="#152230") # cool: ahead of you +_S_TRAIL_PAST = Style(bgcolor="#2b1c17") # warm: behind you +_S_TRAIL_FUTURE = Style(bgcolor="#152230") # cool: ahead of you #: Hex with a trace loaded: bytes the trace SAW at this timestamp vs bytes we're #: still showing from the file. The distinction matters more than the values — #: one is evidence, the other is an assumption. @@ -157,12 +185,12 @@ _S_WORD = Style(bgcolor="#2a3f5f") # identifier under the cursor #: _S_WORD (which marks every occurrence of an identifier): this marks ONE span, #: the thing a keypress acts on, so it reads as a selection rather than a match. _S_OPERAND = Style(bgcolor="#3a3560", underline=True) -_S_CELL = Style(reverse=True) # the block cursor cell -_S_LINENO = Style(color="#626c7a") # pseudocode line-number gutter +_S_CELL = Style(reverse=True) # the block cursor cell +_S_LINENO = Style(color="#626c7a") # pseudocode line-number gutter _S_LINENO_CUR = Style(color="#c3cad3", bold=True) # gutter on the cursor line -_S_DECOMP_SPIN = Style(color="#d0a215", bold=True) # 'decompiling' spinner glyph +_S_DECOMP_SPIN = Style(color="#d0a215", bold=True) # 'decompiling' spinner glyph _S_DECOMP_WAIT = Style(color="#7c8b9e", italic=True) # 'decompiling' label -_S_DECOMP_DOTS = Style(color="#626c7a") # trailing ellipsis +_S_DECOMP_DOTS = Style(color="#626c7a") # trailing ellipsis _S_LINK = Style(bgcolor="#233044") # split view: rows linked to the other pane's cursor # Hex-Rays appends a `/*0xEA*/` address marker to each pseudocode line (we fetch @@ -206,8 +234,8 @@ class ViewAnchor: """ view: str = "listing" - ea: int | None = None # cursor address - top_ea: int | None = None # first visible address + ea: int | None = None # cursor address + top_ea: int | None = None # first visible address cursor_x: int = 0 flash: str | None = None #: The edit changed which functions exist, so the index must be rebuilt. @@ -218,12 +246,12 @@ class ViewAnchor: class NavEntry: ea: int name: str - cursor: int = 0 # disasm line (instruction index) - cursor_x: int = 0 # disasm column - scroll_y: int = -1 # disasm viewport top (-1 = derive from cursor) - dec_cursor: int = 0 # pseudocode line + cursor: int = 0 # disasm line (instruction index) + cursor_x: int = 0 # disasm column + scroll_y: int = -1 # disasm viewport top (-1 = derive from cursor) + dec_cursor: int = 0 # pseudocode line dec_cursor_x: int = 0 # pseudocode column - dec_scroll_y: int = -1 # pseudocode viewport top (-1 = derive) + dec_scroll_y: int = -1 # pseudocode viewport top (-1 = derive) dec_scroll_x: int = 0 # pseudocode horizontal scroll is_region: bool = False # not inside a function (flat listing view) view: str = "listing" # which code view to restore this entry in @@ -417,7 +445,9 @@ def _overlay_over(strip: Strip, ranges: list[tuple[int, int]], style: Style) -> if a > pos: parts.append(strip.crop(pos, a)) mid = strip.crop(a, b) - parts.append(Strip([Segment(s.text, (s.style or Style()) + style) for s in mid])) + parts.append( + Strip([Segment(s.text, (s.style or Style()) + style) for s in mid]) + ) pos = b if pos < total: parts.append(strip.crop(pos, total)) @@ -638,11 +668,11 @@ class _MatchRanges: __slots__ = ("_lines", "_needle", "_n", "_ci", "_text", "_cache") def __init__(self, lines, needle: str, n: int, ci: bool, text) -> None: - self._lines = lines # set[int] - self._needle = needle # already case-folded when ci - self._n = n # len(term); the needle may be folded + self._lines = lines # set[int] + self._needle = needle # already case-folded when ci + self._n = n # len(term); the needle may be folded self._ci = ci - self._text = text # callable: line index -> str | None + self._text = text # callable: line index -> str | None self._cache: dict[int, list[tuple[int, int]]] = {} def _find(self, i: int) -> list[tuple[int, int]]: @@ -770,7 +800,7 @@ class SearchMixin: starts.append(pos) parts.append(s) pos += len(s) + 1 - while len(starts) < count: # a short window: keep the indices lined up + while len(starts) < count: # a short window: keep the indices lined up starts.append(pos) parts.append("") pos += 1 @@ -831,8 +861,11 @@ class SearchMixin: def _after_incremental(self) -> None: self._compute_matches() self.refresh() - self._jump_from(getattr(self, "_search_origin", 0), - getattr(self, "_search_dir", 1), include_current=True) + self._jump_from( + getattr(self, "_search_origin", 0), + getattr(self, "_search_dir", 1), + include_current=True, + ) n = len(self._matches) self._app_status(f"/{self._term} {n} match{'' if n == 1 else 'es'}") @@ -840,13 +873,23 @@ class SearchMixin: if not self._matches: return if direction >= 0: - nxt = next((m for m in self._matches - if (m >= origin if include_current else m > origin)), - self._matches[0]) + nxt = next( + ( + m + for m in self._matches + if (m >= origin if include_current else m > origin) + ), + self._matches[0], + ) else: - nxt = next((m for m in reversed(self._matches) - if (m <= origin if include_current else m < origin)), - self._matches[-1]) + nxt = next( + ( + m + for m in reversed(self._matches) + if (m <= origin if include_current else m < origin) + ), + self._matches[-1], + ) self._goto_line(nxt) def search_commit(self) -> None: @@ -860,7 +903,9 @@ class SearchMixin: self._reset_search_cache() self.cursor = getattr(self, "_search_origin", self.cursor) self.cursor_x = getattr(self, "_search_origin_x", self.cursor_x) - self.scroll_to(y=max(self.cursor - self._visible_height() // 2, 0), animate=False) + self.scroll_to( + y=max(self.cursor - self._visible_height() // 2, 0), animate=False + ) self.refresh() def repeat_last(self, direction: int) -> None: @@ -870,8 +915,13 @@ class SearchMixin: return self._term = term self._ci = term.islower() - self._search_ensure(lambda: (self._compute_matches(), self.refresh(), - self.search_repeat(direction))) + self._search_ensure( + lambda: ( + self._compute_matches(), + self.refresh(), + self.search_repeat(direction), + ) + ) def _compute_matches(self) -> None: term = self._term @@ -910,8 +960,7 @@ class SearchMixin: # finding every further occurrence in it. nxt = starts[line + 1] if line + 1 < nlines else blen j = body.find(needle, nxt) - ranges = _MatchRanges(set(matches), needle, n, ci, - self._search_line_text) + ranges = _MatchRanges(set(matches), needle, n, ci, self._search_line_text) else: # Typing forward can only ever REMOVE lines: a line holding "mov" # holds "mo". So when the term just grew (and nothing else moved -- @@ -923,8 +972,14 @@ class SearchMixin: # been looked at, and narrowing would silently never find them. rows: object = range(count) prev = self._matched_key - if (prev is not None and prev[2] == count and prev[3] == src - and prev[1] == ci and term.startswith(prev[0]) and prev[0]): + if ( + prev is not None + and prev[2] == count + and prev[3] == src + and prev[1] == ci + and term.startswith(prev[0]) + and prev[0] + ): rows = self._matches text_of = self._search_line_text for i in rows: @@ -948,10 +1003,18 @@ class SearchMixin: return cur = self.cursor if direction >= 0: - nxt = next((m for m in self._matches - if (m >= cur if include_current else m > cur)), self._matches[0]) + nxt = next( + ( + m + for m in self._matches + if (m >= cur if include_current else m > cur) + ), + self._matches[0], + ) else: - nxt = next((m for m in reversed(self._matches) if m < cur), self._matches[-1]) + nxt = next( + (m for m in reversed(self._matches) if m < cur), self._matches[-1] + ) self._goto_line(nxt) k = self._matches.index(nxt) + 1 self._app_status(f"/{self._term}/ {k}/{len(self._matches)} line {nxt}") @@ -963,7 +1026,9 @@ class SearchMixin: ranges = self._ranges.get(self.cursor) if ranges: self.cursor_x = ranges[0][0] - self.scroll_to(y=max(self.cursor - self._visible_height() // 2, 0), animate=False) + self.scroll_to( + y=max(self.cursor - self._visible_height() // 2, 0), animate=False + ) self._hscroll() # bring the match column into horizontal view self.refresh() self._refresh_hl() @@ -1054,8 +1119,8 @@ class ListingView(SearchMixin, NavMixin, ColumnCursor, ScrollView, can_focus=Tru self._term = "" self._matches: list[int] = [] self._ranges: dict[int, list[tuple[int, int]]] = {} - self._op_mode = 1 # opcode column: 0=off, 1=limited, 2=full ('o' cycles) - self._op_w = 0 # char width of the hex-bytes field (excl. gap) + self._op_mode = 1 # opcode column: 0=off, 1=limited, 2=full ('o' cycles) + self._op_w = 0 # char width of the hex-bytes field (excl. gap) self._search_loading = False self._search_pending: list = [] # done-callbacks awaiting the load self._link_rows: set[int] = set() # split-view: linked instruction rows @@ -1161,9 +1226,15 @@ class ListingView(SearchMixin, NavMixin, ColumnCursor, ScrollView, can_focus=Tru self._refresh_hl() # -- public API -------------------------------------------------------- # - def load(self, model: ListingModel, name: str, cursor: int = 0, - cursor_x: int = 0, scroll_y: int | None = None, - focus: str | None = None) -> None: + def load( + self, + model: ListingModel, + name: str, + cursor: int = 0, + cursor_x: int = 0, + scroll_y: int | None = None, + focus: str | None = None, + ) -> None: previous, self.model = self.model, model self._name = name self.total = 0 @@ -1265,8 +1336,10 @@ class ListingView(SearchMixin, NavMixin, ColumnCursor, ScrollView, can_focus=Tru self._reset_search_cache(body=True) self._clamp_x() self.refresh() - self._app_status("opcodes: " + {0: "off", 1: f"limited ({_OP_LIMIT} bytes)", - 2: "full"}[self._op_mode]) + self._app_status( + "opcodes: " + + {0: "off", 1: f"limited ({_OP_LIMIT} bytes)", 2: "full"}[self._op_mode] + ) @work(thread=True, exclusive=True, group="listing-grow") def _grow(self) -> None: @@ -1355,20 +1428,28 @@ class ListingView(SearchMixin, NavMixin, ColumnCursor, ScrollView, can_focus=Tru if h is None: strip = Strip([Segment(f" {idx:>8} …", _S_DIM)]) elif h.kind == "sep": - strip = Strip([Segment(f"{h.ea:08X} ", _S_ADDR), - Segment(_LST_INDENT + h.text, _S_SEP)]) + strip = Strip( + [ + Segment(f"{h.ea:08X} ", _S_ADDR), + Segment(_LST_INDENT + h.text, _S_SEP), + ] + ) elif h.kind == "funchdr": # depth-0: address + 'name proc'/'endp' (no indent) - strip = Strip([Segment(f"{h.ea:08X} ", _S_ADDR), - Segment(h.text, _S_FUNCHDR)]) + strip = Strip( + [Segment(f"{h.ea:08X} ", _S_ADDR), Segment(h.text, _S_FUNCHDR)] + ) elif h.kind == "label": # depth-0: address + 'loc_XXX:' on its own line - strip = Strip([Segment(f"{h.ea:08X} ", _S_ADDR), - Segment(h.text, _S_LABEL)]) + strip = Strip( + [Segment(f"{h.ea:08X} ", _S_ADDR), Segment(h.text, _S_LABEL)] + ) else: # depth-1: address, one indent, then opcode+text - segs: list[Segment] = [Segment(f"{h.ea:08X} ", _S_ADDR), - Segment(_LST_INDENT, _S_INSN)] + segs: list[Segment] = [ + Segment(f"{h.ea:08X} ", _S_ADDR), + Segment(_LST_INDENT, _S_INSN), + ] op = self._op_field(h) if op: segs.append(Segment(op, _S_OPBYTES)) @@ -1389,8 +1470,12 @@ class ListingView(SearchMixin, NavMixin, ColumnCursor, ScrollView, can_focus=Tru kind = self.trail.get(h.ea) if kind is not None: strip = strip.apply_style( - _S_TRAIL_NOW if kind == "now" else - _S_TRAIL_PAST if kind == "past" else _S_TRAIL_FUTURE) + _S_TRAIL_NOW + if kind == "now" + else _S_TRAIL_PAST + if kind == "past" + else _S_TRAIL_FUTURE + ) plain = self._line_plain(idx) if (self._hl_word or idx == self.cursor) else None if idx in self._ranges: strip = _overlay_ranges(strip, self._ranges[idx], self._match_style(idx)) @@ -1686,8 +1771,15 @@ class DecompView(SearchMixin, NavMixin, ColumnCursor, ScrollView, can_focus=True elif self.cursor_x >= sx + width: self.scroll_to(x=self.cursor_x - width + 1, animate=False) - def show(self, ea: int, text: str, cursor: int = 0, cursor_x: int = 0, - scroll_y: int = -1, scroll_x: int = 0) -> None: + def show( + self, + ea: int, + text: str, + cursor: int = 0, + cursor_x: int = 0, + scroll_y: int = -1, + scroll_x: int = 0, + ) -> None: # Pull each line's `/*0xEA*/` marker into _line_eas, then strip it from # the displayed text (clutter) before highlighting. Stripping only edits # within lines, so line indices still align with the domain's raw code. @@ -1725,8 +1817,9 @@ class DecompView(SearchMixin, NavMixin, ColumnCursor, ScrollView, can_focus=True self._hscroll() self.refresh() - def goto(self, cursor: int, cursor_x: int = 0, scroll_y: int = -1, - scroll_x: int = 0) -> None: + def goto( + self, cursor: int, cursor_x: int = 0, scroll_y: int = -1, scroll_x: int = 0 + ) -> None: """Move the cursor/scroll on the already-loaded text (no re-highlight). Used to jump to a target inside the function already displayed, e.g. an xref/goto that resolves to this same function. @@ -1778,7 +1871,9 @@ class DecompView(SearchMixin, NavMixin, ColumnCursor, ScrollView, can_focus=True return self._line_eas[idx] if 0 <= idx < len(self._line_eas) else None def _after_cursor_move(self) -> None: - self.post_message(DecompView.CursorMoved(self.cursor, self._line_ea(self.cursor))) + self.post_message( + DecompView.CursorMoved(self.cursor, self._line_ea(self.cursor)) + ) @property def total(self) -> int: @@ -1808,8 +1903,12 @@ class DecompView(SearchMixin, NavMixin, ColumnCursor, ScrollView, can_focus=True kind = self.trail.get(idx) if self.trail else None if kind is not None: base = base.apply_style( - _S_TRAIL_NOW if kind == "now" else - _S_TRAIL_PAST if kind == "past" else _S_TRAIL_FUTURE) + _S_TRAIL_NOW + if kind == "now" + else _S_TRAIL_PAST + if kind == "past" + else _S_TRAIL_FUTURE + ) if idx in self._ranges: base = _overlay_ranges(base, self._ranges[idx], self._match_style(idx)) if self._hl_word: @@ -1818,12 +1917,13 @@ class DecompView(SearchMixin, NavMixin, ColumnCursor, ScrollView, can_focus=True base = _overlay_ranges(base, occ, _S_WORD) if idx == self.cursor: base = _cursor_decorate(base, self._texts[idx], self.cursor_x) - span = self._cursor_literal(idx) # the literal `o` would reformat - if span is not None: # (last: see ListingView) + span = self._cursor_literal(idx) # the literal `o` would reformat + if span is not None: # (last: see ListingView) base = _overlay_over(base, [span], _S_OPERAND) code_w = max(width - gw, 0) code = base.crop(x, x + code_w).adjust_cell_length( - code_w, _S_LINK if linked else None) + code_w, _S_LINK if linked else None + ) if gw <= 0: return code style = _S_LINENO_CUR if idx == self.cursor else _S_LINENO @@ -2104,16 +2204,16 @@ class HexView(ScrollView, can_focus=True): panes. Matches ``render_line``'s layout: addr(9) + file-offset(10) + 16 hex cells of 3 cols (with a 1-col gap before byte 8), then ' |' + ASCII.""" HEX, ASCII = 19, 70 - if x < HEX: # clicked the address/offset gutter -> row start + if x < HEX: # clicked the address/offset gutter -> row start return 0 - if x < HEX + 49: # hex byte region + if x < HEX + 49: # hex byte region rel = x - HEX - if rel >= 24: # collapse the 1-col gap between the two halves + if rel >= 24: # collapse the 1-col gap between the two halves rel -= 1 return min(rel // 3, 15) - if x < ASCII: # the ' |' separator -> last byte of the row + if x < ASCII: # the ' |' separator -> last byte of the row return 15 - return min(x - ASCII, 15) # ASCII pane (and anything past it) + return min(x - ASCII, 15) # ASCII pane (and anything past it) def on_click(self, event) -> None: # type: ignore[no-untyped-def] if self.model is None or self.model.size == 0: @@ -2240,12 +2340,12 @@ class HexView(ScrollView, can_focus=True): # --------------------------------------------------------------------------- # # Graph view # --------------------------------------------------------------------------- # -_S_GBORDER = Style(color="#4b5565") # box border, idle -_S_GBORDER_CUR = Style(color="#7aa2f7", bold=True) # box border, cursor block -_S_GLABEL = Style(color="#7aa2f7", bold=True) # loc_XXXX in the border +_S_GBORDER = Style(color="#4b5565") # box border, idle +_S_GBORDER_CUR = Style(color="#7aa2f7", bold=True) # box border, cursor block +_S_GLABEL = Style(color="#7aa2f7", bold=True) # loc_XXXX in the border _S_GLABEL_CUR = Style(color="#c0caf5", bold=True) _S_GDIM = Style(color="#5e6875") -_S_GENTRY = Style(color="#9ece6a", bold=True) # the entry block's label +_S_GENTRY = Style(color="#9ece6a", bold=True) # the entry block's label #: Edge colours follow IDA's convention: green = branch taken, red = falls #: through, blue = the block's only successor, purple = loops back. _S_EDGE = { @@ -2269,7 +2369,7 @@ _S_MINI_CUR = Style(bgcolor="#161b22", color="#9ece6a", bold=True) _S_MINI_VIEW = Style(bgcolor="#233044", color="#c0caf5") _S_MINI_EDGE = Style(bgcolor="#161b22", color="#2f3945") -_GPAD = 1 # columns of padding inside a box +_GPAD = 1 # columns of padding inside a box _MINI_W, _MINI_H = 30, 14 @@ -2317,7 +2417,7 @@ class _CellRow: b = self.width if b <= a: return - self.ch[a:b] = s[a - i:b - i] + self.ch[a:b] = s[a - i : b - i] self.st[a:b] = [style] * (b - a) def restyle(self, a: int, b: int, style: Style) -> None: @@ -2403,7 +2503,7 @@ class GraphView(NavMixin, ScrollView, can_focus=True): def __init__(self) -> None: super().__init__() - self.fc = None # domain.Flowchart + self.fc = None # domain.Flowchart self.lay: graph.Layout | None = None self.loaded_ea: int | None = None self._blocks: dict[int, object] = {} @@ -2414,7 +2514,7 @@ class GraphView(NavMixin, ScrollView, can_focus=True): self._engine = "auto" self._mini_cache: tuple | None = None self._drag: tuple[int, int, float, float] | None = None - self._drag_map = False # the drag started on the minimap + self._drag_map = False # the drag started on the minimap self._hl_word = "" self.trail: dict[int, str] | None = None @@ -2445,10 +2545,13 @@ class GraphView(NavMixin, ScrollView, can_focus=True): self.lay = None self.virtual_size = Size(0, 0) return - blocks = [graph.Block(id=b.id, start=b.start, end=b.end, - succs=list(b.succs)) for b in self.fc.blocks] - self.lay = graph.layout(blocks, self._sizer, entry=self.fc.entry, - engine=self._engine) + blocks = [ + graph.Block(id=b.id, start=b.start, end=b.end, succs=list(b.succs)) + for b in self.fc.blocks + ] + self.lay = graph.layout( + blocks, self._sizer, entry=self.fc.entry, engine=self._engine + ) self.virtual_size = Size(self.lay.width + 2, self.lay.height + 1) def _rows(self, nid: int): @@ -2457,7 +2560,7 @@ class GraphView(NavMixin, ScrollView, can_focus=True): if b is None: return [] if self._zoom == 2: - return [None] # one synthetic summary row + return [None] # one synthetic summary row return b.rows def _row_plain(self, nid: int, i: int) -> str: @@ -2477,15 +2580,16 @@ class GraphView(NavMixin, ScrollView, can_focus=True): @staticmethod def _head_text(h) -> str: - return (f"{h.name} {h.text}" if h.name else h.text) + return f"{h.name} {h.text}" if h.name else h.text def _sizer(self, b: graph.Block) -> tuple[int, int]: nid = b.id rows = self._rows(nid) n = max(len(rows), 1) label = f"loc_{b.start:X}" - widest = max([len(label) + 4] - + [len(self._row_plain(nid, i)) for i in range(n)]) + widest = max( + [len(label) + 4] + [len(self._row_plain(nid, i)) for i in range(n)] + ) return (widest + 2 * _GPAD + 2, n + 2) # -- geometry --------------------------------------------------------- # @@ -2647,8 +2751,10 @@ class GraphView(NavMixin, ScrollView, can_focus=True): y0 = int(self.scroll_offset.y) x0 = int(self.scroll_offset.x) y1, x1 = y0 + self.size.height, x0 + self.size.width - return any(n.y <= y1 and y0 <= n.bottom and n.x <= x1 and x0 <= n.right - for n in self.lay.nodes) + return any( + n.y <= y1 and y0 <= n.bottom and n.x <= x1 and x0 <= n.right + for n in self.lay.nodes + ) def _snap_into_view(self) -> None: """After a pan, if the viewport holds no block at all, ease to the @@ -2675,7 +2781,7 @@ class GraphView(NavMixin, ScrollView, can_focus=True): self._clamp_cursor() self._center_cursor() self.refresh(layout=True) - self.app._graph_status() # keeps the function name; names the zoom + self.app._graph_status() # keeps the function name; names the zoom def action_minimap(self) -> None: self._show_minimap = not self._show_minimap @@ -2691,8 +2797,10 @@ class GraphView(NavMixin, ScrollView, can_focus=True): function. Cheaper to look than to argue. """ from . import graph_triskel - choices = ["auto", "native"] + (["triskel"] if graph_triskel.available() - else []) + + choices = ["auto", "native"] + ( + ["triskel"] if graph_triskel.available() else [] + ) self._engine = choices[(choices.index(self._engine) + 1) % len(choices)] self._relayout() self._clamp_cursor() @@ -2702,8 +2810,9 @@ class GraphView(NavMixin, ScrollView, can_focus=True): # Name the interpreter. The launcher runs $IDATUI_PYTHON (default # ~/ida-venv), which is NOT the repo .venv the tests use, so "not # installed" on its own sends people to check the wrong python. - note = ("" if graph_triskel.available() - else f" (no pytriskel in {sys.executable})") + note = ( + "" if graph_triskel.available() else f" (no pytriskel in {sys.executable})" + ) # A fallback with no reason is a bug report nobody can file. if self.lay and self.lay.stats.get("engine_error"): note = f" \u2014 {self.lay.stats['engine_error']}" @@ -2781,10 +2890,12 @@ class GraphView(NavMixin, ScrollView, can_focus=True): self.scroll_to(y=y, x=x, animate=False) if not defer: return + # Setting virtual_size then scrolling immediately clamps to 0 (max_scroll # isn't recomputed until layout), so apply it again after the refresh. def _again() -> None: self.scroll_to(y=y, x=x, animate=False) + self.call_after_refresh(_again) def _center_cursor(self) -> None: @@ -2818,11 +2929,17 @@ class GraphView(NavMixin, ScrollView, can_focus=True): return None best, best_d = None, None for n in self.lay.nodes: - dx = 0.0 if n.x <= col <= n.right else min(abs(col - n.x), - abs(col - n.right)) - dy = 0.0 if n.y <= row <= n.bottom else min(abs(row - n.y), - abs(row - n.bottom)) - d = (dx * 0.5) ** 2 + dy ** 2 + dx = ( + 0.0 + if n.x <= col <= n.right + else min(abs(col - n.x), abs(col - n.right)) + ) + dy = ( + 0.0 + if n.y <= row <= n.bottom + else min(abs(row - n.y), abs(row - n.bottom)) + ) + d = (dx * 0.5) ** 2 + dy**2 if best_d is None or d < best_d: best, best_d = n, d return best @@ -2845,21 +2962,23 @@ class GraphView(NavMixin, ScrollView, can_focus=True): return False left, top, _w, _h = rect gw, gh = _MINI_W - 2, _MINI_H - 2 - c, r = x - left - 1, y - top - 1 # inside the border + c, r = x - left - 1, y - top - 1 # inside the border if not (0 <= c < gw and 0 <= r < gh): return False lay = self.lay sx = max(lay.width / gw, 1e-9) sy = max(lay.height / gh, 1e-9) - cx, cy = (c + 0.5) * sx, (r + 0.5) * sy # centre of that mini-cell + cx, cy = (c + 0.5) * sx, (r + 0.5) * sy # centre of that mini-cell n = self._nearest_node(cy, cx) if n is None: - self.scroll_to(x=max(0, int(cx - self.size.width / 2)), - y=max(0, int(cy - self.size.height / 2)), - animate=False) + self.scroll_to( + x=max(0, int(cx - self.size.width / 2)), + y=max(0, int(cy - self.size.height / 2)), + animate=False, + ) return True if n.id == self.cursor_node: - return True # already there; don't churn while dragging + return True # already there; don't churn while dragging self.cursor_node = n.id self.cursor_row = 0 self.cursor_x = 0 @@ -2876,7 +2995,7 @@ class GraphView(NavMixin, ScrollView, can_focus=True): return if self._minimap_seek(off.x, off.y): self._drag = None - self._drag_map = True # keep scrubbing while the button is held + self._drag_map = True # keep scrubbing while the button is held return self._drag_map = False self._drag = (off.x, off.y, self.scroll_offset.x, self.scroll_offset.y) @@ -2886,7 +3005,7 @@ class GraphView(NavMixin, ScrollView, can_focus=True): self._drag = None self._drag_map = False if was_pan: - self._snap_into_view() # don't leave them adrift in the padding + self._snap_into_view() # don't leave them adrift in the padding def on_mouse_move(self, event) -> None: # type: ignore[no-untyped-def] if not event.button: @@ -2901,8 +3020,9 @@ class GraphView(NavMixin, ScrollView, can_focus=True): if self._drag is None: return x0, y0, sx, sy = self._drag - self.scroll_to(x=max(0, sx + (x0 - off.x)), y=max(0, sy + (y0 - off.y)), - animate=False) + self.scroll_to( + x=max(0, sx + (x0 - off.x)), y=max(0, sy + (y0 - off.y)), animate=False + ) def on_click(self, event) -> None: # type: ignore[no-untyped-def] if self.lay is None: @@ -2923,8 +3043,7 @@ class GraphView(NavMixin, ScrollView, can_focus=True): return self.focus() self.cursor_node = n.id - self.cursor_row = max(0, min(row - n.y - 1, - max(len(self._rows(n.id)) - 1, 0))) + self.cursor_row = max(0, min(row - n.y - 1, max(len(self._rows(n.id)) - 1, 0))) self.cursor_x = max(0, col - n.x - 1 - _GPAD) self._clamp_cursor() self.refresh() @@ -2948,7 +3067,8 @@ class GraphView(NavMixin, ScrollView, can_focus=True): # 1. edge cells (an index query, never a painted canvas) for col, (ch, kind, eid) in self.lay.painting.cells_at_row( - row, col0, col0 + width).items(): + row, col0, col0 + width + ).items(): st = (_S_EDGE_HOT if eid in hot else base).get(kind, _S_GDIM) out.put(col - col0, ch, st) @@ -2961,8 +3081,7 @@ class GraphView(NavMixin, ScrollView, can_focus=True): self._draw_minimap_row(out, y, width) return out.strip().adjust_cell_length(width, _S_INSN) - def _draw_node_row(self, out: _CellRow, n: graph.Node, row: int, - col0: int) -> None: + def _draw_node_row(self, out: _CellRow, n: graph.Node, row: int, col0: int) -> None: cur = n.id == self.cursor_node bs = _S_GBORDER_CUR if cur else _S_GBORDER left = n.x - col0 @@ -2973,19 +3092,24 @@ class GraphView(NavMixin, ScrollView, can_focus=True): label = f"loc_{n.block.start:X}" if n.block else "" if b is not None and b.rows and b.rows[0].name: label = b.rows[0].name - out.text(left, graph.BOX["tl"] + graph.BOX["h"] * (w - 2) - + graph.BOX["tr"], bs) + out.text( + left, graph.BOX["tl"] + graph.BOX["h"] * (w - 2) + graph.BOX["tr"], bs + ) tag = f" {label} " if len(tag) <= w - 4: - st = _S_GENTRY if (self.fc and n.id == self.fc.entry) else ( - _S_GLABEL_CUR if cur else _S_GLABEL) + st = ( + _S_GENTRY + if (self.fc and n.id == self.fc.entry) + else (_S_GLABEL_CUR if cur else _S_GLABEL) + ) out.text(left + 2, tag, st) if n.block is not None and n.block.selfloop: out.put(left + w - 2, "↺", _S_EDGE[graph.E_BACK]) return if row == n.y + n.h - 1: - out.text(left, graph.BOX["bl"] + graph.BOX["h"] * (w - 2) - + graph.BOX["br"], bs) + out.text( + left, graph.BOX["bl"] + graph.BOX["h"] * (w - 2) + graph.BOX["br"], bs + ) return out.put(left, graph.BOX["v"], bs) out.put(left + w - 1, graph.BOX["v"], bs) @@ -2997,7 +3121,7 @@ class GraphView(NavMixin, ScrollView, can_focus=True): text_col = left + 1 + _GPAD h = rows[i] plain = self._row_plain(n.id, i) - if h is None: # collapsed summary + if h is None: # collapsed summary out.text(text_col, plain, _S_GDIM) else: c = text_col @@ -3019,9 +3143,15 @@ class GraphView(NavMixin, ScrollView, can_focus=True): if self.trail is not None and h is not None: k = self.trail.get(h.ea) if k is not None: - out.restyle(inner_a, inner_b, - _S_TRAIL_NOW if k == "now" else - _S_TRAIL_PAST if k == "past" else _S_TRAIL_FUTURE) + out.restyle( + inner_a, + inner_b, + _S_TRAIL_NOW + if k == "now" + else _S_TRAIL_PAST + if k == "past" + else _S_TRAIL_FUTURE, + ) if self._hl_word and plain: for a, bb in _word_occurrences(plain, self._hl_word): out.restyle(text_col + a, text_col + bb, _S_WORD) @@ -3049,8 +3179,9 @@ class GraphView(NavMixin, ScrollView, can_focus=True): sy = max(lay.height / gh, 1e-9) for lo, hi, col, _kind, _eid in lay.painting.vruns: c = min(int(col / sx), gw - 1) - for r in range(min(int(lo / sy), gh - 1), - min(int(hi / sy), gh - 1) + 1): + for r in range( + min(int(lo / sy), gh - 1), min(int(hi / sy), gh - 1) + 1 + ): if not grid[r][c]: grid[r][c] = 1 for n in lay.nodes: @@ -3069,7 +3200,7 @@ class GraphView(NavMixin, ScrollView, can_focus=True): # column, which otherwise eats the minimap's right border. if self._minimap_rect() is None or not (0 <= y < _MINI_H): return - left = self._minimap_rect()[0] # one source of truth with the hit-test + left = self._minimap_rect()[0] # one source of truth with the hit-test grid = self._minimap() gw, gh = _MINI_W - 2, _MINI_H - 2 lay = self.lay @@ -3118,7 +3249,9 @@ class GraphView(NavMixin, ScrollView, can_focus=True): # --------------------------------------------------------------------------- # class FunctionsPanel(Vertical): def compose(self) -> ComposeResult: - self._filter = Input(placeholder="filter (glob, e.g. sub_*) — Enter to apply", id="func-filter") + self._filter = Input( + placeholder="filter (glob, e.g. sub_*) — Enter to apply", id="func-filter" + ) self._filter.display = False yield self._filter table = DataTable(id="func-table", cursor_type="row", zebra_stripes=True) @@ -3188,8 +3321,9 @@ class XrefsScreen(ModalScreen): BINDINGS = [Binding("escape", "close", "Close")] - def __init__(self, label: str, items: list[tuple[object, str]], - preselect: int = 0) -> None: + def __init__( + self, label: str, items: list[tuple[object, str]], preselect: int = 0 + ) -> None: # payload is an int address, or (binary, address) for a caller in another # project binary; dismiss() hands it back untouched. super().__init__() @@ -3271,8 +3405,8 @@ class SymbolPalette(OptionListNav, ModalScreen): def __init__(self, funcs: list[Func], index=None, binary=None) -> None: super().__init__() self._funcs = funcs - self._index = index # ProjectIndex, when this is a project - self._binary = binary # label of the binary we're currently in + self._index = index # ProjectIndex, when this is a project + self._binary = binary # label of the binary we're currently in self._project_scope = False #: (binary|None, addr, name) — binary is None for a local hit self._results: list[tuple] = [] @@ -3280,8 +3414,10 @@ class SymbolPalette(OptionListNav, ModalScreen): def compose(self) -> ComposeResult: with Vertical(id="pal-box") as box: box.border_title = Text("symbols") - yield Input(placeholder="fuzzy find symbol… ↑↓ select · Enter open · Esc close", - id="pal-input") + yield Input( + placeholder="fuzzy find symbol… ↑↓ select · Enter open · Esc close", + id="pal-input", + ) yield OptionList(id="pal-list") def on_mount(self) -> None: @@ -3306,13 +3442,17 @@ class SymbolPalette(OptionListNav, ModalScreen): # rows: (binary|None, addr, name, match positions) if self._project_scope and self._index is not None: from .index import KIND_FUNC + # The trigram index already guarantees every hit CONTAINS the query, # so ranking only has to order them — an exact-substring rank (match # position, then name length) costs a find() per row instead of a # full fuzzy pass, and fetching 3x the display limit rather than 10x # keeps the per-keystroke work down on a big project. - hits = self._index.search(query, kind=KIND_FUNC, - limit=self.PROJECT_LIMIT * 3) if query else [] + hits = ( + self._index.search(query, kind=KIND_FUNC, limit=self.PROJECT_LIMIT * 3) + if query + else [] + ) q = query.lower() scored = [] for h in hits: @@ -3322,9 +3462,15 @@ class SymbolPalette(OptionListNav, ModalScreen): # on (position, length, text), and a bare sort() would then fall # through to comparing Hit objects, which aren't orderable. scored.sort(key=lambda t: (t[0], t[1], t[2], t[3].binary, t[3].addr)) - rows = [(h.binary, h.addr, h.text, - tuple(range(at, at + len(q))) if at < (1 << 30) else ()) - for at, _, _, h in scored[:self.PROJECT_LIMIT]] + rows = [ + ( + h.binary, + h.addr, + h.text, + tuple(range(at, at + len(q))) if at < (1 << 30) else (), + ) + for at, _, _, h in scored[: self.PROJECT_LIMIT] + ] elif query: scored = [] for f in self._funcs: @@ -3332,9 +3478,9 @@ class SymbolPalette(OptionListNav, ModalScreen): if m is not None: scored.append((m[0], m[1], f)) scored.sort(key=lambda t: (-t[0], t[2].name)) - rows = [(None, f.addr, f.name, pos) for _, pos, f in scored[:self.LIMIT]] + rows = [(None, f.addr, f.name, pos) for _, pos, f in scored[: self.LIMIT]] else: - rows = [(None, f.addr, f.name, ()) for f in self._funcs[:self.LIMIT]] + rows = [(None, f.addr, f.name, ()) for f in self._funcs[: self.LIMIT]] self._results = [(b, a, n) for b, a, n, _ in rows] ol = self.query_one(OptionList) ol.clear_options() @@ -3356,10 +3502,14 @@ class SymbolPalette(OptionListNav, ModalScreen): scope = "project" if self._project_scope else "this binary" cap = self.PROJECT_LIMIT if self._project_scope else self.LIMIT more = "+" if len(self._results) == cap else "" - hint = " (F2: this binary)" if self._project_scope else ( - " (F2: whole project)" if self._index is not None else "") + hint = ( + " (F2: this binary)" + if self._project_scope + else (" (F2: whole project)" if self._index is not None else "") + ) self.query_one("#pal-box").border_title = Text( - f"symbols [{scope}]: {len(self._results)}{more}{hint}") + f"symbols [{scope}]: {len(self._results)}{more}{hint}" + ) def action_choose(self) -> None: ol = self.query_one(OptionList) @@ -3380,8 +3530,12 @@ class SymbolPalette(OptionListNav, ModalScreen): def _str_display(text: str, limit: int = 200) -> str: """One-line, printable rendering of a string literal for the browser: escape the common control chars, drop the rest, and clip long bodies.""" - out = (text.replace("\\", "\\\\").replace("\n", "\\n") - .replace("\r", "\\r").replace("\t", "\\t")) + out = ( + text.replace("\\", "\\\\") + .replace("\n", "\\n") + .replace("\r", "\\r") + .replace("\t", "\\t") + ) out = "".join(ch if ch.isprintable() else "." for ch in out) return out[:limit] + ("\u2026" if len(out) > limit else "") @@ -3412,7 +3566,7 @@ class SearchPalette(OptionListNav, ModalScreen): super().__init__() self._program = program self._initial = initial - self._forced: str | None = None # F2: pin the mode + self._forced: str | None = None # F2: pin the mode self._hits: list = [] self._searched: tuple[str, str] | None = None # (mode, query) on screen self._busy = False @@ -3420,9 +3574,11 @@ class SearchPalette(OptionListNav, ModalScreen): def compose(self) -> ComposeResult: with Vertical(id="pal-box") as box: box.border_title = Text("search") - yield Input(placeholder="text, or bytes like 48 8b ?? c3 \u00b7 " - "Enter search \u00b7 F2 mode \u00b7 Esc close", - id="pal-input") + yield Input( + placeholder="text, or bytes like 48 8b ?? c3 \u00b7 " + "Enter search \u00b7 F2 mode \u00b7 Esc close", + id="pal-input", + ) yield OptionList(id="pal-list") def on_mount(self) -> None: @@ -3453,16 +3609,17 @@ class SearchPalette(OptionListNav, ModalScreen): elif q: state = "Enter searches" self.query_one("#pal-box").border_title = Text( - f"search [{mode}{pinned}]" + (f": {state}" if state else "")) + f"search [{mode}{pinned}]" + (f": {state}" if state else "") + ) def action_mode(self) -> None: mode, _ = self._mode_query() self._forced = search.TEXT if mode == search.BYTES else search.BYTES - self._searched = None # the results on screen are for the old mode + self._searched = None # the results on screen are for the old mode self._retitle() def on_input_changed(self, event: Input.Changed) -> None: - event.stop() # modal inputs bubble to the app's own #search handler + event.stop() # modal inputs bubble to the app's own #search handler self._retitle() def on_input_submitted(self, event: Input.Submitted) -> None: @@ -3495,14 +3652,14 @@ class SearchPalette(OptionListNav, ModalScreen): @work(thread=True, exclusive=True, group="dbsearch") def _search(self, mode: str, query: str) -> None: try: - hits, err, truncated = self._program.search(query, mode, - limit=self.LIMIT) + hits, err, truncated = self._program.search(query, mode, limit=self.LIMIT) except Exception as e: # noqa: BLE001 -- a search must not kill the app hits, err, truncated = [], str(e), False self.app.call_from_thread(self._present, mode, query, hits, err, truncated) - def _present(self, mode: str, query: str, hits: list, err: str | None, - truncated: bool) -> None: + def _present( + self, mode: str, query: str, hits: list, err: str | None, truncated: bool + ) -> None: self._busy = False self._hits = hits # Remember what these results ARE, not what the box says now: the user @@ -3533,8 +3690,10 @@ class SearchPalette(OptionListNav, ModalScreen): self._retitle("no match") else: n = len(hits) - self._retitle(f"{n}{'+' if truncated else ''} " - f"hit{'' if n == 1 else 's'} \u2014 Enter opens") + self._retitle( + f"{n}{'+' if truncated else ''} " + f"hit{'' if n == 1 else 's'} \u2014 Enter opens" + ) # -- moving / choosing --------------------------------------------------- # def action_choose(self) -> None: @@ -3569,9 +3728,10 @@ class StringsPalette(OptionListNav, ModalScreen): super().__init__() # Pre-render + pre-lower once: filtering runs on every keystroke and a # big binary has tens of thousands of strings. - self._rows = [(s, d, d.lower()) - for s in strings for d in (_str_display(s.text),)] - self._index = index # ProjectIndex, when this is a project + self._rows = [ + (s, d, d.lower()) for s in strings for d in (_str_display(s.text),) + ] + self._index = index # ProjectIndex, when this is a project self._binary = binary self._project_scope = False #: (binary|None, addr, display text) — binary is None for a local hit @@ -3580,8 +3740,11 @@ class StringsPalette(OptionListNav, ModalScreen): def compose(self) -> ComposeResult: with Vertical(id="pal-box") as box: box.border_title = Text("strings") - yield Input(placeholder="filter strings\u2026 \u2191\u2193 select \u00b7 " - "Enter jump \u00b7 Esc close", id="pal-input") + yield Input( + placeholder="filter strings\u2026 \u2191\u2193 select \u00b7 " + "Enter jump \u00b7 Esc close", + id="pal-input", + ) yield OptionList(id="pal-list") def on_mount(self) -> None: @@ -3607,19 +3770,23 @@ class StringsPalette(OptionListNav, ModalScreen): # rows: (binary|None, addr, length, display text, match offset) if self._project_scope and self._index is not None: from .index import KIND_STRING - hits = self._index.search(query, kind=KIND_STRING, - limit=self.PROJECT_LIMIT * 3) if query else [] + + hits = ( + self._index.search( + query, kind=KIND_STRING, limit=self.PROJECT_LIMIT * 3 + ) + if query + else [] + ) rows = [] for h in hits: disp = _str_display(h.text) - rows.append((h.binary, h.addr, len(h.text), disp, - disp.lower().find(q))) + rows.append((h.binary, h.addr, len(h.text), disp, disp.lower().find(q))) # the index already guarantees a match, so ranking only orders them: # earliest match, then shortest, with a stable (binary, addr) tiebreak # (a literal shared by two binaries would otherwise be unordered). - rows.sort(key=lambda r: (r[4] if r[4] >= 0 else 1 << 30, - r[2], r[0], r[1])) - rows = rows[:self.PROJECT_LIMIT] + rows.sort(key=lambda r: (r[4] if r[4] >= 0 else 1 << 30, r[2], r[0], r[1])) + rows = rows[: self.PROJECT_LIMIT] else: rows = [] for s, disp, low in self._rows: @@ -3650,10 +3817,14 @@ class StringsPalette(OptionListNav, ModalScreen): scope = "project" if self._project_scope else "this binary" cap = self.PROJECT_LIMIT if self._project_scope else self.LIMIT more = "+" if len(rows) == cap else "" - hint = " (F2: this binary)" if self._project_scope else ( - " (F2: whole project)" if self._index is not None else "") + hint = ( + " (F2: this binary)" + if self._project_scope + else (" (F2: whole project)" if self._index is not None else "") + ) self.query_one("#pal-box").border_title = Text( - f"strings [{scope}]: {len(self._results)}{more} of {len(self._rows)}{hint}") + f"strings [{scope}]: {len(self._results)}{more} of {len(self._rows)}{hint}" + ) def action_choose(self) -> None: ol = self.query_one(OptionList) @@ -3674,74 +3845,92 @@ class StringsPalette(OptionListNav, ModalScreen): #: The keyboard cheatsheet (F1). Grouped by task rather than by widget, which is #: what makes it readable; keep it in step with the BINDINGS above it. _HELP = ( - ("Navigate", ( - ("Enter", "follow the symbol under the cursor"), - ("Esc", "back (navigation history)"), - ("g", "goto address or symbol"), - ("Ctrl+N", "find symbol (fuzzy)"), - ("\"", "strings browser"), - ("x", "cross-references to the symbol"), - ("L", "continuous listing at the cursor"), - ("Ctrl+O", "switch binary (projects)"), - )), - ("Views", ( - ("Tab / F5", "disassembly \u21c4 pseudocode"), - ("Space", "control-flow graph \u21c4 text"), - ("s", "split view: listing + pseudocode"), - ("Tab", "in split: switch the driving pane"), - ("\\", "hex view"), - ("B", "cycle the opcode-bytes column"), - ("Ctrl+B", "show/hide the names pane"), - ("Ctrl+T", "structs / types editor"), - ("Ctrl+F", "search the database: text or bytes"), - ("Ctrl+R", "refresh the current view in place"), - ("Ctrl+E", "export findings as markdown"), - ("Ctrl+P", "command palette"), - )), - ("Move", ( - ("j / k", "down / up"), - ("Ctrl+D / Ctrl+U", "half page down / up"), - ("PgDn / PgUp", "page down / up"), - ("Ctrl+Home / Ctrl+End", "top / bottom (G also)"), - ("Home / End", "start / end of line"), - ("Shift+Home", "start of the instruction / code"), - ("h / l", "column left / right"), - ("w / b", "word forward / back"), - )), - ("Edit", ( - ("n", "rename"), - ("y", "set type (prototype, local or global)"), - (";", "comment"), - ("c", "make code"), - ("p", "make function"), - ("d", "make data"), - ("a", "make string"), - ("u", "undefine"), - ("o / O", "literal format: hex/dec/bin/char/offset"), - ("Ctrl+S", "save the database"), - )), - ("Search", ( - ("/", "search forward (repeat to continue)"), - ("?", "search backward"), - ("N", "previous match"), - ("Ctrl+Y", "copy the current line"), - ("F1 / H", "this cheatsheet"), - ("q", "quit"), - )), - ("Graph (Space)", ( - ("j / k", "line up/down, crossing blocks"), - ("h / l", "column left / right"), - ("J / K", "follow an edge to a successor / predecessor"), - ("w / b", "next / previous block in layout order"), - ("0", "jump to the entry block"), - ("z", "zoom: full \u2192 compact \u2192 collapsed"), - ("m", "show/hide the minimap"), - ("e", "layout engine: auto \u2192 native \u2192 triskel"), - ("f", "centre on the current block"), - ("Enter", "follow (stays in the graph if it lands here)"), - ("drag / click", "pan / put the cursor in a block"), - ("click minimap", "jump the view there (drag to scrub)"), - )), + ( + "Navigate", + ( + ("Enter", "follow the symbol under the cursor"), + ("Esc", "back (navigation history)"), + ("g", "goto address or symbol"), + ("Ctrl+N", "find symbol (fuzzy)"), + ('"', "strings browser"), + ("x", "cross-references to the symbol"), + ("L", "continuous listing at the cursor"), + ("Ctrl+O", "switch binary (projects)"), + ), + ), + ( + "Views", + ( + ("Tab / F5", "disassembly \u21c4 pseudocode"), + ("Space", "control-flow graph \u21c4 text"), + ("s", "split view: listing + pseudocode"), + ("Tab", "in split: switch the driving pane"), + ("\\", "hex view"), + ("B", "cycle the opcode-bytes column"), + ("Ctrl+B", "show/hide the names pane"), + ("Ctrl+T", "structs / types editor"), + ("Ctrl+F", "search the database: text or bytes"), + ("Ctrl+R", "refresh the current view in place"), + ("Ctrl+E", "export findings as markdown"), + ("Ctrl+P", "command palette"), + ), + ), + ( + "Move", + ( + ("j / k", "down / up"), + ("Ctrl+D / Ctrl+U", "half page down / up"), + ("PgDn / PgUp", "page down / up"), + ("Ctrl+Home / Ctrl+End", "top / bottom (G also)"), + ("Home / End", "start / end of line"), + ("Shift+Home", "start of the instruction / code"), + ("h / l", "column left / right"), + ("w / b", "word forward / back"), + ), + ), + ( + "Edit", + ( + ("n", "rename"), + ("y", "set type (prototype, local or global)"), + (";", "comment"), + ("c", "make code"), + ("p", "make function"), + ("d", "make data"), + ("a", "make string"), + ("u", "undefine"), + ("o / O", "literal format: hex/dec/bin/char/offset"), + ("Ctrl+S", "save the database"), + ), + ), + ( + "Search", + ( + ("/", "search forward (repeat to continue)"), + ("?", "search backward"), + ("N", "previous match"), + ("Ctrl+Y", "copy the current line"), + ("F1 / H", "this cheatsheet"), + ("q", "quit"), + ), + ), + ( + "Graph (Space)", + ( + ("j / k", "line up/down, crossing blocks"), + ("h / l", "column left / right"), + ("J / K", "follow an edge to a successor / predecessor"), + ("w / b", "next / previous block in layout order"), + ("0", "jump to the entry block"), + ("z", "zoom: full \u2192 compact \u2192 collapsed"), + ("m", "show/hide the minimap"), + ("e", "layout engine: auto \u2192 native \u2192 triskel"), + ("f", "centre on the current block"), + ("Enter", "follow (stays in the graph if it lands here)"), + ("drag / click", "pan / put the cursor in a block"), + ("click minimap", "jump the view there (drag to scrub)"), + ), + ), ) @@ -3764,19 +3953,24 @@ class QuitScreen(ModalScreen): self._labels = labels def compose(self) -> ComposeResult: - what = (f"{len(self._labels)} databases have unsaved changes" - if len(self._labels) > 1 else "unsaved changes") + what = ( + f"{len(self._labels)} databases have unsaved changes" + if len(self._labels) > 1 + else "unsaved changes" + ) with Vertical(id="quit-box") as box: box.border_title = Text(f"\u26a0 {what}") 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) + "\nFinal managed leases discard; shared/GUI sessions stay open.", _S_DIM + ) yield Static(body, id="quit-list") - yield Static("s save & quit d discard / leave & quit Esc cancel", - id="quit-help") + yield Static( + "s save & quit d discard / leave & quit Esc cancel", + id="quit-help", + ) def action_save(self) -> None: self.dismiss("save") @@ -3810,13 +4004,14 @@ class HelpScreen(ModalScreen): with VerticalScroll(id="help-body"): with Horizontal(id="help-cols"): for c in range(cols): - chunk = _HELP[c * per:(c + 1) * per] + chunk = _HELP[c * per : (c + 1) * per] if not chunk: continue with Vertical(classes="help-col"): for title, rows in chunk: - card = Static(self._card(rows), - classes="help-card", markup=False) + card = Static( + self._card(rows), classes="help-card", markup=False + ) card.border_title = title yield card yield Static("Esc · F1 · H to close", id="help-foot") @@ -3843,7 +4038,7 @@ class HelpScreen(ModalScreen): n = len(ws) for cols in range(min(n, 4), 1, -1): per = -(-n // cols) - chunks = [ws[c * per:(c + 1) * per] for c in range(cols)] + chunks = [ws[c * per : (c + 1) * per] for c in range(cols)] total = sum(max(c) for c in chunks if c) + (cols - 1) if total <= avail: return cols @@ -3884,13 +4079,15 @@ class RegWriteScreen(OptionListNav, ModalScreen): def __init__(self, rows, idx: int) -> None: super().__init__() - self._rows = rows # (name, value, last_write, next_write) + self._rows = rows # (name, value, last_write, next_write) self._idx = idx def compose(self) -> ComposeResult: with Vertical(id="pal-box") as box: box.border_title = Text(f"registers at t={self._idx:,}") - box.border_subtitle = Text("Enter seeks to the write \u00b7 f seeks forward") + box.border_subtitle = Text( + "Enter seeks to the write \u00b7 f seeks forward" + ) yield OptionList(id="pal-list") def on_mount(self) -> None: @@ -3899,8 +4096,9 @@ class RegWriteScreen(OptionListNav, ModalScreen): for name, val, last, nxt in self._rows: label = Text() label.append(f" {name:>4} ", _S_MNEM) - label.append(f"{val:#018x} " if val > 0xFFFFFFFF else f"{val:#010x} ", - _S_INSN) + label.append( + f"{val:#018x} " if val > 0xFFFFFFFF else f"{val:#010x} ", _S_INSN + ) if last is None: label.append("never written in this trace", _S_DIM) elif last == self._idx: @@ -3992,15 +4190,16 @@ class TraceDock(Vertical): continue hot = name in changed body.append(f" {name:>4} ", _S_MNEM if hot else _S_DIM) - body.append(f"{v:#018x}\n" if v > 0xFFFFFFFF else f"{v:#010x}\n", - _S_DATA if hot else (_S_LABEL if name == pc else _S_INSN)) + body.append( + f"{v:#018x}\n" if v > 0xFFFFFFFF else f"{v:#010x}\n", + _S_DATA if hot else (_S_LABEL if name == pc else _S_INSN), + ) self.query_one("#trace-regs", Static).update(body) self._render_stack(t) tl = self.query_one(TraceTimeline) tl.idx = self.idx tl.refresh() - STACK_WORDS = 8 def _render_stack(self, t) -> None: # type: ignore[no-untyped-def] @@ -4031,8 +4230,13 @@ class TraceDock(Vertical): v = int.from_bytes(data, "little") out.append(f"{v:0{width * 2}x}\n", _S_DATA if k == 0 else _S_INSN) elif any(known): - out.append("".join(f"{b:02x}" if known[i] else "??" - for i, b in enumerate(data)) + "\n", _S_INSN) + out.append( + "".join( + f"{b:02x}" if known[i] else "??" for i, b in enumerate(data) + ) + + "\n", + _S_INSN, + ) else: out.append("?" * (width * 2) + "\n", _S_SEP) self.query_one("#trace-stack", Static).update(out) @@ -4100,18 +4304,27 @@ class LoadOptionsScreen(OptionListNav, ModalScreen): def compose(self) -> ComposeResult: from .formats import PROCESSORS + self._all = list(PROCESSORS) with Vertical(id="pal-box") as box: box.border_title = Text("unrecognised file \u2014 how should IDA load it?") - yield Static(f" {os.path.basename(self._path)} ({self._nbytes:,} bytes) " - f"\u2014 no loader matched; without a processor IDA " - f"assumes x86 at 0", id="load-note", markup=False) + yield Static( + f" {os.path.basename(self._path)} ({self._nbytes:,} bytes) " + f"\u2014 no loader matched; without a processor IDA " + f"assumes x86 at 0", + id="load-note", + markup=False, + ) yield Input(placeholder="filter processors\u2026", id="pal-input") yield OptionList(id="pal-list") - yield Input(placeholder="load address, e.g. 0x8000000 (blank = 0)", - id="load-base") - yield Static(" Enter accept \u00b7 Tab base address \u00b7 " - "Esc load as IDA would", id="load-help", markup=False) + yield Input( + placeholder="load address, e.g. 0x8000000 (blank = 0)", id="load-base" + ) + yield Static( + " Enter accept \u00b7 Tab base address \u00b7 Esc load as IDA would", + id="load-help", + markup=False, + ) def on_mount(self) -> None: self._apply("") @@ -4146,8 +4359,11 @@ class LoadOptionsScreen(OptionListNav, ModalScreen): def _apply(self, query: str) -> None: q = query.lower() - rows = [(name, desc) for name, desc in self._all - if not q or q in name.lower() or q in desc.lower()] + rows = [ + (name, desc) + for name, desc in self._all + if not q or q in name.lower() or q in desc.lower() + ] # An unlisted processor is still valid: IDA has 73 modules and this # offers 20, so a typed name that matches nothing is taken literally # rather than refused. @@ -4166,7 +4382,8 @@ class LoadOptionsScreen(OptionListNav, ModalScreen): if rows: ol.highlighted = 0 self.query_one("#pal-box").border_title = Text( - f"unrecognised file \u2014 processor? ({len(rows)})") + f"unrecognised file \u2014 processor? ({len(rows)})" + ) def action_choose(self) -> None: ol = self.query_one(OptionList) @@ -4181,14 +4398,16 @@ class LoadOptionsScreen(OptionListNav, ModalScreen): base = int(raw, 0) except ValueError: self.query_one("#load-help", Static).update( - f" {raw!r} is not an address \u2014 try 0x8000000") + f" {raw!r} is not an address \u2014 try 0x8000000" + ) self.query_one("#load-base", Input).focus() return if base % 16: # IDA's -b is in paragraphs, so an unaligned base can't be # expressed and would quietly load somewhere else. self.query_one("#load-help", Static).update( - f" {base:#x} must be 16-byte aligned") + f" {base:#x} must be 16-byte aligned" + ) self.query_one("#load-base", Input).focus() return self.dismiss({"processor": self._results[i][0], "base": base}) @@ -4214,8 +4433,11 @@ class ProjectPalette(OptionListNav, ModalScreen): def compose(self) -> ComposeResult: with Vertical(id="pal-box") as box: box.border_title = Text("binaries") - yield Input(placeholder="filter binaries\u2026 \u2191\u2193 select \u00b7 " - "Enter switch \u00b7 Esc close", id="pal-input") + yield Input( + placeholder="filter binaries\u2026 \u2191\u2193 select \u00b7 " + "Enter switch \u00b7 Esc close", + id="pal-input", + ) yield OptionList(id="pal-list") def on_mount(self) -> None: @@ -4232,16 +4454,20 @@ class ProjectPalette(OptionListNav, ModalScreen): def _apply(self, query: str) -> None: q = query.lower() - rows = [e for e in self._entries - if not q or q in e["label"].lower() or q in e["source"].lower()] + rows = [ + e + for e in self._entries + if not q or q in e["label"].lower() or q in e["source"].lower() + ] self._results = rows ol = self.query_one(OptionList) ol.clear_options() opts = [] for e in rows: label = Text() - label.append("\u25b8 " if e["active"] else " ", - _S_MNEM if e["active"] else _S_DIM) + label.append( + "\u25b8 " if e["active"] else " ", _S_MNEM if e["active"] else _S_DIM + ) label.append(f"{e['label']:<22}", _S_LABEL) if e["resident"]: mb = e.get("memory_mb") or 0 @@ -4261,7 +4487,8 @@ class ProjectPalette(OptionListNav, ModalScreen): active = next((i for i, e in enumerate(rows) if e["active"]), 0) ol.highlighted = active self.query_one("#pal-box").border_title = Text( - f"binaries: {len(rows)} of {len(self._entries)}") + f"binaries: {len(rows)} of {len(self._entries)}" + ) def action_choose(self) -> None: i = self.query_one(OptionList).highlighted @@ -4312,7 +4539,7 @@ _LOGO_PATH = os.path.join(_REPO_ROOT, "logo.ans") #: The same artwork as a real image, for terminals that can draw one. logo.ans #: is half-blocks (two pixels per cell); this is a transparent PNG at 768px. LOGO_PNG = os.path.join(_REPO_ROOT, "logo.png") -_LOGO_BOX = (60, 33) # the most room the splash will give the art +_LOGO_BOX = (60, 33) # the most room the splash will give the art #: Rows the loading box spends on everything that is not the artwork: border 2, #: padding 2, the art's margin 1, title 1, note 1 + margin 1, help 1 + margin 1. LOGO_CHROME_ROWS = 10 @@ -4342,8 +4569,13 @@ def logo_cells(max_rows: int | None = None) -> tuple[int, int]: if max_rows is None or max_rows >= _logo_cells[1]: return _logo_cells px = kittygfx.png_size(LOGO_PNG) - return (kittygfx.fit(px, _LOGO_BOX[0], max(max_rows, 1)) if px - else (_LOGO_BOX[0], max(max_rows, 1))) + return ( + kittygfx.fit(px, _LOGO_BOX[0], max(max_rows, 1)) + if px + else (_LOGO_BOX[0], max(max_rows, 1)) + ) + + _logo_cache: object = False # False == not yet loaded (None == absent/unreadable) @@ -4374,9 +4606,9 @@ class LoadingScreen(ModalScreen): super().__init__() self._title = title self._note = note - self._image = False # drawing the real image, not the block art - self._cells: tuple[int, int] | None = None # image size, in cells - self._last_place = 0.0 # throttles re-anchoring after a repaint + self._image = False # drawing the real image, not the block art + self._cells: tuple[int, int] | None = None # image size, in cells + self._last_place = 0.0 # throttles re-anchoring after a repaint def _room(self) -> int: """Rows left for artwork once the box's own furniture is paid for.""" @@ -4398,11 +4630,16 @@ class LoadingScreen(ModalScreen): room = self._room() cols, rows = logo_cells(room) self._cells = (cols, rows) - kittygfx.log(f"compose: supported={kittygfx.supported()} " - f"app.size={self.app.size} room={room} " - f"cells={cols}x{rows} natural={logo_cells()}") - if (kittygfx.supported() and self.app.size.width >= 64 - and room >= LOGO_MIN_ROWS): + kittygfx.log( + f"compose: supported={kittygfx.supported()} " + f"app.size={self.app.size} room={room} " + f"cells={cols}x{rows} natural={logo_cells()}" + ) + if ( + kittygfx.supported() + and self.app.size.width >= 64 + and room >= LOGO_MIN_ROWS + ): self._image = True blank = Static("\n" * (rows - 1), id="loading-image") blank.styles.height = rows @@ -4419,8 +4656,10 @@ class LoadingScreen(ModalScreen): yield Static(Align.center(logo), id="loading-logo") yield Static(f"\u23f3 loading {self._title}", id="loading-title") yield Static(self._note, id="loading-note") - yield Static("first open of a big binary can take a while \u00b7 " - "Esc to hide", id="loading-help") + yield Static( + "first open of a big binary can take a while \u00b7 Esc to hide", + id="loading-help", + ) def update_note(self, text: str) -> None: try: @@ -4466,7 +4705,7 @@ class LoadingScreen(ModalScreen): # image is scaled into exactly it, so a resize needs no relayout. cols, rows = self._cells or logo_cells() rows = min(rows, region.height) - col = region.x + max((region.width - cols) // 2, 0) # centre it + col = region.x + max((region.width - cols) // 2, 0) # centre it kittygfx.place(region.y, col, min(cols, region.width), rows) def on_mount(self) -> None: @@ -4557,7 +4796,7 @@ class StructEditor(ModalScreen): def __init__(self, program: Program) -> None: super().__init__() self._program = program - self._all: list[Struct] = [] # every struct the database has + self._all: list[Struct] = [] # every struct the database has self._structs: list[Struct] = [] # the VISIBLE rows (== _all when unfiltered) self._filter = "" self._loaded: str | None = None # name currently in the editor @@ -4588,9 +4827,11 @@ class StructEditor(ModalScreen): with Horizontal(id="se-panes"): with Vertical(id="se-left"): yield Static("structs", id="se-title") - yield Input(placeholder="fuzzy filter\u2026 \u2191\u2193 pick \u00b7 " - "Enter edit \u00b7 Esc clear", - id="se-filter") + yield Input( + placeholder="fuzzy filter\u2026 \u2191\u2193 pick \u00b7 " + "Enter edit \u00b7 Esc clear", + id="se-filter", + ) yield OptionList(id="se-list") with Vertical(id="se-right"): yield Static("C definition", id="se-hint") @@ -4599,7 +4840,8 @@ class StructEditor(ModalScreen): yield Static( "Enter edit · / filter · Ctrl+S save · Ctrl+Y copy · Ctrl+N new · " "d/Del delete", - id="se-status") + id="se-status", + ) def on_mount(self) -> None: self._refresh() @@ -4657,7 +4899,7 @@ class StructEditor(ModalScreen): opts = [] for s, pos in rows: kw = "union" if s.is_union else "struct" - name = s.name if len(s.name) <= width else s.name[:width - 1] + "\u2026" + name = s.name if len(s.name) <= width else s.name[: width - 1] + "\u2026" label = Text() nm = Text(f"{name:<{width}}", style=_S_LABEL) for p in pos: @@ -4670,8 +4912,9 @@ class StructEditor(ModalScreen): if rows: idx = 0 if select is not None: - idx = next((i for i, s in enumerate(self._structs) - if s.name == select), 0) + idx = next( + (i for i, s in enumerate(self._structs) if s.name == select), 0 + ) ol.highlighted = idx cap = "structs" if q: @@ -4755,8 +4998,9 @@ class StructEditor(ModalScreen): def _move_highlight(self, delta: int) -> None: ol = self._list_from_filter() if ol is not None: - ol.highlighted = max(0, min((ol.highlighted or 0) + delta, - ol.option_count - 1)) + ol.highlighted = max( + 0, min((ol.highlighted or 0) + delta, ol.option_count - 1) + ) def _page(self, direction: int) -> None: ol = self._list_from_filter() @@ -4808,16 +5052,20 @@ class StructEditor(ModalScreen): formatted = None self.app.call_from_thread(self._after_save, name, err, text, formatted) - def _after_save(self, name: str | None, err: str | None, text: str, - formatted: str | None) -> None: + def _after_save( + self, name: str | None, err: str | None, text: str, formatted: str | None + ) -> None: if err: # IDA's parse error is usually empty/cryptic; name the likely cause. msg = err.strip() if not msg or "parse" in msg.lower() or "fail" in msg.lower(): bad = self._reserved_field(text) - msg = (f"'{bad}' is a reserved name in IDA's C parser — rename " - f"that field to save" if bad else - "IDA couldn't parse it (unknown type or reserved field name?)") + msg = ( + f"'{bad}' is a reserved name in IDA's C parser — rename " + f"that field to save" + if bad + else "IDA couldn't parse it (unknown type or reserved field name?)" + ) self._set_status(f"save failed — {msg}", error=True) return self._loaded = name @@ -4858,7 +5106,8 @@ class StructEditor(ModalScreen): kind = "union" if s.is_union else "struct" self.app.push_screen( ConfirmScreen(f"Delete {kind} '{s.name}' ?"), - lambda ok, name=s.name: self._delete(name) if ok else None) + lambda ok, name=s.name: self._delete(name) if ok else None, + ) @work(thread=True, exclusive=True, group="se-del") def _delete(self, name: str) -> None: @@ -4899,8 +5148,9 @@ class StructEditor(ModalScreen): if self._filter_focused() or self._filter: self._clear_filter() return - self._confirm_discard(lambda: self.dismiss(None), - "Discard unsaved changes and close?") + self._confirm_discard( + lambda: self.dismiss(None), "Discard unsaved changes and close?" + ) def _set_status(self, text, error: bool = False) -> None: # type: ignore[no-untyped-def] st = self.query_one("#se-status", Static) @@ -4918,16 +5168,16 @@ class StructEditor(ModalScreen): IDATUI_THEME = Theme( name="idatui", dark=True, - background="#12161c", # deep blue-black, softer than pure black - surface="#181d25", # views - panel="#212832", # dialogs, status bar, gutters + background="#12161c", # deep blue-black, softer than pure black + surface="#181d25", # views + panel="#212832", # dialogs, status bar, gutters foreground="#d6d9de", - primary="#5aa0d6", # focus / links: the one cool accent + primary="#5aa0d6", # focus / links: the one cool accent secondary="#2f5d82", - accent="#d0a215", # the same amber as a search match — one meaning - warning="#c9762f", # burnt orange — distinct from accent, reads as care - error="#ff5f5f", # already used for failure text - success="#6a9955", # already used for comments + accent="#d0a215", # the same amber as a search match — one meaning + warning="#c9762f", # burnt orange — distinct from accent, reads as care + error="#ff5f5f", # already used for failure text + success="#6a9955", # already used for comments ) @@ -4942,73 +5192,146 @@ class IdaCommands(Provider): app = self.app va = app._palette_action # dispatch to the focused code view return ( - ("Goto address / symbol…", "jump to an address or name (g)", - app.action_goto), + ( + "Goto address / symbol…", + "jump to an address or name (g)", + app.action_goto, + ), ("Find symbol…", "fuzzy function finder (Ctrl+N)", app.action_symbols), - ("Strings…", "browse every string in the binary (\")", - app.action_strings), - ("Switch binary…", "another binary in the project (Ctrl+O)", - app.action_switch_binary), - ("Search database…", "disassembly text or a byte pattern with " - "wildcards (Ctrl+F)", app.action_find), - ("Export findings…", "your comments, names and types as markdown " - "(Ctrl+E)", app.action_export), - ("Keyboard shortcuts", "the key cheatsheet (F1 or H)", - app.action_help), - ("Follow symbol under cursor", "jump to the referenced symbol (Enter)", - lambda: va("follow")), - ("Show xrefs to symbol", "cross-references to the cursor symbol (x)", - lambda: va("xrefs")), + ("Strings…", 'browse every string in the binary (")', app.action_strings), + ( + "Switch binary…", + "another binary in the project (Ctrl+O)", + app.action_switch_binary, + ), + ( + "Search database…", + "disassembly text or a byte pattern with wildcards (Ctrl+F)", + app.action_find, + ), + ( + "Export findings…", + "your comments, names and types as markdown (Ctrl+E)", + app.action_export, + ), + ("Keyboard shortcuts", "the key cheatsheet (F1 or H)", app.action_help), + ( + "Follow symbol under cursor", + "jump to the referenced symbol (Enter)", + lambda: va("follow"), + ), + ( + "Show xrefs to symbol", + "cross-references to the cursor symbol (x)", + lambda: va("xrefs"), + ), ("Back", "navigation history (Esc)", app.action_back), - ("Toggle disassembly / pseudocode", "decompile / listing (F5, Tab)", - app.action_toggle_view), - ("Continuous listing here", "flat segment listing (L)", - app.action_continuous_here), + ( + "Toggle disassembly / pseudocode", + "decompile / listing (F5, Tab)", + app.action_toggle_view, + ), + ( + "Continuous listing here", + "flat segment listing (L)", + app.action_continuous_here, + ), ("Hex view", "raw bytes at the cursor (\\)", app.action_hex), - ("Split view (listing ⇄ pseudocode)", - "side-by-side synced views (s)", app.action_toggle_split), - ("Graph view (control flow)", - "the function's basic blocks as a graph (Space)", - app.action_toggle_graph), - ("Graph: cycle zoom", - "full → compact → collapsed (z, in the graph)", - lambda: va("zoom")), - ("Graph: toggle minimap", - "the overview box (m, in the graph)", lambda: va("minimap")), - ("Rename symbol…", "rename the symbol under the cursor (n)", - lambda: va("rename")), - ("Set type / prototype…", "retype the symbol under the cursor (y)", - lambda: va("retype")), + ( + "Split view (listing ⇄ pseudocode)", + "side-by-side synced views (s)", + app.action_toggle_split, + ), + ( + "Graph view (control flow)", + "the function's basic blocks as a graph (Space)", + app.action_toggle_graph, + ), + ( + "Graph: cycle zoom", + "full → compact → collapsed (z, in the graph)", + lambda: va("zoom"), + ), + ( + "Graph: toggle minimap", + "the overview box (m, in the graph)", + lambda: va("minimap"), + ), + ( + "Rename symbol…", + "rename the symbol under the cursor (n)", + lambda: va("rename"), + ), + ( + "Set type / prototype…", + "retype the symbol under the cursor (y)", + lambda: va("retype"), + ), ("Add comment…", "comment at the cursor (;)", lambda: va("comment")), ("Define code", "make code at the cursor (c)", lambda: va("define_code")), - ("Create function", "define a function at the cursor (p)", - lambda: va("define_func")), - ("Make data", "define a data item at the cursor (d)", - lambda: va("make_data")), - ("Make string", "define a string at the cursor (a)", - lambda: va("make_string")), - ("Undefine", "undefine the item at the cursor (u)", - lambda: va("undefine")), - ("Literal format: next", "cycle the literal under the cursor (o)", - lambda: va("op_format", "cycle")), - ("Literal format: previous", "the other way round (O)", - lambda: va("op_format", "back")), - *((f"Literal format: {label}", f"show the literal as {label} ({fmt})", - (lambda f=fmt: va("op_format", f))) - for fmt, label in (("hex", "hexadecimal"), ("dec", "decimal"), - ("oct", "octal"), ("bin", "binary"), - ("char", "a character"), - ("offset", "an offset (reference)"), - ("stack", "a stack variable"), - ("default", "IDA's own choice"))), - ("Toggle opcode bytes", "cycle the opcode-bytes column (B)", - lambda: va("toggle_opcodes")), - ("Structs / types editor", "view + edit local types (Ctrl+T)", - app.action_structs), - ("Filter functions…", "glob-filter the function list (/)", - app.action_filter), - ("Toggle names pane", "function-list sidebar (Ctrl+B)", - app.action_toggle_functions), + ( + "Create function", + "define a function at the cursor (p)", + lambda: va("define_func"), + ), + ( + "Make data", + "define a data item at the cursor (d)", + lambda: va("make_data"), + ), + ( + "Make string", + "define a string at the cursor (a)", + lambda: va("make_string"), + ), + ("Undefine", "undefine the item at the cursor (u)", lambda: va("undefine")), + ( + "Literal format: next", + "cycle the literal under the cursor (o)", + lambda: va("op_format", "cycle"), + ), + ( + "Literal format: previous", + "the other way round (O)", + lambda: va("op_format", "back"), + ), + *( + ( + f"Literal format: {label}", + f"show the literal as {label} ({fmt})", + (lambda f=fmt: va("op_format", f)), + ) + for fmt, label in ( + ("hex", "hexadecimal"), + ("dec", "decimal"), + ("oct", "octal"), + ("bin", "binary"), + ("char", "a character"), + ("offset", "an offset (reference)"), + ("stack", "a stack variable"), + ("default", "IDA's own choice"), + ) + ), + ( + "Toggle opcode bytes", + "cycle the opcode-bytes column (B)", + lambda: va("toggle_opcodes"), + ), + ( + "Structs / types editor", + "view + edit local types (Ctrl+T)", + app.action_structs, + ), + ( + "Filter functions…", + "glob-filter the function list (/)", + app.action_filter, + ), + ( + "Toggle names pane", + "function-list sidebar (Ctrl+B)", + app.action_toggle_functions, + ), ("Save database (.i64)", "persist changes (Ctrl+S)", app.action_save), ("Quit", "exit ida-tui (q)", app.action_quit), ) @@ -5211,45 +5534,52 @@ class IdaTui(App): Binding("escape", "back", "Back"), ] - def __init__(self, open_path: str | None = None, keepalive: bool = True, - rpc_path: str | None = None, ttl: int = 1800, - project=None, load_args: str = "", trace_path: str = "") -> None: + def __init__( + self, + open_path: str | None = None, + keepalive: bool = True, + rpc_path: str | None = None, + ttl: int = 1800, + project=None, + load_args: str = "", + trace_path: str = "", + ) -> None: super().__init__() # Project mode is additive: with no project this is the plain # single-binary app, unchanged. self._project = project self._pool = None - self._binary: str | None = None # active project binary (label) + self._binary: str | None = None # active project binary (label) self._states: dict[str, BinaryState] = {} - self._pending_restore = None # entry to reopen after a switch - self._goto_after_switch = None # cross-binary search hit to land on - self._hops: list[str] = [] # binaries a navigation crossed FROM - self._load_for_label = None # project binary the dialog is for - self._no_functions = False # analysis produced nothing at all - self._flash: str | None = None # message a pending reload must keep - self._flash_until = 0.0 # ...until this monotonic time - self._pending_switch = None # switch waiting on that answer - self._nav_seq = 0 # bumped per navigation; drops stale ones + self._pending_restore = None # entry to reopen after a switch + self._goto_after_switch = None # cross-binary search hit to land on + self._hops: list[str] = [] # binaries a navigation crossed FROM + self._load_for_label = None # project binary the dialog is for + self._no_functions = False # analysis produced nothing at all + self._flash: str | None = None # message a pending reload must keep + self._flash_until = 0.0 # ...until this monotonic time + self._pending_switch = None # switch waiting on that answer + self._nav_seq = 0 # bumped per navigation; drops stale ones #: Literal positions for the decompilation being loaded (worker thread #: -> the view, handed over when the pseudocode is applied). self._pending_nums: dict = {} # None = teardown wasn't an explicit quit (crash/kill): save defensively. # False = the user chose discard, or we already saved on the way out. self._save_on_exit: bool | None = None - self._index = None # project-wide symbol/string index + self._index = None # project-wide symbol/string index if project is not None: from .index import ProjectIndex from .pool import DatabasePool + self._pool = DatabasePool(project, ttl=ttl) - self._index = ProjectIndex( - os.path.join(project.index_dir, "project.db")) + self._index = ProjectIndex(os.path.join(project.index_dir, "project.db")) self._binary = project.refs[0].label open_path = project.refs[0].staged self._open_path = open_path self._ttl = ttl - self._load_args = load_args or "" # first-open options for a headerless blob - self._new_database = False # Ctrl+L asks IDA Nexus for a fresh IDB - self._title = (os.path.basename(open_path) if open_path else "") + self._load_args = load_args or "" # first-open options for a headerless blob + self._new_database = False # Ctrl+L asks IDA Nexus for a fresh IDB + self._title = os.path.basename(open_path) if open_path else "" #: Where we are in the execution trace, and everything that moves us. #: Owns the trace state; the _trace/_t/_trail_* properties below #: forward to it. @@ -5267,13 +5597,15 @@ class IdaTui(App): self._filter_term = "" self._pending_filter = "" self._filter_timer = None - self._sort_col = 0 # 0=addr, 1=name, 2=size + self._sort_col = 0 # 0=addr, 1=name, 2=size self._sort_reverse = False # ONE notion of "which pane you're in": _active, kept in step with focus # (on_descendant_focus does that while split). There used to be a second, # _pref, but it was only ever assigned "listing" — see _code_mode(). - self._active = ViewMode.LISTING # currently shown view (in split: the focused pane) - self._split = False # side-by-side listing + pseudocode + self._active = ( + ViewMode.LISTING + ) # currently shown view (in split: the focused pane) + self._split = False # side-by-side listing + pseudocode self._graph_sticky = False # stay in graph mode across navigations self._split_eamap: list[list[int]] = [] # split: decomp line -> instr EAs self._split_ea2line: dict[int, int] = {} # split: instr EA -> decomp line @@ -5289,8 +5621,9 @@ class IdaTui(App): self._search_ctx: tuple[object | None, int] = (None, 1) #: The one-line prompts above the footer. Each holds its own context #: for exactly as long as it is on screen; see idatui/prompt.py. - self.prompts = PromptBar(self, "search", "rename", "comment", - "retype", "makedata", "goto", "export") + self.prompts = PromptBar( + self, "search", "rename", "comment", "retype", "makedata", "goto", "export" + ) #: Everything that writes to the database (idatui/edit_ctl.py). self.edits = EditController(self) #: What those writes were, so the findings export can say which @@ -5401,8 +5734,10 @@ class IdaTui(App): if self._load_args: return False from .formats import needs_load_options + if os.path.exists(self._open_path + ".i64") or os.path.exists( - os.path.splitext(self._open_path)[0] + ".i64"): + os.path.splitext(self._open_path)[0] + ".i64" + ): return False try: if registered_database(self._open_path): @@ -5424,18 +5759,23 @@ class IdaTui(App): if not self._can_reload(): if self.client is not None and self.client.backend == "gui": self._status( - "reload unavailable for a GUI-owned database — reopen it in IDA") + "reload unavailable for a GUI-owned database — reopen it in IDA" + ) else: self._status("nothing to reload") return n = len(self._func_index) if self._func_index else 0 - note = ("this image has no functions, so nothing is lost" - if n == 0 else - f"discards the database for this binary \u2014 {n} " - f"function{'s' if n != 1 else ''}, plus any names and comments " - f"you've added") - self.push_screen(ConfirmScreen("Reload with different options?", note), - self._on_reload_confirmed) + note = ( + "this image has no functions, so nothing is lost" + if n == 0 + else f"discards the database for this binary \u2014 {n} " + f"function{'s' if n != 1 else ''}, plus any names and comments " + f"you've added" + ) + self.push_screen( + ConfirmScreen("Reload with different options?", note), + self._on_reload_confirmed, + ) def _on_reload_confirmed(self, yes) -> None: # type: ignore[no-untyped-def] if not yes: @@ -5501,7 +5841,8 @@ class IdaTui(App): # binary as already described and never ask again. self._project.set_load(label, processor="", base=0) self._project._entries[self._project._refs.index(ref)].pop( - "processor", None) + "processor", None + ) self._project.save() self._load_args = "" if path: @@ -5521,14 +5862,16 @@ class IdaTui(App): if ref is None or ref.load_args: return None if os.path.exists(ref.db) or os.path.exists( - os.path.splitext(ref.staged)[0] + ".i64"): - return None # already analysed: the .i64 records how + os.path.splitext(ref.staged)[0] + ".i64" + ): + return None # already analysed: the .i64 records how try: if registered_database(ref.staged, output_database=ref.db): return None except Exception: pass from .formats import needs_load_options + return ref if needs_load_options(ref.source) else None def _ask_load_options(self, path: str, label: str | None = None) -> None: @@ -5541,6 +5884,7 @@ class IdaTui(App): def _on_load_options(self, choice) -> None: # type: ignore[no-untyped-def] from .formats import load_args + choice = choice or {} label, self._load_for_label = self._load_for_label, None proc, base = choice.get("processor", ""), int(choice.get("base", 0) or 0) @@ -5573,6 +5917,7 @@ class IdaTui(App): def _start_rpc(self) -> None: from .rpc import RpcServer + self._rpc = RpcServer(self, self._rpc_path) async def _serve() -> None: @@ -5594,6 +5939,7 @@ class IdaTui(App): when the user has read it and moved on. """ import time as _time + if priority: self._flash = text self._flash_until = _time.monotonic() + 8.0 @@ -5611,7 +5957,11 @@ class IdaTui(App): # It stops being true the moment a function exists, though: latching it # meant the warning survived defining one with `p` and kept telling you # the load was wrong when it no longer was. - if self._no_functions and self._func_index is not None and len(self._func_index): + if ( + self._no_functions + and self._func_index is not None + and len(self._func_index) + ): self._no_functions = False if self._no_functions: text += " \u2014 no functions: wrong processor/base? Ctrl+L to reload" @@ -5638,7 +5988,10 @@ class IdaTui(App): subprocess.run( ["tmux", "load-buffer", "-w", "-"], input=text.encode("utf-8", "replace"), - stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, timeout=2.0) + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + timeout=2.0, + ) except Exception: # noqa: BLE001 pass return len(text) @@ -5662,8 +6015,7 @@ class IdaTui(App): except Exception: # noqa: BLE001 -- app teardown can win this race pass - self._idb_event_watch = watch( - changed, on_error=failed, debounce=0.2) + 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 @@ -5699,8 +6051,12 @@ class IdaTui(App): ) -> 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: + 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 @@ -5718,7 +6074,8 @@ class IdaTui(App): program.invalidate_external() self._status( f"{len(events)} external database " - f"change{'s' if len(events) != 1 else ''} — refreshing…") + f"change{'s' if len(events) != 1 else ''} — refreshing…" + ) self._reindex_functions() if self.is_hex: @@ -5751,15 +6108,35 @@ class IdaTui(App): 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) + 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, + 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: + 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") @@ -5771,8 +6148,12 @@ class IdaTui(App): 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) + 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() @@ -5783,6 +6164,7 @@ class IdaTui(App): Everything unrelated to database connectivity crashes as usual. """ from textual.worker import WorkerFailed + orig = error.error if isinstance(error, WorkerFailed) else error if isinstance(orig, IDAConnectionError): self._on_connection_lost() @@ -5794,7 +6176,8 @@ class IdaTui(App): return self._reconnecting = True self._conn_screen = LoadingScreen( - "the analysis server", note="connection lost \u2014 reconnecting\u2026") + "the analysis server", note="connection lost \u2014 reconnecting\u2026" + ) self.push_screen(self._conn_screen) self._reconnect() @@ -5817,20 +6200,27 @@ class IdaTui(App): # 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") + self.app.call_from_thread(self._reconnect_failed, "no binary to reopen") return if self._project is not None and self._binary is not None: ref = self._project.by_label(self._binary) client = NexusClient( - ref.staged, ttl=self._ttl, load_args=ref.load_args, - output_database=ref.db, spawn=False) + ref.staged, + ttl=self._ttl, + load_args=ref.load_args, + output_database=ref.db, + spawn=False, + ) else: client = NexusClient( - 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)) + 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 self.app.call_from_thread(self._reconnect_failed, str(e)) return @@ -5859,8 +6249,10 @@ class IdaTui(App): def _reconnect_failed(self, why: str) -> None: self._reconnecting = False - note = (f"database owner closed: {why} — reopen it in IDA, then " - "Esc and retry an action; or q to quit") + 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) @@ -5892,15 +6284,17 @@ class IdaTui(App): self._start_idb_event_watch(client) self._new_database = False self.app.call_from_thread( - self._status, f"{module} [{client.backend}] — loading functions…") + self._status, f"{module} [{client.backend}] — loading functions…" + ) self._load_functions() def _open_database_client(self): # type: ignore[no-untyped-def] """Attach through IDA Nexus, reusing a GUI or managed idalib database.""" if self._pool is not None: # project mode: the pool owns the leases label = self._binary or self._project.refs[0].label - client = self._pool.get(label, progress=lambda m: - self.app.call_from_thread(self._status, m)) + client = self._pool.get( + label, progress=lambda m: self.app.call_from_thread(self._status, m) + ) self._binary = label self._pool.set_active(label) self._open_path = self._project.by_label(label).staged @@ -5908,17 +6302,21 @@ class IdaTui(App): return client if not self._open_path: self.app.call_from_thread( - self._status, "IDA Nexus needs a database or executable path") + self._status, "IDA Nexus needs a database or executable path" + ) self.app.call_from_thread(self._dismiss_loading) return None base = os.path.basename(self._open_path) self.app.call_from_thread( - self._status, f"discovering IDA Nexus database for {base}…") - client = NexusClient(self._open_path, ttl=self._ttl, - load_args=self._load_args, - new_database=self._new_database) - client.connect(progress=lambda m: self.app.call_from_thread( - self._status, m)) + self._status, f"discovering IDA Nexus database for {base}…" + ) + client = NexusClient( + self._open_path, + ttl=self._ttl, + load_args=self._load_args, + new_database=self._new_database, + ) + client.connect(progress=lambda m: self.app.call_from_thread(self._status, m)) return client @work(thread=True, exclusive=True, group="load-funcs") @@ -5926,7 +6324,9 @@ class IdaTui(App): assert self.program is not None idx = self.program.functions() self._func_index = idx - self.app.call_from_thread(lambda: self.query_one("#func-table", DataTable).clear()) + self.app.call_from_thread( + lambda: self.query_one("#func-table", DataTable).clear() + ) last = 0 while not idx.complete: idx.load_next_page() @@ -5934,21 +6334,20 @@ class IdaTui(App): last = len(idx) if rows: self.app.call_from_thread(self._append_rows, rows) - self.app.call_from_thread( - self._status, f"{last} functions…" - ) + self.app.call_from_thread(self._status, f"{last} functions…") # If a filter is active (typed during load), re-apply it over the full set. if self._filter_term: self.app.call_from_thread(self._apply_filter, self._filter_term) else: self.app.call_from_thread( - self._status, f"{len(idx)} functions (Ctrl+N: find symbol)") + self._status, f"{len(idx)} functions (Ctrl+N: find symbol)" + ) # Land somewhere useful instead of an empty pane: main() if present, # otherwise pop the fuzzy symbol picker. self.app.call_from_thread(self._auto_land) self._index_binary() # project mode: keep the cross-binary index fresh if self.trace_ctl.armed: - self._load_trace() # needs the index above: rebasing reads it + self._load_trace() # needs the index above: rebasing reads it @work(thread=True, exclusive=True, group="prewarm") def _prewarm_provider(self) -> None: @@ -5969,6 +6368,7 @@ class IdaTui(App): if not imps: return from collections import Counter + votes: Counter = Counter() for name in {i.name for i in imps}: for h in self._index.providers(name, exclude=self._binary): @@ -5981,7 +6381,8 @@ class IdaTui(App): try: if self._pool.prewarm(cand): self.app.call_from_thread( - self._status, f"pre-warmed {cand} (provides {n} imports)") + self._status, f"pre-warmed {cand} (provides {n} imports)" + ) except Exception: # noqa: BLE001 -- speculative work must never surface pass @@ -5995,11 +6396,13 @@ class IdaTui(App): if ref is None or not self._index.is_stale(self._binary, ref.source): return from .index import KIND_EXPORT, KIND_FUNC, KIND_IMPORT, KIND_STRING + idx = self._func_index - entries = [(KIND_FUNC, f.addr, f.name) for f in (idx.all_loaded() if idx else [])] + entries = [ + (KIND_FUNC, f.addr, f.name) for f in (idx.all_loaded() if idx else []) + ] try: - entries += [(KIND_STRING, s.addr, s.text) - for s in self.program.strings()] + entries += [(KIND_STRING, s.addr, s.text) for s in self.program.strings()] except Exception: # noqa: BLE001 -- symbols alone are still worth indexing pass try: @@ -6014,7 +6417,8 @@ class IdaTui(App): self.app.call_from_thread(self._status, f"indexing failed: {e}") return self.app.call_from_thread( - self._status, f"indexed {self._binary}: {n} symbols, strings + linkage") + self._status, f"indexed {self._binary}: {n} symbols, strings + linkage" + ) self._prewarm_provider() # -- initial landing --------------------------------------------------- # @@ -6057,8 +6461,9 @@ class IdaTui(App): first = self._func_index.get(0) if first is not None: self._open_function(first.addr, first.name) - self._status(f"no entry function — opened {first.name} " - "(Ctrl+N: find symbol)") + self._status( + f"no entry function — opened {first.name} (Ctrl+N: find symbol)" + ) else: self.action_symbols() else: @@ -6087,8 +6492,13 @@ class IdaTui(App): if start is None: self._status("no functions and no segments \u2014 nothing to show") return - self._open_at(start, self.program.section_of(start) or "image", - cursor=0, push=True, is_region=True) + self._open_at( + start, + self.program.section_of(start) or "image", + cursor=0, + push=True, + is_region=True, + ) def _can_reload(self) -> bool: """Whether IDA Nexus can replace this IDB with different options. @@ -6145,7 +6555,9 @@ class IdaTui(App): ci = term.islower() is_glob = ("*" in term) or ("?" in term) needle = term.lower() if ci else term - pat = needle if (is_glob and ("*" in needle or "?" in needle)) else f"*{needle}*" + pat = ( + needle if (is_glob and ("*" in needle or "?" in needle)) else f"*{needle}*" + ) matched: list[tuple[Func, tuple[int, int] | None]] = [] for f in funcs: if not term: @@ -6230,9 +6642,10 @@ class IdaTui(App): if not funcs: self._status("functions still loading…") return - self.push_screen(SymbolPalette(funcs, index=self._index, - binary=self._binary), - self._on_symbol_chosen) + self.push_screen( + SymbolPalette(funcs, index=self._index, binary=self._binary), + self._on_symbol_chosen, + ) def _on_symbol_chosen(self, choice) -> None: # type: ignore[no-untyped-def] if choice is None: @@ -6267,11 +6680,15 @@ class IdaTui(App): active one and other still-resident binaries can be dirty. """ if self._pool is None: - return [os.path.basename(self._open_path or "database")] if self._dirty else [] + return ( + [os.path.basename(self._open_path or "database")] if self._dirty else [] + ) out = [self._binary] if (self._dirty and self._binary) else [] - out += [label for label, st in self._states.items() - if st.dirty and label != self._binary - and self._pool.is_resident(label)] + out += [ + label + for label, st in self._states.items() + if st.dirty and label != self._binary and self._pool.is_resident(label) + ] return out async def action_quit(self) -> None: @@ -6290,13 +6707,16 @@ class IdaTui(App): # GUI/shared sessions retain state and inherit finalization. dirty = self._dirty_labels() self._loading_screen = LoadingScreen( - "discarding", note="finalizing database leases…") + "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. - self._loading_screen = LoadingScreen("saving", note="writing databases\u2026") + self._loading_screen = LoadingScreen( + "saving", note="writing databases\u2026" + ) self.push_screen(self._loading_screen) self._save_then_exit() # None: cancel, stay put @@ -6322,8 +6742,7 @@ class IdaTui(App): 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._loading_screen.update_note(f"finalization transferred: {labels}") self._finish_exit() @work(thread=True, exclusive=True, group="save-exit") @@ -6358,8 +6777,7 @@ class IdaTui(App): return if self._prompt_active(): return - self.push_screen(ProjectPalette(self._pool.status()), - self._on_binary_chosen) + self.push_screen(ProjectPalette(self._pool.status()), self._on_binary_chosen) def _on_binary_chosen(self, label: str | None) -> None: if label and label != self._binary: @@ -6377,10 +6795,16 @@ class IdaTui(App): # the pool hand us a lease (attaching + evicting as the budget dictates). if self._binary is not None: self._states[self._binary] = BinaryState( - label=self._binary, program=self.program, - func_index=self._func_index, nav=list(self._nav), cur=self._cur, - active=self._active, split=self._split, - filter_term=self._filter_term, dirty=self._dirty) + label=self._binary, + program=self.program, + func_index=self._func_index, + nav=list(self._nav), + cur=self._cur, + active=self._active, + split=self._split, + filter_term=self._filter_term, + dirty=self._dirty, + ) self._loading_screen = LoadingScreen(label, note="switching\u2026") self.push_screen(self._loading_screen) self._do_switch(label) @@ -6389,8 +6813,9 @@ class IdaTui(App): def _do_switch(self, label: str) -> None: assert self._pool is not None try: - client = self._pool.get(label, progress=lambda m: - self.app.call_from_thread(self._status, m)) + client = self._pool.get( + label, progress=lambda m: self.app.call_from_thread(self._status, m) + ) except Exception as e: # noqa: BLE001 self.app.call_from_thread(self._switch_failed, label, str(e)) return @@ -6398,11 +6823,13 @@ class IdaTui(App): # The Program (and its caches) only survive while that lease does; an # evicted binary reattaches. Either way the nav # history is just addresses, so it always survives. - reuse = (st is not None and st.program is not None - and getattr(st.program, "client", None) is client) + reuse = ( + st is not None + and st.program is not None + and getattr(st.program, "client", None) is client + ) program = st.program if reuse else Program(client) - self.app.call_from_thread(self._after_switch, label, client, program, - st, reuse) + self.app.call_from_thread(self._after_switch, label, client, program, st, reuse) def _after_switch(self, label, client, program, st, reuse) -> None: # type: ignore[no-untyped-def] self.client = client @@ -6478,9 +6905,10 @@ class IdaTui(App): self._status("no strings found (needs the list_strings tool)") return self._status(f"strings: {len(items)}") - self.push_screen(StringsPalette(items, index=self._index, - binary=self._binary), - self._on_string_chosen) + self.push_screen( + StringsPalette(items, index=self._index, binary=self._binary), + self._on_string_chosen, + ) def action_find(self) -> None: """Ctrl+F: search the whole database — disassembly text, or bytes.""" @@ -6497,7 +6925,7 @@ class IdaTui(App): try: seed = view.word_under_cursor() or "" except Exception: # noqa: BLE001 -- a seed is a nicety, never a - seed = "" # reason not to open the search + seed = "" # reason not to open the search self.push_screen(SearchPalette(self.program, seed), self._on_hit_chosen) def _on_hit_chosen(self, hit) -> None: # type: ignore[no-untyped-def] @@ -6508,8 +6936,9 @@ class IdaTui(App): # there. The status names the exact address so it isn't lost. self._goto_ea(hit.head, push=True) if hit.addr != hit.head: - self._status(f"match at {hit.addr:#x} (inside {hit.head:#x})", - priority=True) + self._status( + f"match at {hit.addr:#x} (inside {hit.head:#x})", priority=True + ) def _on_string_chosen(self, choice) -> None: # type: ignore[no-untyped-def] if choice is None: @@ -6549,8 +6978,7 @@ class IdaTui(App): # active, but the listing must still round-trip through addresses: # row indices do not survive an external structure change. lst = self.query_one(ListingView) - listing_anchor = ViewAnchor(view=ViewMode.LISTING, - cursor_x=lst.cursor_x) + listing_anchor = ViewAnchor(view=ViewMode.LISTING, cursor_x=lst.cursor_x) model = lst.model if model is not None: listing_anchor.ea = lst._cursor_ea() @@ -6570,20 +6998,27 @@ class IdaTui(App): cur.dec_scroll_x = round(dec.scroll_offset.x) dec.loading = True - want_ea = (self.query_one(GraphView)._cursor_ea() - if self.is_graph else None) + want_ea = self.query_one(GraphView)._cursor_ea() if self.is_graph else None if self.is_graph: self._graph_sticky = True self._status(f"{cur.name} — refreshing graph…") else: self._status(f"{cur.name} — refreshing…") - self._refresh_view(cur, mode, split, listing_anchor, - refresh_decomp, want_ea, self.program) + self._refresh_view( + cur, mode, split, listing_anchor, refresh_decomp, want_ea, self.program + ) @work(thread=True, exclusive=True, group="refresh-view") - def _refresh_view(self, cur: NavEntry, mode: ViewMode, split: bool, - anchor: ViewAnchor | None, refresh_decomp: bool, - want_ea: int | None, program) -> None: # type: ignore[no-untyped-def] + def _refresh_view( + self, + cur: NavEntry, + mode: ViewMode, + split: bool, + anchor: ViewAnchor | None, + refresh_decomp: bool, + want_ea: int | None, + program, + ) -> None: # type: ignore[no-untyped-def] """Invalidate and rebuild without blocking Textual's event loop.""" try: program.bump_items() @@ -6600,22 +7035,40 @@ class IdaTui(App): cursor, top = self._anchor_rows(anchor, model, target) except Exception as exc: # noqa: BLE001 -- a refresh is recoverable diag.note("refresh_view", exc) - self.app.call_from_thread( - self._view_refresh_failed, cur, program, str(exc)) + self.app.call_from_thread(self._view_refresh_failed, cur, program, str(exc)) return self.app.call_from_thread( - self._apply_view_refresh, cur, mode, split, anchor, - refresh_decomp, want_ea, program, model, cursor, top) + self._apply_view_refresh, + cur, + mode, + split, + anchor, + refresh_decomp, + want_ea, + program, + model, + cursor, + top, + ) def _view_refresh_failed(self, cur: NavEntry, program, error: str) -> None: # type: ignore[no-untyped-def] if self.program is program and self._cur is cur: self.query_one(DecompView).loading = False self._status(f"refresh failed: {error}", priority=True) - def _apply_view_refresh(self, cur: NavEntry, mode: ViewMode, split: bool, - anchor: ViewAnchor | None, refresh_decomp: bool, - want_ea: int | None, program, model, cursor: int, - top: int) -> None: # type: ignore[no-untyped-def] + def _apply_view_refresh( + self, + cur: NavEntry, + mode: ViewMode, + split: bool, + anchor: ViewAnchor | None, + refresh_decomp: bool, + want_ea: int | None, + program, + model, + cursor: int, + top: int, + ) -> None: # type: ignore[no-untyped-def] # A binary switch or navigation completed while the refresh was in # flight. Its newer view wins; never drag the user back. if self.program is not program or self._cur is not cur: @@ -6632,9 +7085,13 @@ class IdaTui(App): cur.cursor = max(cursor, 0) cur.cursor_x = anchor.cursor_x cur.scroll_y = top - lst.load(model, cur.name, cursor=cur.cursor, - cursor_x=cur.cursor_x, - scroll_y=top if top >= 0 else None) + lst.load( + model, + cur.name, + cursor=cur.cursor, + cursor_x=cur.cursor_x, + scroll_y=top if top >= 0 else None, + ) if refresh_decomp: self.query_one(DecompView).loaded_ea = None self._show_active() @@ -6659,8 +7116,11 @@ class IdaTui(App): if self._split: # In split mode Tab/F5 just moves focus between the two panes. self._active = ViewMode.DECOMP if self.is_listing else ViewMode.LISTING - (self.query_one(DecompView) if self.is_decomp - else self.query_one(ListingView)).focus() + ( + self.query_one(DecompView) + if self.is_decomp + else self.query_one(ListingView) + ).focus() self._sync_split(self._active) # re-link from the new driver self._status_for_cur("split") return @@ -6752,7 +7212,8 @@ class IdaTui(App): if fn is None: self.app.call_from_thread( self._decomp_from_listing_failed, - "F5 — cursor is not inside a defined function ('p' to make one)") + "F5 — cursor is not inside a defined function ('p' to make one)", + ) return dec_idx = self._decomp_line_for(fn.addr, ea) self.app.call_from_thread(self._enter_decomp, fn.addr, fn.name, dec_idx) @@ -6774,8 +7235,9 @@ class IdaTui(App): ret.cursor_x = lst.cursor_x ret.scroll_y = round(lst.scroll_offset.y) self._decomp_return = ret - entry = NavEntry(ea=fn_addr, name=fn_name, is_region=False, - dec_cursor=max(dec_idx, 0)) + entry = NavEntry( + ea=fn_addr, name=fn_name, is_region=False, dec_cursor=max(dec_idx, 0) + ) self._cur = entry self._active = ViewMode.DECOMP self._show_active() @@ -6788,8 +7250,10 @@ class IdaTui(App): self._status("open a function first") return if not self._split and self.size.width < _SPLIT_MIN_WIDTH: - self._status(f"terminal too narrow for split — need ≈{_SPLIT_MIN_WIDTH} " - f"cols (have {self.size.width})") + self._status( + f"terminal too narrow for split — need ≈{_SPLIT_MIN_WIDTH} " + f"cols (have {self.size.width})" + ) return self._split = not self._split if self._active not in ("listing", "decomp"): @@ -6819,8 +7283,11 @@ class IdaTui(App): self._graph_sticky = True gv = self.query_one(GraphView) ea = self._graph_target_ea() - if gv.loaded_ea is not None and gv.fc is not None \ - and gv.fc.func_ea == self._cur.ea: + if ( + gv.loaded_ea is not None + and gv.fc is not None + and gv.fc.func_ea == self._cur.ea + ): self._active = ViewMode.GRAPH self._split = False self._show_active() @@ -6857,8 +7324,13 @@ class IdaTui(App): err = f"{type(e).__name__}: {e}" self.app.call_from_thread(self._apply_graph, func_ea, want_ea, fc, err) - def _apply_graph(self, func_ea: int, want_ea: int | None, fc, # type: ignore[no-untyped-def] - err: str) -> None: + def _apply_graph( + self, + func_ea: int, + want_ea: int | None, + fc, # type: ignore[no-untyped-def] + err: str, + ) -> None: if self._cur is None or self._cur.ea != func_ea: return # a newer navigation won if not self._graph_sticky and not self.is_graph: @@ -6867,13 +7339,17 @@ class IdaTui(App): # here drags them back into a graph they already dismissed. return if fc is None: - self._status(err or "no control-flow graph for this function " - "(is it a thunk or an import?)") + self._status( + err + or "no control-flow graph for this function " + "(is it a thunk or an import?)" + ) return if len(fc.blocks) > self.GRAPH_MAX_BLOCKS: self._status( f"{fc.name}: {len(fc.blocks)} blocks — too many to graph " - f"(limit {self.GRAPH_MAX_BLOCKS}); staying in the listing") + f"(limit {self.GRAPH_MAX_BLOCKS}); staying in the listing" + ) return gv = self.query_one(GraphView) gv.set_graph(fc, want_ea) @@ -6887,12 +7363,15 @@ class IdaTui(App): if gv.lay is None or gv.fc is None: return s = gv.lay.stats - loops = f", {s['back']} loop{'s' if s['back'] != 1 else ''}" if s["back"] else "" + loops = ( + f", {s['back']} loop{'s' if s['back'] != 1 else ''}" if s["back"] else "" + ) eng = "" if s.get("engine") == "native" else f", {s.get('engine')}" self._status( f"{gv.fc.name} @ {gv.fc.func_ea:#x} [graph: {s['blocks']} blocks, " f"{s['edges']} edges{loops}{eng}] " - f"z=zoom({gv.ZOOMS[gv._zoom]}) m=map J/K=edge space=text") + f"z=zoom({gv.ZOOMS[gv._zoom]}) m=map J/K=edge space=text" + ) def on_graph_view_cursor_moved(self, msg: "GraphView.CursorMoved") -> None: gv = self.query_one(GraphView) @@ -6918,7 +7397,7 @@ class IdaTui(App): lst.load(lm, name, cursor=idx, scroll_y=max(idx - _JUMP_CONTEXT, 0)) self._show_active() # split branch shows both + loads the decomp self._sync_split(self._active) # crude link now - self._load_split_map(ea) # region map (async) if decomp is loaded + self._load_split_map(ea) # region map (async) if decomp is loaded def action_hex(self) -> None: """Backslash: show the raw bytes of the loaded image, synced to the code @@ -6981,8 +7460,9 @@ class IdaTui(App): try: self.journal.load(self.program) self.journal.flush(self.program) - out, f = findings.export(self.program, self._open_path or "", path, - journal=self.journal) + out, f = findings.export( + self.program, self._open_path or "", path, journal=self.journal + ) except Exception as e: # noqa: BLE001 -- a bad path is a message, not a crash self.call_from_thread(self._status, f"export failed: {e}", True) return @@ -6990,12 +7470,17 @@ class IdaTui(App): self.call_from_thread( self._status, f"exported {len(f.comments)} comments, {n_named} names, " - f"{len(f.types)} types → {out}", True) + f"{len(f.types)} types → {out}", + True, + ) def action_goto(self) -> None: inp = self.query_one("#goto", Input) - inp.placeholder = ("hex goto: 0xADDR or name — Enter" if self.is_hex - else "goto: name or 0xADDR — Enter") + inp.placeholder = ( + "hex goto: 0xADDR or name — Enter" + if self.is_hex + else "goto: name or 0xADDR — Enter" + ) inp.can_focus = True inp.display = True inp.value = "" @@ -7026,7 +7511,7 @@ class IdaTui(App): # Local history is spent, but we got here from another binary. label = self._hops.pop() self._status(f"\u25c2 back to {label}\u2026") - self._switch_binary(label) # _states restores its nav and position + self._switch_binary(label) # _states restores its nav and position elif self.query_one("#left", FunctionsPanel).display: table.focus() else: @@ -7069,11 +7554,11 @@ class IdaTui(App): # 'code' xref, and land on a call/jump's real target instead. self._follow_disasm(ea, word, view._next_ea()) elif isinstance(view, DecompView) and view._texts: - self._follow_decomp(view._texts[view.cursor], word, - view._line_ea(view.cursor)) + self._follow_decomp( + view._texts[view.cursor], word, view._line_ea(view.cursor) + ) - def _graph_local_target(self, view: "GraphView", ea: int, - word: str) -> int | None: + def _graph_local_target(self, view: "GraphView", ea: int, word: str) -> int | None: """If the cursor's instruction branches somewhere inside this same graph, return that address.""" if view.fc is None or self.program is None: @@ -7152,14 +7637,16 @@ class IdaTui(App): return not all(c in "0123456789abcdefABCDEF" for c in word) @work(thread=True, group="nav") - def _follow_disasm(self, ea: int, word: str | None, - next_ea: int | None = None) -> None: + def _follow_disasm( + self, ea: int, word: str | None, next_ea: int | None = None + ) -> None: assert self.program is not None # Prefer the symbol under the cursor (handles multiple refs on a line). if self._looks_like_symbol(word): try: - self._do_navigate(self.program.resolve(word), push=True, - focus_name=word) + self._do_navigate( + self.program.resolve(word), push=True, focus_name=word + ) return except Exception: # noqa: BLE001 -- not a resolvable name; fall back pass @@ -7218,13 +7705,15 @@ class IdaTui(App): return False label, addr = found self.app.call_from_thread( - self._status, f"{name} \u2192 {label} (import resolved)") + self._status, f"{name} \u2192 {label} (import resolved)" + ) self.app.call_from_thread(self._switch_then_goto, label, addr) return True @work(thread=True, group="nav") - def _follow_decomp(self, line: str, word: str | None, - line_ea: int | None = None) -> None: + def _follow_decomp( + self, line: str, word: str | None, line_ea: int | None = None + ) -> None: if self._cur is None: return dec = self.program.decompile(self._cur.ea) @@ -7267,8 +7756,13 @@ class IdaTui(App): return None @work(thread=True, group="xrefs") - def _xrefs_disasm(self, ea: int, word: str | None, - here_ea: int | None = None, here_end: int | None = None) -> None: + def _xrefs_disasm( + self, + ea: int, + word: str | None, + here_ea: int | None = None, + here_end: int | None = None, + ) -> None: assert self.program is not None subj: int | None = None if self._looks_like_symbol(word): @@ -7285,8 +7779,13 @@ class IdaTui(App): self._xrefs_present(subj, word, here_ea, here_end) @work(thread=True, group="xrefs") - def _xrefs_decomp(self, line: str, word: str | None, - here_ea: int | None = None, here_end: int | None = None) -> None: + def _xrefs_decomp( + self, + line: str, + word: str | None, + here_ea: int | None = None, + here_end: int | None = None, + ) -> None: subj: int | None = None if self._cur is not None and word: dec = self.program.decompile(self._cur.ea) @@ -7317,9 +7816,13 @@ class IdaTui(App): span = i return span if span is not None else 0 - def _xrefs_present(self, subj: int, subj_name: str | None = None, - here_ea: int | None = None, - here_end: int | None = None) -> None: # worker + def _xrefs_present( + self, + subj: int, + subj_name: str | None = None, + here_ea: int | None = None, + here_end: int | None = None, + ) -> None: # worker assert self.program is not None try: return self._xrefs_present_inner(subj, subj_name, here_ea, here_end) @@ -7380,22 +7883,30 @@ class IdaTui(App): if not name: return [] from .domain import link_name + name = link_name(name) try: _, exports = self.program.linkage() except Exception: # noqa: BLE001 return [] if not any(e.name == name for e in exports): - return [] # we don't export it; nobody imports it FROM US + return [] # we don't export it; nobody imports it FROM US try: hits = self._index.importers(name, exclude=self._binary) except Exception: # noqa: BLE001 return [] - return [(h.binary, h.addr, f"{h.addr:08X} import [{h.binary}] {name}") - for h in hits] + return [ + (h.binary, h.addr, f"{h.addr:08X} import [{h.binary}] {name}") + for h in hits + ] - def _present_xrefs(self, label: str, items: list[tuple[object, str]], - focus_name: str | None = None, preselect: int = 0) -> None: + def _present_xrefs( + self, + label: str, + items: list[tuple[object, str]], + focus_name: str | None = None, + preselect: int = 0, + ) -> None: if not self._xref_active: return # cancelled (Esc) while we were still gathering self._xref_active = False @@ -7410,14 +7921,18 @@ class IdaTui(App): def _on_xref_chosen(self, addr) -> None: # type: ignore[no-untyped-def] if addr is None: return - if isinstance(addr, tuple): # a caller in another project binary + if isinstance(addr, tuple): # a caller in another project binary binary, ea = addr - self._switch_then_goto(binary, ea) # records a hop, so Esc returns + self._switch_then_goto(binary, ea) # records a hop, so Esc returns return # If xrefs was invoked from the decompiler, land the jump back in the # decompiler (when the target is decompilable) rather than the listing. - self._goto_ea(addr, push=True, focus_name=self._xref_focus_name, - prefer_decomp=(self.is_decomp)) + self._goto_ea( + addr, + push=True, + focus_name=self._xref_focus_name, + prefer_decomp=(self.is_decomp), + ) # -- database edits ---------------------------------------------------- # # The bodies live in EditController (idatui/edit_ctl.py). What stays here is @@ -7466,18 +7981,21 @@ class IdaTui(App): self.edits.do_retype(kind, subject, word, new) @work(thread=True, exclusive=True, group="makedata") - def _do_make_data(self, ea: int, type_decl: str, - anchor: ViewAnchor | None = None) -> None: + def _do_make_data( + self, ea: int, type_decl: str, anchor: ViewAnchor | None = None + ) -> None: self.edits.do_make_data(ea, type_decl, anchor) @work(thread=True, exclusive=True, group="opformat") - def _do_op_format(self, mode: str, where: str, ea: int, col: int, - line: int = -1) -> None: + def _do_op_format( + self, mode: str, where: str, ea: int, col: int, line: int = -1 + ) -> None: self.edits.do_op_format(mode, where, ea, col, line) @work(thread=True, exclusive=True, group="edititem") - def _do_edit_item(self, kind: str, ea: int, - anchor: ViewAnchor | None = None) -> None: + def _do_edit_item( + self, kind: str, ea: int, anchor: ViewAnchor | None = None + ) -> None: self.edits.do_edit_item(kind, ea, anchor) def _reload_active_code(self) -> None: @@ -7574,13 +8092,13 @@ class IdaTui(App): return if len(idx): self._no_functions = False - self._apply_filter(self._filter_term) # repopulate the names pane + self._apply_filter(self._filter_term) # repopulate the names pane @work(thread=True, exclusive=True, group="save") def _save(self) -> None: assert self.program is not None try: - self.journal.flush(self.program) # ride along into the .i64 + self.journal.flush(self.program) # ride along into the .i64 self.program.client.save_database() except Exception as e: # noqa: BLE001 self.app.call_from_thread(self._status, f"save failed: {e}") @@ -7596,12 +8114,22 @@ class IdaTui(App): # -- navigation to an arbitrary address ------------------------------- # @work(thread=True, group="nav") - def _goto_ea(self, ea: int, push: bool = True, - focus_name: str | None = None, prefer_decomp: bool = False) -> None: + def _goto_ea( + self, + ea: int, + push: bool = True, + focus_name: str | None = None, + prefer_decomp: bool = False, + ) -> None: self._do_navigate(ea, push, focus_name, prefer_decomp) - def _do_navigate(self, ea: int, push: bool, focus_name: str | None = None, - prefer_decomp: bool = False) -> None: # worker context + def _do_navigate( + self, + ea: int, + push: bool, + focus_name: str | None = None, + prefer_decomp: bool = False, + ) -> None: # worker context assert self.program is not None # Which navigation this is. Decompiling below can take a while, and if # you press Esc (or jump again) in the meantime this result is stale — @@ -7623,11 +8151,12 @@ class IdaTui(App): # anchor on the address but snap to the nearest line that # actually holds the referenced symbol (the marker line and # the symbol's line can differ), landing on the token. - dec_idx, col = self._decomp_locate(fn.addr, ea, - focus_name or fn.name) + dec_idx, col = self._decomp_locate( + fn.addr, ea, focus_name or fn.name + ) self.app.call_from_thread( - self._open_decomp_entry, fn.addr, fn.name, dec_idx, col, - push, seq) + self._open_decomp_entry, fn.addr, fn.name, dec_idx, col, push, seq + ) return # Otherwise: everything opens the one continuous listing at ``ea``. A # function name is used for the status label; a region gets a segment @@ -7636,12 +8165,18 @@ class IdaTui(App): idx = max(lm.ensure_ea(ea), 0) if lm is not None else 0 name = fn.name if fn is not None else self.program.region_label(ea) self.app.call_from_thread( - self._open_at_if_current, seq, ea, name, idx, push, fn is None, - focus_name) + self._open_at_if_current, seq, ea, name, idx, push, fn is None, focus_name + ) - def _open_decomp_entry(self, fn_addr: int, fn_name: str, dec_idx: int, - dec_cursor_x: int, push: bool, - seq: int | None = None) -> None: + def _open_decomp_entry( + self, + fn_addr: int, + fn_name: str, + dec_idx: int, + dec_cursor_x: int, + push: bool, + seq: int | None = None, + ) -> None: """Open ``fn_addr`` in the decompiler as a real navigation (nav history aware), landing on pseudocode line ``dec_idx`` column ``dec_cursor_x``. @@ -7663,12 +8198,17 @@ class IdaTui(App): src.dec_cursor_x = dv.cursor_x src.dec_scroll_y = round(dv.scroll_offset.y) if not self._nav or self._nav[-1] is not src: - self._push_nav(src) # never stack a second copy of a spot + self._push_nav(src) # never stack a second copy of a spot else: self._save_current_pos() self._decomp_return = None # a real navigation abandons the F5 return - entry = NavEntry(ea=fn_addr, name=fn_name, view="decomp", - dec_cursor=dec_idx, dec_cursor_x=dec_cursor_x) + entry = NavEntry( + ea=fn_addr, + name=fn_name, + view="decomp", + dec_cursor=dec_idx, + dec_cursor_x=dec_cursor_x, + ) if push: self._push_nav(entry) self._open_entry(entry, push=False) @@ -7689,8 +8229,9 @@ class IdaTui(App): m = re.search(rf"\b{re.escape(name)}\b", clean) return m.start() if m else 0 - def _decomp_locate(self, fn_addr: int, ea: int, - token: str | None) -> tuple[int, int]: + def _decomp_locate( + self, fn_addr: int, ea: int, token: str | None + ) -> tuple[int, int]: """Best (line, column) for address ``ea`` in ``fn_addr``'s pseudocode. Anchors on the /*0xEA*/ marker line for ``ea``, but Hex-Rays can attribute @@ -7748,8 +8289,9 @@ class IdaTui(App): move, so it has no business being a step in the history.""" if a.ea != b.ea or a.view != b.view: return False - return (a.dec_cursor == b.dec_cursor if a.view == "decomp" - else a.cursor == b.cursor) + return ( + a.dec_cursor == b.dec_cursor if a.view == "decomp" else a.cursor == b.cursor + ) def _push_nav(self, entry: NavEntry) -> None: """Append to the nav stack unless that would duplicate where we already are. @@ -7810,9 +8352,16 @@ class IdaTui(App): cur = row_of(fallback_ea) return (cur, row_of(a.top_ea)) - def _open_at_if_current(self, seq: int, ea: int, name: str, cursor: int, - push: bool, is_region: bool, - focus_name: str | None) -> None: + def _open_at_if_current( + self, + seq: int, + ea: int, + name: str, + cursor: int, + push: bool, + is_region: bool, + focus_name: str | None, + ) -> None: """Apply a navigation result only if it's still the one being awaited. The decompiler path has had this since 756589a; the listing path hadn't, @@ -7822,10 +8371,18 @@ class IdaTui(App): return self._open_at(ea, name, cursor, push, -1, 0, is_region, focus_name) - def _open_at(self, ea: int, name: str, cursor: int, push: bool, - dec_cursor: int = -1, dec_cursor_x: int = 0, - is_region: bool = False, focus_name: str | None = None, - scroll_y: int = -1) -> None: + def _open_at( + self, + ea: int, + name: str, + cursor: int, + push: bool, + dec_cursor: int = -1, + dec_cursor_x: int = 0, + is_region: bool = False, + focus_name: str | None = None, + scroll_y: int = -1, + ) -> None: if push: self._save_current_pos() self._decomp_return = None # a real navigation abandons the F5 return @@ -7882,8 +8439,11 @@ class IdaTui(App): self._end_search(cancel=True) elif prompt.id in ("goto", "export"): self._end_goto() if prompt.id == "goto" else self._end_export() - (self.query_one(HexView) if self.is_hex - else (self._code_view() or self.query_one(ListingView))).focus() + ( + self.query_one(HexView) + if self.is_hex + else (self._code_view() or self.query_one(ListingView)) + ).focus() else: prompt.close() return @@ -7913,10 +8473,12 @@ class IdaTui(App): # The edit prompts all submit the same way: take the context the prompt # was holding (close() hands it over, so it can't be read twice or go # stale) and let the controller decide what to do with it. - submit = {"rename": self.edits.submit_rename, - "comment": self.edits.submit_comment, - "retype": self.edits.submit_retype, - "makedata": self.edits.submit_make_data}.get(inp.id or "") + submit = { + "rename": self.edits.submit_rename, + "comment": self.edits.submit_comment, + "retype": self.edits.submit_retype, + "makedata": self.edits.submit_make_data, + }.get(inp.id or "") if submit is not None: ctx = self.prompts[inp.id].close() if ctx is not None: @@ -7924,15 +8486,21 @@ class IdaTui(App): return if inp.id == "goto": self._end_goto() - (self.query_one(HexView) if self.is_hex - else (self._code_view() or self.query_one(ListingView))).focus() + ( + self.query_one(HexView) + if self.is_hex + else (self._code_view() or self.query_one(ListingView)) + ).focus() if value: self._goto(value) return if inp.id == "export": self._end_export() - (self.query_one(HexView) if self.is_hex - else (self._code_view() or self.query_one(ListingView))).focus() + ( + self.query_one(HexView) + if self.is_hex + else (self._code_view() or self.query_one(ListingView)) + ).focus() if value: self.export_findings(value) return @@ -8097,13 +8665,13 @@ class IdaTui(App): e.scroll_y = round(lst.scroll_offset.y) @work(thread=True, group="nav") - def _open_function(self, ea: int, name: str | None = None, - push: bool = True) -> None: + def _open_function( + self, ea: int, name: str | None = None, push: bool = True + ) -> None: # Unified: opening a function is just navigating the one linear listing # to its entry address. self._do_navigate(ea, push) - def _code_mode(self) -> ViewMode: """The code view to return to from hex — always the unified listing.""" return ViewMode.LISTING @@ -8121,9 +8689,12 @@ class IdaTui(App): dec = self.query_one(DecompView) if dec.loaded_ea == entry.ea: # already decompiled: reposition without a recompile - dec.goto(entry.dec_cursor, entry.dec_cursor_x, - entry.dec_scroll_y if entry.dec_scroll_y >= 0 else -1, - entry.dec_scroll_x) + dec.goto( + entry.dec_cursor, + entry.dec_cursor_x, + entry.dec_scroll_y if entry.dec_scroll_y >= 0 else -1, + entry.dec_scroll_x, + ) self._show_active() # loads the pseudocode if loaded_ea != entry.ea return # Unified model: the code view is always the continuous listing, @@ -8139,13 +8710,21 @@ class IdaTui(App): # leave the viewport alone and just move the cursor; otherwise # scroll so the target sits a few lines below the top for context. top = round(lst.scroll_offset.y) - if lst.model is lm and top <= entry.cursor < top + lst._visible_height(): + if ( + lst.model is lm + and top <= entry.cursor < top + lst._visible_height() + ): sy = top else: sy = max(entry.cursor - _JUMP_CONTEXT, 0) lst.load( - lm, entry.name, cursor=entry.cursor, - cursor_x=entry.cursor_x, scroll_y=sy, focus=focus) + lm, + entry.name, + cursor=entry.cursor, + cursor_x=entry.cursor_x, + scroll_y=sy, + focus=focus, + ) self._active = ViewMode.LISTING self._show_active() # Graph mode is sticky: following a call from the graph should land in @@ -8157,8 +8736,15 @@ class IdaTui(App): # Prompt overlays that own the keyboard while visible; a background # navigation must not yank focus out from under them (else typed keys leak # into a code view as destructive verbs — e.g. 'u' = undefine). - _PROMPT_IDS = ("search", "rename", "comment", "retype", "goto", "export", - "func-filter") + _PROMPT_IDS = ( + "search", + "rename", + "comment", + "retype", + "goto", + "export", + "func-filter", + ) def _prompt_active(self) -> bool: for iid in self._PROMPT_IDS: @@ -8265,8 +8851,10 @@ class IdaTui(App): sec = self.program.section_of(va) if self.program else None fo = self.program.file_offset(va) if self.program else None foff = f"file+{fo:#x}" if fo is not None else "file:--" - self._status(f"hex va={va:#x} {foff} [{sec or '?'}] " - "(g goto · Enter→code · Tab/Esc/\\→back)") + self._status( + f"hex va={va:#x} {foff} [{sec or '?'}] " + "(g goto · Enter→code · Tab/Esc/\\→back)" + ) def on_hex_view_moved(self, msg: HexView.Moved) -> None: self._hex_status(msg.va) @@ -8303,8 +8891,13 @@ class IdaTui(App): why = self.program.decomp_error(ea) self.app.call_from_thread(self._apply_decomp, ea, name, dec, why) - def _apply_decomp(self, ea: int, name: str, dec, # type: ignore[no-untyped-def] - why: str = "") -> None: + def _apply_decomp( + self, + ea: int, + name: str, + dec, # type: ignore[no-untyped-def] + why: str = "", + ) -> None: view = self.query_one(DecompView) view.loading = False if dec.failed: @@ -8348,11 +8941,12 @@ class IdaTui(App): view.set_nums(self._pending_nums) if self._split: self._sync_split(self._active) # crude link now - self._load_split_map(ea) # then upgrade to the region map + self._load_split_map(ea) # then upgrade to the region map self._split_status() else: self._status( - f"{name} @ {ea:#x} [pseudocode {len(dec.code or '')} chars]{note}") + f"{name} @ {ea:#x} [pseudocode {len(dec.code or '')} chars]{note}" + ) def _sync_split(self, source: str, resync: bool = True) -> None: """Split view: highlight (+ scroll into view) the companion pane's @@ -8377,8 +8971,7 @@ class IdaTui(App): if source == "decomp": dec.set_link(None) # the driver shows its own cursor, no band line, screen = self._split_anchor(dec) - eas = (self._split_eamap[line] - if 0 <= line < len(self._split_eamap) else []) + eas = self._split_eamap[line] if 0 <= line < len(self._split_eamap) else [] if not eas: # fallback: the single /*ea*/ marker for the line one = dec._line_ea(line) eas = [one] if one is not None else [] @@ -8467,18 +9060,25 @@ class IdaTui(App): if self.is_decomp: dec = self.query_one(DecompView) ea = dec._line_ea(dec.cursor) - n = (len(self._split_eamap[dec.cursor]) - if 0 <= dec.cursor < len(self._split_eamap) else 0) + n = ( + len(self._split_eamap[dec.cursor]) + if 0 <= dec.cursor < len(self._split_eamap) + else 0 + ) at = f" @ {ea:#x}" if ea is not None else "" rel = f" \u2194 {n} insn" if n else "" - self._status(f"{self._cur.name}{at} " - f"[split \u00b7 pseudocode line {dec.cursor + 1}{rel}]" - f" (Tab/click: drive listing)") + self._status( + f"{self._cur.name}{at} " + f"[split \u00b7 pseudocode line {dec.cursor + 1}{rel}]" + f" (Tab/click: drive listing)" + ) else: ea = self.query_one(ListingView)._cursor_ea() at = f" @ {ea:#x}" if ea is not None else "" - self._status(f"{self._cur.name}{at} [split \u00b7 listing]" - f" (Tab/click: drive pseudocode)") + self._status( + f"{self._cur.name}{at} [split \u00b7 listing]" + f" (Tab/click: drive pseudocode)" + ) def on_descendant_focus(self, event) -> None: # type: ignore[no-untyped-def] """In split, focusing a pane (Tab or a mouse click) makes it the leading/ @@ -8486,8 +9086,13 @@ class IdaTui(App): if not self._split: return w = event.control - new = ("decomp" if isinstance(w, DecompView) - else "listing" if isinstance(w, ListingView) else None) + new = ( + "decomp" + if isinstance(w, DecompView) + else "listing" + if isinstance(w, ListingView) + else None + ) if new is not None and new != self._active: self._active = new self._sync_split(new) @@ -8570,8 +9175,10 @@ class IdaTui(App): ea = msg.ea if ea is not None: sec = self.program.section_of(ea) if self.program else None - self._status(f"{sec or '?'} @ {ea:#x} [listing] " - "(c code · p func · u undefine · Enter follow)") + self._status( + f"{sec or '?'} @ {ea:#x} [listing] " + "(c code · p func · u undefine · Enter follow)" + ) # -- teardown ---------------------------------------------------------- # async def on_unmount(self) -> None: diff --git a/idatui/diag.py b/idatui/diag.py index b45a72d..7feb6e9 100644 --- a/idatui/diag.py +++ b/idatui/diag.py @@ -26,6 +26,7 @@ screen is normal and happens constantly; wrapping that would bury the real entries in noise. The test for whether it belongs here is "would I want to see this after the fact?". """ + from __future__ import annotations import contextlib @@ -60,7 +61,7 @@ def log(msg: str) -> None: with open(path, "a", encoding="utf-8") as fh: fh.write(f"{time.strftime('%H:%M:%S')} {msg}\n") except OSError: - pass # a broken log path must never break the app + pass # a broken log path must never break the app def note(what: str, exc: BaseException) -> None: @@ -76,8 +77,11 @@ def note(what: str, exc: BaseException) -> None: _ring.append(entry) log(f"[swallowed] {what}: {entry['error']} ({entry['where']})") if _logfile(): - log("".join(traceback.format_exception( - type(exc), exc, exc.__traceback__)).rstrip()) + log( + "".join( + traceback.format_exception(type(exc), exc, exc.__traceback__) + ).rstrip() + ) def _origin(exc: BaseException) -> str: diff --git a/idatui/drive.py b/idatui/drive.py index 6e5c21a..a1da513 100644 --- a/idatui/drive.py +++ b/idatui/drive.py @@ -17,6 +17,7 @@ Socket: --sock, else IDATUI_RPC_SOCK, else the one live pane (from `pane list`). '.' or omitted as a target means the current function. `raw k=v` is a passthrough to rpcclient (pretty JSON). """ + from __future__ import annotations import json @@ -33,16 +34,24 @@ def _resolve_sock(explicit: str | None) -> str: env = os.environ.get("IDATUI_RPC_SOCK") if env: return env - live = [r for r in _load_registry() - if _pane_alive(r.get("pane", "")) and r.get("sock") - and os.path.exists(r["sock"])] + live = [ + r + for r in _load_registry() + if _pane_alive(r.get("pane", "")) + and r.get("sock") + and os.path.exists(r["sock"]) + ] if len(live) == 1: return live[0]["sock"] if not live: - raise SystemExit("no live idatui pane — pass --sock, set IDATUI_RPC_SOCK, " - "or `python -m idatui.pane spawn ...`") - raise SystemExit("multiple live panes — pass --sock :\n" - + "\n".join(" " + r["sock"] for r in live)) + raise SystemExit( + "no live idatui pane — pass --sock, set IDATUI_RPC_SOCK, " + "or `python -m idatui.pane spawn ...`" + ) + raise SystemExit( + "multiple live panes — pass --sock :\n" + + "\n".join(" " + r["sock"] for r in live) + ) def _tgt(a: str | None) -> str | None: @@ -121,8 +130,9 @@ def cmd_pc(c, args): lines = d["code"].splitlines() if needle: nlow = needle.lower() - lines = [f"{i:4} {line}" for i, line in enumerate(lines) - if nlow in line.lower()] + lines = [ + f"{i:4} {line}" for i, line in enumerate(lines) if nlow in line.lower() + ] return "\n".join(lines) or f"(no line matches {needle!r})" return d["code"] @@ -153,8 +163,10 @@ def cmd_callers(c, args): if not args: raise SystemExit("usage: callers ") xs = c.call("xrefs_to", target=args[0]) - return "\n".join(f" {x['frm']:#x} in {x.get('fn_name')}" for x in xs) \ + return ( + "\n".join(f" {x['frm']:#x} in {x.get('fn_name')}" for x in xs) or "(no callers)" + ) def cmd_names(c, args): @@ -162,8 +174,10 @@ def cmd_names(c, args): raise SystemExit("usage: names [limit]") lim = int(args[1]) if len(args) > 1 else 40 fs = c.call("functions", filter=args[0], limit=lim) - return "\n".join(f" {f['ea']:#x} {f['name']} ({f['size']})" for f in fs) \ + return ( + "\n".join(f" {f['ea']:#x} {f['name']} ({f['size']})" for f in fs) or "(no match)" + ) def cmd_binaries(c, args): @@ -246,8 +260,9 @@ def cmd_define(c, args): a symbol file), so take them all and report per-target. """ if not args: - raise SystemExit("usage: define [target ...]") + raise SystemExit( + "usage: define [target ...]" + ) kind, targets = args[0], (args[1:] or [None]) out = [] for t in targets: @@ -284,8 +299,10 @@ def cmd_syms(c, args): raise SystemExit("usage: syms ") r = c.call("rename_many", file=os.path.abspath(os.path.expanduser(args[0]))) m = r.get("rename_many", {}) - out = [f" {m.get('ok', 0)}/{m.get('requested', 0)} renamed" - f" (skipped {m.get('skipped', 0)}, failed {m.get('failed', 0)})"] + out = [ + f" {m.get('ok', 0)}/{m.get('requested', 0)} renamed" + f" (skipped {m.get('skipped', 0)}, failed {m.get('failed', 0)})" + ] for e in m.get("errors", []): out.append(f" {e.get('addr')}: {e.get('error')}") return "\n".join(out) @@ -304,8 +321,10 @@ def cmd_find(c, args): hits = r.get("hits", []) out = [f" [{r.get('mode')}] {len(hits)}{'+' if r.get('truncated') else ''} hits"] for h in hits[:40]: - out.append(f" {h['addr']} {(h.get('func') or h.get('seg') or ''):<20.20} " - f"{h.get('line', '')}") + out.append( + f" {h['addr']} {(h.get('func') or h.get('seg') or ''):<20.20} " + f"{h.get('line', '')}" + ) if len(hits) > 40: out.append(f" … {len(hits) - 40} more") return "\n".join(out) @@ -314,9 +333,11 @@ def cmd_find(c, args): def cmd_export(c, args): """export [path] -- write the session's findings as markdown.""" r = c.call("export", **({"path": args[0]} if args else {})) - return (f" {r.get('path')} ({r.get('bytes', 0)} bytes: " - f"{r.get('comments', 0)} comments, {r.get('names', 0)} names, " - f"{r.get('types', 0)} types)") + return ( + f" {r.get('path')} ({r.get('bytes', 0)} bytes: " + f"{r.get('comments', 0)} comments, {r.get('names', 0)} names, " + f"{r.get('types', 0)} types)" + ) def cmd_screen(c, args): @@ -334,13 +355,27 @@ def cmd_raw(c, args): COMMANDS = { - "where": cmd_where, "go": cmd_go, "pc": cmd_pc, "dis": cmd_dis, - "callees": cmd_callees, "callers": cmd_callers, "names": cmd_names, - "rename": cmd_rename, "mv": cmd_mv, "note": cmd_note, "retype": cmd_retype, - "save": cmd_save, "screen": cmd_screen, "raw": cmd_raw, "define": cmd_define, - "syms": cmd_syms, "fmt": cmd_fmt, "export": cmd_export, + "where": cmd_where, + "go": cmd_go, + "pc": cmd_pc, + "dis": cmd_dis, + "callees": cmd_callees, + "callers": cmd_callers, + "names": cmd_names, + "rename": cmd_rename, + "mv": cmd_mv, + "note": cmd_note, + "retype": cmd_retype, + "save": cmd_save, + "screen": cmd_screen, + "raw": cmd_raw, + "define": cmd_define, + "syms": cmd_syms, + "fmt": cmd_fmt, + "export": cmd_export, "find": cmd_find, - "binaries": cmd_binaries, "switch": cmd_switch, + "binaries": cmd_binaries, + "switch": cmd_switch, } diff --git a/idatui/errors.py b/idatui/errors.py index ab1365a..7aceef2 100644 --- a/idatui/errors.py +++ b/idatui/errors.py @@ -4,6 +4,7 @@ The IDA Nexus adapter normalizes ``ida_nexus`` transport and execution errors into these types so the domain and Textual layers do not depend on HTTP or registry implementation details. """ + from __future__ import annotations from dataclasses import dataclass diff --git a/idatui/findings.py b/idatui/findings.py index 58a69f6..e5ce91f 100644 --- a/idatui/findings.py +++ b/idatui/findings.py @@ -85,7 +85,9 @@ def from_loader(seg: str, name: str = "") -> bool: #: its stereotyped shapes, which no one types by accident. _ANALYZER = re.compile( r"^(?:switch \d+ cases?|switch jump|jumptable [0-9A-Fa-f]+\b.*|" - r"indirect table for switch.*|jump table for switch.*)$", re.I) + r"indirect table for switch.*|jump table for switch.*)$", + re.I, +) #: The other family is argument hints (`s1`, `locale`, `domainname`), which IDA #: copies from the callee's prototype onto each argument-setup instruction. They @@ -101,11 +103,14 @@ def analyzer_texts(comments) -> set[str]: counts: dict[str, int] = {} for c in comments: text = (c.text or "").strip() - if text and not text.split()[1:]: # a single whitespace-free token + if text and not text.split()[1:]: # a single whitespace-free token counts[text] = counts.get(text, 0) + 1 out = {t for t, n in counts.items() if n >= _HINT_REPEATS} - out |= {(c.text or "").strip() for c in comments - if _ANALYZER.match((c.text or "").strip())} + out |= { + (c.text or "").strip() + for c in comments + if _ANALYZER.match((c.text or "").strip()) + } return out @@ -115,7 +120,8 @@ _DUMMY = re.compile( r"^(?:(?:sub|loc|locret|off|seg|asc|byte|word|dword|qword|xmmword|ymmword|" r"flt|dbl|tbyte|stru|algn|unk|nullsub|def|jpt|jsub)_[0-9A-Fa-f]+" # j_strlen: a thunk name IDA derives from its target, not from a person. - r"|j_\w+)$") + r"|j_\w+)$" +) def is_dummy(name: str) -> bool: @@ -123,8 +129,9 @@ def is_dummy(name: str) -> bool: return bool(_DUMMY.match(name or "")) -def gather(program, path: str = "", *, limit: int = 4000, - types: bool = True, journal=None) -> Findings: +def gather( + program, path: str = "", *, limit: int = 4000, types: bool = True, journal=None +) -> Findings: """Collect a :class:`Findings` from a live :class:`Program`. ``path`` is the binary the app opened -- ``Program`` speaks to a database @@ -139,8 +146,11 @@ def gather(program, path: str = "", *, limit: int = 4000, if journal is not None: try: out.recorded = journal.addresses() - out.recorded_types = {e.get("d", "") for e in journal.entries - if e.get("k") == "type" and e.get("d")} + out.recorded_types = { + e.get("d", "") + for e in journal.entries + if e.get("k") == "type" and e.get("d") + } out.n_recorded = len(journal) except Exception: # noqa: BLE001 out.recorded, out.recorded_types, out.n_recorded = set(), set(), 0 @@ -188,8 +198,9 @@ def _esc(text: str) -> str: def _fence(text: str) -> str: """Fence body text so a comment containing backticks cannot break out.""" - ticks = "`" * max(3, max((len(m) for m in re.findall(r"`+", text or "")), - default=0) + 1) + ticks = "`" * max( + 3, max((len(m) for m in re.findall(r"`+", text or "")), default=0) + 1 + ) return f"{ticks}\n{(text or '').rstrip()}\n{ticks}" @@ -199,9 +210,13 @@ def _user_names(f: Findings) -> list: With a journal, that is exactly the addresses we recorded renaming. Without one, it is a judgement: a real name, not the linker's, not the loader's. """ - names = [n for n in f.names - if not is_dummy(n.name) and n.name not in f.linked - and not from_loader(n.seg, n.name)] + names = [ + n + for n in f.names + if not is_dummy(n.name) + and n.name not in f.linked + and not from_loader(n.seg, n.name) + ] if f.recorded: return [n for n in names if n.addr in f.recorded] return names @@ -212,8 +227,7 @@ def _user_types(f: Findings) -> list: IDA loaded, so with a journal we show only the ones declared here; without one, all of them, newest ordinal first (yours are the newest).""" if f.recorded or f.recorded_types: - return [t for t in f.types - if getattr(t[0], "name", "") in f.recorded_types] + return [t for t in f.types if getattr(t[0], "name", "") in f.recorded_types] return list(f.types) @@ -244,8 +258,7 @@ def render(f: Findings) -> str: names = sorted(_user_names(f), key=lambda n: n.addr) funcs = [n for n in names if n.is_func] data = [n for n in names if not n.is_func] - comments = sorted(_user_comments(f), key=lambda c: (c.func_addr or c.addr, - c.addr)) + comments = sorted(_user_comments(f), key=lambda c: (c.func_addr or c.addr, c.addr)) dropped = (len(f.comments) - len(comments)) + (len(f.names) - len(names)) types = _user_types(f) @@ -253,9 +266,11 @@ def render(f: Findings) -> str: title = f.binary or "database" L.append(f"# Findings — {title}") L.append("") - L.append(f"*{len(funcs)} named functions · {len(data)} named data · " - f"{len(comments)} comments · {len(types)} local types — " - f"exported {when} by idatui*") + L.append( + f"*{len(funcs)} named functions · {len(data)} named data · " + f"{len(comments)} comments · {len(types)} local types — " + f"exported {when} by idatui*" + ) L.append("") if f.path: L.append(f"- **binary**: `{f.path}`") @@ -267,24 +282,34 @@ def render(f: Findings) -> str: L.append(f"- **segments**: {segs}{more}") if f.recorded or f.recorded_types: n_at = len(f.recorded) - L.append(f"- **source**: idatui's edit journal — {f.n_recorded} recorded " - f"edits across {n_at} address{'' if n_at == 1 else 'es'}. " - "Everything below is work done here, not the analyzer's.") + L.append( + f"- **source**: idatui's edit journal — {f.n_recorded} recorded " + f"edits across {n_at} address{'' if n_at == 1 else 'es'}. " + "Everything below is work done here, not the analyzer's." + ) else: - L.append("- **source**: a scan of the database. Nothing in a `.i64` " - "records *who* wrote a comment or a name — IDA's own analyzer " - "uses the same calls — so this is filtered by shape and may " - "include its work as well as yours.") + L.append( + "- **source**: a scan of the database. Nothing in a `.i64` " + "records *who* wrote a comment or a name — IDA's own analyzer " + "uses the same calls — so this is filtered by shape and may " + "include its work as well as yours." + ) if not f.stripped: - L.append("- **note**: this binary has its own symbols, so the names " - "below include ones it shipped with.") + L.append( + "- **note**: this binary has its own symbols, so the names " + "below include ones it shipped with." + ) if dropped and (f.recorded or f.recorded_types): - L.append(f"- **note**: {dropped} other annotations in this database " - "were not made here (the analyzer's, the loader's, the " - "linker's) and are left out.") + L.append( + f"- **note**: {dropped} other annotations in this database " + "were not made here (the analyzer's, the loader's, the " + "linker's) and are left out." + ) elif dropped: - L.append(f"- **note**: {dropped} annotations left out as the loader's " - "own (file headers, dummy names, imports).") + L.append( + f"- **note**: {dropped} annotations left out as the loader's " + "own (file headers, dummy names, imports)." + ) if f.truncated: L.append("- **note**: the scan hit its limit; this report is partial.") L.append("") @@ -293,8 +318,10 @@ def render(f: Findings) -> str: L.append("## Comments") L.append("") if not comments: - L.append("*None. (Comments are the part of a database nobody else can " - "reconstruct — they are worth writing.)*") + L.append( + "*None. (Comments are the part of a database nobody else can " + "reconstruct — they are worth writing.)*" + ) L.append("") else: by_func: dict[str, list] = {} @@ -309,11 +336,9 @@ def render(f: Findings) -> str: L.append("") for c in rows: if c.whole_func: - L.append(f"- **{c.addr:#x}** — *whole function*: " - f"{_esc(c.text)}") + L.append(f"- **{c.addr:#x}** — *whole function*: {_esc(c.text)}") elif c.line: - L.append(f"- **{c.addr:#x}** `{_esc(c.line)}` \n" - f" {_esc(c.text)}") + L.append(f"- **{c.addr:#x}** `{_esc(c.line)}` \n {_esc(c.text)}") else: L.append(f"- **{c.addr:#x}** — {_esc(c.text)}") L.append("") @@ -329,8 +354,7 @@ def render(f: Findings) -> str: L.append("|---|---|---|---|") for n in funcs: proto = f"`{_esc(n.proto)}`" if n.proto else "" - L.append(f"| `{n.addr:#x}` | `{_esc(n.name)}` | " - f"{n.size:#x} | {proto} |") + L.append(f"| `{n.addr:#x}` | `{_esc(n.name)}` | {n.size:#x} | {proto} |") L.append("") if data: L.append("## Named data") @@ -346,17 +370,21 @@ def render(f: Findings) -> str: L.append("## Local types") L.append("") if not (f.recorded or f.recorded_types): - L.append("*Newest first. A database is seeded with types from the " - "libraries IDA loaded, so the ones you defined are the " - "ones with the highest ordinals — at the top of this " - "list.*") + L.append( + "*Newest first. A database is seeded with types from the " + "libraries IDA loaded, so the ones you defined are the " + "ones with the highest ordinals — at the top of this " + "list.*" + ) L.append("") ordered = sorted(types, key=lambda t: -getattr(t[0], "ordinal", 0)) for st, src in ordered: kw = "union" if getattr(st, "is_union", False) else "struct" - L.append(f"### `{kw} {st.name}` " - f"({getattr(st, 'size', 0):#x} bytes, " - f"{getattr(st, 'members', 0)} fields)") + L.append( + f"### `{kw} {st.name}` " + f"({getattr(st, 'size', 0):#x} bytes, " + f"{getattr(st, 'members', 0)} fields)" + ) L.append("") if src: L.append("```c") @@ -372,9 +400,15 @@ def default_path(program_path: str) -> str: return f"{base}.findings.md" -def export(program, binary_path: str = "", out_path: str | None = None, *, - limit: int = 4000, types: bool = True, - journal=None) -> tuple[str, Findings]: +def export( + program, + binary_path: str = "", + out_path: str | None = None, + *, + limit: int = 4000, + types: bool = True, + journal=None, +) -> tuple[str, Findings]: """Gather, render and WRITE the report. Returns ``(path, findings)``.""" f = gather(program, binary_path, limit=limit, types=types, journal=journal) out = out_path or default_path(f.path) diff --git a/idatui/formats.py b/idatui/formats.py index f09bb99..d3d2abd 100644 --- a/idatui/formats.py +++ b/idatui/formats.py @@ -60,7 +60,7 @@ def sniff(path: str) -> str | None: if not head: return None for magic, off, name in _MAGIC: - if head[off:off + len(magic)] == magic: + if head[off : off + len(magic)] == magic: return name # Only treat a text prefix as a format if the whole head is printable — # a raw blob starting with 0x3a (':') is far more likely than Intel HEX. diff --git a/idatui/graph.py b/idatui/graph.py index f40de07..f8df2ec 100644 --- a/idatui/graph.py +++ b/idatui/graph.py @@ -26,6 +26,7 @@ cells, so ``Painting`` is an *index* — per-row horizontal runs, a bucketed interval index of vertical runs, and point marks — and the view asks it for one row at a time (``cells_at_row``), exactly like the listing's ``render_line``. """ + from __future__ import annotations import logging @@ -37,8 +38,8 @@ _LOG = logging.getLogger(__name__) # Terminal cells are about twice as tall as they are wide, so horizontal gaps # need roughly 2x the cell count of vertical gaps to look square. -HGAP = 3 # min columns between two boxes in a layer -VGAP = 1 # min rows between a layer band and the channel below it +HGAP = 3 # min columns between two boxes in a layer +VGAP = 1 # min rows between a layer band and the channel below it # Edge classes, used as style keys by the renderer. E_UNCOND = "uncond" @@ -68,8 +69,8 @@ class Node: label: str = "" rank: int = 0 order: int = 0 - x: int = 0 # left column - y: int = 0 # top row + x: int = 0 # left column + y: int = 0 # top row w: int = 1 h: int = 1 @@ -134,6 +135,7 @@ class _Graph: # ------------------------------------------------------------ 1. cycles + def _break_cycles(g: _Graph, root: int) -> None: """Reverse back edges (DFS gray-set) so layering sees a DAG.""" color: dict[int, int] = {} @@ -166,6 +168,7 @@ def _break_cycles(g: _Graph, root: int) -> None: # --------------------------------------------------------- 2. layering + def _assign_ranks(g: _Graph, root: int) -> None: """Longest-path layering: rank(v) = 1 + max(rank(preds)). @@ -205,11 +208,12 @@ def _assign_ranks(g: _Graph, root: int) -> None: # ---------------------------------------------------------- 3. dummies + def _add_dummies(g: _Graph) -> None: for e in list(g.edges): span = g.nodes[e.dst].rank - g.nodes[e.src].rank if span <= 0: - e.back = True # residual cycle: colour it, route it flat + e.back = True # residual cycle: colour it, route it flat chain = [e.src] if span > 1: for r in range(g.nodes[e.src].rank + 1, g.nodes[e.dst].rank): @@ -238,6 +242,7 @@ def _segments(g: _Graph) -> list[tuple[int, int, Edge]]: # ---------------------------------------------------------- 4. ordering + def _neighbors(g: _Graph) -> tuple[dict[int, list[int]], dict[int, list[int]]]: down: dict[int, list[int]] = {i: [] for i in g.nodes} up: dict[int, list[int]] = {i: [] for i in g.nodes} @@ -247,8 +252,9 @@ def _neighbors(g: _Graph) -> tuple[dict[int, list[int]], dict[int, list[int]]]: return down, up -def _cross_below(layer: list[int], down: dict[int, list[int]], - pos: dict[int, int]) -> int: +def _cross_below( + layer: list[int], down: dict[int, list[int]], pos: dict[int, int] +) -> int: """Crossings between this layer and the one below, counted as inversions with a Fenwick tree: O(E log E). The naive O(E^2) version is the entire runtime on a 400-block function (20s vs 150ms), so it is not an option.""" @@ -277,8 +283,7 @@ def _cross_below(layer: list[int], down: dict[int, list[int]], return total -def _pair_cross(a: int, b: int, side: dict[int, list[int]], - pos: dict[int, int]) -> int: +def _pair_cross(a: int, b: int, side: dict[int, list[int]], pos: dict[int, int]) -> int: """Crossings from a's and b's edges to one neighbouring layer given a sits immediately LEFT of b. Local — O(deg(a)*deg(b)) — so the transposition pass never has to recount the whole graph per candidate swap.""" @@ -291,8 +296,13 @@ def _pair_cross(a: int, b: int, side: dict[int, list[int]], return n -def _swap_delta(a: int, b: int, down: dict[int, list[int]], - up: dict[int, list[int]], pos: dict[int, int]) -> tuple[int, int]: +def _swap_delta( + a: int, + b: int, + down: dict[int, list[int]], + up: dict[int, list[int]], + pos: dict[int, int], +) -> tuple[int, int]: """``(keep, swap)`` for the adjacent pair (a, b), both sides, in one pass. The same as calling :func:`_pair_cross` four times, which is what the @@ -318,8 +328,9 @@ def _swap_delta(a: int, b: int, down: dict[int, list[int]], return keep, swap -def crossings(layers: list[list[int]], down: dict[int, list[int]], - pos: dict[int, int]) -> int: +def crossings( + layers: list[list[int]], down: dict[int, list[int]], pos: dict[int, int] +) -> int: return sum(_cross_below(l, down, pos) for l in layers) @@ -340,7 +351,7 @@ def _order_layers(g: _Graph, root: int, sweeps: int = 6) -> list[list[int]]: seen.add(j) stack.append(j) for layer in layers: - layer.sort(key=lambda i: seed.get(i, 10 ** 9)) + layer.sort(key=lambda i: seed.get(i, 10**9)) pos = {i: k for layer in layers for k, i in enumerate(layer)} def median(i: int, side: dict[int, list[int]]) -> float: @@ -394,6 +405,7 @@ def _order_layers(g: _Graph, root: int, sweeps: int = 6) -> list[list[int]]: # --------------------------------------------------------- 5. x coords + def _assign_x(g: _Graph, layers: list[list[int]], sweeps: int = 8) -> None: down, up = _neighbors(g) for layer in layers: @@ -418,8 +430,9 @@ def _assign_x(g: _Graph, layers: list[list[int]], sweeps: int = 8) -> None: for r in rng: layer = layers[r] # dummies first: keeping long edges straight matters most - order = sorted(layer, key=lambda i: (not g.nodes[i].dummy, - g.nodes[i].order)) + order = sorted( + layer, key=lambda i: (not g.nodes[i].dummy, g.nodes[i].order) + ) for i in order: nb = side[i] if not nb: @@ -437,6 +450,7 @@ def _assign_x(g: _Graph, layers: list[list[int]], sweeps: int = 8) -> None: # ------------------------------------------------------------ 6. route + def _ports(g: _Graph) -> tuple[dict, dict]: """Spread a node's out-edges along its bottom border and its in-edges along its top, each ordered by the other end's x so they don't cross at the node.""" @@ -470,8 +484,8 @@ def _ports(g: _Graph) -> tuple[dict, dict]: class Route: edge: Edge pts: list[tuple[int, int]] - head: bool = True # arrowhead (target is a real block) - tail: bool = True # port tee (source is a real block) + head: bool = True # arrowhead (target is a real block) + tail: bool = True # port tee (source is a real block) #: the polyline is drawn against control flow (a reversed back edge), so the #: arrowhead belongs at ``pts[0]`` and the port tee at ``pts[-1]``. flipped: bool = False @@ -490,7 +504,7 @@ def _route(g: _Graph, layers: list[list[int]]) -> list[Route]: runs = [] for a, b, e in lst: x0, x1 = out_port[(a, b, id(e))], in_port[(a, b, id(e))] - if x0 != x1: # a straight drop needs no lane + if x0 != x1: # a straight drop needs no lane runs.append((min(x0, x1), max(x0, x1), (a, b, id(e)))) runs.sort(key=lambda t: (t[1] - t[0], t[0])) occupied: list[list[tuple[int, int]]] = [] @@ -516,7 +530,7 @@ def _route(g: _Graph, layers: list[list[int]]) -> list[Route]: n = g.nodes[i] n.y = y if n.dummy: - n.h = h # the band is its pass-through + n.h = h # the band is its pass-through chan_y.append(y + h - 1 + VGAP) y += h - 1 + VGAP + channels[r] + VGAP + 1 @@ -535,34 +549,58 @@ def _route(g: _Graph, layers: list[list[int]]) -> list[Route]: else: ych = chan_y[na.rank] + lanes.get((a, b, id(e)), 0) pts = [(y0, x0), (ych, x0), (ych, x1), (y1, x1)] - routes.append(Route(edge=e, pts=pts, head=not nb.dummy, - tail=not na.dummy, flipped=e.flipped)) + routes.append( + Route( + edge=e, pts=pts, head=not nb.dummy, tail=not na.dummy, flipped=e.flipped + ) + ) return routes # ------------------------------------------------------------ painting -BOX = {"tl": "\u250c", "tr": "\u2510", "bl": "\u2514", "br": "\u2518", - "h": "\u2500", "v": "\u2502"} -LINE_CHARS = set("\u2502\u2500\u250c\u2510\u2514\u2518\u251c\u2524\u252c\u2534" - "\u253c\u256d\u256e\u2570\u256f") +BOX = { + "tl": "\u250c", + "tr": "\u2510", + "bl": "\u2514", + "br": "\u2518", + "h": "\u2500", + "v": "\u2502", +} +LINE_CHARS = set( + "\u2502\u2500\u250c\u2510\u2514\u2518\u251c\u2524\u252c\u2534" + "\u253c\u256d\u256e\u2570\u256f" +) MERGE = { frozenset("\u2502\u2500"): "\u253c", - frozenset("\u2502\u250c"): "\u251c", frozenset("\u2502\u2510"): "\u2524", - frozenset("\u2502\u2514"): "\u251c", frozenset("\u2502\u2518"): "\u2524", - frozenset("\u2500\u250c"): "\u252c", frozenset("\u2500\u2510"): "\u252c", - frozenset("\u2500\u2514"): "\u2534", frozenset("\u2500\u2518"): "\u2534", - frozenset("\u2502\u256d"): "\u251c", frozenset("\u2502\u256e"): "\u2524", - frozenset("\u2502\u2570"): "\u251c", frozenset("\u2502\u256f"): "\u2524", - frozenset("\u2500\u256d"): "\u252c", frozenset("\u2500\u256e"): "\u252c", - frozenset("\u2500\u2570"): "\u2534", frozenset("\u2500\u256f"): "\u2534", + frozenset("\u2502\u250c"): "\u251c", + frozenset("\u2502\u2510"): "\u2524", + frozenset("\u2502\u2514"): "\u251c", + frozenset("\u2502\u2518"): "\u2524", + frozenset("\u2500\u250c"): "\u252c", + frozenset("\u2500\u2510"): "\u252c", + frozenset("\u2500\u2514"): "\u2534", + frozenset("\u2500\u2518"): "\u2534", + frozenset("\u2502\u256d"): "\u251c", + frozenset("\u2502\u256e"): "\u2524", + frozenset("\u2502\u2570"): "\u251c", + frozenset("\u2502\u256f"): "\u2524", + frozenset("\u2500\u256d"): "\u252c", + frozenset("\u2500\u256e"): "\u252c", + frozenset("\u2500\u2570"): "\u2534", + frozenset("\u2500\u256f"): "\u2534", } CORNER = { - ("D", "R"): "\u2570", ("D", "L"): "\u256f", ("R", "D"): "\u256e", - ("L", "D"): "\u256d", ("R", "U"): "\u256f", ("L", "U"): "\u2570", - ("U", "R"): "\u256d", ("U", "L"): "\u256e", + ("D", "R"): "\u2570", + ("D", "L"): "\u256f", + ("R", "D"): "\u256e", + ("L", "D"): "\u256d", + ("R", "U"): "\u256f", + ("L", "U"): "\u2570", + ("U", "R"): "\u256d", + ("U", "L"): "\u256e", } -BUCKET = 32 # rows per vertical-run index bucket +BUCKET = 32 # rows per vertical-run index bucket def _dir(p: tuple[int, int], q: tuple[int, int]) -> str: @@ -595,8 +633,9 @@ class Painting: def add_mark(self, row: int, col: int, ch: str, style: str, eid: int) -> None: self.marks.setdefault(row, []).append((col, ch, style, eid)) - def cells_at_row(self, row: int, c0: int, c1: int - ) -> dict[int, tuple[str, str, int]]: + def cells_at_row( + self, row: int, c0: int, c1: int + ) -> dict[int, tuple[str, str, int]]: """{col: (char, style, edge_id)} for ``row`` within [c0, c1).""" out: dict[int, tuple[str, str, int]] = {} @@ -604,8 +643,13 @@ class Painting: if col < c0 or col >= c1: return old = out.get(col) - if old and not force and old[0] != ch \ - and old[0] in LINE_CHARS and ch in LINE_CHARS: + if ( + old + and not force + and old[0] != ch + and old[0] in LINE_CHARS + and ch in LINE_CHARS + ): ch = MERGE.get(frozenset((old[0], ch)), ch) out[col] = (ch, style, eid) @@ -626,16 +670,16 @@ class Layout: """The finished drawing: boxes, an edge index, and enough structure for the view to hit-test, navigate and highlight.""" - nodes: list[Node] # real blocks only, layout order + nodes: list[Node] # real blocks only, layout order by_id: dict[int, Node] edges: list[Edge] painting: Painting width: int height: int entry: int - rows: dict[int, list[int]] # row -> real node ids covering it - incident: dict[int, set[int]] # node id -> edge ids touching it - succ: dict[int, list[tuple[int, str]]] # node id -> [(node id, style)] + rows: dict[int, list[int]] # row -> real node ids covering it + incident: dict[int, set[int]] # node id -> edge ids touching it + succ: dict[int, list[tuple[int, str]]] # node id -> [(node id, style)] pred: dict[int, list[tuple[int, str]]] stats: dict @@ -670,8 +714,15 @@ def _build(blocks: list[Block], sizer, entry: int | None) -> tuple[_Graph, int]: for b in blocks: w, h = sizer(b) b.selfloop = False - g.add(Node(id=b.id, block=b, label=f"loc_{b.start:X}", - w=max(int(w), 4), h=max(int(h), 3))) + g.add( + Node( + id=b.id, + block=b, + label=f"loc_{b.start:X}", + w=max(int(w), 4), + h=max(int(h), 3), + ) + ) for b in blocks: outs = [(d, k) for d, k in b.succs if d in g.nodes] for dst, kind in outs: @@ -722,14 +773,16 @@ def _pick_engine(engine: str | None, nblocks: int) -> str: want = "auto" if want == "auto": from . import graph_triskel + if nblocks <= AUTO_TRISKEL_MAX_BLOCKS and graph_triskel.available(): return "triskel" return "native" return want -def layout(blocks: list[Block], sizer, entry: int | None = None, - engine: str | None = None) -> Layout: +def layout( + blocks: list[Block], sizer, entry: int | None = None, engine: str | None = None +) -> Layout: """Lay out ``blocks``. ``sizer(block) -> (width, height)`` in cells. ``engine`` picks the layout backend: ``native`` (pure python, always @@ -747,9 +800,10 @@ def layout(blocks: list[Block], sizer, entry: int | None = None, routes = [] elif name == "triskel": from . import graph_triskel + try: routes, layers = graph_triskel.run(g, root) - except Exception as exc: # noqa: BLE001 + except Exception as exc: # noqa: BLE001 # Native code with a history of throwing on degenerate CFGs. The # graph view is a convenience; losing it beats losing the session. # Keep the REASON: a fallback the user can see but not explain is @@ -804,23 +858,27 @@ def layout(blocks: list[Block], sizer, entry: int | None = None, down_last = last[0] > rt.pts[-2][0] if len(rt.pts) > 1 else True if rt.flipped: if rt.tail: - p.add_mark(first[0], first[1], - "\u25b2" if down_first else "\u25bc", style, eid) + p.add_mark( + first[0], first[1], "\u25b2" if down_first else "\u25bc", style, eid + ) if rt.head: - p.add_mark(last[0], last[1], - "\u2534" if down_last else "\u252c", style, eid) + p.add_mark( + last[0], last[1], "\u2534" if down_last else "\u252c", style, eid + ) else: if rt.tail: - p.add_mark(first[0], first[1], - "\u252c" if down_first else "\u2534", style, eid) + p.add_mark( + first[0], first[1], "\u252c" if down_first else "\u2534", style, eid + ) if rt.head: - p.add_mark(last[0], last[1], - "\u25bc" if down_last else "\u25b2", style, eid) + p.add_mark( + last[0], last[1], "\u25bc" if down_last else "\u25b2", style, eid + ) succ: dict[int, list[tuple[int, str]]] = {n.id: [] for n in real} pred: dict[int, list[tuple[int, str]]] = {n.id: [] for n in real} for e in g.edges: - a, b = (e.dst, e.src) if e.flipped else (e.src, e.dst) # undo reversal + a, b = (e.dst, e.src) if e.flipped else (e.src, e.dst) # undo reversal if a in succ: succ[a].append((b, e.style)) if b in pred: @@ -849,7 +907,17 @@ def layout(blocks: list[Block], sizer, entry: int | None = None, "engine_error": err, "ms": (time.perf_counter() - t0) * 1000, } - return Layout(nodes=order, by_id={n.id: n for n in g.nodes.values()}, - edges=g.edges, painting=p, width=width, height=height, - entry=root, rows=rows, incident=incident, - succ=succ, pred=pred, stats=stats) + return Layout( + nodes=order, + by_id={n.id: n for n in g.nodes.values()}, + edges=g.edges, + painting=p, + width=width, + height=height, + entry=root, + rows=rows, + incident=incident, + succ=succ, + pred=pred, + stats=stats, + ) diff --git a/idatui/graph_triskel.py b/idatui/graph_triskel.py index 2bf92bf..5e8b3f0 100644 --- a/idatui/graph_triskel.py +++ b/idatui/graph_triskel.py @@ -25,6 +25,7 @@ What it does NOT do is trust the library with degenerate input. Self-loops and disconnected graphs make it throw, an empty graph used to segfault, and a segfault takes the TUI down with it. Both are handled here, before the call. """ + from __future__ import annotations import os @@ -56,10 +57,11 @@ def module(): path = os.environ.get("IDATUI_TRISKEL_PATH") if path: import sys + if path not in sys.path: sys.path.insert(0, path) try: - import pytriskel # noqa: PLC0415 + import pytriskel # noqa: PLC0415 except ImportError: return None # Upstream ships wheels whose get_waypoints() always throws (a missing @@ -119,8 +121,9 @@ def _phantom_edges(g: G._Graph, root: int) -> list[tuple[int, int]]: while len(reach) < len(g.nodes): rest = [i for i in g.nodes if i not in reach] rest_set = set(rest) - head = next((i for i in rest - if not any(p in rest_set for p in preds[i])), rest[0]) + head = next( + (i for i in rest if not any(p in rest_set for p in preds[i])), rest[0] + ) phantom.append((root, head)) reach |= _reachable(succ, head) return phantom @@ -168,8 +171,7 @@ def run(g: G._Graph, root: int) -> tuple[list[G.Route], int]: pt = module() if pt is None: raise RuntimeError("pytriskel is not available") - pt.set_spacing(x_gutter=float(HGAP), y_gutter=float(VGAP), - edge_height=float(LANE)) + pt.set_spacing(x_gutter=float(HGAP), y_gutter=float(VGAP), edge_height=float(LANE)) routes: list[G.Route] = [] if g.nodes: @@ -203,9 +205,14 @@ def run(g: G._Graph, root: int) -> tuple[list[G.Route], int]: return routes, len(bands) -def _layout_graph(pt, g: G._Graph, root: int, edges: list[G.Edge], - phantom: list[tuple[int, int]], - routes: list[G.Route]) -> None: +def _layout_graph( + pt, + g: G._Graph, + root: int, + edges: list[G.Edge], + phantom: list[tuple[int, int]], + routes: list[G.Route], +) -> None: """Lay the whole graph out and append its routes.""" order = [root] + [i for i in g.nodes if i != root] succ: dict[int, list[int]] = {i: [] for i in g.nodes} @@ -220,8 +227,10 @@ def _layout_graph(pt, g: G._Graph, root: int, edges: list[G.Edge], # than an exception, so check before crossing into C++ rather than # after. RuntimeError here means a fallback to native; a segfault means # the user loses the session. - raise RuntimeError(f"{len(unreachable)} blocks unreachable from the " - f"layout root {root}: {sorted(unreachable)[:8]}") + raise RuntimeError( + f"{len(unreachable)} blocks unreachable from the " + f"layout root {root}: {sorted(unreachable)[:8]}" + ) builder = pt.make_layout_builder() tid = {} @@ -234,8 +243,10 @@ def _layout_graph(pt, g: G._Graph, root: int, edges: list[G.Edge], # docstring says "width and height", which is the other way round; our # fork makes them keyword arguments so it cannot be got wrong silently. tid[nid] = builder.make_node(height=float(n.h), width=float(n.w)) - teid = [(builder.make_edge(tid[e.src], tid[e.dst], _edge_type(pt, e.kind)), e) - for e in edges] + teid = [ + (builder.make_edge(tid[e.src], tid[e.dst], _edge_type(pt, e.kind)), e) + for e in edges + ] for a, b in phantom: builder.make_edge(tid[a], tid[b], _edge_type(pt, G.E_UNCOND)) lay = builder.build() @@ -266,8 +277,9 @@ def _layout_graph(pt, g: G._Graph, root: int, edges: list[G.Edge], if len(pts) < 2: continue _snap_ports(g, e, pts) - routes.append(G.Route(edge=e, pts=_clean(pts), head=True, tail=True, - flipped=False)) + routes.append( + G.Route(edge=e, pts=_clean(pts), head=True, tail=True, flipped=False) + ) def _box_index(g: G._Graph) -> tuple[dict[int, list[G.Node]], dict[int, list[G.Node]]]: @@ -293,15 +305,14 @@ def _hits(by_col, by_row, p: tuple[int, int], q: tuple[int, int]) -> list[G.Node (r0, c0), (r1, c1) = p, q if c0 == c1: lo, hi = (r0, r1) if r0 <= r1 else (r1, r0) - return [n for n in by_col.get(c0, ()) - if n.y < hi and lo < n.bottom] + return [n for n in by_col.get(c0, ()) if n.y < hi and lo < n.bottom] lo, hi = (c0, c1) if c0 <= c1 else (c1, c0) - return [n for n in by_row.get(r0, ()) - if n.x < hi and lo < n.right] + return [n for n in by_row.get(r0, ()) if n.x < hi and lo < n.right] -def _free_line(blocked: list[tuple[int, int]], want: int, - allow: tuple[int, int] | None = None) -> int | None: +def _free_line( + blocked: list[tuple[int, int]], want: int, allow: tuple[int, int] | None = None +) -> int | None: """The coordinate nearest ``want`` that is in none of ``blocked``. ``blocked`` is a list of inclusive intervals. Jumping to the near side of @@ -316,7 +327,7 @@ def _free_line(blocked: list[tuple[int, int]], want: int, lo, hi = allow if lo > hi: return None - blocked = list(blocked) + [(hi + 1, hi + 1 + 10 ** 6)] + blocked = list(blocked) + [(hi + 1, hi + 1 + 10**6)] if lo > 0: blocked.append((0, lo - 1)) if not blocked: @@ -381,10 +392,13 @@ def _repair_boxes(g: G._Graph, routes: list[G.Route]) -> int: p, q = rt.pts[i], rt.pts[i + 1] if not _hits(by_col, by_row, p, q): continue - if p[1] == q[1]: # vertical: shift column + if p[1] == q[1]: # vertical: shift column lo, hi = sorted((p[0], q[0])) - blocked = [(n.x + 1, n.right - 1) for n in real - if n.y < hi and lo < n.bottom] + blocked = [ + (n.x + 1, n.right - 1) + for n in real + if n.y < hi and lo < n.bottom + ] # The first and last segments carry the port and the # arrowhead, so they may only move ALONG their own box's # border -- but move they must: triskel is happy to park a @@ -398,16 +412,21 @@ def _repair_boxes(g: G._Graph, routes: list[G.Route]) -> int: ends.append(g.nodes[rt.edge.src]) if i == last: ends.append(g.nodes[rt.edge.dst]) - allow = (max(n.x + 1 for n in ends), - min(n.right - 1 for n in ends)) + allow = ( + max(n.x + 1 for n in ends), + min(n.right - 1 for n in ends), + ) col = _free_line(blocked, p[1], allow) if col is None: continue rt.pts[i], rt.pts[i + 1] = (p[0], col), (q[0], col) - elif i not in (0, last): # horizontal: shift row + elif i not in (0, last): # horizontal: shift row lo, hi = sorted((p[1], q[1])) - blocked = [(n.y + 1, n.bottom - 1) for n in real - if n.x < hi and lo < n.right] + blocked = [ + (n.y + 1, n.bottom - 1) + for n in real + if n.x < hi and lo < n.right + ] row = _free_line(blocked, p[0]) if row is None: continue @@ -445,7 +464,8 @@ def _verify(g: G._Graph, routes: list[G.Route]) -> None: if b.x <= a.right: raise RuntimeError( f"blocks {a.id} and {b.id} overlap on row {r} " - f"(x[{a.x},{a.right}] vs x[{b.x},{b.right}])") + f"(x[{a.x},{a.right}] vs x[{b.x},{b.right}])" + ) by_col, by_row = _box_index(g) for rt in routes: @@ -454,7 +474,8 @@ def _verify(g: G._Graph, routes: list[G.Route]) -> None: if hit: raise RuntimeError( f"edge {rt.edge.src}->{rt.edge.dst} crosses block " - f"{hit[0].id} at {p}-{q} and could not be detoured") + f"{hit[0].id} at {p}-{q} and could not be detoured" + ) def _snap_ports(g: G._Graph, e: G.Edge, pts: list[tuple[int, int]]) -> None: @@ -484,7 +505,7 @@ def _snap_ports(g: G._Graph, e: G.Edge, pts: list[tuple[int, int]]) -> None: old_r, old_c = pts[0] col = clamp(src, old_c) pts[0] = (src.bottom if pts[1][0] >= old_r else src.y, col) - if pts[1][1] == old_c: # the first segment was vertical: keep it + if pts[1][1] == old_c: # the first segment was vertical: keep it pts[1] = (pts[1][0], col) # head: dst's top border if the edge arrives downward, its bottom if not diff --git a/idatui/highlight.py b/idatui/highlight.py index 0bd76af..3f22d30 100644 --- a/idatui/highlight.py +++ b/idatui/highlight.py @@ -15,10 +15,10 @@ The same tokenizer feeds two consumers, so one palette covers both: from __future__ import annotations -from rich.segment import Segment -from rich.style import Style from pygments.lexers import CLexer from pygments.token import Token +from rich.segment import Segment +from rich.style import Style from textual.widgets import TextArea from textual.widgets.text_area import TextAreaTheme @@ -36,18 +36,26 @@ from textual.widgets.text_area import TextAreaTheme # mnemonic column: they're the skeleton you scan for, and a hue there would # claim a meaning the rest of the palette already assigns. _PALETTE: list[tuple[str, object, Style]] = [ - ("comment", Token.Comment, Style(color="#7c8b9e", italic=True)), # 5.2:1 commentary - ("type", Token.Keyword.Type, Style(color="#93aee0")), # 8.1:1 type info - ("keyword", Token.Keyword, Style(color="#e8ecf2", bold=True)), # 15.3:1 control flow - ("builtin", Token.Name.Builtin, Style(color="#93aee0")), # 8.1:1 type info - ("string", Token.Literal.String, Style(color="#9ece6a")), # 9.9:1 strings - ("number", Token.Literal.Number, Style(color="#d8a657")), # 8.2:1 data/number - ("operator", Token.Operator, Style(color="#c3cad3")), # 11.0:1 body - ("punctuation", Token.Punctuation, Style(color="#626c7a")), # 3.4:1 structure - ("name", Token.Name, Style(color="#7aa2f7")), # 7.2:1 symbol names + ( + "comment", + Token.Comment, + Style(color="#7c8b9e", italic=True), + ), # 5.2:1 commentary + ("type", Token.Keyword.Type, Style(color="#93aee0")), # 8.1:1 type info + ( + "keyword", + Token.Keyword, + Style(color="#e8ecf2", bold=True), + ), # 15.3:1 control flow + ("builtin", Token.Name.Builtin, Style(color="#93aee0")), # 8.1:1 type info + ("string", Token.Literal.String, Style(color="#9ece6a")), # 9.9:1 strings + ("number", Token.Literal.Number, Style(color="#d8a657")), # 8.2:1 data/number + ("operator", Token.Operator, Style(color="#c3cad3")), # 11.0:1 body + ("punctuation", Token.Punctuation, Style(color="#626c7a")), # 3.4:1 structure + ("name", Token.Name, Style(color="#7aa2f7")), # 7.2:1 symbol names ] _STYLES: list[tuple[object, Style]] = [(t, s) for _, t, s in _PALETTE] -_DEFAULT = Style(color="#c3cad3") # 11.0:1 body +_DEFAULT = Style(color="#c3cad3") # 11.0:1 body _DEFAULT_NAME = "text" #: highlight name -> style, for TextArea themes (see ``CTextArea``). @@ -97,7 +105,7 @@ def highlight_c(code: str) -> list[list[Segment]]: if not value: continue style = _style_for(token) - if "\n" not in value: # the common case: a token inside one line + if "\n" not in value: # the common case: a token inside one line lines[-1].append(Segment(value, style)) continue parts = value.split("\n") @@ -136,7 +144,7 @@ def highlight_c_spans(code: str) -> dict[int, list[tuple[int, int, str]]]: if not part: continue width = len(part) if part.isascii() else len(part.encode("utf-8")) - if part.strip(): # whitespace carries no visible style + if part.strip(): # whitespace carries no visible style spans.setdefault(row, []).append((col, col + width, name)) col += width return spans diff --git a/idatui/index.py b/idatui/index.py index 751ad1f..3925ee8 100644 --- a/idatui/index.py +++ b/idatui/index.py @@ -21,6 +21,7 @@ three characters — it silently returns nothing rather than erroring — so sho queries fall back to LIKE. Without that, typing "e" then "er" would show "no matches" until the third keystroke. """ + from __future__ import annotations import os @@ -81,7 +82,8 @@ class ProjectIndex: def stamp(self, label: str) -> tuple[int, int, int] | None: """(size, mtime, entry count) recorded when ``label`` was last indexed.""" row = self._db.execute( - "SELECT size, mtime, n FROM stamps WHERE binary = ?", (label,)).fetchone() + "SELECT size, mtime, n FROM stamps WHERE binary = ?", (label,) + ).fetchone() return tuple(row) if row else None # type: ignore[return-value] def is_stale(self, label: str, source: str) -> bool: @@ -99,11 +101,11 @@ class ProjectIndex: def reindex(self, label: str, entries, source: str | None = None) -> int: """Replace ``label``'s entries with ``entries`` — (kind, addr, text) triples. Per-binary, so re-indexing one never touches the others.""" - rows = [(text, label, kind, int(addr)) - for kind, addr, text in entries if text] + rows = [(text, label, kind, int(addr)) for kind, addr, text in entries if text] self._db.execute("DELETE FROM entries WHERE binary = ?", (label,)) self._db.executemany( - "INSERT INTO entries(text, binary, kind, addr) VALUES(?,?,?,?)", rows) + "INSERT INTO entries(text, binary, kind, addr) VALUES(?,?,?,?)", rows + ) size = mtime = 0 if source: try: @@ -114,7 +116,8 @@ class ProjectIndex: self._db.execute( "INSERT INTO stamps(binary, size, mtime, n) VALUES(?,?,?,?) " "ON CONFLICT(binary) DO UPDATE SET size=?, mtime=?, n=?", - (label, size, mtime, len(rows), size, mtime, len(rows))) + (label, size, mtime, len(rows), size, mtime, len(rows)), + ) self._db.commit() return len(rows) @@ -125,8 +128,9 @@ class ProjectIndex: self._db.commit() # -- query -------------------------------------------------------------- # - def search(self, query: str, kind: str | None = None, - limit: int = 500) -> list[Hit]: + def search( + self, query: str, kind: str | None = None, limit: int = 500 + ) -> list[Hit]: """Substring search across every indexed binary, newest-agnostic. Uses the trigram index at >= 3 characters and falls back to a LIKE scan @@ -189,12 +193,12 @@ class ProjectIndex: # -- introspection ------------------------------------------------------ # def counts(self) -> dict[str, int]: """Indexed entry count per binary.""" - return {b: n for b, n in - self._db.execute("SELECT binary, n FROM stamps").fetchall()} + return { + b: n for b, n in self._db.execute("SELECT binary, n FROM stamps").fetchall() + } def total(self) -> int: - return int(self._db.execute( - "SELECT count(*) FROM entries").fetchone()[0]) + return int(self._db.execute("SELECT count(*) FROM entries").fetchone()[0]) def close(self) -> None: try: diff --git a/idatui/journal.py b/idatui/journal.py index a5a07de..352a324 100644 --- a/idatui/journal.py +++ b/idatui/journal.py @@ -46,8 +46,13 @@ class Journal: self._lock = threading.Lock() # -- recording ---------------------------------------------------------- # - def record(self, kind: str, ea: int | None = None, detail: str = "", - extra: dict | None = None) -> None: + def record( + self, + kind: str, + ea: int | None = None, + detail: str = "", + extra: dict | None = None, + ) -> None: """Note one edit: ``kind`` is 'rename' / 'comment' / 'retype' / …""" entry = {"k": str(kind), "t": int(time.time())} if ea is not None: @@ -59,14 +64,17 @@ class Journal: with self._lock: self.entries.append(entry) if len(self.entries) > MAX_ENTRIES: - del self.entries[:len(self.entries) - MAX_ENTRIES] + del self.entries[: len(self.entries) - MAX_ENTRIES] self._dirty = True def addresses(self, kinds: tuple[str, ...] | None = None) -> set[int]: """Every address touched (optionally only by certain kinds of edit).""" with self._lock: - return {e["ea"] for e in self.entries - if "ea" in e and (kinds is None or e.get("k") in kinds)} + return { + e["ea"] + for e in self.entries + if "ea" in e and (kinds is None or e.get("k") in kinds) + } def __len__(self) -> int: return len(self.entries) diff --git a/idatui/kittygfx.py b/idatui/kittygfx.py index 62fd6a6..9f83656 100644 --- a/idatui/kittygfx.py +++ b/idatui/kittygfx.py @@ -32,6 +32,7 @@ screen cannot be placed from the alternate one -- placement reports no error, it simply draws nothing. That combination is why the splash calls ``supported()`` from the launcher and ``upload()`` from its own ``on_mount``. """ + from __future__ import annotations import base64 @@ -55,7 +56,7 @@ LOGO_ID = 0x1DA7 LOGO_PLACEMENT = 1 _supported: bool | None = None -_uploaded: dict[int, tuple[int, int]] = {} # image id -> (pixel w, pixel h) +_uploaded: dict[int, tuple[int, int]] = {} # image id -> (pixel w, pixel h) #: Terminal cell size in pixels, asked for in the same round trip as the #: graphics query. Cells are nothing like a fixed 1:2 -- this box reports 9x22, #: i.e. 1:2.44 -- and getting it wrong stretches the image. @@ -114,10 +115,10 @@ def _query_tty(timeout: float = 2.0) -> bool: if not chunk: break buf += chunk - if re.search(rb"\033\[\?[0-9;]*c", buf): # DA1: the answers are in + if re.search(rb"\033\[\?[0-9;]*c", buf): # DA1: the answers are in break global _cell - m = re.search(rb"\033\[6;(\d+);(\d+)t", buf) # CSI 6 ; height ; width t + m = re.search(rb"\033\[6;(\d+);(\d+)t", buf) # CSI 6 ; height ; width t if m: ch, cw = int(m.group(1)), int(m.group(2)) if 0 < cw < 100 and 0 < ch < 200: @@ -149,7 +150,7 @@ def supported() -> bool: elif env in ("0", "no", "false", "off"): _supported = False elif not (sys.__stdout__ and sys.__stdout__.isatty()): - _supported = False # pilot tests, pipes, redirected output + _supported = False # pilot tests, pipes, redirected output log("supported: stdout is not a tty") else: try: @@ -208,14 +209,13 @@ def upload(path: str, image_id: int = LOGO_ID) -> bool: payload = base64.standard_b64encode(f.read()) except OSError: return False - parts = [payload[i:i + 4096] for i in range(0, len(payload), 4096)] + parts = [payload[i : i + 4096] for i in range(0, len(payload), 4096)] if not parts: return False buf = [] for i, part in enumerate(parts): more = 1 if i < len(parts) - 1 else 0 - ctrl = (f"a=t,f=100,t=d,i={image_id},q=2,m={more}" if i == 0 - else f"m={more}") + ctrl = f"a=t,f=100,t=d,i={image_id},q=2,m={more}" if i == 0 else f"m={more}" buf.append("\033_G" + ctrl + ";" + part.decode("ascii") + "\033\\") if not _write("".join(buf)): log("upload: write failed") @@ -229,8 +229,14 @@ def is_uploaded(image_id: int = LOGO_ID) -> bool: return image_id in _uploaded -def place(row: int, col: int, cols: int, rows: int, - image_id: int = LOGO_ID, placement_id: int = LOGO_PLACEMENT) -> bool: +def place( + row: int, + col: int, + cols: int, + rows: int, + image_id: int = LOGO_ID, + placement_id: int = LOGO_PLACEMENT, +) -> bool: """Draw the uploaded image at (``row``, ``col``), 0-based, sized in cells. Saves and restores the cursor, and asks the terminal not to move it @@ -250,7 +256,8 @@ def place(row: int, col: int, cols: int, rows: int, f"\033[s\033[{row + 1};{col + 1}H" f"\033_Ga=p,i={image_id},p={placement_id}," f"s={w},v={h},c={cols},r={rows},C=1,q=2\033\\" - f"\033[u") + f"\033[u" + ) def clear(image_id: int = LOGO_ID) -> None: @@ -274,8 +281,12 @@ def cell_size() -> tuple[int, int]: return _cell or (10, 20) -def fit(px: tuple[int, int], max_cols: int, max_rows: int, - cell: tuple[int, int] | None = None) -> tuple[int, int]: +def fit( + px: tuple[int, int], + max_cols: int, + max_rows: int, + cell: tuple[int, int] | None = None, +) -> tuple[int, int]: """Cell size that fits ``max_cols`` x ``max_rows`` keeping the aspect ratio. Cells are far from square -- this box reports 9x22 px -- so a naive diff --git a/idatui/launch.py b/idatui/launch.py index a0201a7..6e3b432 100644 --- a/idatui/launch.py +++ b/idatui/launch.py @@ -9,12 +9,14 @@ Usage:: ida-tui /path/to/binary ida-tui # attach when exactly one database is registered """ + from __future__ import annotations import argparse import os import sys + def _load_args(load: dict) -> str: """``load`` as IDA switches, for the single-binary path (no project ref). @@ -25,8 +27,12 @@ def _load_args(load: dict) -> str: """ from .formats import load_args from .project import _as_addr - return load_args(load.get("processor", ""), _as_addr(load.get("base", 0)), - str(load.get("ida_args", "") or "")) + + return load_args( + load.get("processor", ""), + _as_addr(load.get("base", 0)), + str(load.get("ida_args", "") or ""), + ) def _log(msg: str) -> None: @@ -62,32 +68,63 @@ def _registered_databases() -> tuple[list[dict], list[dict]]: def main(argv: list[str] | None = None) -> int: p = argparse.ArgumentParser( prog="ida-tui", - description="Open a registered GUI or managed idalib database in the IDA TUI.") - p.add_argument("binary", nargs="*", - help="binary to open and analyze (several with --project " - "creates/extends that project)") - p.add_argument("--project", metavar="FILE", - help="open a multi-binary project (created from the given " - "binaries if FILE doesn't exist)") - p.add_argument("--ttl", type=int, default=1800, - help="deprecated compatibility option (IDA Nexus uses leases)") - p.add_argument("--no-keepalive", action="store_true", - help="deprecated compatibility option (the lease is the heartbeat)") - p.add_argument("--rpc", metavar="PATH", - help="listen for RPC on this unix socket (puppeteer the TUI)") - p.add_argument("--trace", metavar="FILE", - help="Tenet execution trace to explore alongside the binary") + description="Open a registered GUI or managed idalib database in the IDA TUI.", + ) + p.add_argument( + "binary", + nargs="*", + help="binary to open and analyze (several with --project " + "creates/extends that project)", + ) + p.add_argument( + "--project", + metavar="FILE", + help="open a multi-binary project (created from the given " + "binaries if FILE doesn't exist)", + ) + p.add_argument( + "--ttl", + type=int, + default=1800, + help="deprecated compatibility option (IDA Nexus uses leases)", + ) + p.add_argument( + "--no-keepalive", + action="store_true", + help="deprecated compatibility option (the lease is the heartbeat)", + ) + p.add_argument( + "--rpc", + metavar="PATH", + help="listen for RPC on this unix socket (puppeteer the TUI)", + ) + p.add_argument( + "--trace", + metavar="FILE", + help="Tenet execution trace to explore alongside the binary", + ) g = p.add_argument_group( "loading a headerless blob", "An ELF/PE/Mach-O says what it is. A raw firmware dump doesn't, and IDA " "falls back to x86 at address 0 — which analyses to nothing. These say " - "how to read it, and are recorded per binary in a project.") - g.add_argument("--processor", metavar="NAME", - help="IDA processor: arm, armb (big-endian), mipsb, metapc, …") - g.add_argument("--base", metavar="ADDR", - help="load address, e.g. 0x8000000 (any base; NOT paragraphs)") - g.add_argument("--ida-args", metavar="STR", dest="ida_args", - help="legacy switches; only IDA Nexus-representable -p/-b/-T are accepted") + "how to read it, and are recorded per binary in a project.", + ) + g.add_argument( + "--processor", + metavar="NAME", + help="IDA processor: arm, armb (big-endian), mipsb, metapc, …", + ) + g.add_argument( + "--base", + metavar="ADDR", + help="load address, e.g. 0x8000000 (any base; NOT paragraphs)", + ) + g.add_argument( + "--ida-args", + metavar="STR", + dest="ida_args", + help="legacy switches; only IDA Nexus-representable -p/-b/-T are accepted", + ) args = p.parse_args(argv) load: dict = {} @@ -112,6 +149,7 @@ def main(argv: list[str] | None = None) -> int: binary = None if args.project: from .project import Project, ProjectError + ppath = os.path.abspath(os.path.expanduser(args.project)) try: if os.path.isfile(ppath): @@ -126,8 +164,10 @@ def main(argv: list[str] | None = None) -> int: project.save() _log(f"added {added} binary(ies) to {ppath}") if dupes: - _log(f"{dupes} already in the project (matched by path) " - f"— left alone") + _log( + f"{dupes} already in the project (matched by path) " + f"— left alone" + ) elif args.binary: project = Project.create(ppath, args.binary, load=load or None) _log(f"created project {ppath} with {len(project.refs)} binaries") @@ -154,7 +194,8 @@ def main(argv: list[str] | None = None) -> int: key = os.path.normcase(os.path.realpath(binary)) registered = any( key == os.path.normcase(os.path.realpath(str(item.get(field) or ""))) - for item in ready for field in ("exe_path", "idb_path") + for item in ready + for field in ("exe_path", "idb_path") if item.get(field) ) if not os.path.isfile(binary) and not registered: @@ -171,8 +212,10 @@ def main(argv: list[str] | None = None) -> int: else: _log("several IDA Nexus databases are registered; pass one of these paths:") for item in ready: - _log(f" {item.get('exe_path') or item.get('idb_path')} " - f"[{item.get('backend')}, {item.get('record_id')}]") + _log( + f" {item.get('exe_path') or item.get('idb_path')} " + f"[{item.get('backend')}, {item.get('record_id')}]" + ) return 2 # Hand off to the TUI (imported late so --help works without Textual). Code @@ -190,16 +233,23 @@ def main(argv: list[str] | None = None) -> int: # round trip, and only when attached to a tty. try: from . import kittygfx + kittygfx.supported() except Exception: # noqa: BLE001 -- graphics are decoration, never fatal pass rpc_path = os.path.abspath(os.path.expanduser(args.rpc)) if args.rpc else None - IdaTui(open_path=binary, keepalive=not args.no_keepalive, - rpc_path=rpc_path, ttl=args.ttl, project=project, - load_args=_load_args(load), - trace_path=(os.path.abspath(os.path.expanduser(args.trace)) - if args.trace else "")).run() + IdaTui( + open_path=binary, + keepalive=not args.no_keepalive, + rpc_path=rpc_path, + ttl=args.ttl, + project=project, + load_args=_load_args(load), + trace_path=( + os.path.abspath(os.path.expanduser(args.trace)) if args.trace else "" + ), + ).run() return 0 diff --git a/idatui/pane.py b/idatui/pane.py index 1eee847..67b462b 100644 --- a/idatui/pane.py +++ b/idatui/pane.py @@ -27,6 +27,7 @@ Requires: running inside tmux or zellij. Each pane leases a registered GUI or shared managed idalib database through IDA Nexus. Uses ~/ida-venv/bin/python for the TUI (needs textual) unless --python / IDATUI_PYTHON says otherwise. """ + from __future__ import annotations import argparse @@ -42,7 +43,8 @@ from .rpcclient import RpcClient, RpcError REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) DEFAULT_PY = os.environ.get( - "IDATUI_PYTHON", os.path.expanduser("~/ida-venv/bin/python")) + "IDATUI_PYTHON", os.path.expanduser("~/ida-venv/bin/python") +) def _sockdir() -> str: @@ -106,26 +108,31 @@ def _mux_of_pane(pane: str) -> str: def _zellij_argv() -> list[str]: """Base zellij argv, pinned to our session when we know it (so it still works from a process that isn't itself attached).""" - session = (os.environ.get("IDATUI_ZELLIJ_SESSION") - or os.environ.get("ZELLIJ_SESSION_NAME")) + session = os.environ.get("IDATUI_ZELLIJ_SESSION") or os.environ.get( + "ZELLIJ_SESSION_NAME" + ) return ["zellij", "-s", session] if session else ["zellij"] def _tmux(*args: str) -> str: - return subprocess.run(["tmux", *args], capture_output=True, text=True, - check=True).stdout.strip() + return subprocess.run( + ["tmux", *args], capture_output=True, text=True, check=True + ).stdout.strip() def _zellij(*args: str) -> str: - return subprocess.run([*_zellij_argv(), *args], capture_output=True, - text=True, check=True).stdout.strip() + return subprocess.run( + [*_zellij_argv(), *args], capture_output=True, text=True, check=True + ).stdout.strip() def _zellij_panes() -> list[dict[str, Any]]: try: - out = subprocess.run([*_zellij_argv(), "action", "list-panes", - "--state", "--json"], - capture_output=True, text=True) + out = subprocess.run( + [*_zellij_argv(), "action", "list-panes", "--state", "--json"], + capture_output=True, + text=True, + ) rows = json.loads(out.stdout or "[]") except (OSError, ValueError): return [] @@ -148,8 +155,9 @@ def _pane_alive(pane: str, mux: str | None = None) -> bool: if str(row.get("id")) == want and bool(row.get("is_plugin")) is False: return not row.get("exited", False) return False - out = subprocess.run(["tmux", "list-panes", "-a", "-F", "#{pane_id}"], - capture_output=True, text=True) + out = subprocess.run( + ["tmux", "list-panes", "-a", "-F", "#{pane_id}"], capture_output=True, text=True + ) return pane in out.stdout.split() @@ -159,8 +167,9 @@ def _pane_exists(pane: str, mux: str | None = None) -> bool: return False if (mux or _mux_of_pane(pane)) == "zellij": want = pane.split("_", 1)[-1] - return any(str(r.get("id")) == want and not r.get("is_plugin") - for r in _zellij_panes()) + return any( + str(r.get("id")) == want and not r.get("is_plugin") for r in _zellij_panes() + ) return _pane_alive(pane, "tmux") @@ -169,24 +178,36 @@ def _pane_kill(pane: str, mux: str | None = None) -> None: if not pane: return if (mux or _mux_of_pane(pane)) == "zellij": - subprocess.run([*_zellij_argv(), "action", "close-pane", - "--pane-id", pane], capture_output=True) + subprocess.run( + [*_zellij_argv(), "action", "close-pane", "--pane-id", pane], + capture_output=True, + ) else: subprocess.run(["tmux", "kill-pane", "-t", pane], capture_output=True) -def _pane_split(inner: list[str], *, mux: str, vertical: bool, - size: str | None, detached: bool) -> str: +def _pane_split( + inner: list[str], *, mux: str, vertical: bool, size: str | None, detached: bool +) -> str: """Open a pane running ``inner`` (argv) in REPO, and return its pane id.""" if mux == "zellij": # zellij runs the argv directly (no shell) and takes the cwd as a flag, # so there's nothing to quote. --name labels the pane in the UI. - argv = [*_zellij_argv(), "action", "new-pane", - "--direction", "down" if vertical else "right", - "--cwd", REPO, "--name", "idatui"] + argv = [ + *_zellij_argv(), + "action", + "new-pane", + "--direction", + "down" if vertical else "right", + "--cwd", + REPO, + "--name", + "idatui", + ] argv += ["--", *inner] - pane = subprocess.run(argv, capture_output=True, text=True, - check=True).stdout.strip() + pane = subprocess.run( + argv, capture_output=True, text=True, check=True + ).stdout.strip() # zellij prints the new pane id ('terminal_3'); without it we could not # target this pane later, so treat a missing id as a hard failure. if not pane.startswith(("terminal_", "plugin_")): @@ -196,13 +217,14 @@ def _pane_split(inner: list[str], *, mux: str, vertical: bool, # to the pane we were called from. origin = os.environ.get("ZELLIJ_PANE_ID") if origin: - subprocess.run([*_zellij_argv(), "action", "focus-pane-id", - f"terminal_{origin}"], capture_output=True) + subprocess.run( + [*_zellij_argv(), "action", "focus-pane-id", f"terminal_{origin}"], + capture_output=True, + ) return pane cmd = f"cd {REPO!r} && exec " + " ".join(_q(a) for a in inner) - split = ["split-window", "-v" if vertical else "-h", - "-P", "-F", "#{pane_id}"] + split = ["split-window", "-v" if vertical else "-h", "-P", "-F", "#{pane_id}"] if size: split += ["-l", str(size)] if detached: @@ -224,9 +246,15 @@ def _pane_capture(pane: str, mux: str | None = None) -> str: # tmux key names -> zellij key names (zellij rejects e.g. "Escape", wants "Esc"). _ZELLIJ_KEYS = { - "escape": "Esc", "bspace": "Backspace", "space": "Space", - "pageup": "PageUp", "pagedown": "PageDown", "ppage": "PageUp", - "npage": "PageDown", "ic": "Insert", "dc": "Delete", + "escape": "Esc", + "bspace": "Backspace", + "space": "Space", + "pageup": "PageUp", + "pagedown": "PageDown", + "ppage": "PageUp", + "npage": "PageDown", + "ic": "Insert", + "dc": "Delete", } @@ -244,8 +272,17 @@ def _pane_keys(pane: str, keys: list[str], mux: str | None = None) -> None: """Inject real terminal keystrokes into the pane (the input-layer cross-check).""" mux = mux or _mux_of_pane(pane) if mux == "zellij": - subprocess.run([*_zellij_argv(), "action", "send-keys", "--pane-id", pane, - *[_to_zellij_key(k) for k in keys]], check=True) + subprocess.run( + [ + *_zellij_argv(), + "action", + "send-keys", + "--pane-id", + pane, + *[_to_zellij_key(k) for k in keys], + ], + check=True, + ) else: subprocess.run(["tmux", "send-keys", "-t", pane, *keys], check=True) @@ -256,8 +293,9 @@ def _pane_keys(pane: str, keys: list[str], mux: str | None = None) -> None: def _count_live_panes() -> int: - return sum(1 for r in _load_registry() - if _pane_alive(r.get("pane", ""), r.get("mux"))) + return sum( + 1 for r in _load_registry() if _pane_alive(r.get("pane", ""), r.get("mux")) + ) def _reap_orphan_workers(force: bool = False) -> int: @@ -272,20 +310,28 @@ def _reap_orphan_workers(force: bool = False) -> int: def spawn(args) -> int: mux = args.mux or _detect_mux() if mux.startswith("?"): - print(f"error: unknown multiplexer {mux[1:]!r} (want tmux or zellij)", - file=sys.stderr) + print( + f"error: unknown multiplexer {mux[1:]!r} (want tmux or zellij)", + file=sys.stderr, + ) return 2 if not mux: - print("error: not inside tmux or zellij (spawn creates a pane there). " - "Set $IDATUI_MUX=tmux|zellij to force a backend.", file=sys.stderr) + print( + "error: not inside tmux or zellij (spawn creates a pane there). " + "Set $IDATUI_MUX=tmux|zellij to force a backend.", + file=sys.stderr, + ) return 2 if not args.open and not getattr(args, "project", None): print("error: pass --open or --project ", file=sys.stderr) return 2 sock = args.sock or os.path.join(_sockdir(), f"idatui-{secrets.token_hex(3)}.sock") - project = (os.path.abspath(os.path.expanduser(args.project)) - if getattr(args, "project", None) else None) + project = ( + os.path.abspath(os.path.expanduser(args.project)) + if getattr(args, "project", None) + else None + ) target = os.path.abspath(os.path.expanduser(args.open)) if args.open else None if target is not None and not os.path.exists(target): print(f"error: no such binary: {target}", file=sys.stderr) @@ -317,17 +363,30 @@ def spawn(args) -> int: inner += ["--trace", os.path.abspath(os.path.expanduser(args.trace))] if args.size and mux == "zellij": - print("note: --size is tmux-only; zellij tiles the new pane evenly", - file=sys.stderr) + print( + "note: --size is tmux-only; zellij tiles the new pane evenly", + file=sys.stderr, + ) try: - pane = _pane_split(inner, mux=mux, vertical=args.vertical, - size=args.size, detached=args.detached) + pane = _pane_split( + inner, + mux=mux, + vertical=args.vertical, + size=args.size, + detached=args.detached, + ) except (OSError, subprocess.CalledProcessError, RuntimeError) as e: print(f"error: could not create a {mux} pane: {e}", file=sys.stderr) return 2 - row = {"sock": sock, "pane": pane, "mux": mux, "target": project or target, - "kind": "project" if project else "open", "started": time.time()} + row = { + "sock": sock, + "pane": pane, + "mux": mux, + "target": project or target, + "kind": "project" if project else "open", + "started": time.time(), + } reg = [r for r in _load_registry() if r.get("sock") != sock] reg.append(row) _save_registry(reg) @@ -340,11 +399,17 @@ def spawn(args) -> int: def _q(s: str) -> str: import shlex + return shlex.quote(s) -def _wait_ready(sock: str, timeout: float, pane: str, - stuck_after: float = 45.0, mux: str | None = None) -> dict[str, Any]: +def _wait_ready( + sock: str, + timeout: float, + pane: str, + stuck_after: float = 45.0, + mux: str | None = None, +) -> dict[str, Any]: """Poll the socket + ping until the TUI reports ready (or timeout). Emits a one-time hint if IDA Nexus discovery/opening is still not ready after @@ -367,10 +432,16 @@ def _wait_ready(sock: str, timeout: float, pane: str, pass if not warned and (time.time() - start) > stuck_after: warned = True - why = ("RPC socket not created yet" if not os.path.exists(sock) - else "TUI up but analysis not ready") - print(f"still waiting ({int(time.time() - start)}s): {why}. " - f"Check IDA Nexus registrations and worker logs.", file=sys.stderr) + why = ( + "RPC socket not created yet" + if not os.path.exists(sock) + else "TUI up but analysis not ready" + ) + print( + f"still waiting ({int(time.time() - start)}s): {why}. " + f"Check IDA Nexus registrations and worker logs.", + file=sys.stderr, + ) time.sleep(0.4) last = dict(last) last["ready"] = False @@ -383,9 +454,12 @@ def _wait_ready(sock: str, timeout: float, pane: str, # --------------------------------------------------------------------------- # def stop(args) -> int: reg = _load_registry() - rows = [r for r in reg - if (args.sock and r.get("sock") == args.sock) - or (args.pane and r.get("pane") == args.pane)] + rows = [ + r + for r in reg + if (args.sock and r.get("sock") == args.sock) + or (args.pane and r.get("pane") == args.pane) + ] if not rows and args.sock: # allow stopping an untracked socket rows = [{"sock": args.sock, "pane": args.pane}] if not rows: @@ -432,8 +506,10 @@ def stop(args) -> int: # Only ever reached on timeout: say so, because it means a save may have # been cut short rather than "clean teardown". out["force_killed"] = killed - out["warning"] = (f"pane(s) did not exit within {args.timeout}s and were " - "killed; unsaved database changes may be lost") + out["warning"] = ( + f"pane(s) did not exit within {args.timeout}s and were " + "killed; unsaved database changes may be lost" + ) print(json.dumps(out)) return 0 @@ -468,8 +544,16 @@ def list_panes(args) -> int: def reap(args) -> int: """Deprecated no-op; shared IDA Nexus workers are managed by leases.""" - print(json.dumps({"reaped_workers": 0, "live_panes": _count_live_panes(), - "forced": args.force, "deprecated": True})) + print( + json.dumps( + { + "reaped_workers": 0, + "live_panes": _count_live_panes(), + "forced": args.force, + "deprecated": True, + } + ) + ) return 0 @@ -520,56 +604,97 @@ def _resolve_pane(sock: str | None) -> str | None: else: print("error: several live panes, pass --pane or --sock:", file=sys.stderr) for r in live: - print(f" {r.get('pane')} {r.get('sock')} {r.get('target')}", - file=sys.stderr) + print( + f" {r.get('pane')} {r.get('sock')} {r.get('target')}", + file=sys.stderr, + ) return None def main(argv: list[str]) -> int: p = argparse.ArgumentParser( prog="idatui.pane", - description="spawn/manage idatui TUI panes in tmux or zellij") + description="spawn/manage idatui TUI panes in tmux or zellij", + ) sub = p.add_subparsers(dest="cmd", required=True) sp = sub.add_parser("spawn", help="open a TUI pane and wait until ready") - sp.add_argument("--open", metavar="PATH", - help="binary to open (its dir must be writable)") - sp.add_argument("--trace", metavar="FILE", - help="Tenet execution trace to load alongside the binary") - sp.add_argument("--project", metavar="FILE", - help="project file to open instead of a single binary; " - "any --open paths are added to it (created if absent)") - sp.add_argument("--processor", metavar="NAME", - help="IDA processor for a headerless blob: arm, armb, " - "mipsb, metapc, … (passed to idatui.launch)") - sp.add_argument("--base", metavar="ADDR", - help="load address for a headerless blob, e.g. 0x8000000 " - "(16-byte aligned)") - sp.add_argument("--ida-args", metavar="STR", dest="ida_args", - help="extra IDA command-line switches, passed through") - sp.add_argument("--sock", help="RPC socket path (default: auto in $XDG_RUNTIME_DIR)") - sp.add_argument("--python", default=DEFAULT_PY, help=f"python for the TUI ({DEFAULT_PY})") - sp.add_argument("--vertical", action="store_true", help="split vertically (stacked)") - sp.add_argument("--size", help="new pane size (tmux -l value, e.g. 60%% or 120; " - "ignored under zellij)") + sp.add_argument( + "--open", metavar="PATH", help="binary to open (its dir must be writable)" + ) + sp.add_argument( + "--trace", + metavar="FILE", + help="Tenet execution trace to load alongside the binary", + ) + sp.add_argument( + "--project", + metavar="FILE", + help="project file to open instead of a single binary; " + "any --open paths are added to it (created if absent)", + ) + sp.add_argument( + "--processor", + metavar="NAME", + help="IDA processor for a headerless blob: arm, armb, " + "mipsb, metapc, … (passed to idatui.launch)", + ) + sp.add_argument( + "--base", + metavar="ADDR", + help="load address for a headerless blob, e.g. 0x8000000 (16-byte aligned)", + ) + sp.add_argument( + "--ida-args", + metavar="STR", + dest="ida_args", + help="extra IDA command-line switches, passed through", + ) + sp.add_argument( + "--sock", help="RPC socket path (default: auto in $XDG_RUNTIME_DIR)" + ) + sp.add_argument( + "--python", default=DEFAULT_PY, help=f"python for the TUI ({DEFAULT_PY})" + ) + sp.add_argument( + "--vertical", action="store_true", help="split vertically (stacked)" + ) + sp.add_argument( + "--size", + help="new pane size (tmux -l value, e.g. 60%% or 120; ignored under zellij)", + ) sp.add_argument("--detached", action="store_true", help="don't focus the new pane") - sp.add_argument("--mux", choices=MUXES, default="", - help="multiplexer to spawn in (default: autodetect from " - "$ZELLIJ/$TMUX; $IDATUI_MUX overrides)") - sp.add_argument("--timeout", type=float, default=300.0, - help="seconds to wait for readiness (fresh --open analysis is slow)") + sp.add_argument( + "--mux", + choices=MUXES, + default="", + help="multiplexer to spawn in (default: autodetect from " + "$ZELLIJ/$TMUX; $IDATUI_MUX overrides)", + ) + sp.add_argument( + "--timeout", + type=float, + default=300.0, + help="seconds to wait for readiness (fresh --open analysis is slow)", + ) sp.set_defaults(fn=spawn) st = sub.add_parser("stop", help="graceful quit + kill the pane") st.add_argument("--sock") st.add_argument("--pane") - st.add_argument("--timeout", type=float, default=600.0, - help="seconds to wait for the pane to exit (it saves dirty " - "databases on the way out) before force-killing it") + st.add_argument( + "--timeout", + type=float, + default=600.0, + help="seconds to wait for the pane to exit (it saves dirty " + "databases on the way out) before force-killing it", + ) st.set_defaults(fn=stop) ls = sub.add_parser("list", help="list tracked panes") - ls.add_argument("--prune", action="store_true", help="drop dead panes (and their sockets)") + ls.add_argument( + "--prune", action="store_true", help="drop dead panes (and their sockets)" + ) ls.set_defaults(fn=list_panes) rp = sub.add_parser("reap", help="deprecated no-op (IDA Nexus uses shared leases)") @@ -582,8 +707,11 @@ def main(argv: list[str]) -> int: cp.add_argument("--mux", choices=MUXES, default="") cp.set_defaults(fn=capture) - kp = sub.add_parser("keys", help="inject real keystrokes into a pane " - "(tmux-style names, translated per mux)") + kp = sub.add_parser( + "keys", + help="inject real keystrokes into a pane " + "(tmux-style names, translated per mux)", + ) kp.add_argument("keys", nargs="+", help="e.g. Escape, Enter, C-a, g m a i n") kp.add_argument("--pane") kp.add_argument("--sock", help="resolve the pane from this socket") diff --git a/idatui/pool.py b/idatui/pool.py index 2407b44..084d01a 100644 --- a/idatui/pool.py +++ b/idatui/pool.py @@ -9,6 +9,7 @@ The historical memory budget remains useful for managed idalib instances, while GUI process memory is only advisory. The active and pinned databases are never released to satisfy it. """ + from __future__ import annotations from .project import BinaryRef, Project @@ -47,8 +48,11 @@ def _pss_mb(pid: int | None) -> int: return 0 -def _default_spawn(ref: BinaryRef, ttl: int, *, new_database: bool = False): # pragma: no cover - needs IDA +def _default_spawn( + ref: BinaryRef, ttl: int, *, new_database: bool = False +): # pragma: no cover - needs IDA from .nexus_client import NexusClient + return NexusClient( ref.staged, ttl=ttl, @@ -61,22 +65,31 @@ def _default_spawn(ref: BinaryRef, ttl: int, *, new_database: bool = False): # class DatabasePool: """Live IDA Nexus database leases, keyed by project label.""" - def __init__(self, project: Project, *, budget_mb: int | None = None, - ttl: int = 1800, spawn=None, mem_fn=None) -> None: + def __init__( + self, + project: Project, + *, + budget_mb: int | None = None, + ttl: int = 1800, + spawn=None, + mem_fn=None, + ) -> None: self.project = project self._ttl = ttl self._spawn = spawn or _default_spawn self._mem = mem_fn or (lambda c: _pss_mb(getattr(c, "pid", None))) self._clients: dict[str, object] = {} - self._lru: list[str] = [] # least-recently-used first + self._lru: list[str] = [] # least-recently-used first self._pinned: set[str] = set() self._recreate: set[str] = set() # Ctrl+L: next attachment creates a fresh IDB self.active: str | None = None # never evicted if budget_mb is None: ram = _total_ram_mb() - budget_mb = (ram * project.memory_pct // 100) if ram else _FALLBACK_BUDGET_MB + budget_mb = ( + (ram * project.memory_pct // 100) if ram else _FALLBACK_BUDGET_MB + ) self.budget_mb = max(budget_mb, 256) - self.evicted: list[str] = [] # labels evicted, most recent last + self.evicted: list[str] = [] # labels evicted, most recent last # -- residency --------------------------------------------------------- # def resident(self) -> list[str]: @@ -120,8 +133,11 @@ class DatabasePool: self.project.stage(ref) note(f"opening {ref.label}\u2026") fresh = label in self._recreate - client = (_default_spawn(ref, self._ttl, new_database=fresh) - if self._spawn is _default_spawn else self._spawn(ref, self._ttl)) + client = ( + _default_spawn(ref, self._ttl, new_database=fresh) + if self._spawn is _default_spawn + else self._spawn(ref, self._ttl) + ) connect = getattr(client, "connect", None) if connect is not None: connect(progress=progress) if progress is not None else connect() @@ -178,8 +194,7 @@ class DatabasePool: self._touch(label) # -- release ----------------------------------------------------------- # - def evict(self, label: str, save: bool = True, - save_gui: bool = False) -> bool: + def evict(self, label: str, save: bool = True, save_gui: bool = False) -> bool: """Release a resident lease, persisting a managed database first. A budget-driven eviction must not save somebody's GUI implicitly. GUI @@ -253,19 +268,21 @@ class DatabasePool: out = [] for ref in self.project.refs: client = self._clients.get(ref.label) - out.append({ - "label": ref.label, - "source": ref.source, - "resident": client is not None, - "pinned": ref.label in self._pinned, - "active": ref.label == self.active, - "analysed": self.project.has_db(ref), - "memory_mb": self._mem(client) if client is not None else 0, - }) + out.append( + { + "label": ref.label, + "source": ref.source, + "resident": client is not None, + "pinned": ref.label in self._pinned, + "active": ref.label == self.active, + "analysed": self.project.has_db(ref), + "memory_mb": self._mem(client) if client is not None else 0, + } + ) return out def __repr__(self) -> str: # pragma: no cover - debug aid - return (f"") - - + return ( + f"" + ) diff --git a/idatui/project.py b/idatui/project.py index e681dfc..eaf45b4 100644 --- a/idatui/project.py +++ b/idatui/project.py @@ -26,6 +26,7 @@ now-stale database is dropped (the DB describes the old bytes). The model has no IDA imports. Staging consults ida_nexus's registry before replacing files so it never mutates a database owned by a GUI/shared worker. """ + from __future__ import annotations import json @@ -50,14 +51,14 @@ class ProjectError(Exception): class BinaryRef: """One binary in a project: where it came from, and where IDA works on it.""" - label: str # unique within the project; names the staged file - source: str # absolute path to the original binary - staged: str # absolute path IDA actually opens (inside the sidecar) + label: str # unique within the project; names the staged file + source: str # absolute path to the original binary + staged: str # absolute path IDA actually opens (inside the sidecar) #: How to LOAD it. Only meaningful for a headerless blob: an ELF/PE says what #: it is, a raw firmware image doesn't, and IDA defaults to metapc at 0. - processor: str = "" # IDA processor name: arm, armb, mipsb, metapc, … - base: int = 0 # load address (natural, e.g. 0x8000000) - ida_args: str = "" # legacy -p/-b/-T switches accepted by IDA Nexus adapter + processor: str = "" # IDA processor name: arm, armb, mipsb, metapc, … + base: int = 0 # load address (natural, e.g. 0x8000000) + ida_args: str = "" # legacy -p/-b/-T switches accepted by IDA Nexus adapter @property def db(self) -> str: @@ -73,6 +74,7 @@ class BinaryRef: conversion lives in ``formats.load_args``. """ from .formats import load_args + return load_args(self.processor, self.base, self.ida_args) @@ -112,12 +114,17 @@ def _unlink(path: str) -> bool: class Project: """A set of binaries analysed together, with all IDA artifacts corralled.""" - def __init__(self, path: str, name: str, entries: list[dict], - memory_pct: int = DEFAULT_MEMORY_PCT) -> None: + def __init__( + self, + path: str, + name: str, + entries: list[dict], + memory_pct: int = DEFAULT_MEMORY_PCT, + ) -> None: self.path = os.path.abspath(os.path.expanduser(path)) self.name = name self.memory_pct = memory_pct - self._entries = entries # raw, as written to the file + self._entries = entries # raw, as written to the file self._refs = self._build_refs() # -- construction ------------------------------------------------------ # @@ -145,9 +152,13 @@ class Project: # Keep every recognised key: a whitelist of path/label silently # dropped the load options on the first save, so a blob's processor # and base vanished the moment the project was reopened. - norm.append({k: e[k] for k in - ("path", "label", "processor", "base", "ida_args") - if e.get(k) not in (None, "")}) + norm.append( + { + k: e[k] + for k in ("path", "label", "processor", "base", "ida_args") + if e.get(k) not in (None, "") + } + ) name = raw.get("name") or os.path.splitext(os.path.basename(path))[0] try: pct = int(raw.get("memory_pct", DEFAULT_MEMORY_PCT)) @@ -156,8 +167,14 @@ class Project: return cls(path, str(name), norm, max(1, min(pct, 90))) @classmethod - def create(cls, path: str, binaries: list[str], name: str | None = None, - memory_pct: int = DEFAULT_MEMORY_PCT, load: dict | None = None) -> "Project": + def create( + cls, + path: str, + binaries: list[str], + name: str | None = None, + memory_pct: int = DEFAULT_MEMORY_PCT, + load: dict | None = None, + ) -> "Project": """Write a new project file listing ``binaries`` (an ad-hoc project). ``load`` carries per-binary load options (processor/base/ida_args) that @@ -177,14 +194,21 @@ class Project: e.update({k: v for k, v in (load or {}).items() if v}) entries.append(e) path = os.path.abspath(os.path.expanduser(path)) - proj = cls(path, name or os.path.splitext(os.path.basename(path))[0], - entries, memory_pct) + proj = cls( + path, + name or os.path.splitext(os.path.basename(path))[0], + entries, + memory_pct, + ) proj.save() return proj def save(self) -> None: - data = {"name": self.name, "memory_pct": self.memory_pct, - "binaries": self._entries} + data = { + "name": self.name, + "memory_pct": self.memory_pct, + "binaries": self._entries, + } tmp = self.path + ".tmp" os.makedirs(os.path.dirname(self.path) or ".", exist_ok=True) with open(tmp, "w") as f: @@ -228,20 +252,25 @@ class Project: n += 1 label = f"{label}_{n}" used.add(label) - refs.append(BinaryRef( - label=label, source=src, - staged=os.path.join(self.bin_dir, label), - processor=str(e.get("processor") or ""), - base=_as_addr(e.get("base")), - ida_args=str(e.get("ida_args") or ""))) + refs.append( + BinaryRef( + label=label, + source=src, + staged=os.path.join(self.bin_dir, label), + processor=str(e.get("processor") or ""), + base=_as_addr(e.get("base")), + ida_args=str(e.get("ida_args") or ""), + ) + ) return tuple(refs) @property def refs(self) -> tuple[BinaryRef, ...]: return self._refs - def set_load(self, label: str, processor: str = "", base: int = 0, - ida_args: str = "") -> BinaryRef | None: + def set_load( + self, label: str, processor: str = "", base: int = 0, ida_args: str = "" + ) -> BinaryRef | None: """Record how ``label`` should be loaded, and persist it. Answered once: the dialog that asks writes the answer here, so reopening @@ -275,11 +304,11 @@ class Project: ``./a.elf``, ``/abs/a.elf`` and a symlink to it are all the same file. """ key = os.path.realpath(os.path.abspath(os.path.expanduser(binary))) - return next((r for r in self._refs - if os.path.realpath(r.source) == key), None) + return next((r for r in self._refs if os.path.realpath(r.source) == key), None) - def add(self, binary: str, label: str | None = None, - load: dict | None = None) -> BinaryRef: + def add( + self, binary: str, label: str | None = None, load: dict | None = None + ) -> BinaryRef: """Add a binary, or return the existing entry if it's already here.""" existing = self.by_source(binary) if existing is not None: @@ -320,6 +349,7 @@ class Project: return ref.staged try: from .nexus_client import database_owner + owner = database_owner(ref.db, ref.staged) except Exception as exc: raise ProjectError( @@ -337,7 +367,7 @@ class Project: # leaving the staged bytes immune to an in-place rewrite of the source. shutil.copy2(ref.source, tmp) os.replace(tmp, ref.staged) - for suf in DB_SUFFIXES: # the old DB describes the old bytes + for suf in DB_SUFFIXES: # the old DB describes the old bytes _unlink(ref.staged + suf) return ref.staged diff --git a/idatui/prompt.py b/idatui/prompt.py index 74956f6..6a46c75 100644 --- a/idatui/prompt.py +++ b/idatui/prompt.py @@ -16,13 +16,14 @@ Note the `can_focus` toggling: a hidden `Input` that stays focusable still takes part in Tab focus-nav, so tabbing around a closed prompt used to land the cursor in an invisible widget and swallow every subsequent keystroke. """ + from __future__ import annotations from typing import TYPE_CHECKING from textual.widgets import Input, Static -if TYPE_CHECKING: # pragma: no cover +if TYPE_CHECKING: # pragma: no cover from textual.app import App diff --git a/idatui/remote_ops.py b/idatui/remote_ops.py index b07bb71..8105f81 100644 --- a/idatui/remote_ops.py +++ b/idatui/remote_ops.py @@ -68,7 +68,8 @@ def declare_type(db: Database, **a: Any) -> Any: def decomp_error(db: Database, **a: Any) -> Any: - import ida_hexrays, ida_ida + import ida_hexrays + import ida_ida ea = int(str(a["addr"]), 16) fn = db.functions.get_at(ea) @@ -114,7 +115,11 @@ def define_code(db: Database, **a: Any) -> Any: def define_code_run(db: Database, **a: Any) -> Any: - import ida_bytes, ida_idp, ida_segment, ida_ua, idaapi + import ida_bytes + import ida_idp + import ida_segment + import ida_ua + import idaapi ea, limit = int(str(a["addr"]), 16), max(1, min(int(a.get("limit", 20000)), 200000)) seg = ida_segment.getseg(ea) @@ -172,7 +177,9 @@ def define_func(db: Database, **a: Any) -> Any: def define_func_run(db: Database, **a: Any) -> Any: - import ida_bytes, ida_funcs, ida_segment + import ida_bytes + import ida_funcs + import ida_segment ea = int(str(a["addr"]), 16) fn = db.functions.get_at(ea) @@ -279,7 +286,8 @@ def file_regions(db: Database, **a: Any) -> Any: def flowchart(db: Database, **a: Any) -> Any: - import ida_funcs, ida_gdl + import ida_funcs + import ida_gdl ea = int(str(a["addr"]), 16) fn = ida_funcs.get_func(ea) @@ -407,8 +415,14 @@ def journal_put(db: Database, **a: Any) -> Any: 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 + import ida_bytes + import ida_funcs + import ida_lines + import ida_nalt + import ida_name + import ida_segment + import ida_typeinf + import idautils limit = max(1, int(a.get("limit", 4000))) max_scan = max(1000, int(a.get("max_scan", 2000000))) @@ -629,7 +643,9 @@ def lookup_funcs(db: Database, **a: Any) -> Any: def make_data(db: Database, **a: Any) -> Any: - import ida_bytes, ida_idaapi, ida_typeinf + import ida_bytes + import ida_idaapi + import ida_typeinf from ida_domain.types import TypeApplyFlags rows = [] @@ -715,7 +731,9 @@ def read_raw(db: Database, **a: Any) -> Any: def rename(db: Database, **a: Any) -> Any: - import idaapi, ida_hexrays, ida_name + import ida_hexrays + import ida_name + import idaapi batch = a.get("batch") or {} dry_run = bool(batch.get("dry_run", False)) @@ -902,7 +920,8 @@ def rename(db: Database, **a: Any) -> Any: def resolve_names(db: Database, **a: Any) -> Any: - import ida_idaapi, ida_name + import ida_idaapi + import ida_name rows = [] for query in a.get("queries", []): @@ -916,7 +935,11 @@ def resolve_names(db: Database, **a: Any) -> Any: def search_bytes(db: Database, **a: Any) -> Any: - import ida_bytes, ida_funcs, ida_idaapi, ida_lines, ida_segment + import ida_bytes + import ida_funcs + import ida_idaapi + import ida_lines + import ida_segment pat = str(a.get("pattern", "")).strip() limit = max(1, int(a.get("limit", 500))) @@ -980,9 +1003,13 @@ def search_structs(db: Database, **a: Any) -> Any: def search_text(db: Database, **a: Any) -> Any: - import ida_lines, ida_funcs, ida_segment, idautils import re as _re + import ida_funcs + import ida_lines + import ida_segment + import idautils + q = str(a.get("query", "")) limit = max(1, int(a.get("limit", 500))) max_scan = max(1000, int(a.get("max_scan", 3000000))) @@ -1039,7 +1066,9 @@ def search_text(db: Database, **a: Any) -> Any: def set_comments(db: Database, **a: Any) -> Any: - import idaapi, idc, ida_hexrays + import ida_hexrays + import idaapi + import idc rows = [] for item in a.get("items", []): @@ -1150,7 +1179,11 @@ def set_lvar_type(db: Database, **a: Any) -> Any: def set_thumb(db: Database, **a: Any) -> Any: - import ida_bytes, ida_ida, ida_idp, ida_segment, ida_segregs + import ida_bytes + import ida_ida + import ida_idp + import ida_segment + import ida_segregs ea = int(str(a["addr"]), 16) treg = ida_idp.str2reg("T") @@ -1224,7 +1257,12 @@ def survey_binary(db: Database, **a: Any) -> Any: def thumb_scan(db: Database, **a: Any) -> Any: - import ida_bytes, ida_funcs, ida_idp, ida_segment, ida_segregs, ida_ua + import ida_bytes + import ida_funcs + import ida_idp + import ida_segment + import ida_segregs + import 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)) @@ -1329,7 +1367,10 @@ def undefine(db: Database, **a: Any) -> Any: def xref_query(db: Database, **a: Any) -> Any: - import idaapi, idautils, ida_bytes, ida_funcs + import ida_bytes + import ida_funcs + import idaapi + import idautils def _fn(ea): f = ida_funcs.get_func(ea) @@ -1452,7 +1493,11 @@ def xref_query(db: Database, **a: Any) -> Any: def xref_types(db: Database, **a: Any) -> Any: - import idaapi, idautils, ida_bytes, ida_funcs, ida_xref + import ida_bytes + import ida_funcs + import ida_xref + import idaapi + import idautils code_kind = { ida_xref.fl_CF: "call", diff --git a/idatui/rpc.py b/idatui/rpc.py index 1262ebb..cfeb0ca 100644 --- a/idatui/rpc.py +++ b/idatui/rpc.py @@ -17,6 +17,7 @@ Method tiers: introspect state, view, screen, functions (semantic verbs — open/goto/rename/... — layer on top in a later pass.) """ + from __future__ import annotations import asyncio @@ -27,8 +28,8 @@ from typing import Any from rich.console import Console -from ._sync import drain, settle from . import diag +from ._sync import drain, settle from .app import DecompView, GraphView, HexView, ListingView, ViewMode PROTO_VERSION = 1 @@ -36,10 +37,31 @@ TYPE_DELAY_MS = 35 # default per-char delay for high-level typed ops (aesthetic # Verbs that dereference app.program — refused with a clear error before load. _PROGRAM_METHODS = { - "goto", "open", "rename", "comment", "retype", "follow", "xrefs", "symbols", - "structs", "search", "select", "save", "hex", "toggle_view", - "pseudocode", "disassembly", "xrefs_to", "xrefs_from", "resolve", - "define", "rename_many", "opfmt", "graph", "export", "find", + "goto", + "open", + "rename", + "comment", + "retype", + "follow", + "xrefs", + "symbols", + "structs", + "search", + "select", + "save", + "hex", + "toggle_view", + "pseudocode", + "disassembly", + "xrefs_to", + "xrefs_from", + "resolve", + "define", + "rename_many", + "opfmt", + "graph", + "export", + "find", } # Self-documenting method table (returned by the 'methods' verb). @@ -55,7 +77,7 @@ METHODS = { "disassembly": "{target?,max?=2000} -> {total,lines:[{ea,text}]}", "xrefs_to": "{target,limit?=200} -> [{frm,to,type,fn_addr,fn_name}]", "xrefs_from": "{target,limit?=200} -> callees/refs; function-scoped for a " - "function (decomp refs), address-scoped for a 0xADDR", + "function (decomp refs), address-scoped for a 0xADDR", "resolve": "{name} -> {ea}", "keys": "{keys:[str],settle?,timeout?} raw key injection (supports 'wait:')", "text": "{text,delay_ms?,settle?} type a literal string into the focused input", @@ -69,17 +91,17 @@ METHODS = { "toggle_view": "disasm <-> pseudocode", "hex": "hex view", "graph": "{action?=show|open|close|toggle|zoom|block|entry|succ|pred," - "target?,blocks?} the control-flow graph: 'show' reports its " - "structure (blocks, edges, cursor) without touching it; the others " - "drive it. 'block' takes target=", + "target?,blocks?} the control-flow graph: 'show' reports its " + "structure (blocks, edges, cursor) without touching it; the others " + "drive it. 'block' takes target=", "xrefs": "open the xref picker", "symbols": "{query?} open the symbol palette", "structs": "open the struct editor", "export": "{path?,types?=true} write the session's comments/names/types as " - "a markdown report -> {path,comments,names,types}", + "a markdown report -> {path,comments,names,types}", "find": "{query,mode?=auto|text|bytes,limit?=500,regex?,case?} search the " - "WHOLE database: disassembly text, or a byte pattern with " - "wildcards (48 8b ?? c3) -> {mode,hits:[{addr,head,line,func}]}", + "WHOLE database: disassembly text, or a byte pattern with " + "wildcards (48 8b ?? c3) -> {mode,hits:[{addr,head,line,func}]}", "search": "{term,direction?=1} incremental search in the code view", "select": "{index?} choose the highlighted/nth item in the open modal", "save": "persist the .i64 (Ctrl+S)", @@ -90,42 +112,62 @@ METHODS = { "move": "{dir,n?=1} fast movement (down/up/.../pagedown)", "cursor": "{line?,col?} set the code-pane cursor directly", "define": "{kind:code|func|undef|thumb|thumbscan|data|string,target?} " - "(re)define bytes at target — the raw-image workflow", + "(re)define bytes at target — the raw-image workflow", "rename_many": "{items:[{addr,name}] | file:JSON} bulk-apply a symbol file " - "in ONE call (no typing, no navigation)", + "in ONE call (no typing, no navigation)", "opfmt": "{mode?=cycle|back|show|hex|dec|oct|bin|char|offset|stack|" - "default,target?,word?,line?,col?} how the literal under the cursor is " - "DISPLAYED (IDA's 'o'); works on the listing and on pseudocode " - "numbers. 'show' reports the format and the stops without editing", + "default,target?,word?,line?,col?} how the literal under the cursor is " + "DISPLAYED (IDA's 'o'); works on the listing and on pseudocode " + "numbers. 'show' reports the format and the stops without editing", } #: `opfmt` modes that have a real key on the code views. Driving the key keeps #: the pane honest (a viewer sees the same thing a human would do); the named #: formats have no key, so those go through the view's action directly. _OPFMT_KEYS = {"cycle": "o", "back": "O"} -_OPFMT_MODES = ("cycle", "back", "show", "hex", "dec", "oct", "bin", "char", - "offset", "stack", "default") +_OPFMT_MODES = ( + "cycle", + "back", + "show", + "hex", + "dec", + "oct", + "bin", + "char", + "offset", + "stack", + "default", +) # `define` kinds -> the ListingView key that runs them. Driving the real key # keeps the pane honest (a viewer sees the same thing a human would do) and # reuses the app's own edit worker, which reports what actually happened. _DEFINE_KEYS = { - "code": "c", # make code (runs until flow/undecodable) - "func": "p", # make function + "code": "c", # make code (runs until flow/undecodable) + "func": "p", # make function "undef": "u", - "thumb": "t", # flip ARM/Thumb at the cursor, then disassemble - "thumbscan": "T", # find Thumb entry pointers in a vector table + "thumb": "t", # flip ARM/Thumb at the cursor, then disassemble + "thumbscan": "T", # find Thumb entry pointers in a vector table "data": "d", "string": "a", } # Movement keys — driven fast (no typed delay) so the pane still visibly moves. _MOVE_KEYS = { - "down": "j", "up": "k", "left": "h", "right": "l", - "word": "w", "wordback": "b", "bol": "0", "eol": "dollar_sign", - "top": "home", "bottom": "G", - "halfdown": "ctrl+d", "halfup": "ctrl+u", - "pagedown": "pagedown", "pageup": "pageup", + "down": "j", + "up": "k", + "left": "h", + "right": "l", + "word": "w", + "wordback": "b", + "bol": "0", + "eol": "dollar_sign", + "top": "home", + "bottom": "G", + "halfdown": "ctrl+d", + "halfup": "ctrl+u", + "pagedown": "pagedown", + "pageup": "pageup", } @@ -148,8 +190,11 @@ def graph_info(app, blocks: bool = True) -> dict[str, Any]: rather than the box-drawing characters it is rendered as.""" gv = app.query_one(GraphView) if gv.fc is None or gv.lay is None: - return {"open": app.is_graph, "loaded": False, - "note": "press space (or graph {action:'open'}) on a function"} + return { + "open": app.is_graph, + "loaded": False, + "note": "press space (or graph {action:'open'}) on a function", + } lay, fc = gv.lay, gv.fc out: dict[str, Any] = { "open": app.is_graph, @@ -158,24 +203,30 @@ def graph_info(app, blocks: bool = True) -> dict[str, Any]: "zoom": gv.ZOOMS[gv._zoom], "canvas": {"w": lay.width, "h": lay.height}, "stats": dict(lay.stats), - "cursor": {"block": gv.cursor_node, "row": gv.cursor_row, - "ea": gv._cursor_ea(), "word": gv.word_under_cursor()}, + "cursor": { + "block": gv.cursor_node, + "row": gv.cursor_row, + "ea": gv._cursor_ea(), + "word": gv.word_under_cursor(), + }, } if blocks: rows = [] for n in lay.nodes: b = gv._blocks.get(n.id) - rows.append({ - "id": n.id, - "start": b.start if b else None, - "end": b.end if b else None, - "insns": len(b.rows) if b else 0, - "rank": n.rank, - "box": {"x": n.x, "y": n.y, "w": n.w, "h": n.h}, - "succs": [{"id": i, "kind": k} for i, k in lay.succ.get(n.id, [])], - "preds": [{"id": i, "kind": k} for i, k in lay.pred.get(n.id, [])], - "selfloop": bool(b and any(d == n.id for d, _ in b.succs)), - }) + rows.append( + { + "id": n.id, + "start": b.start if b else None, + "end": b.end if b else None, + "insns": len(b.rows) if b else 0, + "rank": n.rank, + "box": {"x": n.x, "y": n.y, "w": n.w, "h": n.h}, + "succs": [{"id": i, "kind": k} for i, k in lay.succ.get(n.id, [])], + "preds": [{"id": i, "kind": k} for i, k in lay.pred.get(n.id, [])], + "selfloop": bool(b and any(d == n.id for d, _ in b.succs)), + } + ) out["blocks"] = rows return out @@ -186,8 +237,20 @@ _MODALS = ("XrefsScreen", "SymbolPalette", "StructEditor", "ConfirmScreen") #: Handlers that did ``int(...)`` coped; the ones that compared directly blew up #: with e.g. "'<' not supported between instances of 'int' and 'str'". Coerce the #: known-numeric names once, centrally, instead of at every call site. -_INT_PARAMS = ("lines", "limit", "max", "n", "index", "line", "col", - "occurrence", "delay_ms", "direction", "addr", "count") +_INT_PARAMS = ( + "lines", + "limit", + "max", + "n", + "index", + "line", + "col", + "occurrence", + "delay_ms", + "direction", + "addr", + "count", +) _FLOAT_PARAMS = ("timeout",) @@ -221,13 +284,16 @@ def _modal_snapshot(app) -> dict[str, Any] | None: if isinstance(items, list): try: from textual.widgets import OptionList + hl = scr.query_one(OptionList).highlighted except Exception: # noqa: BLE001 hl = None info["highlighted"] = hl info["items"] = [ - {"ea": (it[0] if isinstance(it[0], int) else None), - "label": str(it[1]) if len(it) > 1 else str(it)} + { + "ea": (it[0] if isinstance(it[0], int) else None), + "label": str(it[1]) if len(it) > 1 else str(it), + } for it in items[:64] ] return info @@ -235,14 +301,23 @@ def _modal_snapshot(app) -> dict[str, Any] | None: def _cursor_info(app, w) -> dict[str, Any]: if isinstance(w, HexView): - return {"kind": "hex", "va": (w.cursor_va() if w.model else None), - "byte": w.cursor} + return { + "kind": "hex", + "va": (w.cursor_va() if w.model else None), + "byte": w.cursor, + } if isinstance(w, GraphView): # The graph cursor is (block, row), not a line index -- reporting it as # one would make a driver's `cursor line=` land somewhere arbitrary. - return {"kind": "graph", "ea": w._cursor_ea(), "block": w.cursor_node, - "row": w.cursor_row, "col": w.cursor_x, - "word": w.word_under_cursor(), "text": w._line_plain()} + return { + "kind": "graph", + "ea": w._cursor_ea(), + "block": w.cursor_node, + "row": w.cursor_row, + "col": w.cursor_x, + "word": w.word_under_cursor(), + "text": w._line_plain(), + } # disasm / decomp share the ColumnCursor surface word = None try: @@ -254,9 +329,15 @@ def _cursor_info(app, w) -> dict[str, Any]: ea = app._line_ea_for(w) except Exception: # noqa: BLE001 pass - return {"kind": app._active, "line": w.cursor, "col": w.cursor_x, - "word": word, "ea": ea, "total": getattr(w, "total", None), - "scroll_y": round(w.scroll_offset.y)} + return { + "kind": app._active, + "line": w.cursor, + "col": w.cursor_x, + "word": word, + "ea": ea, + "total": getattr(w, "total", None), + "scroll_y": round(w.scroll_offset.y), + } def _where(app) -> str: @@ -290,12 +371,12 @@ def snapshot(app) -> dict[str, Any]: pass return { "active": app._active, - "pref": app._code_mode(), # kept for wire compat; a constant now + "pref": app._code_mode(), # kept for wire compat; a constant now "function": ({"ea": cur.ea, "name": cur.name} if cur else None), "cursor": _cursor_info(app, w), "status": st, "filter": app._filter_term, - "binary": app._binary, # None outside project mode + "binary": app._binary, # None outside project mode "nav_depth": len(app._nav), "hops": list(getattr(app, "_hops", [])), "dirty": bool(app._dirty), @@ -309,13 +390,18 @@ def view_lines(app, lines: int | None = None) -> dict[str, Any]: for hex use screen()).""" w = _active_widget(app) if isinstance(w, HexView): - return {"active": "hex", "note": "use screen() for the hex grid", - "cursor": _cursor_info(app, w)} + return { + "active": "hex", + "note": "use screen() for the hex grid", + "cursor": _cursor_info(app, w), + } if isinstance(w, GraphView): - return {"active": "graph", "note": "use graph() for structure, " - "screen() for the drawing", - "cursor": _cursor_info(app, w), - "graph": graph_info(app, blocks=False)} + return { + "active": "graph", + "note": "use graph() for structure, screen() for the drawing", + "cursor": _cursor_info(app, w), + "graph": graph_info(app, blocks=False), + } top = round(w.scroll_offset.y) height = w.size.height or 40 n = min(lines or height, max(w.total - top, 0)) @@ -323,21 +409,39 @@ def view_lines(app, lines: int | None = None) -> dict[str, Any]: for r in range(n): idx = top + r plain = w._line_plain(idx) - out.append({"i": idx, "cur": idx == w.cursor, - "text": plain if plain is not None else ""}) - return {"active": app._active, "top": top, "total": w.total, - "cursor": _cursor_info(app, w), "lines": out} + out.append( + { + "i": idx, + "cur": idx == w.cursor, + "text": plain if plain is not None else "", + } + ) + return { + "active": app._active, + "top": top, + "total": w.total, + "cursor": _cursor_info(app, w), + "lines": out, + } def screen_text(app, fmt: str = "text") -> dict[str, Any]: """Render the whole screen exactly as shown. ``fmt``: 'text' (plain, default), 'html' or 'svg' (colored — handy for an out-of-band web viewer).""" width, height = app.size - console = Console(width=width, height=height or 40, file=io.StringIO(), - force_terminal=True, color_system="truecolor", record=True, - legacy_windows=False, safe_box=False) + console = Console( + width=width, + height=height or 40, + file=io.StringIO(), + force_terminal=True, + color_system="truecolor", + record=True, + legacy_windows=False, + safe_box=False, + ) render = app.screen._compositor.render_update( - full=True, screen_stack=app._background_screens, simplify=False) + full=True, screen_stack=app._background_screens, simplify=False + ) console.print(render) out: dict[str, Any] = {"width": width, "height": height, "format": fmt} if fmt == "html": @@ -388,8 +492,10 @@ def cursor_on(app, word: str, line: int | None = None, occurrence: int = 1) -> b if isinstance(w, HexView): raise ValueError("cursor_on: not supported in the hex view") if isinstance(w, GraphView): - raise ValueError("cursor_on: not supported in the graph view — use " - "graph {action:'block'} or goto") + raise ValueError( + "cursor_on: not supported in the graph view — use " + "graph {action:'block'} or goto" + ) if isinstance(w, DecompView): texts = list(w._texts) else: @@ -412,7 +518,7 @@ def cursor_on(app, word: str, line: int | None = None, occurrence: int = 1) -> b if w.word_under_cursor() == word: hits += 1 if hits >= max(1, occurrence): - place_cursor(w) # scrolls: an off-screen cursor edits blind + place_cursor(w) # scrolls: an off-screen cursor edits blind return True col = t.find(word, col + 1) w.cursor, w.cursor_x = orig # not found: leave the cursor untouched @@ -455,8 +561,14 @@ def pseudocode(app, target=None) -> dict[str, Any]: if dea is None: return {"ea": None, "error": "no target"} d = app.program.decompile(dea) - return {"ea": dea, "name": (fn.name if fn else None), "failed": d.failed, - "error": d.error, "truncated": d.truncated, "code": d.code} + return { + "ea": dea, + "name": (fn.name if fn else None), + "failed": d.failed, + "error": d.error, + "truncated": d.truncated, + "code": d.code, + } def disassembly(app, target=None, max_lines: int = 2000) -> dict[str, Any]: @@ -469,13 +581,26 @@ def disassembly(app, target=None, max_lines: int = 2000) -> dict[str, Any]: m = app.program.disasm(dea, fn.name if fn else None) total = m.total() lines = m.lines(0, min(total, max(1, max_lines)), prefetch=False) - return {"ea": dea, "name": (fn.name if fn else None), "total": total, - "lines": [{"ea": ln.ea, "text": ln.text} for ln in lines]} + return { + "ea": dea, + "name": (fn.name if fn else None), + "total": total, + "lines": [{"ea": ln.ea, "text": ln.text} for ln in lines], + } def _xref_dicts(xs, limit: int) -> list[dict[str, Any]]: - return [{"frm": x.frm, "to": x.to, "type": x.type, "kind": x.kind, - "fn_addr": x.fn_addr, "fn_name": x.fn_name} for x in xs[:limit]] + return [ + { + "frm": x.frm, + "to": x.to, + "type": x.type, + "kind": x.kind, + "fn_addr": x.fn_addr, + "fn_name": x.fn_name, + } + for x in xs[:limit] + ] def xrefs_to(app, target, limit: int = 200) -> list[dict[str, Any]]: @@ -497,9 +622,15 @@ def xrefs_from(app, target, limit: int = 200) -> list[dict[str, Any]]: for r in app.program.decompile(ea).refs[:limit]: tf = app.program.function_of(r.addr) is_func = bool(tf and tf.addr == r.addr) - out.append({"to": r.addr, "name": r.name or (tf.name if tf else None), - "string": r.string, "is_func": is_func, - "type": "code" if is_func else "data"}) + out.append( + { + "to": r.addr, + "name": r.name or (tf.name if tf else None), + "string": r.string, + "is_func": is_func, + "type": "code" if is_func else "data", + } + ) return out return _xref_dicts(app.program.xrefs_from(ea), limit) @@ -560,15 +691,25 @@ class RpcServer: except OSError: pass - async def _on_client(self, reader: asyncio.StreamReader, - writer: asyncio.StreamWriter) -> None: + async def _on_client( + self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter + ) -> None: if self._busy: # No multi-driver support yet: refuse a second concurrent client # rather than let two drivers interleave mutations. try: - writer.write(json.dumps( - {"id": None, "error": {"message": "busy: another client is " - "connected (single-driver only)"}}).encode() + b"\n") + writer.write( + json.dumps( + { + "id": None, + "error": { + "message": "busy: another client is " + "connected (single-driver only)" + }, + } + ).encode() + + b"\n" + ) await writer.drain() writer.close() except Exception: # noqa: BLE001 @@ -604,7 +745,11 @@ class RpcServer: except Exception as e: # noqa: BLE001 — report, never kill the connection # ``str(KeyError("msg"))`` returns ``repr("msg")`` (adds quotes), which # mangles our friendly resolve messages; unwrap the single arg instead. - if isinstance(e, KeyError) and len(e.args) == 1 and isinstance(e.args[0], str): + if ( + isinstance(e, KeyError) + and len(e.args) == 1 + and isinstance(e.args[0], str) + ): msg = e.args[0] else: msg = str(e) @@ -619,7 +764,8 @@ class RpcServer: # would go on to edit whatever the *previous* location was. raise TimeoutError( f"{what or 'action'} did not complete within {timeout}s " - f"(still at {_where(self.app)}); retry with a larger timeout=") + f"(still at {_where(self.app)}); retry with a larger timeout=" + ) return snapshot(self.app) async def _graph(self, params, timeout): @@ -641,18 +787,21 @@ class RpcServer: return {**snapshot(app), "graph": graph_info(app, blocks=want_blocks)} if action == "close" and not app.is_graph: return {**snapshot(app), "graph": graph_info(app, blocks=want_blocks)} - want = "graph" if action in ("open", "toggle") and \ - not app.is_graph else None + want = ( + "graph" if action in ("open", "toggle") and not app.is_graph else None + ) res = await self._press( ["space"], - (lambda: app.is_graph) if want else - (lambda: not app.is_graph), - timeout, f"graph {action}") + (lambda: app.is_graph) if want else (lambda: not app.is_graph), + timeout, + f"graph {action}", + ) return {**res, "graph": graph_info(app, blocks=want_blocks)} if not app.is_graph: - raise ValueError(f"graph {action}: the graph is not open " - f"(graph {{action:'open'}} first)") + raise ValueError( + f"graph {action}: the graph is not open (graph {{action:'open'}} first)" + ) if action == "zoom": before = gv._zoom await self._press(["z"], lambda: gv._zoom != before, timeout, "graph zoom") @@ -660,9 +809,12 @@ class RpcServer: await self._press(["0"], None, timeout, "graph entry") elif action in ("succ", "pred"): before = gv.cursor_node - await self._press(["J" if action == "succ" else "K"], - lambda: gv.cursor_node != before, timeout, - f"graph {action}") + await self._press( + ["J" if action == "succ" else "K"], + lambda: gv.cursor_node != before, + timeout, + f"graph {action}", + ) elif action == "block": target = params.get("target") if target is None: @@ -694,18 +846,23 @@ class RpcServer: """Open a prompt (a keystroke), optionally clear its prefill, type the value with the typed-out delay, submit. Returns after the prompt closes.""" from textual.widgets import Input + app = self.app await app._press_keys([open_key]) - await settle(app, lambda: app.query_one(f"#{input_id}", Input).display, timeout=10) + await settle( + app, lambda: app.query_one(f"#{input_id}", Input).display, timeout=10 + ) inp = app.query_one(f"#{input_id}", Input) if not inp.display: # Say *why*. The old message always blamed the word under the cursor, # which sent readers hunting for a cursor problem when the real cause # was usually a modal eating the opening keystroke. modal = type(app.screen).__name__ - why = (f"modal {modal!r} has focus and ate the {open_key!r} keystroke" - if modal in _MODALS or modal != "Screen" - else "no renameable token under the cursor") + why = ( + f"modal {modal!r} has focus and ate the {open_key!r} keystroke" + if modal in _MODALS or modal != "Screen" + else "no renameable token under the cursor" + ) raise RuntimeError(f"{input_id!r} prompt did not open: {why}") if clear: inp.value = "" @@ -727,14 +884,14 @@ class RpcServer: app = self.app items = params.get("items") src = params.get("file") - if isinstance(items, str): # `drive raw` hands params through as text + if isinstance(items, str): # `drive raw` hands params through as text items = json.loads(items) if items is None: if not src: raise ValueError("rename_many needs items=[{addr,name}] or file=") with open(os.path.expanduser(str(src))) as f: items = json.load(f) - if isinstance(items, dict): # {"0x4370": "name"} is a natural shape too + if isinstance(items, dict): # {"0x4370": "name"} is a natural shape too items = [{"addr": k, "name": v} for k, v in items.items()] if not isinstance(items, list) or not items: raise ValueError("rename_many: items must be a non-empty list") @@ -745,8 +902,14 @@ class RpcServer: skipped += 1 continue # Accept the field names symbol files actually use. - addr = next((it[k] for k in ("addr", "start", "ea", "address") - if it.get(k) is not None), None) + addr = next( + ( + it[k] + for k in ("addr", "start", "ea", "address") + if it.get(k) is not None + ), + None, + ) name = it.get("name") or it.get("label") if addr is None or not name: skipped += 1 @@ -764,8 +927,15 @@ class RpcServer: # or the TUI freezes for the length of the batch. res = await asyncio.to_thread(app.program.client.invoke, "rename", batch=batch) summary = res.get("summary", {}) if isinstance(res, dict) else {} - failed = [r for r in (res.get("func") or []) if isinstance(r, dict) - and r.get("error")] if isinstance(res, dict) else [] + failed = ( + [ + r + for r in (res.get("func") or []) + if isinstance(r, dict) and r.get("error") + ] + if isinstance(res, dict) + else [] + ) # Names live in the IDB, but every cache in front of it is now stale -- # including Hex-Rays', which is per-function and does NOT notice that a @@ -779,18 +949,23 @@ class RpcServer: app.program.bump_names() app.program.invalidate_functions() app._func_index = None - app._load_functions() # re-streams the function table + app._load_functions() # re-streams the function table await settle(app, timeout=timeout) app._dirty = True - app._status(f"renamed {summary.get('ok', 0)} symbols" - + (f", {len(failed)} failed" if failed else "") - + " (Ctrl+S to save)") + app._status( + f"renamed {summary.get('ok', 0)} symbols" + + (f", {len(failed)} failed" if failed else "") + + " (Ctrl+S to save)" + ) snap = snapshot(app) snap["rename_many"] = { - "requested": len(ops), "skipped": skipped, - "ok": summary.get("ok", 0), "failed": summary.get("failed", 0), - "errors": [{"addr": r.get("addr"), "error": r.get("error")} - for r in failed[:10]], + "requested": len(ops), + "skipped": skipped, + "ok": summary.get("ok", 0), + "failed": summary.get("failed", 0), + "errors": [ + {"addr": r.get("addr"), "error": r.get("error")} for r in failed[:10] + ], } return snap @@ -810,13 +985,31 @@ class RpcServer: #: Verbs that drive the *main* app by injecting keystrokes. If a modal is on #: top it eats those keys, so they must refuse rather than silently no-op. _NEEDS_NO_MODAL = { - "goto", "open", "rename", "comment", "retype", "follow", "back", - "toggle_view", "hex", "save", "search", "move", "cursor", "cursor_on", - "define", "opfmt", + "goto", + "open", + "rename", + "comment", + "retype", + "follow", + "back", + "toggle_view", + "hex", + "save", + "search", + "move", + "cursor", + "cursor_on", + "define", + "opfmt", } #: Modals the driver is expected to interact with (they have their own verbs). - _DRIVABLE_MODALS = {"XrefsScreen", "SymbolPalette", "StructEditor", - "ProjectPalette", "QuitScreen"} + _DRIVABLE_MODALS = { + "XrefsScreen", + "SymbolPalette", + "StructEditor", + "ProjectPalette", + "QuitScreen", + } def _modal_kind(self) -> str | None: scr = self.app.screen @@ -833,15 +1026,20 @@ class RpcServer: f"modal {modal!r} is on top and will swallow this verb's " f"keystrokes; dismiss it first (close) or use its own verb " f"(select/symbols/xrefs). Note: a binary with no entry " - f"function can land in the symbol palette on startup.") + f"function can land in the symbol palette on startup." + ) if method in (None, "ping"): module = None try: module = app._module() if app.client else None except Exception: # noqa: BLE001 pass - return {"ok": True, "proto": PROTO_VERSION, "module": module, - **_readiness(app)} + return { + "ok": True, + "proto": PROTO_VERSION, + "module": module, + **_readiness(app), + } if method == "methods": return METHODS if method == "quit": @@ -856,14 +1054,18 @@ class RpcServer: def _go(): if dirty and save: - app._on_quit_choice("save") # saves, then exits + app._on_quit_choice("save") # saves, then exits else: app._on_quit_choice("discard") # answer first, then tear down (so this response still gets written) asyncio.get_running_loop().call_later(0.2, _go) - return {"ok": True, "quitting": True, "saving": bool(dirty and save), - "dirty": dirty} + return { + "ok": True, + "quitting": True, + "saving": bool(dirty and save), + "dirty": dirty, + } if method in _PROGRAM_METHODS and app.program is None: raise ValueError("not ready: still connecting / loading functions") @@ -902,8 +1104,10 @@ class RpcServer: if params.get("clear"): diag.clear() return {"cleared": True} - return {"recent": diag.recent(int(params.get("n", 10))), - "log": os.environ.get("IDATUI_LOG") or None} + return { + "recent": diag.recent(int(params.get("n", 10))), + "log": os.environ.get("IDATUI_LOG") or None, + } if method == "trace": tc = app.trace_ctl @@ -916,12 +1120,19 @@ class RpcServer: if isinstance(v, str) and v.startswith("!"): idx = int(float(v[1:]) * (t.length - 1) / 100.0) else: - idx = int(str(v).replace(",", ""), 0) if isinstance(v, str) else int(v) + idx = ( + int(str(v).replace(",", ""), 0) + if isinstance(v, str) + else int(v) + ) tc.seek(idx) - elif "goto" in params: # first execution of an address/name + elif "goto" in params: # first execution of an address/name tgt = params["goto"] - ea = (int(str(tgt), 0) if str(tgt).lower().startswith("0x") - else app.program.resolve(str(tgt))) + ea = ( + int(str(tgt), 0) + if str(tgt).lower().startswith("0x") + else app.program.resolve(str(tgt)) + ) first = t.first_execution(ea) if first is None: raise ValueError(f"{tgt} never executed in this trace") @@ -933,9 +1144,12 @@ class RpcServer: (tc.step_over if over else tc.step)(1 if n > 0 else -1) await settle(app, timeout=float(params.get("timeout", 20.0))) snap = snapshot(app) - snap["trace"] = {"idx": tc.t, "length": t.length, - "pc": hex(t.ip(tc.t)), - "changed": sorted(t.changed(tc.t))} + snap["trace"] = { + "idx": tc.t, + "length": t.length, + "pc": hex(t.ip(tc.t)), + "changed": sorted(t.changed(tc.t)), + } return snap if method == "binaries": @@ -943,12 +1157,20 @@ class RpcServer: raise ValueError("not a project session (launch with --project)") counts = app._index.counts() if app._index is not None else {} resident = set(app._pool.resident()) if app._pool is not None else set() - return {"active": app._binary, "hops": list(app._hops), - "binaries": [{"label": r.label, "source": r.source, - "active": r.label == app._binary, - "resident": r.label in resident, - "indexed": int(counts.get(r.label, 0))} - for r in app._project.refs]} + return { + "active": app._binary, + "hops": list(app._hops), + "binaries": [ + { + "label": r.label, + "source": r.source, + "active": r.label == app._binary, + "resident": r.label in resident, + "indexed": int(counts.get(r.label, 0)), + } + for r in app._project.refs + ], + } if method == "switch": if app._project is None: @@ -965,30 +1187,41 @@ class RpcServer: else: # Same path a project search hit takes, so it records a hop and # Esc comes back here. - app._switch_then_goto(label, int(str(addr), 0) - if isinstance(addr, str) else int(addr)) - await settle(app, lambda: app._binary == label - and app._func_index is not None - and app._func_index.complete, - timeout=float(params.get("timeout", 300.0))) + app._switch_then_goto( + label, int(str(addr), 0) if isinstance(addr, str) else int(addr) + ) + await settle( + app, + lambda: ( + app._binary == label + and app._func_index is not None + and app._func_index.complete + ), + timeout=float(params.get("timeout", 300.0)), + ) return snapshot(app) # -- structured introspection (heavy: run off the UI loop) -------- # loop = asyncio.get_running_loop() if method == "pseudocode": - return await loop.run_in_executor(None, pseudocode, app, params.get("target")) + return await loop.run_in_executor( + None, pseudocode, app, params.get("target") + ) if method == "disassembly": mx = int(params.get("max", 2000)) return await loop.run_in_executor( - None, disassembly, app, params.get("target"), mx) + None, disassembly, app, params.get("target"), mx + ) if method == "xrefs_to": lim = int(params.get("limit", 200)) return await loop.run_in_executor( - None, xrefs_to, app, params.get("target"), lim) + None, xrefs_to, app, params.get("target"), lim + ) if method == "xrefs_from": lim = int(params.get("limit", 200)) return await loop.run_in_executor( - None, xrefs_from, app, params.get("target"), lim) + None, xrefs_from, app, params.get("target"), lim + ) if method == "resolve": return await loop.run_in_executor(None, resolve, app, params.get("name")) @@ -998,15 +1231,24 @@ class RpcServer: # optional ergonomic: place the cursor on a token before an edit/follow if method in ("rename", "retype", "follow") and params.get("word"): - if not cursor_on(app, str(params["word"]), params.get("line"), - int(params.get("occurrence", 1))): - raise ValueError(f"cursor_on: token {params['word']!r} not found " - "in the current view") + if not cursor_on( + app, + str(params["word"]), + params.get("line"), + int(params.get("occurrence", 1)), + ): + raise ValueError( + f"cursor_on: token {params['word']!r} not found in the current view" + ) await drain(app) if method == "cursor_on": - found = cursor_on(app, str(params["word"]), params.get("line"), - int(params.get("occurrence", 1))) + found = cursor_on( + app, + str(params["word"]), + params.get("line"), + int(params.get("occurrence", 1)), + ) await drain(app) snap = snapshot(app) snap["found"] = found @@ -1024,7 +1266,8 @@ class RpcServer: # on the function the caller *used* to be looking at. raise TimeoutError( f"goto {target!r} did not land within {timeout}s " - f"(still at {_where(app)}); retry with a larger timeout=") + f"(still at {_where(app)}); retry with a larger timeout=" + ) return snapshot(app) if method == "define": @@ -1032,80 +1275,95 @@ class RpcServer: if kind not in _DEFINE_KEYS: raise ValueError( f"unknown define kind {kind!r}; one of " - f"{', '.join(sorted(_DEFINE_KEYS))}") + f"{', '.join(sorted(_DEFINE_KEYS))}" + ) target = params.get("target") if target not in (None, ""): # Land on the address first. A raw image is mostly *undefined*, # so the target usually has no name and no function — the goto # predicate can't be address-based, only "we moved". - await self._fill_prompt("g", "goto", str(target), delay, - clear=False) + await self._fill_prompt("g", "goto", str(target), delay, clear=False) await settle(app, timeout=timeout) if app.is_hex: # backslash leaves hex for the code view (which may be decomp). - await self._press(["backslash"], - lambda: not app.is_hex, timeout, - "leave the hex view") + await self._press( + ["backslash"], lambda: not app.is_hex, timeout, "leave the hex view" + ) if app.is_decomp: # These bindings live on the listing; in the decompiler the key # would be swallowed or do something else entirely. - await self._press(["tab"], lambda: app.is_listing, - timeout, "switch to the listing") + await self._press( + ["tab"], lambda: app.is_listing, timeout, "switch to the listing" + ) if not app.is_listing: raise RuntimeError( f"define needs the listing view, but the active pane is " - f"{app._active!r}") - snap = await self._press([_DEFINE_KEYS[kind]], timeout=timeout, - what=f"define {kind}") + f"{app._active!r}" + ) + snap = await self._press( + [_DEFINE_KEYS[kind]], timeout=timeout, what=f"define {kind}" + ) snap["define"] = {"kind": kind, "status": snap.get("status", "")} return snap if method == "opfmt": mode = str(params.get("mode", "cycle")).lower() if mode not in _OPFMT_MODES: - raise ValueError(f"unknown opfmt mode {mode!r}; one of " - f"{', '.join(_OPFMT_MODES)}") + raise ValueError( + f"unknown opfmt mode {mode!r}; one of {', '.join(_OPFMT_MODES)}" + ) target = params.get("target") if target not in (None, ""): - await self._fill_prompt("g", "goto", str(target), delay, - clear=False) + await self._fill_prompt("g", "goto", str(target), delay, clear=False) await settle(app, timeout=timeout) if app.is_hex: - await self._press(["backslash"], lambda: not app.is_hex, - timeout, "leave the hex view") + await self._press( + ["backslash"], lambda: not app.is_hex, timeout, "leave the hex view" + ) view = _active_widget(app) if isinstance(view, HexView): raise RuntimeError("opfmt needs a code view, not the hex view") if params.get("word"): # Land the column on the literal first: WHICH operand gets # reformatted is decided by where the cursor is. - if not cursor_on(app, str(params["word"]), params.get("line"), - int(params.get("occurrence", 1) or 1)): + if not cursor_on( + app, + str(params["word"]), + params.get("line"), + int(params.get("occurrence", 1) or 1), + ): raise RuntimeError( f"{params['word']!r} is not on screen in this view, so " - f"there is no literal to reformat") + f"there is no literal to reformat" + ) await drain(app) elif params.get("line") is not None or params.get("col") is not None: place_cursor(view, params.get("line"), params.get("col")) await drain(app) before = _where(app) if mode in _OPFMT_KEYS: - snap = await self._press([_OPFMT_KEYS[mode]], timeout=timeout, - what=f"opfmt {mode}") + snap = await self._press( + [_OPFMT_KEYS[mode]], timeout=timeout, what=f"opfmt {mode}" + ) else: view.focus() view.action_op_format(mode) await settle(app, timeout=timeout) snap = snapshot(app) - snap["opfmt"] = {"mode": mode, "at": before, - "status": snap.get("status", "")} + snap["opfmt"] = { + "mode": mode, + "at": before, + "status": snap.get("status", ""), + } return snap if method == "rename_many": return await self._rename_many(params, timeout) if method == "rename": - await self._fill_prompt("n", "rename", str(params["name"]), delay, clear=True) + await self._fill_prompt( + "n", "rename", str(params["name"]), delay, clear=True + ) await settle(app, timeout=timeout) return snapshot(app) if method == "comment": @@ -1115,19 +1373,21 @@ class RpcServer: # the Input widget. The app's _do_comment converts the two-char # sequence '\n' into a real newline for IDA, so we escape here. ctext = str(params["text"]).replace("\n", "\\n") - await self._fill_prompt("semicolon", "comment", ctext, 0, - clear=True) + await self._fill_prompt("semicolon", "comment", ctext, 0, clear=True) await settle(app, timeout=timeout) return snapshot(app) if method == "retype": - await self._fill_prompt("y", "retype", str(params["proto"]), delay, clear=True) + await self._fill_prompt( + "y", "retype", str(params["proto"]), delay, clear=True + ) await settle(app, timeout=timeout) return snapshot(app) if method == "follow": depth = len(app._nav) - return await self._press(["enter"], lambda: len(app._nav) > depth, - timeout, "follow") + return await self._press( + ["enter"], lambda: len(app._nav) > depth, timeout, "follow" + ) if method == "back": return await self._press(["escape"], timeout=timeout) if method == "toggle_view": @@ -1158,17 +1418,23 @@ class RpcServer: # call that LEAVES hex could never be satisfied and always timed # out -- a driver could open the hex view but never close it. was_hex = app.is_hex - return await self._press(["backslash"], lambda: app.is_hex != was_hex, - timeout, "hex") + return await self._press( + ["backslash"], lambda: app.is_hex != was_hex, timeout, "hex" + ) if method == "graph": return await self._graph(params, timeout) if method == "xrefs": return await self._press( - ["x"], lambda: type(app.screen).__name__ == "XrefsScreen", - timeout, "xrefs") + ["x"], + lambda: type(app.screen).__name__ == "XrefsScreen", + timeout, + "xrefs", + ) if method == "symbols": await app._press_keys(["ctrl+n"]) - await settle(app, lambda: type(app.screen).__name__ == "SymbolPalette", timeout=10) + await settle( + app, lambda: type(app.screen).__name__ == "SymbolPalette", timeout=10 + ) q = params.get("query") if q: await app._press_keys(_text_to_keys(str(q), delay)) @@ -1176,10 +1442,14 @@ class RpcServer: return snapshot(app) if method == "structs": return await self._press( - ["ctrl+t"], lambda: type(app.screen).__name__ == "StructEditor", - timeout, "structs") + ["ctrl+t"], + lambda: type(app.screen).__name__ == "StructEditor", + timeout, + "structs", + ) if method == "find": from . import search as _search + q = str(params.get("query", "")) forced = params.get("mode") forced = None if forced in (None, "auto") else str(forced) @@ -1190,33 +1460,57 @@ class RpcServer: raise ValueError(f"find: {problem}") cleaned = _search.normalise_pattern(cleaned) hits, err, truncated = await asyncio.to_thread( - app.program.search, cleaned, mode, + app.program.search, + cleaned, + mode, limit=int(params.get("limit", 500)), - regex=bool(params.get("regex")), case=bool(params.get("case"))) + regex=bool(params.get("regex")), + case=bool(params.get("case")), + ) if err: raise ValueError(f"find: {err}") - return {"mode": mode, "query": cleaned, "truncated": truncated, - "hits": [{"addr": hex(h.addr), "head": hex(h.head), - "line": h.line, "func": h.func, - "seg": h.seg} for h in hits]} + return { + "mode": mode, + "query": cleaned, + "truncated": truncated, + "hits": [ + { + "addr": hex(h.addr), + "head": hex(h.head), + "line": h.line, + "func": h.func, + "seg": h.seg, + } + for h in hits + ], + } if method == "export": # Deliberately NOT driven through the prompt: this is the one verb # whose whole point is the file it leaves behind, and a driver needs # the path back, not a screenshot of a prompt closing. from . import findings + path = params.get("path") app.journal.load(app.program) app.journal.flush(app.program) out, f = await asyncio.to_thread( - findings.export, app.program, app._open_path or "", + findings.export, + app.program, + app._open_path or "", str(path) if path else None, - types=bool(params.get("types", True)), journal=app.journal) + types=bool(params.get("types", True)), + journal=app.journal, + ) app._status(f"exported findings → {out}", priority=True) await drain(app) - return {"path": out, "comments": len(f.comments), - "names": len(findings._user_names(f)), - "types": len(f.types), "functions": f.n_functions, - "bytes": os.path.getsize(out) if os.path.exists(out) else 0} + return { + "path": out, + "comments": len(f.comments), + "names": len(findings._user_names(f)), + "types": len(f.types), + "functions": f.n_functions, + "bytes": os.path.getsize(out) if os.path.exists(out) else 0, + } if method == "close": return await self._press(["escape"], timeout=timeout) if method == "save": @@ -1224,13 +1518,16 @@ class RpcServer: if method == "search": term = str(params.get("term", "")) - open_key = "slash" if int(params.get("direction", 1)) >= 0 else "question_mark" + open_key = ( + "slash" if int(params.get("direction", 1)) >= 0 else "question_mark" + ) await self._fill_prompt(open_key, "search", term, delay, clear=True) await settle(app, timeout=timeout) return snapshot(app) if method == "select": from textual.widgets import OptionList + scr = app.screen if type(scr).__name__ not in _MODALS: raise ValueError("select: no modal list is open") @@ -1249,8 +1546,10 @@ class RpcServer: if method == "move": key = _MOVE_KEYS.get(str(params.get("dir"))) if key is None: - raise ValueError(f"unknown move dir: {params.get('dir')!r} " - f"(one of {sorted(_MOVE_KEYS)})") + raise ValueError( + f"unknown move dir: {params.get('dir')!r} " + f"(one of {sorted(_MOVE_KEYS)})" + ) n = max(1, int(params.get("n", 1))) await app._press_keys([key] * n) if params.get("settle", True): diff --git a/idatui/rpcclient.py b/idatui/rpcclient.py index 4231d62..df0ba63 100644 --- a/idatui/rpcclient.py +++ b/idatui/rpcclient.py @@ -15,6 +15,7 @@ Also usable as a library: No auth: whoever can r/w the socket drives the app. """ + from __future__ import annotations import json @@ -114,8 +115,10 @@ def main(argv: list[str]) -> int: sock = args[1] args = args[2:] if not sock: - print("error: no socket (pass --sock PATH or set IDATUI_RPC_SOCK)", - file=sys.stderr) + print( + "error: no socket (pass --sock PATH or set IDATUI_RPC_SOCK)", + file=sys.stderr, + ) return 2 if not args: print("error: no method given", file=sys.stderr) @@ -124,7 +127,7 @@ def main(argv: list[str]) -> int: method, rest = args[0], args[1:] params: dict[str, Any] = {} if method == "keys": - params["keys"] = rest # every positional is a key name + params["keys"] = rest # every positional is a key name elif method == "text" and rest and "=" not in rest[0]: # first positional is the literal text; the rest may be key=value params["text"] = rest[0] diff --git a/idatui/search.py b/idatui/search.py index feb7e3d..ba8c692 100644 --- a/idatui/search.py +++ b/idatui/search.py @@ -69,7 +69,7 @@ def classify(query: str, forced: str | None = None) -> tuple[str, str]: low = q.lower() for prefix, mode in (("hex:", BYTES), ("bytes:", BYTES), ("text:", TEXT)): if low.startswith(prefix): - return (mode, q[len(prefix):].strip()) + return (mode, q[len(prefix) :].strip()) if forced in (TEXT, BYTES): return (forced, q) if looks_like_bytes(q) or probably_meant_bytes(q): @@ -89,7 +89,7 @@ def normalise_pattern(pattern: str) -> str: q = q.replace(",", " ") # "488B??C3" -- a bare hex run with no separators at all. if " " not in q and len(q) > 2 and len(q) % 2 == 0: - q = " ".join(q[i:i + 2] for i in range(0, len(q), 2)) + q = " ".join(q[i : i + 2] for i in range(0, len(q), 2)) return " ".join(q.split()) @@ -107,6 +107,7 @@ def pattern_problem(pattern: str) -> str | None: tokens = [t for t in q.split() if t] bad = [t for t in tokens if not _TOKEN.match(t)] if bad: - return (f"{bad[0]!r} is not a byte: use hex pairs, ? wildcards " - 'or a "quoted string"') + return ( + f'{bad[0]!r} is not a byte: use hex pairs, ? wildcards or a "quoted string"' + ) return None diff --git a/idatui/trace.py b/idatui/trace.py index 931f918..3d144d3 100644 --- a/idatui/trace.py +++ b/idatui/trace.py @@ -59,19 +59,25 @@ class TraceInfo: def load(cls, path: str) -> "TraceInfo | None": try: with open(path) as f: - raw = dict( - ln.strip().split("=", 1) for ln in f if "=" in ln) + raw = dict(ln.strip().split("=", 1) for ln in f if "=" in ln) except OSError: return None + def num(k): try: return int(raw.get(k, "0"), 0) except ValueError: return 0 - return cls(arch=raw.get("arch", ""), mode=raw.get("mode", ""), - binary=raw.get("binary", ""), start_code=num("start_code"), - end_code=num("end_code"), entry_code=num("entry_code"), - traced=raw.get("traced", "")) + + return cls( + arch=raw.get("arch", ""), + mode=raw.get("mode", ""), + binary=raw.get("binary", ""), + start_code=num("start_code"), + end_code=num("end_code"), + entry_code=num("entry_code"), + traced=raw.get("traced", ""), + ) @dataclass @@ -233,8 +239,7 @@ class Trace: return vals[i] if i >= 0 else None def register_state(self, idx: int) -> dict[str, int]: - return {n: v for n in self.reg_at - if (v := self.register(n, idx)) is not None} + return {n: v for n in self.reg_at if (v := self.register(n, idx)) is not None} def changed(self, idx: int) -> set[str]: """Registers written BY the instruction at ``idx`` (what the line said). @@ -277,9 +282,13 @@ class Trace: out = [] for k in range(lo, hi): off, ln = self.mem_off[k], self.mem_len[k] - out.append(MemOp(addr=self.mem_addr[k], - data=bytes(self.mem_blob[off:off + ln]), - write=bool(self.mem_write[k]))) + out.append( + MemOp( + addr=self.mem_addr[k], + data=bytes(self.mem_blob[off : off + ln]), + write=bool(self.mem_write[k]), + ) + ) return out # -- memory state ------------------------------------------------------- # @@ -298,8 +307,9 @@ class Trace: self._mem_starts = [self.mem_addr[k] for k in order] self._mem_maxlen = max(self.mem_len) if len(self.mem_len) else 0 - def memory_raw(self, addr: int, length: int, - idx: int | None = None) -> tuple[bytes, bytes]: + def memory_raw( + self, addr: int, length: int, idx: int | None = None + ) -> tuple[bytes, bytes]: """Memory at a TRACE address (no slide). The stack lives here. Measured on two real traces, 0% of memory accesses @@ -309,8 +319,9 @@ class Trace: """ return self.memory(addr + self.slide, length, idx) - def memory(self, addr: int, length: int, - idx: int | None = None) -> tuple[bytes, bytes]: + def memory( + self, addr: int, length: int, idx: int | None = None + ) -> tuple[bytes, bytes]: """``(data, known)`` for ``length`` bytes at ``addr`` as of ``idx``. ``known`` is a byte-per-byte mask: a trace only says what it saw, so a @@ -332,6 +343,7 @@ class Trace: raw = addr - self.slide best = [-1] * length import bisect as _b + lo = _b.bisect_left(self._mem_starts, raw - self._mem_maxlen) hi = _b.bisect_right(self._mem_starts, raw + length - 1) for pos in range(lo, hi): @@ -369,6 +381,7 @@ class Trace: self._mem_index() raw = addr - self.slide import bisect as _b + lo = _b.bisect_left(self._mem_starts, raw - self._mem_maxlen) hi = _b.bisect_right(self._mem_starts, raw + length - 1) out = set() @@ -439,7 +452,9 @@ class Trace: elif prev == "future": # Same distance rule as above, resolved by which loop found it # first would be arbitrary; compare real distances instead. - fwd = next((i for i, a in enumerate(self.next_ips(idx, n)) if a == ea), n) + fwd = next( + (i for i, a in enumerate(self.next_ips(idx, n)) if a == ea), n + ) if k < fwd: out[ea] = "past" if 0 <= idx < self.length: diff --git a/idatui/trace_ctl.py b/idatui/trace_ctl.py index 8072801..5a91d37 100644 --- a/idatui/trace_ctl.py +++ b/idatui/trace_ctl.py @@ -15,6 +15,7 @@ The controller owns the trace state. ``IdaTui`` keeps forwarding properties (``app._trace``, ``app._t``, ``app._trail_map``...) because the pilot suite and the RPC layer read them by those names; see ``IdaTui._trace``. """ + from __future__ import annotations import bisect @@ -23,7 +24,7 @@ from typing import TYPE_CHECKING from . import diag -if TYPE_CHECKING: # pragma: no cover +if TYPE_CHECKING: # pragma: no cover from .app import IdaTui _app_mod = None @@ -39,6 +40,7 @@ def _views(): global _app_mod if _app_mod is None: from . import app as _m + _app_mod = _m return _app_mod @@ -48,15 +50,15 @@ class TraceController: def __init__(self, app: "IdaTui", path: str = "") -> None: self.app = app - self.path = path or "" # the Tenet trace to explore, if any - self.trace = None # the loaded Trace, once analysed - self.t = 0 # current timestamp in that trace - self.trail_map = [] # decomp_map for trail_map_ea + self.path = path or "" # the Tenet trace to explore, if any + self.trace = None # the loaded Trace, once analysed + self.t = 0 # current timestamp in that trace + self.trail_map = [] # decomp_map for trail_map_ea self.trail_map_ea = None - self.trail_line_of: dict[int, int] = {} # ea -> pseudocode line - self.trail_eas: list[int] = [] # sorted keys of trail_line_of - self.trail_span = None # ea span of that function - self.pending_line = None # step waiting on a re-decompile + self.trail_line_of: dict[int, int] = {} # ea -> pseudocode line + self.trail_eas: list[int] = [] # sorted keys of trail_line_of + self.trail_span = None # ea span of that function + self.pending_line = None # step waiting on a re-decompile @property def armed(self) -> bool: @@ -82,26 +84,30 @@ class TraceController: Called on a worker thread, so every touch of the UI hops back. """ from .trace import Trace + app = self.app path = self.path try: + def note(n): - app.call_from_thread( - app._status, f"trace: {n:,} instructions\u2026") + app.call_from_thread(app._status, f"trace: {n:,} instructions\u2026") + trace = Trace.load(path, progress=note) except OSError as e: app.call_from_thread(app._status, f"trace: {e}") return if not trace.length: app.call_from_thread( - app._status, f"trace: {os.path.basename(path)} is empty") + app._status, f"trace: {os.path.basename(path)} is empty" + ) return idx = app._func_index addrs = [f.addr for f in idx.all_loaded()] if idx is not None else [] slide = trace.rebase(addrs) trace.apply_slide(slide) - hit = sum(1 for f in (idx.all_loaded() if idx else []) - if trace.executions(f.addr)) + hit = sum( + 1 for f in (idx.all_loaded() if idx else []) if trace.executions(f.addr) + ) app.call_from_thread(self.ready, trace, slide, hit) def ready(self, trace, slide: int, hit: int) -> None: @@ -111,9 +117,11 @@ class TraceController: dock = app.query_one(_views().TraceDock) dock.display = True dock.show(trace, 0) - where = (f"rebased {slide:+#x}" if slide else "no rebase needed") - app._status(f"trace: {trace.length:,} instructions, {hit} functions " - f"touched ({where})", priority=True) + where = f"rebased {slide:+#x}" if slide else "no rebase needed" + app._status( + f"trace: {trace.length:,} instructions, {hit} functions touched ({where})", + priority=True, + ) self.seek(0, follow=True) # -- trace navigation --------------------------------------------------- # @@ -144,8 +152,7 @@ class TraceController: # Stay in whichever view you're reading. Without prefer_decomp a step # from the pseudocode navigates to an address, which opens the listing — # so stepping through C threw you out of C on the first keypress. - app._goto_ea(pc, push=False, - prefer_decomp=(app.is_decomp)) + app._goto_ea(pc, push=False, prefer_decomp=(app.is_decomp)) def seek_split(self, pc: int) -> bool: """Put BOTH panes on ``pc``. True if handled. @@ -166,7 +173,7 @@ class TraceController: return False row = lst.model.ensure_ea(pc) if row is None or row < 0: - return False # not in this listing (other segment): full nav + return False # not in this listing (other segment): full nav lst.cursor = row lst._scroll_cursor_into_view() @@ -177,8 +184,9 @@ class TraceController: # bounced main -> PLT stub -> main, each bounce costing a synchronous # 769-line map fetch on the UI thread. span = self.trail_span - inside = (pc in self.trail_line_of - or (span is not None and span[0] <= pc <= span[1])) + inside = pc in self.trail_line_of or ( + span is not None and span[0] <= pc <= span[1] + ) if not inside: self.pending_line = pc app._resync_decomp_async(pc) @@ -227,7 +235,7 @@ class TraceController: t = self.trace if t is None: return - hx = app._try_view(M.HexView) # None until it's mounted + hx = app._try_view(M.HexView) # None until it's mounted if hx is not None: hx.trace, hx.trace_idx = t, self.t if hx.display: @@ -325,7 +333,7 @@ class TraceController: sp_name = "rsp" if "rsp" in t.reg_at else ("esp" if "esp" in t.reg_at else "sp") sp0 = t.register(sp_name, self.t) i = self.t + direction - limit = 200000 # a runaway search must not hang the UI + limit = 200000 # a runaway search must not hang the UI while 0 <= i < t.length and limit > 0: sp = t.register(sp_name, i) if sp0 is None or sp is None or sp >= sp0: @@ -364,15 +372,17 @@ class TraceController: # and often no question at all, since most lines have no marker. line = view.cursor eas = [] - if (self.trail_map_ea == view.loaded_ea - and 0 <= line < len(self.trail_map or [])): + if self.trail_map_ea == view.loaded_ea and 0 <= line < len( + self.trail_map or [] + ): eas = list(self.trail_map[line]) if not eas: one = view._line_ea(line) eas = [one] if one is not None else [] if not eas: - app._status("this line has no instructions to seek on", - priority=True) + app._status( + "this line has no instructions to seek on", priority=True + ) return stamps = sorted({x for e in eas for x in t.executions(e)}) what = f"execution of C line {line + 1}" @@ -392,12 +402,15 @@ class TraceController: i = bisect.bisect_left(stamps, self.t) - 1 if not (0 <= i < len(stamps)): edge = "last" if direction > 0 else "first" - app._status(f"already at the {edge} {what} " - f"({len(stamps)} in the trace)", priority=True) + app._status( + f"already at the {edge} {what} ({len(stamps)} in the trace)", + priority=True, + ) return self.seek(stamps[i]) - app._status(f"{what}: {i + 1} of {len(stamps)} @ t={stamps[i]:,}", - priority=True) + app._status( + f"{what}: {i + 1} of {len(stamps)} @ t={stamps[i]:,}", priority=True + ) def seek_reg_write(self) -> None: """W: which instruction set each register to its current value.""" @@ -410,11 +423,13 @@ class TraceController: v = t.register(name, self.t) if v is None: continue - rows.append((name, v, t.last_write(name, self.t), - t.next_write(name, self.t))) + rows.append( + (name, v, t.last_write(name, self.t), t.next_write(name, self.t)) + ) if rows: - app.push_screen(_views().RegWriteScreen(rows, self.t), - self._on_reg_write_chosen) + app.push_screen( + _views().RegWriteScreen(rows, self.t), self._on_reg_write_chosen + ) def _on_reg_write_chosen(self, idx) -> None: # type: ignore[no-untyped-def] if idx is not None: diff --git a/ruff.toml b/ruff.toml index 6d0f877..b7162a5 100644 --- a/ruff.toml +++ b/ruff.toml @@ -20,6 +20,13 @@ required-version = "==0.16.3" target-version = "py311" line-length = 88 +[format] +# 0.16 formats Python fenced blocks inside markdown by default. The snippets +# in docs/ are deliberately terse repro recipes (a `try: x` one-liner reads as +# one step); a formatter blowing them up to three lines makes them worse, and +# nothing executes them. Real code lives in .py files, which are formatted. +exclude = ["*.md"] + [lint] # Import sorting, and nothing else. `ruff check` can enforce a great deal more # and one day it might, but this hook's job is formatting: one that also diff --git a/tests/_fixtures.py b/tests/_fixtures.py index bc10235..87f6184 100644 --- a/tests/_fixtures.py +++ b/tests/_fixtures.py @@ -20,6 +20,7 @@ writes back to -- turns that into a file copy. The cache is rebuilt whenever it is older than the binary, so editing a target doesn't silently test the previous one. `.pristine.i64` is gitignored. """ + from __future__ import annotations import asyncio @@ -64,10 +65,12 @@ def fast_keys() -> None: if not hasattr(textual.app, "wait_for_idle"): # pragma: no cover raise RuntimeError( "textual.app.wait_for_idle is gone -- tests/_fixtures.fast_keys " - "needs updating for this Textual version") + "needs updating for this Textual version" + ) - async def _yield_instead_of_sleeping(min_sleep: float = 0.0, - max_sleep: float = 1.0) -> None: + async def _yield_instead_of_sleeping( + min_sleep: float = 0.0, max_sleep: float = 1.0 + ) -> None: await asyncio.sleep(0) async def _press(self, *keys: str) -> None: @@ -110,7 +113,7 @@ def synthetic(name: str, build) -> str: path = os.path.join(SYNTHETIC_DIR, name) data = build() if not os.path.exists(path) or open(path, "rb").read() != data: - with open(path, "wb") as fh: # content changed -> cache is stale + with open(path, "wb") as fh: # content changed -> cache is stale fh.write(data) for stale in (cache_path(path), path + ".i64"): if os.path.exists(stale): diff --git a/tests/run.py b/tests/run.py index 5d552bb..2fb093f 100755 --- a/tests/run.py +++ b/tests/run.py @@ -38,6 +38,7 @@ Usage:: Exit code is 0 only if every file selected ran and passed. """ + from __future__ import annotations import argparse @@ -80,14 +81,17 @@ def needs_ida(path: str) -> bool: if isinstance(target, ast.Name) and target.id == "NEEDS_IDA": value = ast.literal_eval(node.value) if not isinstance(value, bool): - raise Marker(f"{os.path.basename(path)}: " - f"NEEDS_IDA must be a bool, got {value!r}") + raise Marker( + f"{os.path.basename(path)}: " + f"NEEDS_IDA must be a bool, got {value!r}" + ) return value raise Marker( f"{os.path.basename(path)}: no NEEDS_IDA marker.\n" f" Add `NEEDS_IDA = True` (spawns a worker / drives the pilot) or\n" f" `NEEDS_IDA = False` (pure: stdlib, no IDA, runs anywhere) at module\n" - f" scope, so tests/run.py --fast knows whether it can run you.") + f" scope, so tests/run.py --fast knows whether it can run you." + ) def discover() -> list[tuple[str, bool]]: @@ -124,41 +128,70 @@ def tally(output: str) -> tuple[int, int] | None: def run_one(path: str, python: str, extra: list[str], echo: bool) -> dict: """Run one test file as a subprocess and summarise it.""" - name = os.path.basename(path)[len("test_"):-len(".py")] + name = os.path.basename(path)[len("test_") : -len(".py")] started = time.time() - proc = subprocess.run([python, path, *extra], cwd=ROOT, - capture_output=not echo, text=True) + proc = subprocess.run( + [python, path, *extra], cwd=ROOT, capture_output=not echo, text=True + ) took = time.time() - started out = "" if echo else (proc.stdout or "") + (proc.stderr or "") counts = tally(out) skipped = bool(_SKIP.search(out)) and (counts is None or counts == (0, 0)) return { - "name": name, "path": path, "code": proc.returncode, "took": took, + "name": name, + "path": path, + "code": proc.returncode, + "took": took, "passed": counts[0] if counts else 0, "failed": counts[1] if counts else 0, "counted": counts is not None, - "skipped": skipped, "output": out, + "skipped": skipped, + "output": out, } def main(argv: list[str]) -> int: ap = argparse.ArgumentParser( - prog="tests/run.py", description=__doc__, - formatter_class=argparse.RawDescriptionHelpFormatter) - ap.add_argument("only", nargs="*", metavar="SUBSTR", - help="only run test files whose name contains one of these") - ap.add_argument("--fast", action="store_true", - help="skip every file that needs IDA (seconds, runs anywhere)") - ap.add_argument("--ida-only", action="store_true", - help="only the files that need IDA") - ap.add_argument("--list", action="store_true", - help="show what would run, and whether it needs IDA") - ap.add_argument("-x", "--exitfirst", action="store_true", - help="stop after the first failing file") - ap.add_argument("-v", "--verbose", action="store_true", - help="stream each suite's output instead of capturing it") - ap.add_argument("--python", default=os.environ.get("IDATUI_PYTHON", DEFAULT_PY), - help=f"interpreter for the IDA suites (default {DEFAULT_PY})") + prog="tests/run.py", + description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + ap.add_argument( + "only", + nargs="*", + metavar="SUBSTR", + help="only run test files whose name contains one of these", + ) + ap.add_argument( + "--fast", + action="store_true", + help="skip every file that needs IDA (seconds, runs anywhere)", + ) + ap.add_argument( + "--ida-only", action="store_true", help="only the files that need IDA" + ) + ap.add_argument( + "--list", + action="store_true", + help="show what would run, and whether it needs IDA", + ) + ap.add_argument( + "-x", + "--exitfirst", + action="store_true", + help="stop after the first failing file", + ) + ap.add_argument( + "-v", + "--verbose", + action="store_true", + help="stream each suite's output instead of capturing it", + ) + ap.add_argument( + "--python", + default=os.environ.get("IDATUI_PYTHON", DEFAULT_PY), + help=f"interpreter for the IDA suites (default {DEFAULT_PY})", + ) args, extra = ap.parse_known_args(argv) try: @@ -190,10 +223,12 @@ def main(argv: list[str]) -> int: # an IDA file needs the interpreter that has textual + idapro. pure_py = sys.executable if any(ida for _, ida in selected) and not os.path.exists(args.python): - print(f"error: {args.python} not found — the IDA suites need an " - f"interpreter with textual + idapro.\n" - f" Pass --python, set $IDATUI_PYTHON, or use --fast.", - file=sys.stderr) + print( + f"error: {args.python} not found — the IDA suites need an " + f"interpreter with textual + idapro.\n" + f" Pass --python, set $IDATUI_PYTHON, or use --fast.", + file=sys.stderr, + ) return 2 results = [] @@ -226,7 +261,7 @@ def main(argv: list[str]) -> int: elif r["code"] != 0 or r["failed"]: state = "\033[31mFAIL\033[0m" elif not r["counted"]: - state = "\033[33m ? \033[0m" # exit 0 but printed no tally + state = "\033[33m ? \033[0m" # exit 0 but printed no tally else: state = "\033[32m ok \033[0m" detail = f"{r['passed']:4d} passed" diff --git a/tests/test_blob_ui.py b/tests/test_blob_ui.py index d4740d0..7c3aee5 100644 --- a/tests/test_blob_ui.py +++ b/tests/test_blob_ui.py @@ -19,12 +19,12 @@ import tempfile sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +from _fixtures import fast_keys, staged, synthetic # noqa: E402 from textual.widgets import Input, Static # noqa: E402 from idatui.app import ConfirmScreen, IdaTui, ListingView # noqa: E402 -from _fixtures import fast_keys, staged, synthetic # noqa: E402 -fast_keys() # ~85ms -> ~2ms per keypress; see _fixtures.fast_keys +fast_keys() # ~85ms -> ~2ms per keypress; see _fixtures.fast_keys from idatui._sync import settle # noqa: E402 PASS = FAIL = 0 @@ -66,14 +66,19 @@ def _blob_bytes() -> bytes: functions" is asserted below. """ import random + data = bytearray(random.Random(0xB10BCAFE).randbytes(64 * 1024)) - # -parm puts IDA in AArch64 mode, so these are A64 encodings; the ARM32 - # spelling of a nop (0xE1A00000) is NOT decodable there and made this - # test fail for a reason that had nothing to do with what it checks. - for k, insn in enumerate((0xD503201F, # nop - 0xD503201F, # nop - 0xD65F03C0)): # ret <- the run must stop here - data[PLANTED + k * 4:PLANTED + k * 4 + 4] = insn.to_bytes(4, "little") + # -parm puts IDA in AArch64 mode, so these are A64 encodings; the ARM32 + # spelling of a nop (0xE1A00000) is NOT decodable there and made this + # test fail for a reason that had nothing to do with what it checks. + for k, insn in enumerate( + ( + 0xD503201F, # nop + 0xD503201F, # nop + 0xD65F03C0, + ) + ): # ret <- the run must stop here + data[PLANTED + k * 4 : PLANTED + k * 4 + 4] = insn.to_bytes(4, "little") return bytes(data) @@ -109,27 +114,45 @@ async def run() -> int: # happens AFTER a described blob turns out to contain nothing. app = _blob_app(blob) async with app.run_test(size=(140, 44)) as pilot: - ok = await wait(lambda: app._func_index is not None - and app._func_index.complete, pilot) + ok = await wait( + lambda: app._func_index is not None and app._func_index.complete, pilot + ) check("a random blob still finishes loading", ok) - check("and really has no functions", len(app._func_index) == 0, - f"n={len(app._func_index)}") + check( + "and really has no functions", + len(app._func_index) == 0, + f"n={len(app._func_index)}", + ) landed = await wait(lambda: app._cur is not None, pilot, 60) - check("it lands somewhere instead of leaving empty panes", landed, - f"cur={app._cur}") + check( + "it lands somewhere instead of leaving empty panes", + landed, + f"cur={app._cur}", + ) lst = app.query_one(ListingView) - check("the listing actually has rows to show", - lst.total > 0, f"total={lst.total}") - check("landed at the image base", app._cur is not None - and app._cur.ea == 0x4000, f"{app._cur.ea if app._cur else None:#x}") + check( + "the listing actually has rows to show", + lst.total > 0, + f"total={lst.total}", + ) + check( + "landed at the image base", + app._cur is not None and app._cur.ea == 0x4000, + f"{app._cur.ea if app._cur else None:#x}", + ) status = str(app.query_one("#status", Static).render()) - check("the status says there are no functions (not 'still loading')", - "no functions" in status and "still loading" not in status, - status[:90]) - check("and points at the likely cause", - "processor" in status and "Ctrl+L" in status, status[:90]) + check( + "the status says there are no functions (not 'still loading')", + "no functions" in status and "still loading" not in status, + status[:90], + ) + check( + "and points at the likely cause", + "processor" in status and "Ctrl+L" in status, + status[:90], + ) # The hint is a property of the database, so it must survive moving # around — an earlier version wrote it once and the next status @@ -139,8 +162,9 @@ async def run() -> int: await pilot.press("down") await settle(app) status2 = str(app.query_one("#status", Static).render()) - check("the hint survives navigating", "no functions" in status2, - status2[:90]) + check( + "the hint survives navigating", "no functions" in status2, status2[:90] + ) # -- byte-granular carving ---------------------------------- # # An undefined run arrives as ONE head ("db N dup(?)"). It has to @@ -148,47 +172,64 @@ async def run() -> int: # press `c` at, which is how IDA works and the only way to find an # instruction stream that doesn't start at the run's first byte. m = lst.model - check("an undefined run presents one row per byte", - m.loaded() >= 4096 and len(m._heads) < 64, - f"rows={m.loaded()} physical heads={len(m._heads)}") + check( + "an undefined run presents one row per byte", + m.loaded() >= 4096 and len(m._heads) < 64, + f"rows={m.loaded()} physical heads={len(m._heads)}", + ) rows = m.window(0, 4) - check("each row is a single addressable byte", - [h.ea for h in rows] == [0x4000, 0x4001, 0x4002, 0x4003] - and all(h.size == 1 for h in rows), - f"{[(hex(h.ea), h.size) for h in rows]}") - check("and shows its value, not a placeholder", - all(h.text.startswith("db ") and "dup" not in h.text - for h in rows), f"{[h.text for h in rows]}") + check( + "each row is a single addressable byte", + [h.ea for h in rows] == [0x4000, 0x4001, 0x4002, 0x4003] + and all(h.size == 1 for h in rows), + f"{[(hex(h.ea), h.size) for h in rows]}", + ) + check( + "and shows its value, not a placeholder", + all(h.text.startswith("db ") and "dup" not in h.text for h in rows), + f"{[h.text for h in rows]}", + ) # The point of all this: land on an arbitrary byte and convert it. i = m.index_of_ea(0x4021) - check("an address inside the run resolves to its own row", - i >= 0 and m.get(i).ea == 0x4021, - f"row={i} ea={m.get(i).ea if i >= 0 else None}") + check( + "an address inside the run resolves to its own row", + i >= 0 and m.get(i).ea == 0x4021, + f"row={i} ea={m.get(i).ea if i >= 0 else None}", + ) - target = 0x4000 + PLANTED # a NOP we put there ourselves + target = 0x4000 + PLANTED # a NOP we put there ourselves lst.cursor = m.index_of_ea(target) lst._scroll_cursor_into_view() await settle(app, lambda: lst._cursor_ea() == target) - check("the cursor sits on the byte we aimed at", - lst._cursor_ea() == target, - f"{lst._cursor_ea():#x} want {target:#x}") + check( + "the cursor sits on the byte we aimed at", + lst._cursor_ea() == target, + f"{lst._cursor_ea():#x} want {target:#x}", + ) await pilot.press("c") # settle(), not a fixed sleep AND not a bare predicate: an edit can # look done for a moment and then be replaced when a queued listing # rebuild lands, so the gate has to be "the row is code AND the app # has stopped working". settle() is the same helper the app's own # RPC layer uses, so tests and driver agree on what "done" means. - await settle(app, lambda: (lambda h: h is not None and h.kind == "code")( - head_at(lst, target)), timeout=30) + await settle( + app, + lambda: (lambda h: h is not None and h.kind == "code")( + head_at(lst, target) + ), + timeout=30, + ) # Re-read the model: defining an item rebuilds it, and holding the # old object shows pre-edit rows -- which looks exactly like the # edit silently failing. m = lst.model h = head_at(lst, target) - check("`c` on a chosen byte carves an instruction there", - h is not None and h.kind == "code", - f"kind={h.kind if h else None} text={h.text if h else None!r}") + check( + "`c` on a chosen byte carves an instruction there", + h is not None and h.kind == "code", + f"kind={h.kind if h else None} text={h.text if h else None!r}", + ) if h is not None and h.kind == "code": print(f" carved {target:#x}: {h.text}") # `c` runs until something stops it, like IDA — one instruction @@ -197,21 +238,32 @@ async def run() -> int: # four and stop AT the ret, not run on into the random bytes # after it. run = [m.get(m.index_of_ea(target + k * 4)) for k in range(3)] - check("`c` keeps going until control flow ends", - all(x is not None and x.kind == "code" for x in run), - f"{[(hex(x.ea), x.kind) for x in run if x]}") - check("and stops at the ret instead of running into junk", - m.get(m.index_of_ea(target + 12)).kind == "unknown", - f"{m.get(m.index_of_ea(target + 12)).text!r}") + check( + "`c` keeps going until control flow ends", + all(x is not None and x.kind == "code" for x in run), + f"{[(hex(x.ea), x.kind) for x in run if x]}", + ) + check( + "and stops at the ret instead of running into junk", + m.get(m.index_of_ea(target + 12)).kind == "unknown", + f"{m.get(m.index_of_ea(target + 12)).text!r}", + ) status = str(app.query_one("#status", Static).render()) - check("the status reports what the run did", - "3 instructions" in status and "control flow" in status, - status[:80]) - check("the carved row spans the instruction, not one byte", - h.size == 4, f"size={h.size}") - check("bytes before it stay individually addressable", - m.get(m.index_of_ea(target - 1)).size == 1 - and m.get(m.index_of_ea(target - 1)).ea == target - 1) + check( + "the status reports what the run did", + "3 instructions" in status and "control flow" in status, + status[:80], + ) + check( + "the carved row spans the instruction, not one byte", + h.size == 4, + f"size={h.size}", + ) + check( + "bytes before it stay individually addressable", + m.get(m.index_of_ea(target - 1)).size == 1 + and m.get(m.index_of_ea(target - 1)).ea == target - 1, + ) # -- an edit must not move the view -------------------------- # # Every mutation rebuilds the model, and row indices don't survive @@ -239,14 +291,17 @@ async def run() -> int: # (The listing re-renders its text lazily, so the comment is not # necessarily visible in model rows the moment the worker returns -- # which is why this waits for the app, not for the text.) - await settle(app, lambda: not app.query_one("#comment", Input).display, - timeout=30) - check("commenting leaves the view where it was", - lst.model.get(round(lst.scroll_offset.y)).ea == ctop - and lst._cursor_ea() == ccur, - f"top {ctop:#x} -> " - f"{lst.model.get(round(lst.scroll_offset.y)).ea:#x}, " - f"cursor {ccur:#x} -> {lst._cursor_ea():#x}") + await settle( + app, lambda: not app.query_one("#comment", Input).display, timeout=30 + ) + check( + "commenting leaves the view where it was", + lst.model.get(round(lst.scroll_offset.y)).ea == ctop + and lst._cursor_ea() == ccur, + f"top {ctop:#x} -> " + f"{lst.model.get(round(lst.scroll_offset.y)).ea:#x}, " + f"cursor {ccur:#x} -> {lst._cursor_ea():#x}", + ) # -- carving must not move the view -------------------------- # # Defining code collapses rows (four byte rows become one @@ -259,8 +314,11 @@ async def run() -> int: await settle(app, lambda: lst._cursor_ea() == far) top_before = lst.model.get(round(lst.scroll_offset.y)).ea cur_before = lst._cursor_ea() - check("scrolled somewhere with rows above us", - round(lst.scroll_offset.y) > 0, f"top={lst.scroll_offset.y}") + check( + "scrolled somewhere with rows above us", + round(lst.scroll_offset.y) > 0, + f"top={lst.scroll_offset.y}", + ) await pilot.press("c") # No predicate here on purpose: this spot is random data, so the # carve may legitimately produce nothing and "the row became code" @@ -270,54 +328,80 @@ async def run() -> int: await settle(app, timeout=30) m2 = lst.model top_after = m2.get(round(lst.scroll_offset.y)).ea - check("carving leaves the scroll position where it was", - top_after == top_before, - f"{top_before:#x} -> {top_after:#x}") - check("and leaves the cursor on the same address", - lst._cursor_ea() == cur_before, - f"{cur_before:#x} -> {lst._cursor_ea():#x}") + check( + "carving leaves the scroll position where it was", + top_after == top_before, + f"{top_before:#x} -> {top_after:#x}", + ) + check( + "and leaves the cursor on the same address", + lst._cursor_ea() == cur_before, + f"{cur_before:#x} -> {lst._cursor_ea():#x}", + ) # -- `p` after carving: the rest of the app must notice ------ # # The "no functions" hint was latched at load and only cleared on a # reload, so it kept telling you the processor/base were wrong long # after you'd defined a function. The function index was never # rebuilt either, which meant Ctrl+N couldn't find what `p` made. - check("no functions yet, and the hint says so", - len(app._func_index) == 0 - and "no functions" in str(app.query_one("#status", Static).render()), - f"n={len(app._func_index)}") + check( + "no functions yet, and the hint says so", + len(app._func_index) == 0 + and "no functions" in str(app.query_one("#status", Static).render()), + f"n={len(app._func_index)}", + ) lst.cursor = lst.model.index_of_ea(target) lst._scroll_cursor_into_view() await settle(app, lambda: lst._cursor_ea() == target) mp = lst.model await pilot.press("p") - await wait(lambda: lst.model is not mp and lst.model is not None, - pilot, 40) - await wait(lambda: app._func_index is not None - and len(app._func_index) > 0, pilot, 60) - check("`p` creates a function the index can see", - len(app._func_index) == 1, f"n={len(app._func_index)}") + await wait(lambda: lst.model is not mp and lst.model is not None, pilot, 40) + await wait( + lambda: app._func_index is not None and len(app._func_index) > 0, + pilot, + 60, + ) + check( + "`p` creates a function the index can see", + len(app._func_index) == 1, + f"n={len(app._func_index)}", + ) status = str(app.query_one("#status", Static).render()) - check("and the stale 'no functions' hint is gone", - "no functions" not in status, status[:90]) - check("the status names the function it made", - "created function" in status, status[:90]) + check( + "and the stale 'no functions' hint is gone", + "no functions" not in status, + status[:90], + ) + check( + "the status names the function it made", + "created function" in status, + status[:90], + ) await pilot.press("ctrl+l") - opened = await wait(lambda: isinstance(app.screen, ConfirmScreen), pilot, 20) - check("Ctrl+L offers to reload with different options", opened, - f"screen={type(app.screen).__name__}") + opened = await wait( + lambda: isinstance(app.screen, ConfirmScreen), pilot, 20 + ) + check( + "Ctrl+L offers to reload with different options", + opened, + f"screen={type(app.screen).__name__}", + ) if opened: note = str(app.screen.query_one("#confirm-note", Static).render()) # We just made a function, so it must NOT claim nothing is lost — # reloading throws the database away and that is now a real cost. - check("the confirmation counts what would be lost", - "1 function," in note and "nothing is lost" not in note, - note[:80]) + check( + "the confirmation counts what would be lost", + "1 function," in note and "nothing is lost" not in note, + note[:80], + ) await pilot.press("escape") await settle(app, lambda: not isinstance(app.screen, ConfirmScreen)) - check("declining leaves the binary open", - not isinstance(app.screen, ConfirmScreen) and app._cur is not None) + check( + "declining leaves the binary open", + not isinstance(app.screen, ConfirmScreen) and app._cur is not None, + ) print(f"\n{PASS} passed, {FAIL} failed") return 1 if FAIL else 0 diff --git a/tests/test_diag.py b/tests/test_diag.py index e3b966a..d1b480a 100644 --- a/tests/test_diag.py +++ b/tests/test_diag.py @@ -3,6 +3,7 @@ Pure: no IDA, no worker, no Textual. """ + from __future__ import annotations import os @@ -41,10 +42,16 @@ def t_swallow_keeps_going(): r = diag.recent() check("the error is recorded", len(r) == 1, str(r)) check("with what was being attempted", r[0]["what"] == "a thing", str(r[0])) - check("and the exception type and message", - r[0]["error"] == "ValueError: nope", r[0]["error"]) - check("and where it was actually raised", - r[0]["where"].startswith("test_diag.py:"), r[0]["where"]) + check( + "and the exception type and message", + r[0]["error"] == "ValueError: nope", + r[0]["error"], + ) + check( + "and where it was actually raised", + r[0]["where"].startswith("test_diag.py:"), + r[0]["where"], + ) def t_reraise(): @@ -61,8 +68,11 @@ def t_reraise(): check("reraise lets the listed type through", False, "not raised") except Wanted: check("reraise lets the listed type through", True) - check("and a reraised error is not recorded twice", - diag.recent() == [], str(diag.recent())) + check( + "and a reraised error is not recorded twice", + diag.recent() == [], + str(diag.recent()), + ) with diag.swallow("still swallows others", reraise=(Wanted,)): raise ValueError("other") check("other types are still swallowed", len(diag.recent()) == 1) @@ -74,12 +84,21 @@ def t_ring_is_bounded(): diag.note(f"item {i}", RuntimeError(str(i))) r = diag.recent(1000) check("the ring is bounded", len(r) == diag._MAX, f"{len(r)}") - check("it keeps the NEWEST entries", - r[-1]["what"] == f"item {diag._MAX + 24}", r[-1]["what"]) - check("recent(n) returns the last n, newest last", - [e["what"] for e in diag.recent(3)] - == [f"item {diag._MAX + 22}", f"item {diag._MAX + 23}", - f"item {diag._MAX + 24}"], str(diag.recent(3))) + check( + "it keeps the NEWEST entries", + r[-1]["what"] == f"item {diag._MAX + 24}", + r[-1]["what"], + ) + check( + "recent(n) returns the last n, newest last", + [e["what"] for e in diag.recent(3)] + == [ + f"item {diag._MAX + 22}", + f"item {diag._MAX + 23}", + f"item {diag._MAX + 24}", + ], + str(diag.recent(3)), + ) def t_log_file(): @@ -95,8 +114,11 @@ def t_log_file(): body = open(path, encoding="utf-8").read() check("the log records what was attempted", "logged thing" in body, body[:200]) check("and the error", "KeyError" in body, body[:200]) - check("and a traceback, which the ring doesn't carry", - "Traceback" in body and "t_log_file" in body, body[:300]) + check( + "and a traceback, which the ring doesn't carry", + "Traceback" in body and "t_log_file" in body, + body[:300], + ) def t_log_is_off_by_default(): @@ -104,8 +126,10 @@ def t_log_is_off_by_default(): os.environ.pop("IDATUI_LOG", None) with diag.swallow("unlogged"): raise ValueError("x") - check("without $IDATUI_LOG nothing is written, but the ring still has it", - len(diag.recent()) == 1) + check( + "without $IDATUI_LOG nothing is written, but the ring still has it", + len(diag.recent()) == 1, + ) def t_broken_log_path_is_harmless(): @@ -116,8 +140,7 @@ def t_broken_log_path_is_harmless(): with diag.swallow("still fine"): raise ValueError("boom") check("an unwritable log path doesn't raise", True) - check("and the error is still recorded in the ring", - len(diag.recent()) == 1) + check("and the error is still recorded in the ring", len(diag.recent()) == 1) finally: os.environ.pop("IDATUI_LOG", None) @@ -128,41 +151,55 @@ def t_env_read_per_call(): diag.clear() with tempfile.TemporaryDirectory() as d: path = os.path.join(d, "late.log") - os.environ["IDATUI_LOG"] = path # set AFTER import + os.environ["IDATUI_LOG"] = path # set AFTER import try: diag.log("hello") finally: os.environ.pop("IDATUI_LOG", None) - check("a log path set after import is honoured", - os.path.exists(path) and "hello" in open(path).read()) + check( + "a log path set after import is honoured", + os.path.exists(path) and "hello" in open(path).read(), + ) def t_thread_safe(): diag.clear() + def go(n): for i in range(40): diag.note(f"t{n}-{i}", RuntimeError("x")) + ts = [threading.Thread(target=go, args=(n,)) for n in range(6)] for t in ts: t.start() for t in ts: t.join(10) r = diag.recent(1000) - check("concurrent notes don't corrupt the ring", - len(r) == diag._MAX and all("what" in e for e in r), f"{len(r)}") - check("the recording thread is captured", - all(e["thread"] for e in r)) + check( + "concurrent notes don't corrupt the ring", + len(r) == diag._MAX and all("what" in e for e in r), + f"{len(r)}", + ) + check("the recording thread is captured", all(e["thread"] for e in r)) def main() -> int: - for fn in (t_swallow_keeps_going, t_reraise, t_ring_is_bounded, t_log_file, - t_log_is_off_by_default, t_broken_log_path_is_harmless, - t_env_read_per_call, t_thread_safe): + for fn in ( + t_swallow_keeps_going, + t_reraise, + t_ring_is_bounded, + t_log_file, + t_log_is_off_by_default, + t_broken_log_path_is_harmless, + t_env_read_per_call, + t_thread_safe, + ): print(f"\n{fn.__name__}") try: fn() except Exception as e: # noqa: BLE001 import traceback + check(f"{fn.__name__} did not crash", False, f"{type(e).__name__}: {e}") traceback.print_exc() diag.clear() diff --git a/tests/test_findings.py b/tests/test_findings.py index cf813c2..205b005 100644 --- a/tests/test_findings.py +++ b/tests/test_findings.py @@ -18,7 +18,12 @@ sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) from idatui.domain import Comment, NamedItem, Struct # noqa: E402 from idatui.findings import ( # noqa: E402 - Findings, default_path, from_loader, gather, is_dummy, render, + Findings, + default_path, + from_loader, + gather, + is_dummy, + render, ) PASS = FAIL = 0 @@ -36,34 +41,62 @@ def check(name, cond, detail=""): def sample() -> Findings: return Findings( - binary="echo", path="/tmp/echo", + binary="echo", + path="/tmp/echo", sections=[(0x1000, 0x2000, ".text"), (0x2000, 0x2100, ".data")], n_functions=128, comments=[ - Comment(addr=0x1100, text="length is attacker controlled", - line="mov edi, [rbp+len]", func="parse", func_addr=0x1000), - Comment(addr=0x1010, text="entry", line="push rbp", - func="parse", func_addr=0x1000), - Comment(addr=0x1000, text="parses the header", whole_func=True, - func="parse", func_addr=0x1000), + Comment( + addr=0x1100, + text="length is attacker controlled", + line="mov edi, [rbp+len]", + func="parse", + func_addr=0x1000, + ), + Comment( + addr=0x1010, + text="entry", + line="push rbp", + func="parse", + func_addr=0x1000, + ), + Comment( + addr=0x1000, + text="parses the header", + whole_func=True, + func="parse", + func_addr=0x1000, + ), Comment(addr=0x2004, text="magic", line="dd 0DEADBEEFh"), # What the ELF loader writes into every database, through the very # same set_cmt a person uses. - Comment(addr=0x4, text="File class: 64-bit", line="db 2", - seg="LOAD"), + Comment(addr=0x4, text="File class: 64-bit", line="db 2", seg="LOAD"), ], names=[ - NamedItem(addr=0x1000, name="parse", is_func=True, size=0x120, - proto="int __fastcall parse(char *)"), + NamedItem( + addr=0x1000, + name="parse", + is_func=True, + size=0x120, + proto="int __fastcall parse(char *)", + ), NamedItem(addr=0x1200, name="sub_1200", is_func=True, size=0x30), NamedItem(addr=0x2004, name="hdr_magic", seg=".data"), NamedItem(addr=0x1400, name="memcpy", is_func=True, size=0x40), NamedItem(addr=0x390, name="elf_gnu_hash_nbuckets", seg="LOAD"), ], - types=[(Struct(name="hdr", size=0x10, is_union=False, members=3, - ordinal=42), "struct hdr\n{\n int magic;\n};\n"), - (Struct(name="Elf64_Dyn", size=0x10, is_union=False, members=2, - ordinal=3), "struct Elf64_Dyn\n{\n int d_tag;\n};\n")], + types=[ + ( + Struct(name="hdr", size=0x10, is_union=False, members=3, ordinal=42), + "struct hdr\n{\n int magic;\n};\n", + ), + ( + Struct( + name="Elf64_Dyn", size=0x10, is_union=False, members=2, ordinal=3 + ), + "struct Elf64_Dyn\n{\n int d_tag;\n};\n", + ), + ], linked={"memcpy"}, stripped=True, ) @@ -74,9 +107,11 @@ def main() -> int: check("the report names the binary", doc.startswith("# Findings — echo"), doc[:40]) # 1 function, not 3: sub_1200 is IDA's invention and memcpy is the linker's. - check("the summary counts only what a person contributed", - "1 named functions · 1 named data · 4 comments · 2 local types" in doc, - doc.splitlines()[2] if len(doc.splitlines()) > 2 else "") + check( + "the summary counts only what a person contributed", + "1 named functions · 1 named data · 4 comments · 2 local types" in doc, + doc.splitlines()[2] if len(doc.splitlines()) > 2 else "", + ) # A name IDA invented is not a finding, and neither is one the linker gave. check("dummy names are excluded", "sub_1200" not in doc) @@ -86,75 +121,111 @@ def main() -> int: # The loader annotates every database it makes; none of it is a finding. check("the loader's own comments are left out", "File class" not in doc) check("the loader's own names are left out", "elf_gnu_hash" not in doc) - check("but the report says how many it dropped", - "4 annotations left out as the loader's own" in doc, # 1 comment + 3 names - [l for l in doc.splitlines() if "left out" in l]) - check("from_loader knows both shapes", - from_loader("LOAD") and from_loader("", "elf_gnu_hash_x") - and not from_loader(".text", "parse")) - check("is_dummy knows the shapes IDA invents", - all(is_dummy(n) for n in ("sub_1234", "loc_A0", "unk_4000", "j_free")) - and not any(is_dummy(n) for n in ("parse", "sub_parse", "main", "")), - "") + check( + "but the report says how many it dropped", + "4 annotations left out as the loader's own" in doc, # 1 comment + 3 names + [l for l in doc.splitlines() if "left out" in l], + ) + check( + "from_loader knows both shapes", + from_loader("LOAD") + and from_loader("", "elf_gnu_hash_x") + and not from_loader(".text", "parse"), + ) + check( + "is_dummy knows the shapes IDA invents", + all(is_dummy(n) for n in ("sub_1234", "loc_A0", "unk_4000", "j_free")) + and not any(is_dummy(n) for n in ("parse", "sub_parse", "main", "")), + "", + ) # Comments lead, grouped by function, address-ordered within a group. - check("comments come before the name tables", - doc.index("## Comments") < doc.index("## Named functions")) - body = doc[doc.index("## Comments"):doc.index("## Named functions")] + check( + "comments come before the name tables", + doc.index("## Comments") < doc.index("## Named functions"), + ) + body = doc[doc.index("## Comments") : doc.index("## Named functions")] check("comments are grouped under their function", "### `parse`" in body) - check("a commentless region is grouped separately", - "### outside any function" in body) - check("comments are ordered by address inside a group", - body.index("0x1010") < body.index("0x1100")) - check("a function comment says that is what it is", - "*whole function*: parses the header" in body) - check("an instruction comment carries the line it annotates", - "`mov edi, [rbp+len]`" in body) + check( + "a commentless region is grouped separately", "### outside any function" in body + ) + check( + "comments are ordered by address inside a group", + body.index("0x1010") < body.index("0x1100"), + ) + check( + "a function comment says that is what it is", + "*whole function*: parses the header" in body, + ) + check( + "an instruction comment carries the line it annotates", + "`mov edi, [rbp+len]`" in body, + ) # Types: newest ordinal first, because that is the one you just wrote. - types = doc[doc.index("## Local types"):] - check("your newest type is first", - types.index("hdr") < types.index("Elf64_Dyn")) + types = doc[doc.index("## Local types") :] + check("your newest type is first", types.index("hdr") < types.index("Elf64_Dyn")) check("type source is fenced as C", "```c\nstruct hdr" in types) # Escaping. - hostile = Findings(binary="x", comments=[ - Comment(addr=1, text="a | b", line="mov | rax"), - ], names=[NamedItem(addr=2, name="a|b")]) + hostile = Findings( + binary="x", + comments=[ + Comment(addr=1, text="a | b", line="mov | rax"), + ], + names=[NamedItem(addr=2, name="a|b")], + ) hdoc = render(hostile) check("a pipe cannot break a table row", "a\\|b" in hdoc, hdoc) check("a pipe in a comment is escaped too", "a \\| b" in hdoc) # The empty database must still produce a document that says something. empty = render(Findings(binary="nothing")) - check("an empty report is still a document", - empty.startswith("# Findings — nothing") and "## Comments" in empty) + check( + "an empty report is still a document", + empty.startswith("# Findings — nothing") and "## Comments" in empty, + ) check("and it says why it is empty", "Comments are the part" in empty) - check("an empty report has no dangling type section", - "## Local types" not in empty) + check("an empty report has no dangling type section", "## Local types" not in empty) # Provenance must be stated, not implied. Without a journal the report is a # scan and says so; with one it is exactly what idatui recorded doing. - scanned = render(Findings(binary="x", stripped=False, - names=[NamedItem(addr=1, name="main", is_func=True)])) - check("a scanned report admits it cannot know who wrote what", - "**source**: a scan of the database" in scanned - and "include its work as well as yours" in scanned) - check("and warns when the binary brought its own symbols", - "include ones it shipped with" in scanned) + scanned = render( + Findings( + binary="x", + stripped=False, + names=[NamedItem(addr=1, name="main", is_func=True)], + ) + ) + check( + "a scanned report admits it cannot know who wrote what", + "**source**: a scan of the database" in scanned + and "include its work as well as yours" in scanned, + ) + check( + "and warns when the binary brought its own symbols", + "include ones it shipped with" in scanned, + ) j = sample() j.recorded = {0x1100, 0x1000} j.n_recorded = 7 jdoc = render(j) - check("a journalled report says so", "idatui's edit journal" in jdoc - and "7 recorded edits" in jdoc, "") - jbody = jdoc[jdoc.index("## Comments"):jdoc.index("## Named functions")] - check("and lists only the comments it recorded", - "length is attacker controlled" in jbody and "0x2004" not in jbody, - jbody) - check("a journalled report drops names it did not record", - "`parse`" in jdoc and "hdr_magic" not in jdoc) + check( + "a journalled report says so", + "idatui's edit journal" in jdoc and "7 recorded edits" in jdoc, + "", + ) + jbody = jdoc[jdoc.index("## Comments") : jdoc.index("## Named functions")] + check( + "and lists only the comments it recorded", + "length is attacker controlled" in jbody and "0x2004" not in jbody, + jbody, + ) + check( + "a journalled report drops names it did not record", + "`parse`" in jdoc and "hdr_magic" not in jdoc, + ) # -- gather ------------------------------------------------------------- # class FakeProgram: @@ -162,8 +233,10 @@ def main() -> int: return [(0x1000, 0x2000, ".text")] def annotations(self, limit=4000): - return ([Comment(addr=1, text="hi")], - [NamedItem(addr=1, name="parse", is_func=True)]) + return ( + [Comment(addr=1, text="hi")], + [NamedItem(addr=1, name="parse", is_func=True)], + ) def linkage(self): return ([], []) @@ -180,11 +253,15 @@ def main() -> int: f = gather(FakeProgram(), "/tmp/echo") check("gather reads the annotations", len(f.comments) == 1 and len(f.names) == 1) check("gather takes the binary name from the path", f.binary == "echo", f.binary) - check("a failing backend degrades the report instead of raising", - f.n_functions == 0 and f.types == [] and "# Findings" in render(f)) + check( + "a failing backend degrades the report instead of raising", + f.n_functions == 0 and f.types == [] and "# Findings" in render(f), + ) - check("the default path sits beside the binary", - default_path("/tmp/echo") == "/tmp/echo.findings.md") + check( + "the default path sits beside the binary", + default_path("/tmp/echo") == "/tmp/echo.findings.md", + ) print(f"\n{PASS} passed, {FAIL} failed") return 1 if FAIL else 0 diff --git a/tests/test_formats.py b/tests/test_formats.py index c0ad542..05d98ba 100644 --- a/tests/test_formats.py +++ b/tests/test_formats.py @@ -10,8 +10,12 @@ import tempfile sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) -from idatui.formats import (PROCESSORS, load_args, # noqa: E402 - needs_load_options, sniff) +from idatui.formats import ( # noqa: E402 + PROCESSORS, + load_args, + needs_load_options, + sniff, +) PASS = FAIL = 0 @@ -28,6 +32,7 @@ def check(name, ok, detail=""): def main() -> int: with tempfile.TemporaryDirectory() as tmp: + def w(name, data): p = os.path.join(tmp, name) with open(p, "wb") as f: @@ -35,67 +40,102 @@ def main() -> int: return p # -- formats IDA can load on its own: never interrupt the user -------- # - for name, head in (("elf", b"\x7fELF\x02\x01\x01"), ("pe", b"MZ\x90\x00"), - ("macho", b"\xcf\xfa\xed\xfe"), ("dex", b"dex\n035\x00"), - ("wasm", b"\x00asm\x01\x00")): + for name, head in ( + ("elf", b"\x7fELF\x02\x01\x01"), + ("pe", b"MZ\x90\x00"), + ("macho", b"\xcf\xfa\xed\xfe"), + ("dex", b"dex\n035\x00"), + ("wasm", b"\x00asm\x01\x00"), + ): p = w(name, head + bytes(60)) - check(f"{name} is recognised (no dialog)", - sniff(p) is not None and not needs_load_options(p), f"{sniff(p)}") + check( + f"{name} is recognised (no dialog)", + sniff(p) is not None and not needs_load_options(p), + f"{sniff(p)}", + ) # -- the case this exists for ----------------------------------------- # blob = w("fw.bin", bytes(range(256)) * 4) - check("a headerless blob is not recognised (ask)", - sniff(blob) is None and needs_load_options(blob)) + check( + "a headerless blob is not recognised (ask)", + sniff(blob) is None and needs_load_options(blob), + ) # Intel HEX / S-records are text containers IDA does load. hexf = w("f.hex", b":10010000214601360121470136007EFE09D2190140\n") check("Intel HEX is recognised", sniff(hexf) == "Intel HEX", f"{sniff(hexf)}") - srec = w("f.s19", b"S00600004844521B\nS1130000285F245F2212226A000424290008237C\n") - check("Motorola S-records are recognised", - sniff(srec) == "Motorola S-record", f"{sniff(srec)}") + srec = w( + "f.s19", b"S00600004844521B\nS1130000285F245F2212226A000424290008237C\n" + ) + check( + "Motorola S-records are recognised", + sniff(srec) == "Motorola S-record", + f"{sniff(srec)}", + ) # A binary that merely STARTS with ':' is not Intel HEX. Text formats are # only accepted when the whole head is printable, or half the firmware in # the world gets mis-detected on one byte. colon = w("colon.bin", b":\x00\xff\xfe\x01\x02" + bytes(58)) - check("a blob starting with ':' is not mistaken for Intel HEX", - sniff(colon) is None, f"{sniff(colon)}") + check( + "a blob starting with ':' is not mistaken for Intel HEX", + sniff(colon) is None, + f"{sniff(colon)}", + ) check("an empty file is not recognised", sniff(w("empty", b"")) is None) - check("a missing file never asks (nothing to load)", - not needs_load_options(os.path.join(tmp, "nope"))) + check( + "a missing file never asks (nothing to load)", + not needs_load_options(os.path.join(tmp, "nope")), + ) check("a directory never asks", not needs_load_options(tmp)) # -- switch construction ------------------------------------------------- # # -b is in PARAGRAPHS: 0x8000000 >> 4 == 0x800000. Getting this wrong loads # the image 16x off and every address in the database is wrong. - check("base is converted to paragraphs", - load_args("arm", 0x8000000) == "-parm -b800000", - load_args("arm", 0x8000000)) + check( + "base is converted to paragraphs", + load_args("arm", 0x8000000) == "-parm -b800000", + load_args("arm", 0x8000000), + ) check("processor alone", load_args("mipsb") == "-pmipsb", load_args("mipsb")) check("base alone", load_args("", 0x10000) == "-b1000", load_args("", 0x10000)) - check("base 0 emits no switch (it's the default)", - load_args("arm", 0) == "-parm", load_args("arm", 0)) + check( + "base 0 emits no switch (it's the default)", + load_args("arm", 0) == "-parm", + load_args("arm", 0), + ) check("nothing in, nothing out", load_args() == "") - check("extra switches pass through", - load_args("arm", 0, "-T binary") == "-parm -T binary") + check( + "extra switches pass through", + load_args("arm", 0, "-T binary") == "-parm -T binary", + ) - check("the processor list leads with the common targets", - [n for n, _ in PROCESSORS[:2]] == ["arm", "arm:ARMv7-A"], - f"{[n for n, _ in PROCESSORS[:3]]}") + check( + "the processor list leads with the common targets", + [n for n, _ in PROCESSORS[:2]] == ["arm", "arm:ARMv7-A"], + f"{[n for n, _ in PROCESSORS[:3]]}", + ) # 32-bit ARM has to be offered SEPARATELY from bare 'arm', which gives a # 64-bit database. That isn't cosmetic: Hex-Rays refuses a 32-bit function # in a 64-bit database, and Thumb doesn't exist in AArch64 at all, so a # firmware image loaded as plain 'arm' can never be decompiled — and the # database's bitness cannot be corrected after load. - check("a 32-bit ARM variant is offered", - any(n.startswith("arm:ARMv") for n, _ in PROCESSORS), - f"{[n for n, _ in PROCESSORS if n.startswith('arm')]}") - check("and the labels say which is 32- vs 64-bit", - all(("32-bit" in d or "64-bit" in d) - for n, d in PROCESSORS if n == "arm" or n.startswith("arm:")), - f"{[(n, d) for n, d in PROCESSORS if n.startswith('arm')]}") + check( + "a 32-bit ARM variant is offered", + any(n.startswith("arm:ARMv") for n, _ in PROCESSORS), + f"{[n for n, _ in PROCESSORS if n.startswith('arm')]}", + ) + check( + "and the labels say which is 32- vs 64-bit", + all( + ("32-bit" in d or "64-bit" in d) + for n, d in PROCESSORS + if n == "arm" or n.startswith("arm:") + ), + f"{[(n, d) for n, d in PROCESSORS if n.startswith('arm')]}", + ) # Every offered name must have been checked against a real IDA, because a # wrong one is REJECTED (rc=4) with nothing useful said — handing the user a @@ -103,16 +143,40 @@ def main() -> int: # output of tools/verify_procs.py; adding a processor without re-running it # fails here on purpose. VERIFIED = { - "arm", "armb", "metapc", "mipsl", "mipsb", "ppc", "ppcl", "sh4", "68k", - "riscv", "tricore", "xtensa", "avr", "z80", "tms320c6", "m32r", "arc", - "h8300", "sparcb", "sparcl", "s390", + "arm", + "armb", + "metapc", + "mipsl", + "mipsb", + "ppc", + "ppcl", + "sh4", + "68k", + "riscv", + "tricore", + "xtensa", + "avr", + "z80", + "tms320c6", + "m32r", + "arc", + "h8300", + "sparcb", + "sparcl", + "s390", # variants: tools/verify_procs.py checks these report the base module # AND change the database bitness, which is the reason they exist - "arm:ARMv7-A", "arm:ARMv7-M", "arm:ARMv6-M", "arm:ARMv5TE", + "arm:ARMv7-A", + "arm:ARMv7-M", + "arm:ARMv6-M", + "arm:ARMv5TE", } offered = {n for n, _ in PROCESSORS} - check("every offered processor name is IDA-verified", - offered <= VERIFIED, f"unverified: {sorted(offered - VERIFIED)}") + check( + "every offered processor name is IDA-verified", + offered <= VERIFIED, + f"unverified: {sorted(offered - VERIFIED)}", + ) # These are module FILENAMES or common aliases, not -p names. IDA refuses # them; they were in the list until a real run said otherwise. @@ -123,16 +187,28 @@ def main() -> int: def finds(q): ql = q.lower() return [n for n, d in PROCESSORS if ql in n.lower() or ql in d.lower()] - check("typing 'arm64' still finds ARM", "arm" in finds("arm64"), f"{finds('arm64')}") + + check( + "typing 'arm64' still finds ARM", "arm" in finds("arm64"), f"{finds('arm64')}" + ) check("typing 'aarch64' still finds ARM", "arm" in finds("aarch64")) check("typing 'm68k' still finds 68k", "68k" in finds("m68k"), f"{finds('m68k')}") - check("typing 'mips' finds both endiannesses", - set(finds("mips")) == {"mipsl", "mipsb"}, f"{finds('mips')}") - check("every processor entry has a human label", - all(n and d for n, d in PROCESSORS)) - check("endianness is spelled out where it matters", - all(any(w in d.lower() for w in ("endian",)) - for n, d in PROCESSORS if n in ("arm", "armb", "mipsb", "mipsl"))) + check( + "typing 'mips' finds both endiannesses", + set(finds("mips")) == {"mipsl", "mipsb"}, + f"{finds('mips')}", + ) + check( + "every processor entry has a human label", all(n and d for n, d in PROCESSORS) + ) + check( + "endianness is spelled out where it matters", + all( + any(w in d.lower() for w in ("endian",)) + for n, d in PROCESSORS + if n in ("arm", "armb", "mipsb", "mipsl") + ), + ) print(f"\n{PASS} passed, {FAIL} failed") return 1 if FAIL else 0 diff --git a/tests/test_graph.py b/tests/test_graph.py index a75ff68..e3b9143 100644 --- a/tests/test_graph.py +++ b/tests/test_graph.py @@ -9,6 +9,7 @@ suite runs anywhere, but when present it is the interesting half: real functions are where the degenerate shapes (switch fan-out, irreducible loops, 400-block monsters) actually live. """ + from __future__ import annotations #: the layout engine is pure: no IDA, no Textual. @@ -55,17 +56,29 @@ def mk(edges: dict[int, list[tuple[int, str]]], n: int | None = None) -> list[G. ids = set(edges) | {d for v in edges.values() for d, _ in v} if n: ids |= set(range(n)) - return [G.Block(id=i, start=0x1000 + i * 0x10, end=0x1000 + i * 0x10 + 8, - succs=list(edges.get(i, []))) for i in sorted(ids)] + return [ + G.Block( + id=i, + start=0x1000 + i * 0x10, + end=0x1000 + i * 0x10 + 8, + succs=list(edges.get(i, [])), + ) + for i in sorted(ids) + ] # ------------------------------------------------------------ invariants + def no_box_overlap(lay: G.Layout) -> bool: for i, a in enumerate(lay.nodes): - for b in lay.nodes[i + 1:]: - if (a.x <= b.right and b.x <= a.right - and a.y <= b.y + b.h - 1 and b.y <= a.y + a.h - 1): + for b in lay.nodes[i + 1 :]: + if ( + a.x <= b.right + and b.x <= a.right + and a.y <= b.y + b.h - 1 + and b.y <= a.y + a.h - 1 + ): return False return True @@ -97,13 +110,16 @@ def all_edges_drawn(lay: G.Layout) -> bool: def invariants(lay: G.Layout, name: str) -> None: check(no_box_overlap(lay), f"{name}: boxes must not overlap") check(no_edge_through_box(lay) == 0, f"{name}: no edge may cross a box") - check(all(n.x >= 0 and n.y >= 0 for n in lay.nodes), - f"{name}: no negative coordinates") + check( + all(n.x >= 0 and n.y >= 0 for n in lay.nodes), + f"{name}: no negative coordinates", + ) check(lay.width > 0 and lay.height > 0, f"{name}: canvas has extent") # ---------------------------------------------------------------- cases + def t_linear() -> None: lay = layout(mk({0: [(1, "uncond")], 1: [(2, "uncond")]}), sizer) invariants(lay, "linear") @@ -113,8 +129,10 @@ def t_linear() -> None: def t_diamond() -> None: - lay = layout(mk({0: [(1, "jump"), (2, "fall")], - 1: [(3, "uncond")], 2: [(3, "uncond")]}), sizer) + lay = layout( + mk({0: [(1, "jump"), (2, "fall")], 1: [(3, "uncond")], 2: [(3, "uncond")]}), + sizer, + ) invariants(lay, "diamond") check(lay.by_id[3].rank == 2, "diamond: join sits below both arms") check(lay.by_id[1].rank == lay.by_id[2].rank, "diamond: arms share a rank") @@ -125,8 +143,10 @@ def t_diamond() -> None: def t_selfloop() -> None: """A self-loop must not stall the ranking — the bug that collapsed a whole function into three layers and made the graph 280 columns wide.""" - lay = layout(mk({0: [(1, "uncond")], 1: [(1, "jump"), (2, "fall")], - 2: [(3, "uncond")]}), sizer) + lay = layout( + mk({0: [(1, "uncond")], 1: [(1, "jump"), (2, "fall")], 2: [(3, "uncond")]}), + sizer, + ) invariants(lay, "selfloop") ranks = [lay.by_id[i].rank for i in (0, 1, 2, 3)] check(ranks == [0, 1, 2, 3], f"selfloop: ranking still stacks ({ranks})") @@ -134,23 +154,36 @@ def t_selfloop() -> None: def t_loop() -> None: - lay = layout(mk({0: [(1, "uncond")], 1: [(2, "jump"), (3, "fall")], - 2: [(1, "uncond")]}), sizer) + lay = layout( + mk({0: [(1, "uncond")], 1: [(2, "jump"), (3, "fall")], 2: [(1, "uncond")]}), + sizer, + ) invariants(lay, "loop") check(any(e.back for e in lay.edges), "loop: a back edge is detected") check(lay.by_id[1].rank < lay.by_id[2].rank, "loop: header above the body") back = [e for e in lay.edges if e.back][0] - check((2, G.E_BACK) in [(a, s) for a, s in lay.pred[1]] - or (1, G.E_BACK) in [(a, s) for a, s in lay.succ[2]], - "loop: the back edge reads 2 -> 1 despite being reversed for layout") + check( + (2, G.E_BACK) in [(a, s) for a, s in lay.pred[1]] + or (1, G.E_BACK) in [(a, s) for a, s in lay.succ[2]], + "loop: the back edge reads 2 -> 1 despite being reversed for layout", + ) def t_switch() -> None: - lay = layout(mk({0: [(i, "switch") for i in range(1, 9)], - **{i: [(9, "uncond")] for i in range(1, 9)}}), sizer) + lay = layout( + mk( + { + 0: [(i, "switch") for i in range(1, 9)], + **{i: [(9, "uncond")] for i in range(1, 9)}, + } + ), + sizer, + ) invariants(lay, "switch") - check(len({lay.by_id[i].rank for i in range(1, 9)}) == 1, - "switch: all cases share a rank") + check( + len({lay.by_id[i].rank for i in range(1, 9)}) == 1, + "switch: all cases share a rank", + ) check(lay.by_id[9].rank == 2, "switch: the join is below the cases") @@ -172,14 +205,26 @@ def t_unreachable_entry() -> None: from, so the engine must never be handed one. """ # entry 0 is a sink; 2 and 3 jump INTO it; 1 and 6 self-loop. - lay = layout(mk({0: [], 1: [(5, "switch"), (1, "fall"), (4, "switch")], - 2: [(5, "jump"), (0, "uncond")], 3: [(0, "switch")], - 4: [], 5: [(4, "jump")], - 6: [(2, "jump"), (6, "switch"), (3, "jump")]}), entry=0) + lay = layout( + mk( + { + 0: [], + 1: [(5, "switch"), (1, "fall"), (4, "switch")], + 2: [(5, "jump"), (0, "uncond")], + 3: [(0, "switch")], + 4: [], + 5: [(4, "jump")], + 6: [(2, "jump"), (6, "switch"), (3, "jump")], + } + ), + entry=0, + ) invariants(lay, "unreachable_entry") check(len(lay.nodes) == 7, "unreachable_entry: every block is placed") - check(lay.stats.get("engine_error") is None, - f"unreachable_entry: no fallback ({lay.stats.get('engine_error')})") + check( + lay.stats.get("engine_error") is None, + f"unreachable_entry: no fallback ({lay.stats.get('engine_error')})", + ) # An entry that reaches nothing at all, with everything hanging off nodes # it cannot see, is the degenerate version of the same thing. @@ -198,8 +243,10 @@ def t_long_edge() -> None: # long edge. Triskel reaches the same end -- an edge that crosses no box, # checked by invariants() above -- without them, so this is engine-specific. if ENGINE == "native": - check(lay.stats["dummies"] >= 4, - f"long_edge: the skip edge is padded ({lay.stats['dummies']} dummies)") + check( + lay.stats["dummies"] >= 4, + f"long_edge: the skip edge is padded ({lay.stats['dummies']} dummies)", + ) check(all_edges_drawn(lay), "long_edge: the long edge is drawn") @@ -212,14 +259,18 @@ def t_empty() -> None: def t_row_query() -> None: """cells_at_row must be windowed: asking for a slice returns only that slice, which is what keeps a 13M-cell graph renderable.""" - lay = layout(mk({0: [(1, "jump"), (2, "fall")], - 1: [(3, "uncond")], 2: [(3, "uncond")]}), sizer) + lay = layout( + mk({0: [(1, "jump"), (2, "fall")], 1: [(3, "uncond")], 2: [(3, "uncond")]}), + sizer, + ) for row in range(lay.height): full = lay.painting.cells_at_row(row, 0, lay.width) part = lay.painting.cells_at_row(row, 5, 12) check(all(5 <= c < 12 for c in part), f"row {row}: window respected") - check(all(full.get(c) == v for c, v in part.items()), - f"row {row}: window agrees with the full row") + check( + all(full.get(c) == v for c, v in part.items()), + f"row {row}: window agrees with the full row", + ) def t_hit_test() -> None: @@ -233,6 +284,7 @@ def t_hit_test() -> None: # ---------------------------------------------------------------- corpus + def t_corpus(path: str) -> None: recs = json.load(open(path)) print(f"\ncorpus: {len(recs)} functions from {path}") @@ -243,9 +295,15 @@ def t_corpus(path: str) -> None: worst_any_name = "" t0 = time.perf_counter() for rec in recs: - blocks = [G.Block(id=b["id"], start=b["start"], end=b["end"], - succs=[(d, k) for d, k in b["succs"]]) - for b in rec["blocks"]] + blocks = [ + G.Block( + id=b["id"], + start=b["start"], + end=b["end"], + succs=[(d, k) for d, k in b["succs"]], + ) + for b in rec["blocks"] + ] lay = layout(blocks, sizer) # A SILENT fallback is the failure mode that matters here: the engine # under test quietly stops being the engine under test, and every @@ -258,46 +316,61 @@ def t_corpus(path: str) -> None: # not that it never happens. if ENGINE != "auto" and lay.stats["engine"] != ENGINE: fellback.append((rec["name"], lay.stats.get("engine_error"))) - check(bool(lay.stats.get("engine_error")), - f"corpus {rec['name']}: a fallback must record its reason") + check( + bool(lay.stats.get("engine_error")), + f"corpus {rec['name']}: a fallback must record its reason", + ) # Time the engine only on the functions it would actually be ASKED for. # `auto` hands anything over AUTO_TRISKEL_MAX_BLOCKS to native, and the # view refuses to draw past 400 blocks at all, so a forced triskel run # on a 495-block monster times a call the app cannot make. - reachable = (ENGINE != "triskel" - or len(blocks) <= G.AUTO_TRISKEL_MAX_BLOCKS) + reachable = ENGINE != "triskel" or len(blocks) <= G.AUTO_TRISKEL_MAX_BLOCKS if reachable and lay.stats["ms"] > worst_ms: worst_ms, worst_name = lay.stats["ms"], rec["name"] if lay.stats["ms"] > worst_any_ms: worst_any_ms, worst_any_name = lay.stats["ms"], rec["name"] check(no_box_overlap(lay), f"corpus {rec['name']}: boxes must not overlap") - check(len(lay.nodes) == len(blocks), - f"corpus {rec['name']}: every block is placed") + check( + len(lay.nodes) == len(blocks), + f"corpus {rec['name']}: every block is placed", + ) # The full cell sweep is O(canvas); only affordable on the small ones, # but that is where a routing bug would show up anyway. if lay.width * lay.height < 400_000: - check(no_edge_through_box(lay) == 0, - f"corpus {rec['name']}: no edge may cross a box") + check( + no_edge_through_box(lay) == 0, + f"corpus {rec['name']}: no edge may cross a box", + ) total = (time.perf_counter() - t0) * 1000 - print(f" laid out {len(recs)} functions in {total:.0f} ms " - f"(worst {worst_ms:.0f} ms: {worst_name})") + print( + f" laid out {len(recs)} functions in {total:.0f} ms " + f"(worst {worst_ms:.0f} ms: {worst_name})" + ) if fellback: print(f" {len(fellback)} fell back to native:") for name, why in fellback: print(f" {name}: {why}") - check(len(fellback) <= max(2, len(recs) // 20), - f"corpus: {ENGINE} fell back on {len(fellback)}/{len(recs)} functions") - check(worst_ms < 2000, f"corpus: worst REACHABLE layout under 2s " - f"({worst_ms:.0f} ms: {worst_name})") + check( + len(fellback) <= max(2, len(recs) // 20), + f"corpus: {ENGINE} fell back on {len(fellback)}/{len(recs)} functions", + ) + check( + worst_ms < 2000, + f"corpus: worst REACHABLE layout under 2s ({worst_ms:.0f} ms: {worst_name})", + ) # Nothing may blow up quadratically even when forced past its own limits. - check(worst_any_ms < 5000, f"corpus: worst layout at any size under 5s " - f"({worst_any_ms:.0f} ms: {worst_any_name})") + check( + worst_any_ms < 5000, + f"corpus: worst layout at any size under 5s " + f"({worst_any_ms:.0f} ms: {worst_any_name})", + ) def main() -> int: global ENGINE print("idatui.graph layout tests") from idatui import graph_triskel + engines = ["native"] if graph_triskel.available(): engines.append("triskel") @@ -306,9 +379,19 @@ def main() -> int: for engine in engines: ENGINE = engine print(f"\nengine: {engine}") - for fn in (t_linear, t_diamond, t_selfloop, t_loop, t_switch, - t_unreachable, t_unreachable_entry, t_long_edge, t_empty, - t_row_query, t_hit_test): + for fn in ( + t_linear, + t_diamond, + t_selfloop, + t_loop, + t_switch, + t_unreachable, + t_unreachable_entry, + t_long_edge, + t_empty, + t_row_query, + t_hit_test, + ): print(f" {fn.__name__}") fn() for path in sys.argv[1:]: diff --git a/tests/test_index.py b/tests/test_index.py index f6761b1..5850821 100644 --- a/tests/test_index.py +++ b/tests/test_index.py @@ -15,8 +15,13 @@ import tempfile import time sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) -from idatui.index import (KIND_EXPORT, KIND_FUNC, KIND_IMPORT, # noqa: E402 - KIND_STRING, ProjectIndex) +from idatui.index import ( # noqa: E402 + KIND_EXPORT, + KIND_FUNC, + KIND_IMPORT, + KIND_STRING, + ProjectIndex, +) PASS = FAIL = 0 @@ -41,70 +46,106 @@ def main() -> int: check("a fresh index is empty", idx.total() == 0 and idx.counts() == {}) check("an unindexed binary is stale", idx.is_stale("libfoo", src)) - n = idx.reindex("libfoo", [ - (KIND_FUNC, 0x1000, "SSL_CTX_new"), - (KIND_FUNC, 0x1100, "SSL_read"), - (KIND_FUNC, 0x1200, "sub_1200"), - (KIND_STRING, 0x8000, "error opening socket"), - (KIND_STRING, 0x8100, "/etc/ssl/certs"), - ], source=src) + n = idx.reindex( + "libfoo", + [ + (KIND_FUNC, 0x1000, "SSL_CTX_new"), + (KIND_FUNC, 0x1100, "SSL_read"), + (KIND_FUNC, 0x1200, "sub_1200"), + (KIND_STRING, 0x8000, "error opening socket"), + (KIND_STRING, 0x8100, "/etc/ssl/certs"), + ], + source=src, + ) check("reindex reports what it stored", n == 5, f"n={n}") check("the entries are there", idx.total() == 5, f"{idx.total()}") check("an indexed binary is fresh", not idx.is_stale("libfoo", src)) # -- substring search (the thing a prefix index can't do) ----------- # hits = idx.search("SSL", kind=KIND_FUNC) - check("finds symbols by substring", {h.text for h in hits} == - {"SSL_CTX_new", "SSL_read"}, f"{[h.text for h in hits]}") - check("matching ignores case across kinds (SSL also hits /etc/ssl)", - {h.text for h in idx.search("SSL")} == - {"SSL_CTX_new", "SSL_read", "/etc/ssl/certs"}, - f"{[h.text for h in idx.search('SSL')]}") + check( + "finds symbols by substring", + {h.text for h in hits} == {"SSL_CTX_new", "SSL_read"}, + f"{[h.text for h in hits]}", + ) + check( + "matching ignores case across kinds (SSL also hits /etc/ssl)", + {h.text for h in idx.search("SSL")} + == {"SSL_CTX_new", "SSL_read", "/etc/ssl/certs"}, + f"{[h.text for h in idx.search('SSL')]}", + ) hits = idx.search("socket") - check("finds strings by substring mid-text", - len(hits) == 1 and hits[0].kind == KIND_STRING - and hits[0].addr == 0x8000, f"{hits}") - check("hits carry the owning binary", - all(h.binary == "libfoo" for h in idx.search("SSL"))) - check("search is case-insensitive", - {h.text for h in idx.search("ssl_read")} == {"SSL_read"}, - f"{[h.text for h in idx.search('ssl_read')]}") - check("kind filter narrows to strings", - [h.text for h in idx.search("ss", kind=KIND_STRING)] == ["/etc/ssl/certs"], - f"{[h.text for h in idx.search('ss', kind=KIND_STRING)]}") + check( + "finds strings by substring mid-text", + len(hits) == 1 and hits[0].kind == KIND_STRING and hits[0].addr == 0x8000, + f"{hits}", + ) + check( + "hits carry the owning binary", + all(h.binary == "libfoo" for h in idx.search("SSL")), + ) + check( + "search is case-insensitive", + {h.text for h in idx.search("ssl_read")} == {"SSL_read"}, + f"{[h.text for h in idx.search('ssl_read')]}", + ) + check( + "kind filter narrows to strings", + [h.text for h in idx.search("ss", kind=KIND_STRING)] == ["/etc/ssl/certs"], + f"{[h.text for h in idx.search('ss', kind=KIND_STRING)]}", + ) # -- the <3 char fallback (trigram silently matches nothing) -------- # - check("2-char query still works (LIKE fallback)", - {h.text for h in idx.search("ss")} == {"SSL_CTX_new", "SSL_read", - "/etc/ssl/certs"}, - f"{[h.text for h in idx.search('ss')]}") - check("1-char query still works", - len(idx.search("/")) == 1, f"{idx.search('/')}") + check( + "2-char query still works (LIKE fallback)", + {h.text for h in idx.search("ss")} + == {"SSL_CTX_new", "SSL_read", "/etc/ssl/certs"}, + f"{[h.text for h in idx.search('ss')]}", + ) + check( + "1-char query still works", len(idx.search("/")) == 1, f"{idx.search('/')}" + ) check("an empty query matches nothing", idx.search(" ") == []) - check("a query with FTS operators is treated literally", - idx.search('SSL OR "') == [] or True) # must not raise + check( + "a query with FTS operators is treated literally", + idx.search('SSL OR "') == [] or True, + ) # must not raise # -- multi-binary: the whole point ---------------------------------- # - idx.reindex("httpd", [ - (KIND_FUNC, 0x2000, "handle_ssl_request"), - (KIND_STRING, 0x9000, "socket bind failed"), - ]) + idx.reindex( + "httpd", + [ + (KIND_FUNC, 0x2000, "handle_ssl_request"), + (KIND_STRING, 0x9000, "socket bind failed"), + ], + ) hits = idx.search("ssl") - check("search spans binaries", - {h.binary for h in hits} == {"libfoo", "httpd"}, - f"{[(h.binary, h.text) for h in hits]}") - check("counts are per binary", - idx.counts() == {"libfoo": 5, "httpd": 2}, f"{idx.counts()}") + check( + "search spans binaries", + {h.binary for h in hits} == {"libfoo", "httpd"}, + f"{[(h.binary, h.text) for h in hits]}", + ) + check( + "counts are per binary", + idx.counts() == {"libfoo": 5, "httpd": 2}, + f"{idx.counts()}", + ) # -- incremental: reindexing one binary leaves the others alone ----- # idx.reindex("libfoo", [(KIND_FUNC, 0x1000, "SSL_CTX_new_v2")], source=src) - check("reindex replaces only that binary's entries", - idx.counts() == {"libfoo": 1, "httpd": 2}, f"{idx.counts()}") - check("the stale entries are gone", - [h.text for h in idx.search("SSL_read")] == [], - f"{idx.search('SSL_read')}") - check("the other binary survived untouched", - len(idx.search("socket bind")) == 1) + check( + "reindex replaces only that binary's entries", + idx.counts() == {"libfoo": 1, "httpd": 2}, + f"{idx.counts()}", + ) + check( + "the stale entries are gone", + [h.text for h in idx.search("SSL_read")] == [], + f"{idx.search('SSL_read')}", + ) + check( + "the other binary survived untouched", len(idx.search("socket bind")) == 1 + ) # -- staleness follows the source ----------------------------------- # time.sleep(0.01) @@ -112,90 +153,128 @@ def main() -> int: f.write(b"\x7fELF binary rebuilt, different size") os.utime(src, (1, 1)) check("a changed source goes stale", idx.is_stale("libfoo", src)) - check("a missing source does NOT wipe the index", - not idx.is_stale("libfoo", os.path.join(tmp, "gone"))) + check( + "a missing source does NOT wipe the index", + not idx.is_stale("libfoo", os.path.join(tmp, "gone")), + ) # -- forget ----------------------------------------------------------- # idx.forget("httpd") - check("forget drops a binary entirely", - idx.counts() == {"libfoo": 1} and idx.search("socket bind") == [], - f"{idx.counts()}") + check( + "forget drops a binary entirely", + idx.counts() == {"libfoo": 1} and idx.search("socket bind") == [], + f"{idx.counts()}", + ) # -- cross-binary linkage join (phase 3) ------------------------------ # - idx.reindex("app", [ - (KIND_FUNC, 0x1000, "main"), - (KIND_IMPORT, 0x2000, "strcmp"), - (KIND_IMPORT, 0x2008, "read"), - (KIND_IMPORT, 0x2010, "SSL_new"), - ]) - idx.reindex("libc", [ - (KIND_EXPORT, 0x8000, "strcmp"), - (KIND_EXPORT, 0x8100, "read"), - (KIND_EXPORT, 0x8200, "pread"), - (KIND_EXPORT, 0x8300, "read_line"), - (KIND_FUNC, 0x8000, "strcmp"), - ]) + idx.reindex( + "app", + [ + (KIND_FUNC, 0x1000, "main"), + (KIND_IMPORT, 0x2000, "strcmp"), + (KIND_IMPORT, 0x2008, "read"), + (KIND_IMPORT, 0x2010, "SSL_new"), + ], + ) + idx.reindex( + "libc", + [ + (KIND_EXPORT, 0x8000, "strcmp"), + (KIND_EXPORT, 0x8100, "read"), + (KIND_EXPORT, 0x8200, "pread"), + (KIND_EXPORT, 0x8300, "read_line"), + (KIND_FUNC, 0x8000, "strcmp"), + ], + ) idx.reindex("libssl", [(KIND_EXPORT, 0x9000, "SSL_new")]) prov = idx.providers("strcmp", exclude="app") - check("an import resolves to the binary that exports it", - [(h.binary, h.addr) for h in prov] == [("libc", 0x8000)], - f"{[(h.binary, hex(h.addr)) for h in prov]}") + check( + "an import resolves to the binary that exports it", + [(h.binary, h.addr) for h in prov] == [("libc", 0x8000)], + f"{[(h.binary, hex(h.addr)) for h in prov]}", + ) # The whole point of exact(): substring search would drag in pread, # read_line and thread_start, and 'read' is also below the trigram floor # for some engines — an import must bind to its exact name or nothing. prov = idx.providers("read", exclude="app") - check("the join is exact, not substring", - [(h.binary, h.addr) for h in prov] == [("libc", 0x8100)], - f"{[(h.binary, h.text) for h in prov]}") + check( + "the join is exact, not substring", + [(h.binary, h.addr) for h in prov] == [("libc", 0x8100)], + f"{[(h.binary, h.text) for h in prov]}", + ) - check("a short name still resolves (below the trigram floor)", - [h.binary for h in idx.providers("SSL_new", exclude="app")] == ["libssl"], - f"{idx.providers('SSL_new')}") + check( + "a short name still resolves (below the trigram floor)", + [h.binary for h in idx.providers("SSL_new", exclude="app")] == ["libssl"], + f"{idx.providers('SSL_new')}", + ) - check("an unprovided import resolves to nothing", - idx.providers("dlopen", exclude="app") == []) + check( + "an unprovided import resolves to nothing", + idx.providers("dlopen", exclude="app") == [], + ) - check("exclude keeps a binary from resolving to itself", - idx.providers("strcmp", exclude="libc") == [], - f"{idx.providers('strcmp', exclude='libc')}") + check( + "exclude keeps a binary from resolving to itself", + idx.providers("strcmp", exclude="libc") == [], + f"{idx.providers('strcmp', exclude='libc')}", + ) imp = idx.importers("strcmp") - check("the reverse join finds who imports an export", - [(h.binary, h.addr) for h in imp] == [("app", 0x2000)], - f"{[(h.binary, hex(h.addr)) for h in imp]}") + check( + "the reverse join finds who imports an export", + [(h.binary, h.addr) for h in imp] == [("app", 0x2000)], + f"{[(h.binary, hex(h.addr)) for h in imp]}", + ) - check("kind keeps functions out of the linkage join", - [h.binary for h in idx.providers("strcmp")] == ["libc"], - "a KIND_FUNC row named strcmp must not answer as an export") + check( + "kind keeps functions out of the linkage join", + [h.binary for h in idx.providers("strcmp")] == ["libc"], + "a KIND_FUNC row named strcmp must not answer as an export", + ) idx.forget("libc") - check("forgetting a provider unresolves its imports", - idx.providers("strcmp", exclude="app") == []) + check( + "forgetting a provider unresolves its imports", + idx.providers("strcmp", exclude="app") == [], + ) # -- ELF symbol versioning -------------------------------------------- # # The importer sees strrchr@@GLIBC_2.2.5 while the provider may export a # different spelling; raw names would resolve almost nothing. link_name # cuts at the first '@' so both sides meet on the bare symbol. from idatui.domain import link_name - check("link_name strips an ELF version suffix", - link_name("strrchr@@GLIBC_2.2.5") == "strrchr", - link_name("strrchr@@GLIBC_2.2.5")) - check("link_name leaves an unversioned name alone", - link_name("strrchr") == "strrchr") - check("link_name handles a single-@ version", - link_name("SSL_new@OPENSSL_3.0.0") == "SSL_new") - check("link_name doesn't eat a leading @", - link_name("@weird") == "@weird", link_name("@weird")) + + check( + "link_name strips an ELF version suffix", + link_name("strrchr@@GLIBC_2.2.5") == "strrchr", + link_name("strrchr@@GLIBC_2.2.5"), + ) + check( + "link_name leaves an unversioned name alone", + link_name("strrchr") == "strrchr", + ) + check( + "link_name handles a single-@ version", + link_name("SSL_new@OPENSSL_3.0.0") == "SSL_new", + ) + check( + "link_name doesn't eat a leading @", + link_name("@weird") == "@weird", + link_name("@weird"), + ) # -- persistence ------------------------------------------------------ # path = idx.path idx.close() idx2 = ProjectIndex(path) - check("the index persists across sessions", - [h.text for h in idx2.search("SSL_CTX")] == ["SSL_CTX_new_v2"], - f"{idx2.search('SSL_CTX')}") + check( + "the index persists across sessions", + [h.text for h in idx2.search("SSL_CTX")] == ["SSL_CTX_new_v2"], + f"{idx2.search('SSL_CTX')}", + ) idx2.close() print(f"\n{PASS} passed, {FAIL} failed") diff --git a/tests/test_kittygfx.py b/tests/test_kittygfx.py index 38d361c..270b4e7 100644 --- a/tests/test_kittygfx.py +++ b/tests/test_kittygfx.py @@ -55,8 +55,11 @@ class Tty: def cmds(self, action): """Every graphics command with the given ``a=`` action.""" - return [c for c in re.findall(r"\x1b_G([^;\x1b]*)", self.blob) - if f"a={action}" in c.split(",")] + return [ + c + for c in re.findall(r"\x1b_G([^;\x1b]*)", self.blob) + if f"a={action}" in c.split(",") + ] def keys(cmd): @@ -76,8 +79,7 @@ def t_no_termios_falls_back(): try: check("missing termios disables graphics", kittygfx._query_tty(0) is False) except Exception as exc: # the original Windows startup crash - check("missing termios does not escape", False, - f"{type(exc).__name__}: {exc}") + check("missing termios does not escape", False, f"{type(exc).__name__}: {exc}") finally: builtins.__import__ = original_import @@ -103,8 +105,7 @@ def t_probe_failure_is_never_fatal(): check("probe exception disables graphics", kittygfx.supported() is False) check("failed result is cached", kittygfx.supported() is False) except Exception as exc: - check("probe exception does not escape", False, - f"{type(exc).__name__}: {exc}") + check("probe exception does not escape", False, f"{type(exc).__name__}: {exc}") finally: kittygfx._query_tty = original_query kittygfx._supported = original_supported @@ -118,7 +119,7 @@ def main() -> int: t_no_termios_falls_back() t_probe_failure_is_never_fatal() - kittygfx._uploaded[kittygfx.LOGO_ID] = (768, 801) # pretend it's uploaded + kittygfx._uploaded[kittygfx.LOGO_ID] = (768, 801) # pretend it's uploaded # -- the bug: anonymous placements STACK ------------------------------- # # A placement is identified by (image id, placement id). With no p key @@ -128,80 +129,126 @@ def main() -> int: for _ in range(50): kittygfx.place(4, 10, 60, 26) placements = tty.cmds("p") - check("place() emits one command per call", len(placements) == 50, - f"{len(placements)}") - check("every placement carries a placement id (replaces, not stacks)", - all("p" in keys(c) for c in placements), - f"{placements[0] if placements else '(none)'}") - check("the placement id is the same every time (one image on screen)", - len({keys(c)["p"] for c in placements}) == 1, - f"{sorted({keys(c).get('p') for c in placements})}") - check("...and it is non-zero (p=0 means anonymous)", - keys(placements[0])["p"] not in ("0", ""), f"{placements[0]}") + check( + "place() emits one command per call", + len(placements) == 50, + f"{len(placements)}", + ) + check( + "every placement carries a placement id (replaces, not stacks)", + all("p" in keys(c) for c in placements), + f"{placements[0] if placements else '(none)'}", + ) + check( + "the placement id is the same every time (one image on screen)", + len({keys(c)["p"] for c in placements}) == 1, + f"{sorted({keys(c).get('p') for c in placements})}", + ) + check( + "...and it is non-zero (p=0 means anonymous)", + keys(placements[0])["p"] not in ("0", ""), + f"{placements[0]}", + ) # -- the rest of the escape still says what it used to ------------------ # with Tty() as tty: ok = kittygfx.place(4, 10, 60, 26) k = keys(tty.cmds("p")[0]) check("place() reports success", ok) - check("image id, source pixels and cell box are unchanged", - (k["i"], k["s"], k["v"], k["c"], k["r"]) - == (str(kittygfx.LOGO_ID), "768", "801", "60", "26"), f"{k}") - check("the terminal is told not to move the cursor (C=1)", - k.get("C") == "1", f"{k}") - check("the cursor is saved and restored around the placement", - tty.blob.startswith("\x1b[s") and tty.blob.endswith("\x1b[u"), - repr(tty.blob[:8] + "..." + tty.blob[-8:])) - check("the placement is positioned 1-based (row 4 -> line 5)", - "\x1b[5;11H" in tty.blob, repr(tty.blob[:24])) + check( + "image id, source pixels and cell box are unchanged", + (k["i"], k["s"], k["v"], k["c"], k["r"]) + == (str(kittygfx.LOGO_ID), "768", "801", "60", "26"), + f"{k}", + ) + check( + "the terminal is told not to move the cursor (C=1)", + k.get("C") == "1", + f"{k}", + ) + check( + "the cursor is saved and restored around the placement", + tty.blob.startswith("\x1b[s") and tty.blob.endswith("\x1b[u"), + repr(tty.blob[:8] + "..." + tty.blob[-8:]), + ) + check( + "the placement is positioned 1-based (row 4 -> line 5)", + "\x1b[5;11H" in tty.blob, + repr(tty.blob[:24]), + ) # -- deleting still removes EVERY placement of the image ---------------- # # d=i is by image id, so it takes the placement with us regardless of p. with Tty() as tty: kittygfx.clear() k = keys(tty.cmds("d")[0]) - check("clear() deletes by image id (d=i), keeping the upload", - k.get("d") == "i" and k.get("i") == str(kittygfx.LOGO_ID), f"{k}") - check("clear() does not free the image data (lowercase d)", - kittygfx.is_uploaded(), "upload was dropped") + check( + "clear() deletes by image id (d=i), keeping the upload", + k.get("d") == "i" and k.get("i") == str(kittygfx.LOGO_ID), + f"{k}", + ) + check( + "clear() does not free the image data (lowercase d)", + kittygfx.is_uploaded(), + "upload was dropped", + ) with Tty() as tty: kittygfx.delete() k = keys(tty.cmds("d")[0]) check("delete() frees the image data too (d=I)", k.get("d") == "I", f"{k}") - check("...and forgets the upload, so the next place() refuses", - not kittygfx.is_uploaded() and kittygfx.place(0, 0, 10, 10) is False) + check( + "...and forgets the upload, so the next place() refuses", + not kittygfx.is_uploaded() and kittygfx.place(0, 0, 10, 10) is False, + ) # -- refusals ----------------------------------------------------------- # kittygfx._uploaded[kittygfx.LOGO_ID] = (768, 801) with Tty() as tty: - check("a zero-sized box is refused, not sent", - kittygfx.place(0, 0, 0, 10) is False - and kittygfx.place(0, 0, 10, 0) is False and not tty.sent, - f"{tty.sent}") + check( + "a zero-sized box is refused, not sent", + kittygfx.place(0, 0, 0, 10) is False + and kittygfx.place(0, 0, 10, 0) is False + and not tty.sent, + f"{tty.sent}", + ) kittygfx._uploaded.pop(kittygfx.LOGO_ID, None) # -- fit(): aspect ratio against non-square cells ----------------------- # - check("fit() keeps the aspect ratio for 9x22 cells", - kittygfx.fit((768, 801), 60, 99, cell=(9, 22)) == (60, 26), - f"{kittygfx.fit((768, 801), 60, 99, cell=(9, 22))}") - check("fit() shrinks to the row budget instead of overflowing", - kittygfx.fit((768, 801), 60, 10, cell=(9, 22))[1] == 10, - f"{kittygfx.fit((768, 801), 60, 10, cell=(9, 22))}") - check("fit() never returns a zero dimension", - all(v >= 1 for v in kittygfx.fit((768, 801), 1, 1, cell=(9, 22)))) - check("fit() survives a degenerate image size", - kittygfx.fit((0, 0), 60, 26) == (60, 26)) + check( + "fit() keeps the aspect ratio for 9x22 cells", + kittygfx.fit((768, 801), 60, 99, cell=(9, 22)) == (60, 26), + f"{kittygfx.fit((768, 801), 60, 99, cell=(9, 22))}", + ) + check( + "fit() shrinks to the row budget instead of overflowing", + kittygfx.fit((768, 801), 60, 10, cell=(9, 22))[1] == 10, + f"{kittygfx.fit((768, 801), 60, 10, cell=(9, 22))}", + ) + check( + "fit() never returns a zero dimension", + all(v >= 1 for v in kittygfx.fit((768, 801), 1, 1, cell=(9, 22))), + ) + check( + "fit() survives a degenerate image size", + kittygfx.fit((0, 0), 60, 26) == (60, 26), + ) # -- png_size() reads the header, not the pixels ------------------------ # - logo = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), - "logo.png") + logo = os.path.join( + os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "logo.png" + ) if os.path.exists(logo): - check("png_size() reads logo.png's IHDR", - kittygfx.png_size(logo) == (768, 801), f"{kittygfx.png_size(logo)}") + check( + "png_size() reads logo.png's IHDR", + kittygfx.png_size(logo) == (768, 801), + f"{kittygfx.png_size(logo)}", + ) check("png_size() returns None for a non-PNG", kittygfx.png_size(__file__) is None) - check("png_size() returns None for a missing file", - kittygfx.png_size("/nonexistent/nope.png") is None) + check( + "png_size() returns None for a missing file", + kittygfx.png_size("/nonexistent/nope.png") is None, + ) print(f"\n{PASS} passed, {FAIL} failed") return 1 if FAIL else 0 diff --git a/tests/test_launch.py b/tests/test_launch.py index 35a0b8d..5c05d08 100644 --- a/tests/test_launch.py +++ b/tests/test_launch.py @@ -10,6 +10,7 @@ is the cheapest guard against someone reintroducing the "helpful" cleanup. Pure: no IDA, no IDA Nexus library, no Textual. """ + from __future__ import annotations import os @@ -53,9 +54,11 @@ def t_no_lock_sweeping(): check("_sweep_locks is gone", not hasattr(launch, "_sweep_locks")) check("the scratch-suffix list is gone", not hasattr(launch, "_LOCK_SUFFIXES")) src = open(launch.__file__, encoding="utf-8").read() - check("the launcher does not remove files at all", - "os.remove" not in src and "shutil.rmtree" not in src, - "launch.py deletes something again") + check( + "the launcher does not remove files at all", + "os.remove" not in src and "shutil.rmtree" not in src, + "launch.py deletes something again", + ) def t_load_args(): @@ -66,10 +69,16 @@ def t_load_args(): # -b is in PARAGRAPHS, not bytes: 0x8000 >> 4 == 0x800. check("a base is converted to paragraphs", "-b800" in a, a) b = _load_args({"base": "0x1000"}) - check("a base given as a hex STRING is accepted (project files write those)", - "-b100" in b, b) - check("no base means no -b switch", "-b" not in _load_args({"processor": "arm"}), - _load_args({"processor": "arm"})) + check( + "a base given as a hex STRING is accepted (project files write those)", + "-b100" in b, + b, + ) + check( + "no base means no -b switch", + "-b" not in _load_args({"processor": "arm"}), + _load_args({"processor": "arm"}), + ) c = _load_args({"ida_args": "-p1"}) check("extra ida_args are passed through", "-p1" in c, c) @@ -81,6 +90,7 @@ def main() -> int: fn() except Exception as e: # noqa: BLE001 import traceback + check(f"{fn.__name__} did not crash", False, f"{type(e).__name__}: {e}") traceback.print_exc() print(f"\n{PASS} passed, {FAIL} failed") diff --git a/tests/test_nexus_client.py b/tests/test_nexus_client.py index ae2bfa7..834a8d6 100644 --- a/tests/test_nexus_client.py +++ b/tests/test_nexus_client.py @@ -5,15 +5,15 @@ from __future__ import annotations import os import queue import sys +import tempfile 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.nexus_client as module # noqa: E402 +from idatui import remote_ops # noqa: E402 +from idatui.errors import IDAConnectionError, IDAToolError # noqa: E402 from idatui.nexus_client import NexusClient, _parse_load_args # noqa: E402 #: Pure: fakes the DatabaseHandle, never touches IDA or the IDA Nexus library. @@ -215,10 +215,7 @@ class FakeRemoteModule: operation_label=label, ) result = response["result"] - if ( - isinstance(result, dict) - and result.get("__remote_ida_status__") == "ok" - ): + if isinstance(result, dict) and result.get("__remote_ida_status__") == "ok": return result.get("__remote_ida_value__") raise module.RemoteError(name, f"remote call failed: {result!r}", 500) @@ -232,6 +229,7 @@ def _open_kwargs_are_real(sent: dict): """ try: import inspect + from ida_nexus import DatabaseHandle as Real except ImportError: return True, "ida_nexus not installed - signature not checked" @@ -249,6 +247,7 @@ def _option_fields_are_real(options): """ try: import dataclasses + from ida_nexus import DatabaseOpenOptions as Real except ImportError: return True, "ida_nexus not installed - fields not checked" diff --git a/tests/test_pool.py b/tests/test_pool.py index 37058bf..203ce95 100644 --- a/tests/test_pool.py +++ b/tests/test_pool.py @@ -33,8 +33,7 @@ def check(name, cond, detail=""): class FakeClient: """Stands in for a NexusClient lease and records saves/closes.""" - def __init__(self, ref, mem=100, backend="idalib", - discardable=True): + def __init__(self, ref, mem=100, backend="idalib", discardable=True): self.ref = ref self.mem = mem self.backend = backend @@ -82,14 +81,18 @@ def main() -> int: made[ref.label] = c return c - pool = DatabasePool(proj, budget_mb=350, spawn=spawn, - mem_fn=lambda c: c.mem) + pool = DatabasePool(proj, budget_mb=350, spawn=spawn, mem_fn=lambda c: c.mem) # -- lazy spawn + reuse -------------------------------------------- # a = pool.get("bin0") - check("get() spawns a database lease on first use", a is made["bin0"] and a.connected) - check("get() stages the binary first", - os.path.isfile(proj.by_label("bin0").staged)) + check( + "get() spawns a database lease on first use", + a is made["bin0"] and a.connected, + ) + check( + "get() stages the binary first", + os.path.isfile(proj.by_label("bin0").staged), + ) check("get() reuses the resident lease", pool.get("bin0") is a) check("resident() reports it", pool.resident() == ["bin0"], pool.resident()) @@ -97,28 +100,40 @@ def main() -> int: pool.get("bin1") pool.get("bin2") pool.get("bin0") # touch: bin0 becomes most-recent - check("LRU order tracks use", pool.resident() == ["bin1", "bin2", "bin0"], - pool.resident()) + check( + "LRU order tracks use", + pool.resident() == ["bin1", "bin2", "bin0"], + pool.resident(), + ) # -- budget eviction -------------------------------------------------- # check("pool reports its memory", pool.memory_mb() == 300, pool.memory_mb()) pool.get("bin3") # 400MB > 350MB budget -> evict LRU (bin1) - check("exceeding the budget evicts the least-recently-used", - pool.evicted == ["bin1"] and not pool.is_resident("bin1"), - f"evicted={pool.evicted} resident={pool.resident()}") + check( + "exceeding the budget evicts the least-recently-used", + pool.evicted == ["bin1"] and not pool.is_resident("bin1"), + f"evicted={pool.evicted} resident={pool.resident()}", + ) check("the just-attached lease is never the victim", pool.is_resident("bin3")) check("eviction saves the database first", made["bin1"].saved == 1) check("eviction closes the lease", made["bin1"].closed) - check("pool is back within budget", pool.memory_mb() <= pool.budget_mb, - f"{pool.memory_mb()}/{pool.budget_mb}") + check( + "pool is back within budget", + pool.memory_mb() <= pool.budget_mb, + f"{pool.memory_mb()}/{pool.budget_mb}", + ) # -- the active binary is never evicted ------------------------------- # pool.set_active("bin2") - check("set_active touches the LRU", pool.resident()[-1] == "bin2", - pool.resident()) + check( + "set_active touches the LRU", pool.resident()[-1] == "bin2", pool.resident() + ) pool.get("bin1") # over budget again -> must evict, but not bin2 - check("the active binary survives eviction", pool.is_resident("bin2"), - f"resident={pool.resident()}") + check( + "the active binary survives eviction", + pool.is_resident("bin2"), + f"resident={pool.resident()}", + ) # -- pinning ---------------------------------------------------------- # pool.close_all() @@ -127,8 +142,11 @@ def main() -> int: pool2.pin("bin0") pool2.get("bin1") pool2.get("bin2") # 300 > 250 -> evict, but bin0 is pinned - check("pinned binaries are never evicted", pool2.is_resident("bin0"), - f"resident={pool2.resident()} evicted={pool2.evicted}") + check( + "pinned binaries are never evicted", + pool2.is_resident("bin0"), + f"resident={pool2.resident()} evicted={pool2.evicted}", + ) check("an unpinned one went instead", "bin1" in pool2.evicted, pool2.evicted) # -- everything pinned/active: stop evicting rather than thrash -------- # @@ -136,21 +154,30 @@ def main() -> int: pool2.set_active("bin2") n_before = len(pool2.evicted) pool2._enforce_budget() - check("nothing evictable -> gives up instead of thrashing", - len(pool2.evicted) == n_before, pool2.evicted) + check( + "nothing evictable -> gives up instead of thrashing", + len(pool2.evicted) == n_before, + pool2.evicted, + ) # -- status for the switcher UI ---------------------------------------- # st = {s["label"]: s for s in pool2.status()} check("status() covers every project binary", len(st) == 4, list(st)) - check("status() marks resident/pinned/active", - st["bin0"]["resident"] and st["bin0"]["pinned"] - and st["bin2"]["active"] and not st["bin3"]["resident"], - f"{st}") + check( + "status() marks resident/pinned/active", + st["bin0"]["resident"] + and st["bin0"]["pinned"] + and st["bin2"]["active"] + and not st["bin3"]["resident"], + f"{st}", + ) # -- teardown ----------------------------------------------------------- # pool2.close_all() - check("close_all() closes every lease", - not pool2.resident() and all(c.closed for c in made.values())) + check( + "close_all() closes every lease", + not pool2.resident() and all(c.closed for c in made.values()), + ) check("close_all() clears the active binary", pool2.active is None) # -- unknown label -------------------------------------------------------- # @@ -164,36 +191,45 @@ def main() -> int: discard_made = {} def spawn_discard(ref, ttl): - client = FakeClient( - ref, discardable=ref.label != "bin1") + 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 = 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) + 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) + 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", - pool3.budget_mb >= 256, pool3.budget_mb) + check( + "default budget is derived, not a fixed lease count", + pool3.budget_mb >= 256, + pool3.budget_mb, + ) # -- prewarm: speculative, and never at the cost of a real binary ------ # with tempfile.TemporaryDirectory() as tmp: @@ -205,24 +241,34 @@ def main() -> int: made2[ref.label] = c return c - pool = DatabasePool(proj, budget_mb=250, spawn=spawn2, - mem_fn=lambda c: c.mem) + pool = DatabasePool(proj, budget_mb=250, spawn=spawn2, mem_fn=lambda c: c.mem) labels = [r.label for r in proj.refs] a, b, c_ = labels[0], labels[1], labels[2] pool.get(a) pool.set_active(a) - check("prewarm warms a binary when the budget has room", - pool.prewarm(b) is True and b in pool.resident(), f"{pool.resident()}") - check("prewarm is a no-op for something already resident", - pool.prewarm(b) is False) + check( + "prewarm warms a binary when the budget has room", + pool.prewarm(b) is True and b in pool.resident(), + f"{pool.resident()}", + ) + check( + "prewarm is a no-op for something already resident", + pool.prewarm(b) is False, + ) # 2 x 100MB resident, estimate 100 more -> 300 > 250: must refuse - check("prewarm refuses rather than making room", - pool.prewarm(c_) is False and c_ not in pool.resident(), - f"resident={pool.resident()} mem={pool.memory_mb()}/{pool.budget_mb}") - check("refusing to prewarm evicts nothing", - set(pool.resident()) == {a, b}, f"{pool.resident()}") - check("prewarm ignores a label outside the project", - pool.prewarm("nope") is False) + check( + "prewarm refuses rather than making room", + pool.prewarm(c_) is False and c_ not in pool.resident(), + f"resident={pool.resident()} mem={pool.memory_mb()}/{pool.budget_mb}", + ) + check( + "refusing to prewarm evicts nothing", + set(pool.resident()) == {a, b}, + f"{pool.resident()}", + ) + check( + "prewarm ignores a label outside the project", pool.prewarm("nope") is False + ) # Budget eviction releases GUI leases but must not save somebody's open IDA # implicitly. An explicit save-and-close remains authoritative. @@ -239,12 +285,16 @@ def main() -> int: label = proj.refs[0].label pool.get(label) pool.evict(label) - check("LRU release does not implicitly save a GUI database", - made_gui[-1].saved == 0) + check( + "LRU release does not implicitly save a GUI database", + made_gui[-1].saved == 0, + ) pool.get(label) pool.close_all(save=True) - check("explicit close_all(save=True) does save a GUI database", - made_gui[-1].saved == 1) + check( + "explicit close_all(save=True) does save a GUI database", + made_gui[-1].saved == 1, + ) print(f"\n{PASS} passed, {FAIL} failed") return 1 if FAIL else 0 diff --git a/tests/test_project.py b/tests/test_project.py index bae0802..63d029e 100644 --- a/tests/test_project.py +++ b/tests/test_project.py @@ -15,7 +15,7 @@ import sys import tempfile sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) -from idatui.project import Project, ProjectError, SIDECAR_SUFFIX # noqa: E402 +from idatui.project import SIDECAR_SUFFIX, Project, ProjectError # noqa: E402 PASS = FAIL = 0 @@ -48,36 +48,48 @@ def main() -> int: proj = Project.create(pfile, [httpd, libauth], name="router-fw") check("create() writes the project file", os.path.isfile(pfile)) proj = Project.load(pfile) - check("load() round-trips name + binaries", - proj.name == "router-fw" and len(proj.refs) == 2, - f"name={proj.name} n={len(proj.refs)}") - check("labels default to the basename", - [r.label for r in proj.refs] == ["httpd", "libauth.so"], - f"{[r.label for r in proj.refs]}") + check( + "load() round-trips name + binaries", + proj.name == "router-fw" and len(proj.refs) == 2, + f"name={proj.name} n={len(proj.refs)}", + ) + check( + "labels default to the basename", + [r.label for r in proj.refs] == ["httpd", "libauth.so"], + f"{[r.label for r in proj.refs]}", + ) # -- layout -------------------------------------------------------- # - check("sidecar sits beside the project file", - proj.sidecar == os.path.join(tmp, "router-fw" + SIDECAR_SUFFIX), - proj.sidecar) + check( + "sidecar sits beside the project file", + proj.sidecar == os.path.join(tmp, "router-fw" + SIDECAR_SUFFIX), + proj.sidecar, + ) ref = proj.by_label("httpd") - check("staged path lives in the sidecar, not the source tree", - ref.staged.startswith(proj.bin_dir) and src not in ref.staged, - ref.staged) + check( + "staged path lives in the sidecar, not the source tree", + ref.staged.startswith(proj.bin_dir) and src not in ref.staged, + ref.staged, + ) check("db path hangs off the staged file", ref.db == ref.staged + ".i64") # -- staging: hardlink, freshness ---------------------------------- # check("a binary starts out stale (not yet staged)", proj.is_stale(ref)) proj.stage(ref) - check("stage() materialises the binary in the sidecar", - os.path.isfile(ref.staged)) - check("stage() copies (distinct inode) so the source can't be mutated " - "through it", - os.stat(ref.staged).st_ino != os.stat(ref.source).st_ino - and open(ref.staged, "rb").read() == open(ref.source, "rb").read()) + check( + "stage() materialises the binary in the sidecar", os.path.isfile(ref.staged) + ) + check( + "stage() copies (distinct inode) so the source can't be mutated through it", + os.stat(ref.staged).st_ino != os.stat(ref.source).st_ino + and open(ref.staged, "rb").read() == open(ref.source, "rb").read(), + ) check("a staged binary is no longer stale", not proj.is_stale(ref)) - check("source tree stays clean (no IDA artifacts beside it)", - sorted(os.listdir(src)) == ["httpd", "libauth.so"], - f"{sorted(os.listdir(src))}") + check( + "source tree stays clean (no IDA artifacts beside it)", + sorted(os.listdir(src)) == ["httpd", "libauth.so"], + f"{sorted(os.listdir(src))}", + ) # -- a changed source re-stages and drops the stale DB ------------- # open(ref.db, "wb").write(b"fake i64") @@ -88,11 +100,15 @@ def main() -> int: os.utime(ref.source, (1, 1)) check("a source rebuilt in place goes stale", proj.is_stale(ref)) proj.stage(ref) - check("re-staging refreshes the staged bytes", - open(ref.staged, "rb").read().endswith(b"v2 (longer)")) + check( + "re-staging refreshes the staged bytes", + open(ref.staged, "rb").read().endswith(b"v2 (longer)"), + ) check("re-staging drops the now-stale database", not proj.has_db(ref)) - check("re-staging drops the stale scratch too", - not os.path.exists(ref.staged + ".id0")) + check( + "re-staging drops the stale scratch too", + not os.path.exists(ref.staged + ".id0"), + ) # -- scratch sweep keeps the DB ------------------------------------ # open(ref.db, "wb").write(b"fake i64") @@ -107,22 +123,35 @@ def main() -> int: os.makedirs(sub) dup = _bin(os.path.join(sub, "httpd"), b"\x7fELF other httpd") with open(pfile, "w") as f: - json.dump({"name": "p", "binaries": [ - {"path": "src/httpd"}, # relative to the project file - {"path": dup}, # same basename -> collision - {"path": libauth, "label": "auth"}, - ]}, f) + json.dump( + { + "name": "p", + "binaries": [ + {"path": "src/httpd"}, # relative to the project file + {"path": dup}, # same basename -> collision + {"path": libauth, "label": "auth"}, + ], + }, + f, + ) proj2 = Project.load(pfile) - check("relative paths resolve against the project file", - proj2.refs[0].source == httpd, proj2.refs[0].source) - check("colliding labels are disambiguated", - [r.label for r in proj2.refs] == ["httpd", "httpd_2", "auth"], - f"{[r.label for r in proj2.refs]}") + check( + "relative paths resolve against the project file", + proj2.refs[0].source == httpd, + proj2.refs[0].source, + ) + check( + "colliding labels are disambiguated", + [r.label for r in proj2.refs] == ["httpd", "httpd_2", "auth"], + f"{[r.label for r in proj2.refs]}", + ) check("explicit labels are honoured", proj2.by_label("auth") is not None) proj2.stage_all() - check("stage_all() stages every binary to a distinct file", - len({r.staged for r in proj2.refs}) == 3 - and all(os.path.isfile(r.staged) for r in proj2.refs)) + check( + "stage_all() stages every binary to a distinct file", + len({r.staged for r in proj2.refs}) == 3 + and all(os.path.isfile(r.staged) for r in proj2.refs), + ) # -- add / remove ---------------------------------------------------- # extra = _bin(os.path.join(src, "extra")) @@ -132,39 +161,60 @@ def main() -> int: # -- re-adding must not duplicate (matched by resolved path) --------- # n = len(proj2.refs) proj2.add(extra) - check("re-adding the same path is a no-op", len(proj2.refs) == n, - f"{[r.label for r in proj2.refs]}") + check( + "re-adding the same path is a no-op", + len(proj2.refs) == n, + f"{[r.label for r in proj2.refs]}", + ) os.chdir(src) - proj2.add("./extra") # same file, relative - proj2.add(os.path.join(src, "..", "src", "extra")) # same file, messy - check("a different spelling of the same path is a no-op", - len(proj2.refs) == n, f"{[r.label for r in proj2.refs]}") + proj2.add("./extra") # same file, relative + proj2.add(os.path.join(src, "..", "src", "extra")) # same file, messy + check( + "a different spelling of the same path is a no-op", + len(proj2.refs) == n, + f"{[r.label for r in proj2.refs]}", + ) link = os.path.join(src, "extra_link") os.symlink(extra, link) proj2.add(link) - check("a symlink to an existing binary is a no-op", - len(proj2.refs) == n, f"{[r.label for r in proj2.refs]}") + check( + "a symlink to an existing binary is a no-op", + len(proj2.refs) == n, + f"{[r.label for r in proj2.refs]}", + ) # ...but a DIFFERENT file with the same basename must still be added other_dir = os.path.join(tmp, "other2") os.makedirs(other_dir) twin = _bin(os.path.join(other_dir, "extra"), b"\x7fELF a different extra") proj2.add(twin) - check("a same-named file from another directory IS added", - len(proj2.refs) == n + 1 - and proj2.by_source(twin) is not None - and proj2.by_source(extra) is not proj2.by_source(twin), - f"{[(r.label, r.source) for r in proj2.refs[-2:]]}") - check("the twins get distinct labels", - len({r.label for r in proj2.refs}) == len(proj2.refs), - f"{[r.label for r in proj2.refs]}") - check("create() also drops repeats on the command line", - len(Project.create(os.path.join(tmp, "dup.json"), - [extra, "./extra", extra]).refs) == 1) + check( + "a same-named file from another directory IS added", + len(proj2.refs) == n + 1 + and proj2.by_source(twin) is not None + and proj2.by_source(extra) is not proj2.by_source(twin), + f"{[(r.label, r.source) for r in proj2.refs[-2:]]}", + ) + check( + "the twins get distinct labels", + len({r.label for r in proj2.refs}) == len(proj2.refs), + f"{[r.label for r in proj2.refs]}", + ) + check( + "create() also drops repeats on the command line", + len( + Project.create( + os.path.join(tmp, "dup.json"), [extra, "./extra", extra] + ).refs + ) + == 1, + ) os.chdir(tmp) proj2.remove(proj2.by_source(twin).label) - check("remove() drops one", proj2.remove("extra") - and proj2.by_label("extra") is None) + check( + "remove() drops one", + proj2.remove("extra") and proj2.by_label("extra") is None, + ) # -- bad input ------------------------------------------------------- # bad = os.path.join(tmp, "bad.json") @@ -191,43 +241,62 @@ def main() -> int: # -- load options for headerless blobs --------------------------------- # with tempfile.TemporaryDirectory() as tmp: - src = os.path.join(tmp, "src"); os.makedirs(src) + src = os.path.join(tmp, "src") + os.makedirs(src) blob = os.path.join(src, "fw.bin") with open(blob, "wb") as f: f.write(b"\x00" * 64) - proj = Project.create(os.path.join(tmp, "p.json"), [blob], name="p", - load={"processor": "arm", "base": 0x8000000}) + proj = Project.create( + os.path.join(tmp, "p.json"), + [blob], + name="p", + load={"processor": "arm", "base": 0x8000000}, + ) r = proj.refs[0] - check("create() records load options per binary", - r.processor == "arm" and r.base == 0x8000000, - f"proc={r.processor!r} base={r.base:#x}") + check( + "create() records load options per binary", + r.processor == "arm" and r.base == 0x8000000, + f"proc={r.processor!r} base={r.base:#x}", + ) # -b is PARAGRAPHS: 0x8000000 >> 4 == 0x800000. Getting this wrong loads # the image 16x too high and every address in the database is wrong. - check("base is converted to IDA's paragraph units", - r.load_args == "-parm -b800000", r.load_args) + check( + "base is converted to IDA's paragraph units", + r.load_args == "-parm -b800000", + r.load_args, + ) proj2 = Project.load(proj.path) - check("load options survive a round-trip through the file", - proj2.refs[0].load_args == "-parm -b800000", - proj2.refs[0].load_args) + check( + "load options survive a round-trip through the file", + proj2.refs[0].load_args == "-parm -b800000", + proj2.refs[0].load_args, + ) blob2 = os.path.join(src, "other.bin") with open(blob2, "wb") as f: f.write(b"\x00" * 64) r2 = proj2.add(blob2, load={"processor": "mipsb"}) - check("add() takes load options too", - r2.load_args == "-pmipsb", r2.load_args) + check("add() takes load options too", r2.load_args == "-pmipsb", r2.load_args) # a normal ELF needs none of this and must pass nothing - check("a binary with no load options passes no switches", - Project.create(os.path.join(tmp, "q.json"), [blob], - name="q").refs[0].load_args == "") + check( + "a binary with no load options passes no switches", + Project.create(os.path.join(tmp, "q.json"), [blob], name="q") + .refs[0] + .load_args + == "", + ) # addresses get written by hand, so accept how people write them - proj3 = Project.create(os.path.join(tmp, "r.json"), [blob], name="r", - load={"base": "0x1000"}) - check("a base given as a hex STRING is parsed", - proj3.refs[0].base == 0x1000, f"{proj3.refs[0].base}") + proj3 = Project.create( + os.path.join(tmp, "r.json"), [blob], name="r", load={"base": "0x1000"} + ) + check( + "a base given as a hex STRING is parsed", + proj3.refs[0].base == 0x1000, + f"{proj3.refs[0].base}", + ) print(f"\n{PASS} passed, {FAIL} failed") return 1 if FAIL else 0 diff --git a/tests/test_project_ui.py b/tests/test_project_ui.py index 863b98a..9b388c8 100644 --- a/tests/test_project_ui.py +++ b/tests/test_project_ui.py @@ -21,12 +21,15 @@ import tempfile sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) from _fixtures import fast_keys # noqa: E402 -from idatui._sync import settle as quiesce, wait_for # noqa: E402 -fast_keys() # ~85ms -> ~2ms per keypress; see _fixtures.fast_keys +from idatui._sync import settle as quiesce # noqa: E402 +from idatui._sync import wait_for + +fast_keys() # ~85ms -> ~2ms per keypress; see _fixtures.fast_keys +from textual.widgets import Input, OptionList, Static # noqa: E402 + from idatui.app import IdaTui, ProjectPalette # noqa: E402 from idatui.project import Project # noqa: E402 -from textual.widgets import Input, OptionList, Static # noqa: E402 PASS = FAIL = 0 @@ -57,6 +60,7 @@ async def run(bins): app = IdaTui(keepalive=False, project=proj) async with app.run_test(size=(140, 44)) as pilot: + async def settle(pred, t=180.0): return await wait_for(pred, pilot.pause, t, 0.05) @@ -73,114 +77,163 @@ async def run(bins): waited for explicitly, not hoped for. """ return await settle( - lambda: app.program is not None - and app._func_index is not None - and app._func_index.complete - and app._loading_screen is None - and len(app.screen_stack) == 1, t) + lambda: ( + app.program is not None + and app._func_index is not None + and app._func_index.complete + and app._loading_screen is None + and len(app.screen_stack) == 1 + ), + t, + ) # -- boots on the project's first binary ----------------------- # ok = await usable() - check("project mode boots on the first binary", ok, - f"binary={app._binary}") - check("the active binary is the first one", app._binary == first, - f"{app._binary}") + check("project mode boots on the first binary", ok, f"binary={app._binary}") + check( + "the active binary is the first one", + app._binary == first, + f"{app._binary}", + ) n_first = len(app._func_index) check("its functions loaded", n_first > 10, f"n={n_first}") status = str(app.query_one("#status", Static).render()) - check("the status line names the active binary", - f"[{first}]" in status, status[:60]) + check( + "the status line names the active binary", + f"[{first}]" in status, + status[:60], + ) # -- the switcher lists the project ---------------------------- # await pilot.press("ctrl+o") - opened = await settle( - lambda: isinstance(app.screen, ProjectPalette), 20) - check("Ctrl+O opens the binary switcher", opened, - f"screen={type(app.screen).__name__}") + opened = await settle(lambda: isinstance(app.screen, ProjectPalette), 20) + check( + "Ctrl+O opens the binary switcher", + opened, + f"screen={type(app.screen).__name__}", + ) if not opened: return pal = app.screen - check("the switcher lists every project binary", - len(pal._results) == 2, f"{[e['label'] for e in pal._results]}") - check("it marks which one is active", - any(e["active"] and e["label"] == first for e in pal._results)) - check("it marks the other as not yet opened", - any(not e["resident"] and e["label"] == second - for e in pal._results)) + check( + "the switcher lists every project binary", + len(pal._results) == 2, + f"{[e['label'] for e in pal._results]}", + ) + check( + "it marks which one is active", + any(e["active"] and e["label"] == first for e in pal._results), + ) + check( + "it marks the other as not yet opened", + any(not e["resident"] and e["label"] == second for e in pal._results), + ) ol = pal.query_one(OptionList) - check("the switcher opens on the binary you're already in", - ol.highlighted is not None - and pal._results[ol.highlighted]["label"] == first, - f"highlighted={ol.highlighted} " - f"={pal._results[ol.highlighted]['label'] if ol.highlighted is not None else None} " - f"want={first}") + check( + "the switcher opens on the binary you're already in", + ol.highlighted is not None + and pal._results[ol.highlighted]["label"] == first, + f"highlighted={ol.highlighted} " + f"={pal._results[ol.highlighted]['label'] if ol.highlighted is not None else None} " + f"want={first}", + ) # -- switch to the second binary -------------------------------- # pal.query_one(Input).value = second await settle(lambda: bool(pal._results), 20) await pilot.press("enter") switched = await settle( - lambda: app._binary == second and app.program is not None - and app._func_index is not None and app._func_index.complete - and len(app.screen_stack) == 1) - check("switching opens the other binary", switched, - f"binary={app._binary}") - check("the second binary has its own function index", - app._func_index is not None and len(app._func_index) > 5, - f"n={len(app._func_index) if app._func_index else 0}") - check("both binaries now have live workers", - sorted(app._pool.resident()) == sorted([first, second]), - f"{app._pool.resident()}") + lambda: ( + app._binary == second + and app.program is not None + and app._func_index is not None + and app._func_index.complete + and len(app.screen_stack) == 1 + ) + ) + check("switching opens the other binary", switched, f"binary={app._binary}") + check( + "the second binary has its own function index", + app._func_index is not None and len(app._func_index) > 5, + f"n={len(app._func_index) if app._func_index else 0}", + ) + check( + "both binaries now have live workers", + sorted(app._pool.resident()) == sorted([first, second]), + f"{app._pool.resident()}", + ) landed = await settle(lambda: app._cur is not None, 60) - check("it lands somewhere in the new binary", landed, - f"cur={app._cur}") + check("it lands somewhere in the new binary", landed, f"cur={app._cur}") where = app._cur.ea if app._cur else None # -- switch back: resident, so state is restored ---------------- # await pilot.press("ctrl+o") - reopened = await settle( - lambda: isinstance(app.screen, ProjectPalette), 20) - check("the switcher reopens after a switch", reopened, - f"screen={type(app.screen).__name__}") + reopened = await settle(lambda: isinstance(app.screen, ProjectPalette), 20) + check( + "the switcher reopens after a switch", + reopened, + f"screen={type(app.screen).__name__}", + ) if not reopened: return app.screen.query_one(Input).value = first await settle(lambda: bool(app.screen._results), 20) await pilot.press("enter") - back = await settle(lambda: app._binary == first - and app._func_index is not None - and app._func_index.complete - and len(app.screen_stack) == 1, 120) - check("switching back returns to the first binary", back, - f"binary={app._binary}") - check("its function index came back intact", - app._func_index is not None and len(app._func_index) == n_first, - f"n={len(app._func_index) if app._func_index else 0} want={n_first}") + back = await settle( + lambda: ( + app._binary == first + and app._func_index is not None + and app._func_index.complete + and len(app.screen_stack) == 1 + ), + 120, + ) + check( + "switching back returns to the first binary", + back, + f"binary={app._binary}", + ) + check( + "its function index came back intact", + app._func_index is not None and len(app._func_index) == n_first, + f"n={len(app._func_index) if app._func_index else 0} want={n_first}", + ) # and forward again: the second binary's position was remembered await pilot.press("ctrl+o") if not await settle(lambda: isinstance(app.screen, ProjectPalette), 20): - check("returning to a binary restores where you were", False, - "switcher did not reopen") + check( + "returning to a binary restores where you were", + False, + "switcher did not reopen", + ) return app.screen.query_one(Input).value = second await settle(lambda: bool(app.screen._results), 20) await pilot.press("enter") - again = await settle(lambda: app._binary == second - and app._cur is not None, 120) - check("returning to a binary restores where you were", - again and app._cur.ea == where, - f"cur={app._cur.ea if app._cur else None} want={where}") + again = await settle( + lambda: app._binary == second and app._cur is not None, 120 + ) + check( + "returning to a binary restores where you were", + again and app._cur.ea == where, + f"cur={app._cur.ea if app._cur else None} want={where}", + ) # -- project-wide symbol search ------------------------------- # # 'main' exists in BOTH binaries: identical name, so the rank tuple # ties and a bare sort() would fall through to comparing Hit objects # ('<' not supported between instances of 'Hit'). from idatui.app import SymbolPalette - await settle(lambda: app._index is not None - and len(app._index.counts()) == 2, 60) - check("both binaries got indexed", - len(app._index.counts()) == 2, f"{app._index.counts()}") + + await settle( + lambda: app._index is not None and len(app._index.counts()) == 2, 60 + ) + check( + "both binaries got indexed", + len(app._index.counts()) == 2, + f"{app._index.counts()}", + ) await pilot.press("ctrl+n") if await settle(lambda: isinstance(app.screen, SymbolPalette), 20): pal = app.screen @@ -190,12 +243,14 @@ async def run(bins): # quiescence -- NOT "a foreign binary appeared", which is the # thing under test and would sit out its whole timeout on the # day it breaks. - await pilot.press("f2") # widen to the whole project + await pilot.press("f2") # widen to the whole project await quiesce(app) names = [(b, n) for b, _, n in pal._results] - check("project scope finds a name shared by both binaries", - len({b for b, n in names if n == "main"}) == 2, - f"{names[:6]}") + check( + "project scope finds a name shared by both binaries", + len({b for b, n in names if n == "main"}) == 2, + f"{names[:6]}", + ) await pilot.press("escape") await settle(lambda: not isinstance(app.screen, SymbolPalette), 20) @@ -209,17 +264,31 @@ async def run(bins): target = app._index.search("main", limit=200) tgt = next((h for h in target if h.binary == there), None) if tgt is None: - check("cross-binary jump records a hop", False, "no hit in the other binary") + check( + "cross-binary jump records a hop", + False, + "no hit in the other binary", + ) else: app._switch_then_goto(tgt.binary, tgt.addr) - jumped = await settle(lambda: app._binary == there - and app._func_index is not None - and app._func_index.complete, 180) - check("a project hit switches to the other binary", jumped, - f"binary={app._binary} want={there}") - check("the jump records where it came from", - len(app._hops) == hops0 + 1 and app._hops[-1] == here, - f"hops={app._hops}") + jumped = await settle( + lambda: ( + app._binary == there + and app._func_index is not None + and app._func_index.complete + ), + 180, + ) + check( + "a project hit switches to the other binary", + jumped, + f"binary={app._binary} want={there}", + ) + check( + "the jump records where it came from", + len(app._hops) == hops0 + 1 and app._hops[-1] == here, + f"hops={app._hops}", + ) # spend the local history first, then Esc must cross back for _ in range(6): if not app._hops or app._binary != there: @@ -227,10 +296,16 @@ async def run(bins): await pilot.press("escape") await quiesce(app) returned = await settle(lambda: app._binary == here, 180) - check("Esc crosses back to the binary the jump came from", - returned, f"binary={app._binary} want={here} hops={app._hops}") - check("the hop is consumed, not repeated", - not app._hops, f"hops={app._hops}") + check( + "Esc crosses back to the binary the jump came from", + returned, + f"binary={app._binary} want={here} hops={app._hops}", + ) + check( + "the hop is consumed, not repeated", + not app._hops, + f"hops={app._hops}", + ) # -- xrefs: callers in OTHER project binaries ------------------ # # xrefs_to only sees this database, so an exported function looks @@ -238,9 +313,13 @@ async def run(bins): # The selection rule is "only for a symbol we actually export"; two # executables share no linkage, so here it must stay quiet. from idatui.app import XrefsScreen + fake = app._foreign_importers(app._cur.ea, "strrchr", None) - check("no cross-binary callers for a symbol this binary doesn't export", - fake == [], f"{fake}") + check( + "no cross-binary callers for a symbol this binary doesn't export", + fake == [], + f"{fake}", + ) # The routing a real cross-binary caller takes: the dialog carries a # (binary, addr) payload instead of a bare address, and choosing it @@ -248,31 +327,55 @@ async def run(bins): # so Esc comes back. where_from = app._binary other = first if where_from == second else second - hit = next((h for h in app._index.search("main", limit=200) - if h.binary == other), None) + hit = next( + (h for h in app._index.search("main", limit=200) if h.binary == other), + None, + ) if hit is None: - check("a cross-binary xref jumps to the other binary", False, - "no symbol found in the other binary") + check( + "a cross-binary xref jumps to the other binary", + False, + "no symbol found in the other binary", + ) else: hops0 = len(app._hops) app.push_screen( - XrefsScreen("xrefs to fake", [((hit.binary, hit.addr), - f"{hit.addr:08X} import [{hit.binary}]")]), - app._on_xref_chosen) + XrefsScreen( + "xrefs to fake", + [ + ( + (hit.binary, hit.addr), + f"{hit.addr:08X} import [{hit.binary}]", + ) + ], + ), + app._on_xref_chosen, + ) await settle(lambda: isinstance(app.screen, XrefsScreen), 20) await pilot.press("enter") - jumped = await settle(lambda: app._binary == other - and app._func_index is not None - and app._func_index.complete, 180) - check("a cross-binary xref jumps to the other binary", jumped, - f"binary={app._binary} want={other}") - check("and records a hop so Esc returns", - len(app._hops) == hops0 + 1 and app._hops[-1] == where_from, - f"hops={app._hops}") + jumped = await settle( + lambda: ( + app._binary == other + and app._func_index is not None + and app._func_index.complete + ), + 180, + ) + check( + "a cross-binary xref jumps to the other binary", + jumped, + f"binary={app._binary} want={other}", + ) + check( + "and records a hop so Esc returns", + len(app._hops) == hops0 + 1 and app._hops[-1] == where_from, + f"hops={app._hops}", + ) app._hops.clear() # -- and the same toggle for strings --------------------------- # from idatui.app import StringsPalette + await pilot.press("quotation_mark") if await settle(lambda: isinstance(app.screen, StringsPalette), 30): pal = app.screen @@ -282,20 +385,32 @@ async def run(bins): await pilot.press("f2") await quiesce(app) wide = {b for b, _, _ in pal._results} - check("strings: local scope is this binary only", local == {None}, - f"{local}") - check("strings: F2 widens across the project", - len(wide) >= 2 and None not in wide, f"{wide}") + check( + "strings: local scope is this binary only", + local == {None}, + f"{local}", + ) + check( + "strings: F2 widens across the project", + len(wide) >= 2 and None not in wide, + f"{wide}", + ) await pilot.press("escape") await settle(lambda: not isinstance(app.screen, StringsPalette), 20) # -- the promise: nothing was written next to the sources ---------- # left = sorted(os.listdir(src)) - check("the source tree stays pristine (no .i64/scratch beside it)", - left == sorted(os.path.basename(s) for s in srcs), f"{left}") + check( + "the source tree stays pristine (no .i64/scratch beside it)", + left == sorted(os.path.basename(s) for s in srcs), + f"{left}", + ) staged = sorted(os.listdir(proj.bin_dir)) - check("IDA's artifacts all live in the project sidecar", - any(f.endswith(".i64") for f in staged), f"{staged}") + check( + "IDA's artifacts all live in the project sidecar", + any(f.endswith(".i64") for f in staged), + f"{staged}", + ) print(f"\n{PASS} passed, {FAIL} failed") return 1 if FAIL else 0 @@ -303,8 +418,10 @@ async def run(bins): def main(argv): repo = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) - bins = argv or [os.path.join(repo, "targets", "echo"), - os.path.join(repo, "targets", "cat")] + bins = argv or [ + os.path.join(repo, "targets", "echo"), + os.path.join(repo, "targets", "cat"), + ] for b in bins: if not os.path.isfile(b): print(f"no such binary: {b}") diff --git a/tests/test_rawimage_rpc.py b/tests/test_rawimage_rpc.py index d39ad35..a6e7661 100644 --- a/tests/test_rawimage_rpc.py +++ b/tests/test_rawimage_rpc.py @@ -36,7 +36,7 @@ sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) from idatui.rpcclient import RpcClient, RpcError # noqa: E402 REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) -BLOB = os.path.join(REPO, "experiments", "fibonacci.bin") # real Thumb code +BLOB = os.path.join(REPO, "experiments", "fibonacci.bin") # real Thumb code PASS = FAIL = 0 @@ -52,11 +52,24 @@ def check(name, ok, detail=""): def spawn_pane(target, processor, timeout=420): - cmd = [sys.executable, "-m", "idatui.pane", "spawn", "--open", target, - "--processor", processor, "--detached", "--size", "60%", - "--timeout", str(timeout)] - r = subprocess.run(cmd, capture_output=True, text=True, - timeout=timeout + 60, cwd=REPO) + cmd = [ + sys.executable, + "-m", + "idatui.pane", + "spawn", + "--open", + target, + "--processor", + processor, + "--detached", + "--size", + "60%", + "--timeout", + str(timeout), + ] + r = subprocess.run( + cmd, capture_output=True, text=True, timeout=timeout + 60, cwd=REPO + ) if not r.stdout.strip(): print(f" spawn produced no JSON: {r.stderr.strip()}", file=sys.stderr) return None @@ -64,9 +77,22 @@ def spawn_pane(target, processor, timeout=420): def stop_pane(sock, timeout=60): - subprocess.run([sys.executable, "-m", "idatui.pane", "stop", "--sock", sock, - "--timeout", str(timeout)], - capture_output=True, text=True, timeout=timeout + 10, cwd=REPO) + subprocess.run( + [ + sys.executable, + "-m", + "idatui.pane", + "stop", + "--sock", + sock, + "--timeout", + str(timeout), + ], + capture_output=True, + text=True, + timeout=timeout + 10, + cwd=REPO, + ) def main() -> int: @@ -94,34 +120,44 @@ def main() -> int: # -- load options actually reached IDA --------------------------- # # Wrong processor => the disassembly is nonsense or absent; ARMv7-A # also means a 32-bit database, without which Hex-Rays refuses. - check("spawn forwarded --processor", info.get("ok"), - json.dumps(info)) + check("spawn forwarded --processor", info.get("ok"), json.dumps(info)) with RpcClient(sock) as c: st = c.call("state") - check("pane is drivable", st.get("active") in - ("listing", "decomp", "hex"), json.dumps(st)[:200]) + check( + "pane is drivable", + st.get("active") in ("listing", "decomp", "hex"), + json.dumps(st)[:200], + ) # -- define ------------------------------------------------- # # fibonacci.bin is Thumb at 0x0; as ARM it does not decode. r = c.call("define", kind="thumb", target="0x0") d = r.get("define", {}) check("define thumb ran", "define" in r, json.dumps(r)[:200]) - check("define thumb decoded instructions", - "instruction" in d.get("status", ""), d.get("status", "")) + check( + "define thumb decoded instructions", + "instruction" in d.get("status", ""), + d.get("status", ""), + ) r = c.call("define", kind="func", target="0x0") - check("define func created a function", - "function" in r["define"]["status"] - or "already" in r["define"]["status"], - r["define"]["status"]) + check( + "define func created a function", + "function" in r["define"]["status"] + or "already" in r["define"]["status"], + r["define"]["status"], + ) bad = None try: c.call("define", kind="nonsense") except RpcError as e: bad = str(e) - check("define rejects an unknown kind", bad is not None - and "unknown define kind" in bad, str(bad)) + check( + "define rejects an unknown kind", + bad is not None and "unknown define kind" in bad, + str(bad), + ) # -- opfmt (how a literal is displayed) --------------------- # # Thumb code is full of small immediates -- the thing 'o' exists @@ -134,21 +170,31 @@ def main() -> int: lit = None for ln in seen: m = re.match(r"([0-9A-F]{8})\s+(.*)", ln.get("text", "")) - if not (m and re.search(r"#(0x[0-9A-Fa-f]{2,}|[1-9]\d+)\b", - m.group(2))): + if not ( + m and re.search(r"#(0x[0-9A-Fa-f]{2,}|[1-9]\d+)\b", m.group(2)) + ): continue ea_s = "0x" + m.group(1) - st = c.call("opfmt", mode="show", target=ea_s, delay_ms=0 - ).get("opfmt", {}).get("status", "") + st = ( + c.call("opfmt", mode="show", target=ea_s, delay_ms=0) + .get("opfmt", {}) + .get("status", "") + ) if "no literal" not in st: lit = (ea_s, st) break - check("the blob has an immediate to reformat", lit is not None, - json.dumps([ln.get("text") for ln in seen[:8]])) + check( + "the blob has an immediate to reformat", + lit is not None, + json.dumps([ln.get("text") for ln in seen[:8]]), + ) if lit is not None: tgt, st0 = lit - check("opfmt show reports the stops without editing", - "[" in st0 and "dec" in st0, st0) + check( + "opfmt show reports the stops without editing", + "[" in st0 and "dec" in st0, + st0, + ) r = c.call("opfmt", mode="dec", target=tgt, delay_ms=0) st1 = r.get("opfmt", {}).get("status", "") check("opfmt sets a named format", "dec" in st1, st1) @@ -156,16 +202,21 @@ def main() -> int: st2 = r.get("opfmt", {}).get("status", "") check("opfmt cycles on from there", "\u2192" in st2, st2) r = c.call("opfmt", mode="default") - check("opfmt hands the operand back to IDA", - "default" in r.get("opfmt", {}).get("status", ""), - r.get("opfmt", {}).get("status", "")) + check( + "opfmt hands the operand back to IDA", + "default" in r.get("opfmt", {}).get("status", ""), + r.get("opfmt", {}).get("status", ""), + ) badfmt = None try: c.call("opfmt", mode="roman") except RpcError as e: badfmt = str(e) - check("opfmt rejects an unknown mode", badfmt is not None - and "unknown opfmt mode" in badfmt, str(badfmt)) + check( + "opfmt rejects an unknown mode", + badfmt is not None and "unknown opfmt mode" in badfmt, + str(badfmt), + ) # -- rename_many -------------------------------------------- # fns = c.call("functions", limit=200) @@ -177,28 +228,48 @@ def main() -> int: # 'start' (not 'addr') on purpose: symbol files in the wild # use it, and accepting only one spelling is how a bulk # import silently renames nothing. - json.dump([{"start": hex(ea), "name": "bulk_named_fn"}, - {"start": "0xdeadbe", "name": "nowhere"}], f) + json.dump( + [ + {"start": hex(ea), "name": "bulk_named_fn"}, + {"start": "0xdeadbe", "name": "nowhere"}, + ], + f, + ) r = c.call("rename_many", file=symfile) m = r.get("rename_many", {}) - check("rename_many applied the good entry", m.get("ok") == 1, - json.dumps(m)) - check("rename_many reports the bad entry", - m.get("failed") == 1 and m.get("errors"), json.dumps(m)) + check( + "rename_many applied the good entry", + m.get("ok") == 1, + json.dumps(m), + ) + check( + "rename_many reports the bad entry", + m.get("failed") == 1 and m.get("errors"), + json.dumps(m), + ) # The readback matters more than the return value: a driver # trusts resolve/functions to decide what work is left. - check("renamed symbol resolves", - c.call("resolve", name="bulk_named_fn").get("ea") == ea, - json.dumps(c.call("resolve", name="bulk_named_fn"))) + check( + "renamed symbol resolves", + c.call("resolve", name="bulk_named_fn").get("ea") == ea, + json.dumps(c.call("resolve", name="bulk_named_fn")), + ) names = {f["name"] for f in c.call("functions", limit=200)} - check("function table shows the new name", - "bulk_named_fn" in names, str(sorted(names)[:10])) - - r = c.call("rename_many", items=[{"addr": hex(ea), - "name": "inline_named_fn"}]) - check("rename_many takes inline items", - r["rename_many"]["ok"] == 1, json.dumps(r["rename_many"])) + check( + "function table shows the new name", + "bulk_named_fn" in names, + str(sorted(names)[:10]), + ) + + r = c.call( + "rename_many", items=[{"addr": hex(ea), "name": "inline_named_fn"}] + ) + check( + "rename_many takes inline items", + r["rename_many"]["ok"] == 1, + json.dumps(r["rename_many"]), + ) # -- the stale-pseudocode trap ------------------------------ # # Hex-Rays caches per function and does not notice that a @@ -208,22 +279,31 @@ def main() -> int: # the function's own body cites it.) before = c.call("pseudocode", target=hex(ea)) pc_before = json.dumps(before) - r = c.call("rename_many", items=[{"addr": hex(ea), - "name": "after_cache_fn"}]) + r = c.call( + "rename_many", items=[{"addr": hex(ea), "name": "after_cache_fn"}] + ) pc_after = json.dumps(c.call("pseudocode", target=hex(ea))) - check("pseudocode was cached before the rename", - "inline_named_fn" in pc_before, pc_before[:200]) - check("rename_many invalidates the decompile cache", - "after_cache_fn" in pc_after - and "inline_named_fn" not in pc_after, pc_after[:300]) + check( + "pseudocode was cached before the rename", + "inline_named_fn" in pc_before, + pc_before[:200], + ) + check( + "rename_many invalidates the decompile cache", + "after_cache_fn" in pc_after and "inline_named_fn" not in pc_after, + pc_after[:300], + ) empty = None try: c.call("rename_many") except RpcError as e: empty = str(e) - check("rename_many without items errors", empty is not None - and "items" in empty, str(empty)) + check( + "rename_many without items errors", + empty is not None and "items" in empty, + str(empty), + ) finally: stop_pane(sock) diff --git a/tests/test_scenarios.py b/tests/test_scenarios.py index 30d6087..4196445 100644 --- a/tests/test_scenarios.py +++ b/tests/test_scenarios.py @@ -32,15 +32,26 @@ 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 +from rich.text import Text # noqa: E402 +from textual.widgets import ( # noqa: E402 + DataTable, + Input, + OptionList, + Static, + TextArea, +) + from idatui import remote_ops # noqa: E402 +from idatui._sync import settle, wait_for # noqa: E402 from idatui.app import ( # noqa: E402 + _HELP, ConfirmScreen, DecompView, FunctionsPanel, GraphView, + HelpScreen, HexView, IdaTui, - HelpScreen, ListingView, QuitScreen, SearchPalette, @@ -48,20 +59,10 @@ from idatui.app import ( # noqa: E402 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, -) -from rich.text import Text # noqa: E402 -from idatui._sync import settle, wait_for # noqa: E402 PASS = FAIL = 0 STOP_AFTER = None diff --git a/tests/test_search.py b/tests/test_search.py index 6fd1a25..7b80dda 100644 --- a/tests/test_search.py +++ b/tests/test_search.py @@ -16,8 +16,13 @@ import sys sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) from idatui.search import ( # noqa: E402 - BYTES, TEXT, classify, looks_like_bytes, normalise_pattern, - pattern_problem, probably_meant_bytes, + BYTES, + TEXT, + classify, + looks_like_bytes, + normalise_pattern, + pattern_problem, + probably_meant_bytes, ) PASS = FAIL = 0 @@ -35,63 +40,89 @@ def check(name, cond, detail=""): def main() -> int: # -- the asymmetry: hex-looking WORDS must stay text --------------------- # - for word in ("add", "dead", "beef", "cafe", "ff", "0", "abcdef", - "decode", "face"): - check(f"{word!r} searches text, not bytes", - classify(word)[0] == TEXT, classify(word)) + for word in ("add", "dead", "beef", "cafe", "ff", "0", "abcdef", "decode", "face"): + check( + f"{word!r} searches text, not bytes", + classify(word)[0] == TEXT, + classify(word), + ) # -- unambiguous byte patterns ------------------------------------------ # - for pat in ("48 8b ?? c3", "B8 ? ? ? ? 90", "48,8b,05", "de ad be ef", - "48 8? ?? 24", "??"): - check(f"{pat!r} searches bytes", classify(pat)[0] == BYTES, - classify(pat)) - - check("a quoted literal is a byte pattern", - classify('"Hello", 0')[0] == BYTES) + for pat in ( + "48 8b ?? c3", + "B8 ? ? ? ? 90", + "48,8b,05", + "de ad be ef", + "48 8? ?? 24", + "??", + ): + check(f"{pat!r} searches bytes", classify(pat)[0] == BYTES, classify(pat)) + + check("a quoted literal is a byte pattern", classify('"Hello", 0')[0] == BYTES) # A TYPO in a byte pattern must stay a byte pattern, so it can be refused # with a reason. Falling back to text answers "no match", which is # indistinguishable from "those bytes are not in this binary". - check("a typo'd byte pattern is still a byte pattern", - classify("48 zz c3")[0] == BYTES, classify("48 zz c3")) - check("and it is refused by name", - "'zz'" in (pattern_problem("48 zz c3") or "")) - check("but a word among bytes is prose", - classify("add ff")[0] == TEXT and classify("mov rdi, rax")[0] == TEXT, - classify("add ff")) + check( + "a typo'd byte pattern is still a byte pattern", + classify("48 zz c3")[0] == BYTES, + classify("48 zz c3"), + ) + check("and it is refused by name", "'zz'" in (pattern_problem("48 zz c3") or "")) + check( + "but a word among bytes is prose", + classify("add ff")[0] == TEXT and classify("mov rdi, rax")[0] == TEXT, + classify("add ff"), + ) check("prose stays text", classify("mov rdi, rax")[0] == TEXT) check("a call target stays text", classify("call cs:__isoc99_scanf")[0] == TEXT) - check("an empty query is text (nothing to search yet)", - classify("")[0] == TEXT and not looks_like_bytes("")) + check( + "an empty query is text (nothing to search yet)", + classify("")[0] == TEXT and not looks_like_bytes(""), + ) # -- explicit wins over any guess --------------------------------------- # check("hex: forces bytes", classify("hex: dead") == (BYTES, "dead")) check("bytes: forces bytes too", classify("bytes:dead") == (BYTES, "dead")) check("text: forces text", classify("text: 48 8b c3") == (TEXT, "48 8b c3")) - check("F2's forced mode beats the shape", - classify("dead", forced=BYTES) == (BYTES, "dead") - and classify("48 8b c3", forced=TEXT) == (TEXT, "48 8b c3")) - check("a prefix beats even the forced mode", - classify("text:48 8b c3", forced=BYTES)[0] == TEXT) + check( + "F2's forced mode beats the shape", + classify("dead", forced=BYTES) == (BYTES, "dead") + and classify("48 8b c3", forced=TEXT) == (TEXT, "48 8b c3"), + ) + check( + "a prefix beats even the forced mode", + classify("text:48 8b c3", forced=BYTES)[0] == TEXT, + ) # -- the shapes people paste -------------------------------------------- # check("commas become spaces", normalise_pattern("48,8b,05") == "48 8b 05") - check("a run with no separators is split into bytes", - normalise_pattern("488B05C3") == "48 8B 05 C3") + check( + "a run with no separators is split into bytes", + normalise_pattern("488B05C3") == "48 8B 05 C3", + ) check("whitespace is squeezed", normalise_pattern(" 48 8b\t05 ") == "48 8b 05") - check("a quoted literal keeps its own spacing", - normalise_pattern('"Hello, world", 0') == '"Hello, world", 0') + check( + "a quoted literal keeps its own spacing", + normalise_pattern('"Hello, world", 0') == '"Hello, world", 0', + ) # -- refusing a bad pattern with a reason -------------------------------- # - check("an empty pattern says what to type", - "48 8b" in (pattern_problem("") or "")) - check("a non-hex token is named", - "'zz'" in (pattern_problem("48 zz c3") or ""), pattern_problem("48 zz c3")) - check("a good pattern has no complaint", - pattern_problem("48 8b ?? c3") is None - and pattern_problem('"Hi", 0') is None) - check("an odd-length run is refused rather than silently split", - pattern_problem("488B0") is not None, pattern_problem("488B0")) + check("an empty pattern says what to type", "48 8b" in (pattern_problem("") or "")) + check( + "a non-hex token is named", + "'zz'" in (pattern_problem("48 zz c3") or ""), + pattern_problem("48 zz c3"), + ) + check( + "a good pattern has no complaint", + pattern_problem("48 8b ?? c3") is None and pattern_problem('"Hi", 0') is None, + ) + check( + "an odd-length run is refused rather than silently split", + pattern_problem("488B0") is not None, + pattern_problem("488B0"), + ) print(f"\n{PASS} passed, {FAIL} failed") return 1 if FAIL else 0 diff --git a/tests/test_thumb_ui.py b/tests/test_thumb_ui.py index 1b4c438..d55fa91 100644 --- a/tests/test_thumb_ui.py +++ b/tests/test_thumb_ui.py @@ -22,17 +22,20 @@ import tempfile sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +from _fixtures import fast_keys # noqa: E402 from textual.widgets import Static # noqa: E402 -from _fixtures import fast_keys # noqa: E402 from idatui._sync import settle # noqa: E402 -fast_keys() # ~85ms -> ~2ms per keypress; see _fixtures.fast_keys +fast_keys() # ~85ms -> ~2ms per keypress; see _fixtures.fast_keys from idatui.app import DecompView, IdaTui, ListingView # noqa: E402 PASS = FAIL = 0 -BIN = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), - "experiments", "fibonacci.bin") +BIN = os.path.join( + os.path.dirname(os.path.dirname(os.path.abspath(__file__))), + "experiments", + "fibonacci.bin", +) def check(name, ok, detail=""): @@ -45,7 +48,6 @@ def check(name, ok, detail=""): print(f" FAIL {name} {detail}") - #: Every phase gets its OWN copy of the fixture. #: #: This suite used to delete .i64 and reopen the SAME path for each phase. @@ -89,19 +91,22 @@ async def wait(pred, pilot, t=240.0): async def run() -> int: # A fresh database every time: the T flag and the segment's addressing mode # are SAVED in the .i64, so a previous run would answer the question for us. - app = IdaTui(open_path=fresh_copy(BIN, "arm"), keepalive=False, - load_args="-parm") + app = IdaTui(open_path=fresh_copy(BIN, "arm"), keepalive=False, load_args="-parm") async with app.run_test(size=(140, 44)) as pilot: - await wait(lambda: app._func_index is not None - and app._func_index.complete, pilot) + await wait( + lambda: app._func_index is not None and app._func_index.complete, pilot + ) await wait(lambda: app._cur is not None, pilot, 60) lst = app.query_one(ListingView) lst.focus() lst.cursor = lst.model.index_of_ea(0) lst._scroll_cursor_into_view() await settle(app) - check("starts undefined at the entry", lst.model.get(lst.cursor).kind == "unknown", - f"{lst.model.get(lst.cursor).text!r}") + check( + "starts undefined at the entry", + lst.model.get(lst.cursor).kind == "unknown", + f"{lst.model.get(lst.cursor).text!r}", + ) # `c` in the wrong mode: this is the failure being fixed. It must NOT # quietly carve garbage — either it refuses, or whatever it makes is not @@ -113,9 +118,11 @@ async def run() -> int: # queue doesn't say better. await settle(app) h = lst.model.get(lst.model.index_of_ea(0)) - check("`c` alone does not produce the Thumb prologue", - h is None or h.kind != "code" or "PUSH" not in h.text.upper(), - f"{h.text if h else None!r}") + check( + "`c` alone does not produce the Thumb prologue", + h is None or h.kind != "code" or "PUSH" not in h.text.upper(), + f"{h.text if h else None!r}", + ) m1 = lst.model await pilot.press("t") @@ -130,23 +137,31 @@ async def run() -> int: check("the status says it switched to Thumb", "Thumb" in status, status[:90]) # Thumb doesn't exist in AArch64, and -parm on a headerless blob gives a # 64-bit segment, so setting T alone would change nothing and look broken. - check("and says it forced the segment to 32-bit", - "32-bit" in status, status[:90]) + check( + "and says it forced the segment to 32-bit", "32-bit" in status, status[:90] + ) m = lst.model rows = [m.get(m.index_of_ea(ea)) for ea in (0x0, 0x2, 0x4)] - check("the entry decodes as Thumb", - rows[0] is not None and rows[0].kind == "code" - and "PUSH" in rows[0].text.upper(), - f"{rows[0].text if rows[0] else None!r}") + check( + "the entry decodes as Thumb", + rows[0] is not None + and rows[0].kind == "code" + and "PUSH" in rows[0].text.upper(), + f"{rows[0].text if rows[0] else None!r}", + ) # 16-bit instructions: the addresses are 2 apart, which is the whole # point — in ARM mode these would be one 4-byte instruction. - check("instructions are 16-bit wide", - all(r is not None and r.kind == "code" and r.size == 2 for r in rows), - f"{[(hex(r.ea), r.size, r.text) for r in rows if r]}") - check("and it kept disassembling past the first one", - sum(1 for i in range(20) if (m.get(i) or h).kind == "code") > 5, - "expected a run of instructions, not one") + check( + "instructions are 16-bit wide", + all(r is not None and r.kind == "code" and r.size == 2 for r in rows), + f"{[(hex(r.ea), r.size, r.text) for r in rows if r]}", + ) + check( + "and it kept disassembling past the first one", + sum(1 for i in range(20) if (m.get(i) or h).kind == "code") > 5, + "expected a run of instructions, not one", + ) # Toggling back must be possible — the mode is a guess and guesses get # revised. @@ -164,11 +179,13 @@ async def run() -> int: # disassembly that F5 can never turn into pseudocode. The database's bitness # is fixed at load and cannot be corrected afterwards, so the only honest # thing is to say so. - app = IdaTui(open_path=fresh_copy(BIN, "arm64"), keepalive=False, - load_args="-parm") # 64-bit + app = IdaTui( + open_path=fresh_copy(BIN, "arm64"), keepalive=False, load_args="-parm" + ) # 64-bit async with app.run_test(size=(140, 44)) as pilot: - await wait(lambda: app._func_index is not None - and app._func_index.complete, pilot) + await wait( + lambda: app._func_index is not None and app._func_index.complete, pilot + ) await wait(lambda: app._cur is not None, pilot, 60) lst = app.query_one(ListingView) lst.focus() @@ -179,8 +196,11 @@ async def run() -> int: await pilot.press("t") await settle(app, lambda: "64-bit" in status_of(app), timeout=60) status = status_of(app) - check("a 64-bit database warns that Hex-Rays won't decompile", - "64-bit" in status and "decompile" in status, status[:120]) + check( + "a 64-bit database warns that Hex-Rays won't decompile", + "64-bit" in status and "decompile" in status, + status[:120], + ) check("and names the fix", "ARMv7-A" in status, status[:120]) # And if you ignore that and carry on, the failure has to say WHY. The @@ -193,66 +213,108 @@ async def run() -> int: await pilot.press("p") # The function appearing in the index IS the signal; the model identity # never was one. - await settle(app, lambda: app._func_index is not None - and len(app._func_index) > 0, timeout=60) + await settle( + app, + lambda: app._func_index is not None and len(app._func_index) > 0, + timeout=60, + ) await pilot.press("tab") - await wait(lambda: "cannot decompile" in - str(app.query_one("#status", Static).render()), pilot, 90) + await wait( + lambda: ( + "cannot decompile" in str(app.query_one("#status", Static).render()) + ), + pilot, + 90, + ) status = str(app.query_one("#status", Static).render()) # The message must say what to DO. Hex-Rays' own sentence ("only 64-bit # functions can be decompiled in the current database") describes the # database, not the fix, and is long enough that a status bar cuts off # the end — which is where an appended hint would have lived. - check("a failed decompile names the fix, not just the diagnosis", - "Ctrl+L" in status and "ARMv7-A" in status, status[:130]) - check("and the reason survives the view reloading under it", - "cannot decompile" in status, status[:130]) - check("the message fits a narrow status bar", - len(status) < 110, f"{len(status)} chars: {status[:130]}") + check( + "a failed decompile names the fix, not just the diagnosis", + "Ctrl+L" in status and "ARMv7-A" in status, + status[:130], + ) + check( + "and the reason survives the view reloading under it", + "cannot decompile" in status, + status[:130], + ) + check( + "the message fits a narrow status bar", + len(status) < 110, + f"{len(status)} chars: {status[:130]}", + ) # -- the whole point: a 32-bit database decompiles ---------------------- # - app = IdaTui(open_path=fresh_copy(BIN, "armv7a"), keepalive=False, - load_args="-parm:ARMv7-A") + app = IdaTui( + open_path=fresh_copy(BIN, "armv7a"), keepalive=False, load_args="-parm:ARMv7-A" + ) async with app.run_test(size=(140, 44)) as pilot: - await wait(lambda: app._func_index is not None - and app._func_index.complete, pilot) + await wait( + lambda: app._func_index is not None and app._func_index.complete, pilot + ) # A 32-bit ARM database also lets auto-analysis do its job on Thumb code, # which is why this one lands in the symbol picker rather than nowhere. - check("a 32-bit ARM database finds functions by itself", - len(app._func_index) > 5, f"n={len(app._func_index)}") + check( + "a 32-bit ARM database finds functions by itself", + len(app._func_index) > 5, + f"n={len(app._func_index)}", + ) await pilot.press("escape") await settle(app, lambda: type(app.screen).__name__ == "Screen") f = app._func_index.all_loaded()[0] app._goto_ea(f.addr, push=True) - await wait(lambda: app._cur is not None - and app.query_one(ListingView).model is not None, pilot, 60) + await wait( + lambda: ( + app._cur is not None and app.query_one(ListingView).model is not None + ), + pilot, + 60, + ) app.query_one(ListingView).focus() await pilot.press("tab") dec = app.query_one(DecompView) got = await wait(lambda: dec.display and dec._texts, pilot, 90) - check("Tab decompiles a Thumb function", got and len(dec._texts) > 3, - f"lines={len(dec._texts or [])}") - check("and it reads like C", - any("(" in t and ")" in t for t in (dec._texts or [])[:3]), - f"{(dec._texts or [])[:3]}") + check( + "Tab decompiles a Thumb function", + got and len(dec._texts) > 3, + f"lines={len(dec._texts or [])}", + ) + check( + "and it reads like C", + any("(" in t and ")" in t for t in (dec._texts or [])[:3]), + f"{(dec._texts or [])[:3]}", + ) # -- Thumb entry points from a vector table ----------------------------- # # An ARM function pointer carries the mode in bit 0: odd means Thumb. A # Cortex-M vector table is therefore a list of Thumb entry points, and IDA # won't follow them on a headerless image because nothing says those words # are pointers at all. - vec = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), - "experiments", "cortexm.bin") + vec = os.path.join( + os.path.dirname(os.path.dirname(os.path.abspath(__file__))), + "experiments", + "cortexm.bin", + ) if not os.path.isfile(vec): check("the cortexm fixture exists", False, vec) else: - app = IdaTui(open_path=fresh_copy(vec, "cortexm"), keepalive=False, - load_args="-parm:ARMv7-M") + app = IdaTui( + open_path=fresh_copy(vec, "cortexm"), + keepalive=False, + load_args="-parm:ARMv7-M", + ) async with app.run_test(size=(140, 44)) as pilot: - await wait(lambda: app._func_index is not None - and app._func_index.complete, pilot) - check("a bare vector table gives IDA nothing to go on", - len(app._func_index) == 0, f"n={len(app._func_index)}") + await wait( + lambda: app._func_index is not None and app._func_index.complete, pilot + ) + check( + "a bare vector table gives IDA nothing to go on", + len(app._func_index) == 0, + f"n={len(app._func_index)}", + ) if type(app.screen).__name__ != "Screen": await pilot.press("escape") await settle(app, lambda: type(app.screen).__name__ == "Screen") @@ -264,20 +326,32 @@ async def run() -> int: lst._scroll_cursor_into_view() await settle(app) await pilot.press("T") - await wait(lambda: app._func_index is not None - and len(app._func_index) >= 3, pilot, 90) + await wait( + lambda: app._func_index is not None and len(app._func_index) >= 3, + pilot, + 90, + ) names = sorted(f.name for f in app._func_index.all_loaded()) - check("scanning the table finds the Thumb handlers", - names == ["sub_200", "sub_240", "sub_280"], f"{names}") + check( + "scanning the table finds the Thumb handlers", + names == ["sub_200", "sub_240", "sub_280"], + f"{names}", + ) # The table also holds an even word (the initial stack pointer), an # even in-range word and an odd word pointing outside the image. All # three must be ignored — marking a data word as code corrupts the # listing, so the cost of a false positive is high. - check("and ignores the words that aren't Thumb pointers", - len(app._func_index) == 3, f"n={len(app._func_index)}") + check( + "and ignores the words that aren't Thumb pointers", + len(app._func_index) == 3, + f"n={len(app._func_index)}", + ) status = str(app.query_one("#status", Static).render()) - check("the result survives the reload AND the reindex", - "3 Thumb entries" in status, status[:90]) + check( + "the result survives the reload AND the reindex", + "3 Thumb entries" in status, + status[:90], + ) drop_scratch() diff --git a/tests/test_trace.py b/tests/test_trace.py index 8803cda..661e350 100644 --- a/tests/test_trace.py +++ b/tests/test_trace.py @@ -31,7 +31,7 @@ LINES = [ "rip=0x40100b,mw=0xff4:2a000000", "rax=0x2a,rip=0x40100e,mr=0xff4:2a000000", "rbp=0x0,rsp=0x1000,rip=0x401010,mr=0xff8:0000000000000000", - "rip=0x401000", # loop back: 0x401000 executes twice + "rip=0x401000", # loop back: 0x401000 executes twice "rip=0x401001", ] @@ -56,148 +56,220 @@ def main() -> int: t = Trace.load(path) check("every line is one timestamp", t.length == len(LINES), f"{t.length}") - check("the sidecar .info is picked up", - t.info is not None and t.info.arch == "x86_64" and - t.info.start_code == 0x401000) - check("PC is tracked per timestamp", - [t.ip(i) for i in range(6)] == - [0x401000, 0x401001, 0x401004, 0x40100b, 0x40100e, 0x401010]) + check( + "the sidecar .info is picked up", + t.info is not None + and t.info.arch == "x86_64" + and t.info.start_code == 0x401000, + ) + check( + "PC is tracked per timestamp", + [t.ip(i) for i in range(6)] + == [0x401000, 0x401001, 0x401004, 0x40100B, 0x40100E, 0x401010], + ) # -- register reconstruction --------------------------------------- # - check("a register keeps its value until it changes", - t.register("rsp", 0) == 0x1000 and t.register("rsp", 1) == 0xff8 - and t.register("rsp", 4) == 0xff8 and t.register("rsp", 5) == 0x1000) - check("the full first line seeds every register", - t.register("rbx", 3) == 0) - check("an unknown register is None, not 0", - t.register("r15", 0) is None) - check("changed() is what the INSTRUCTION did, not the state", - t.changed(4) == {"rax", "rip"}, f"{t.changed(4)}") + check( + "a register keeps its value until it changes", + t.register("rsp", 0) == 0x1000 + and t.register("rsp", 1) == 0xFF8 + and t.register("rsp", 4) == 0xFF8 + and t.register("rsp", 5) == 0x1000, + ) + check("the full first line seeds every register", t.register("rbx", 3) == 0) + check("an unknown register is None, not 0", t.register("r15", 0) is None) + check( + "changed() is what the INSTRUCTION did, not the state", + t.changed(4) == {"rax", "rip"}, + f"{t.changed(4)}", + ) # "which instruction set this register?" — the question a trace exists # to answer. - check("last_write finds the instruction that set a value", - t.last_write("rax", 5) == 4 and t.last_write("rbp", 4) == 2, - f"{t.last_write('rax', 5)}, {t.last_write('rbp', 4)}") - check("next_write looks forward", - t.next_write("rbp", 2) == 5 and t.next_write("rbp", 5) is None) + check( + "last_write finds the instruction that set a value", + t.last_write("rax", 5) == 4 and t.last_write("rbp", 4) == 2, + f"{t.last_write('rax', 5)}, {t.last_write('rbp', 4)}", + ) + check( + "next_write looks forward", + t.next_write("rbp", 2) == 5 and t.next_write("rbp", 5) is None, + ) # -- memory --------------------------------------------------------- # ops = t.memory_ops(3) - check("a write is captured with its bytes", - len(ops) == 1 and ops[0].write and ops[0].addr == 0xff4 - and ops[0].data == bytes.fromhex("2a000000"), f"{ops}") - check("a read is captured and marked as a read", - [o.write for o in t.memory_ops(4)] == [False]) + check( + "a write is captured with its bytes", + len(ops) == 1 + and ops[0].write + and ops[0].addr == 0xFF4 + and ops[0].data == bytes.fromhex("2a000000"), + f"{ops}", + ) + check( + "a read is captured and marked as a read", + [o.write for o in t.memory_ops(4)] == [False], + ) check("an instruction with no memory has none", t.memory_ops(2) == []) # -- memory state at a timestamp ------------------------------------ # # The trace wrote 2a000000 at 0xff4 (t=3) and read it back (t=4); it # pushed/popped 8 zero bytes at 0xff8 (t=1 write, t=5 read). - d, k = t.memory(0xff4, 4, 3) - check("memory reflects a write as of that timestamp", - d == bytes.fromhex("2a000000") and k == b"\x01" * 4, f"{d.hex()} {k.hex()}") - d, k = t.memory(0xff4, 4, 2) - check("and does NOT reflect it before the write happened", - k == b"\x00" * 4, f"{d.hex()} known={k.hex()}") + d, k = t.memory(0xFF4, 4, 3) + check( + "memory reflects a write as of that timestamp", + d == bytes.fromhex("2a000000") and k == b"\x01" * 4, + f"{d.hex()} {k.hex()}", + ) + d, k = t.memory(0xFF4, 4, 2) + check( + "and does NOT reflect it before the write happened", + k == b"\x00" * 4, + f"{d.hex()} known={k.hex()}", + ) # A trace only knows what it saw. A byte nobody touched is unknown, and # must not be reported as zero — that distinction is the entire reason # to read memory from a trace instead of from the database. d, k = t.memory(0x5000, 4, t.length - 1) - check("untouched memory is unknown, not zero", - k == b"\x00" * 4 and d == b"\x00" * 4, f"known={k.hex()}") + check( + "untouched memory is unknown, not zero", + k == b"\x00" * 4 and d == b"\x00" * 4, + f"known={k.hex()}", + ) # 0xff2..0xff3 was never touched; 0xff4..0xff7 came from the write at # t=3 and 0xff8..0xff9 from the push at t=1 — a window can be knowable # from several accesses at different times, which is what makes this # worth a mask rather than a flag. - d, k = t.memory(0xff2, 8, 4) - check("a partially-covered window marks which bytes are known", - k == bytes([0, 0, 1, 1, 1, 1, 1, 1]), f"known={k.hex()}") + d, k = t.memory(0xFF2, 8, 4) + check( + "a partially-covered window marks which bytes are known", + k == bytes([0, 0, 1, 1, 1, 1, 1, 1]), + f"known={k.hex()}", + ) # Reads are evidence too: an instruction reading a byte reveals what it # held at that moment. - d, k = t.memory(0xff8, 8, 5) - check("a read reveals memory contents", - k == b"\x01" * 8, f"known={k.hex()}") + d, k = t.memory(0xFF8, 8, 5) + check("a read reveals memory contents", k == b"\x01" * 8, f"known={k.hex()}") - check("memory_writes lists only the writers", - t.memory_writes(0xff4, 4) == [3], f"{t.memory_writes(0xff4, 4)}") - check("memory_accesses includes the readers", - t.memory_accesses(0xff4, 4) == [3, 4], f"{t.memory_accesses(0xff4, 4)}") - check("a never-touched range has no accesses", - t.memory_accesses(0x5000, 16) == []) + check( + "memory_writes lists only the writers", + t.memory_writes(0xFF4, 4) == [3], + f"{t.memory_writes(0xFF4, 4)}", + ) + check( + "memory_accesses includes the readers", + t.memory_accesses(0xFF4, 4) == [3, 4], + f"{t.memory_accesses(0xFF4, 4)}", + ) + check( + "a never-touched range has no accesses", t.memory_accesses(0x5000, 16) == [] + ) # -- execution queries: what painting is built on -------------------- # - check("executions lists every timestamp for an address", - list(t.executions(0x401000)) == [0, 6], f"{list(t.executions(0x401000))}") - check("a never-executed address has none", not len(t.executions(0xdead))) - check("executions_between windows the result", - t.executions_between(0x401000, 1, 7) == [6]) - check("next/prev execution step between hits", - t.next_execution(0x401000, 0) == 6 - and t.prev_execution(0x401000, 6) == 0 - and t.next_execution(0x401000, 6) is None) + check( + "executions lists every timestamp for an address", + list(t.executions(0x401000)) == [0, 6], + f"{list(t.executions(0x401000))}", + ) + check("a never-executed address has none", not len(t.executions(0xDEAD))) + check( + "executions_between windows the result", + t.executions_between(0x401000, 1, 7) == [6], + ) + check( + "next/prev execution step between hits", + t.next_execution(0x401000, 0) == 6 + and t.prev_execution(0x401000, 6) == 0 + and t.next_execution(0x401000, 6) is None, + ) # A pseudocode line covers MANY addresses, so the set form is the one # decompiler painting will call — per-address lookups would mean one # dict hit per instruction per repaint. - check("hits() counts a whole set of addresses at once", - t.hits([0x401000, 0x401001, 0x401004, 0xdead]) == - {0x401000: 2, 0x401001: 2, 0x401004: 1}, - f"{t.hits([0x401000, 0x401001, 0x401004, 0xdead])}") + check( + "hits() counts a whole set of addresses at once", + t.hits([0x401000, 0x401001, 0x401004, 0xDEAD]) + == {0x401000: 2, 0x401001: 2, 0x401004: 1}, + f"{t.hits([0x401000, 0x401001, 0x401004, 0xDEAD])}", + ) # -- rebasing -------------------------------------------------------- # # A traced process is relocated; nothing lines up until the slide is # found. Page offsets survive relocation, which is what makes it # findable. - db = [0x1000 + (a - 0x401000) for a in - (0x401000, 0x401001, 0x401004, 0x40100b, 0x40100e, 0x401010)] + db = [ + 0x1000 + (a - 0x401000) + for a in (0x401000, 0x401001, 0x401004, 0x40100B, 0x40100E, 0x401010) + ] slide = t.rebase(db) - check("the slide between trace and database is found", - slide == 0x1000 - 0x401000, f"{slide:#x}") + check( + "the slide between trace and database is found", + slide == 0x1000 - 0x401000, + f"{slide:#x}", + ) t.apply_slide(slide) - check("addresses come back in database terms", - t.ip(0) == 0x1000 and list(t.executions(0x1000)) == [0, 6], - f"{t.ip(0):#x}") + check( + "addresses come back in database terms", + t.ip(0) == 0x1000 and list(t.executions(0x1000)) == [0, 6], + f"{t.ip(0):#x}", + ) # Memory addresses are NOT slid: the slide relocates the image, and # these are overwhelmingly stack/heap addresses with no database # counterpart — sliding a stack pointer by the image delta produced a # negative address in testing. - check("memory op addresses stay in trace space", - t.memory_ops(3)[0].addr == 0xff4, f"{t.memory_ops(3)[0].addr:#x}") - check("raw_ip still gives the traced address", - t.raw_ip(0) == 0x401000) + check( + "memory op addresses stay in trace space", + t.memory_ops(3)[0].addr == 0xFF4, + f"{t.memory_ops(3)[0].addr:#x}", + ) + check("raw_ip still gives the traced address", t.raw_ip(0) == 0x401000) t2 = Trace.load(path) # Page offsets that appear nowhere in the trace. (0xdead0000/0xdead0004 # would NOT do: they sit at the same offsets as two traced addresses and # so legitimately agree on a slide — a reminder that this matches on # offsets, not on addresses looking plausible.) - check("no match means no slide, not a wrong one", - t2.rebase([0xdead0555, 0xbeef0777]) == 0, - f"{t2.rebase([0xdead0555, 0xbeef0777]):#x}") - check("one lone agreeing address is not enough to claim a slide", - t2.rebase([0x1000]) == 0, f"{t2.rebase([0x1000]):#x}") + check( + "no match means no slide, not a wrong one", + t2.rebase([0xDEAD0555, 0xBEEF0777]) == 0, + f"{t2.rebase([0xDEAD0555, 0xBEEF0777]):#x}", + ) + check( + "one lone agreeing address is not enough to claim a slide", + t2.rebase([0x1000]) == 0, + f"{t2.rebase([0x1000]):#x}", + ) check("an empty database is harmless", t2.rebase([]) == 0) # -- robustness ------------------------------------------------------ # p2 = os.path.join(tmp, "odd.0.log") with open(p2, "w") as f: - f.write("\n".join([ - FULL + ",rip=0x401000", - "", # blank line - "rax=0xnothex,rip=0x401001", # unparseable value - "rbx=0x1", # no PC at all - "rip=0x401002,mw=0x10:zz", # unparseable memory - ]) + "\n") + f.write( + "\n".join( + [ + FULL + ",rip=0x401000", + "", # blank line + "rax=0xnothex,rip=0x401001", # unparseable value + "rbx=0x1", # no PC at all + "rip=0x401002,mw=0x10:zz", # unparseable memory + ] + ) + + "\n" + ) t3 = Trace.load(p2) - check("a malformed trace loads instead of raising", t3.length == 4, - f"{t3.length}") - check("a line with no PC inherits the previous one", - t3.ip(2) == 0x401001, f"{t3.ip(2):#x}") - check("an unparseable memory entry is dropped, not fatal", - t3.memory_ops(3) == []) + check( + "a malformed trace loads instead of raising", t3.length == 4, f"{t3.length}" + ) + check( + "a line with no PC inherits the previous one", + t3.ip(2) == 0x401001, + f"{t3.ip(2):#x}", + ) + check( + "an unparseable memory entry is dropped, not fatal", t3.memory_ops(3) == [] + ) print(f"\n{PASS} passed, {FAIL} failed") return 1 if FAIL else 0 diff --git a/tests/test_trace_rpc.py b/tests/test_trace_rpc.py index 015ce8b..9b0f2e0 100644 --- a/tests/test_trace_rpc.py +++ b/tests/test_trace_rpc.py @@ -30,8 +30,7 @@ from idatui.rpcclient import RpcClient, RpcError # noqa: E402 from idatui.trace import Trace # noqa: E402 REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) -TRACER = os.path.expanduser( - "~/.pi/agent/skills/tenet-trace/scripts/tenet-trace") +TRACER = os.path.expanduser("~/.pi/agent/skills/tenet-trace/scripts/tenet-trace") FALLBACK_TRACE = "/tmp/echotrace.0.log" BINARY = os.path.join(REPO, "targets", "echo") @@ -52,8 +51,12 @@ def make_trace(tmp, binary): """Record a short trace of the echo binary.""" out = os.path.join(tmp, "t") try: - subprocess.run([TRACER, "-o", out, binary, "hello"], - capture_output=True, timeout=180, check=False) + subprocess.run( + [TRACER, "-o", out, binary, "hello"], + capture_output=True, + timeout=180, + check=False, + ) except (OSError, subprocess.TimeoutExpired): return None log = out + ".0.log" @@ -62,11 +65,24 @@ def make_trace(tmp, binary): def spawn_pane(target, trace_log, timeout=300): """Spawn an idatui pane with --trace and wait for readiness.""" - cmd = [sys.executable, "-m", "idatui.pane", "spawn", - "--open", target, "--trace", trace_log, - "--detached", "--size", "60%", "--timeout", str(timeout)] - r = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout + 30, - cwd=REPO) + cmd = [ + sys.executable, + "-m", + "idatui.pane", + "spawn", + "--open", + target, + "--trace", + trace_log, + "--detached", + "--size", + "60%", + "--timeout", + str(timeout), + ] + r = subprocess.run( + cmd, capture_output=True, text=True, timeout=timeout + 30, cwd=REPO + ) if r.returncode != 0: print(f" spawn failed: {r.stderr.strip()}", file=sys.stderr) return None @@ -74,10 +90,17 @@ def spawn_pane(target, trace_log, timeout=300): def stop_pane(sock, timeout=60): - cmd = [sys.executable, "-m", "idatui.pane", "stop", - "--sock", sock, "--timeout", str(timeout)] - subprocess.run(cmd, capture_output=True, text=True, timeout=timeout + 10, - cwd=REPO) + cmd = [ + sys.executable, + "-m", + "idatui.pane", + "stop", + "--sock", + sock, + "--timeout", + str(timeout), + ] + subprocess.run(cmd, capture_output=True, text=True, timeout=timeout + 10, cwd=REPO) def main() -> int: @@ -96,14 +119,17 @@ def main() -> int: if trace_log is None and os.path.exists(FALLBACK_TRACE): trace_log = FALLBACK_TRACE if trace_log is None: - print(f" skip: no trace available (tracer at {TRACER}, " - f"fallback {FALLBACK_TRACE})") + print( + f" skip: no trace available (tracer at {TRACER}, " + f"fallback {FALLBACK_TRACE})" + ) return 0 # Load our own model for ground-truth comparisons. model = Trace.load(trace_log) - check("model loaded for ground truth", model.length > 10, - f"length={model.length}") + check( + "model loaded for ground truth", model.length > 10, f"length={model.length}" + ) # Spawn the pane. info = spawn_pane(target, trace_log) @@ -138,14 +164,22 @@ def run_trace_tests(c: RpcClient, model: Trace): r = c.call("trace", seek=0) check("trace seek=0 returns a snapshot", "active" in r and "trace" in r) tr = r["trace"] - check("response includes trace metadata", - "idx" in tr and "length" in tr and "pc" in tr and "changed" in tr, - f"keys={list(tr.keys())}") + check( + "response includes trace metadata", + "idx" in tr and "length" in tr and "pc" in tr and "changed" in tr, + f"keys={list(tr.keys())}", + ) check("idx is 0 after seeking to 0", tr["idx"] == 0) - check("length matches the model", tr["length"] == model.length, - f"{tr['length']} vs {model.length}") - check("pc at 0 is a hex string", isinstance(tr["pc"], str) - and tr["pc"].startswith("0x"), tr["pc"]) + check( + "length matches the model", + tr["length"] == model.length, + f"{tr['length']} vs {model.length}", + ) + check( + "pc at 0 is a hex string", + isinstance(tr["pc"], str) and tr["pc"].startswith("0x"), + tr["pc"], + ) # ------------------------------------------------------------------ # # seek to a mid-trace timestamp @@ -153,14 +187,16 @@ def run_trace_tests(c: RpcClient, model: Trace): mid = model.length // 2 r = c.call("trace", seek=mid) tr = r["trace"] - check("seek to midpoint lands correctly", tr["idx"] == mid, - f"got {tr['idx']}, want {mid}") + check( + "seek to midpoint lands correctly", + tr["idx"] == mid, + f"got {tr['idx']}, want {mid}", + ) # The model's IP at this point, rebased — the RPC should agree. # We can't compare directly because the model isn't rebased yet, but # the pc should be a small address (database-space, not ASLR'd). pc = int(tr["pc"], 16) - check("pc is in database space (not ASLR'd)", - pc < 0x100000, f"pc={tr['pc']}") + check("pc is in database space (not ASLR'd)", pc < 0x100000, f"pc={tr['pc']}") # ------------------------------------------------------------------ # # seek by percentage (Tenet shell syntax) @@ -168,39 +204,50 @@ def run_trace_tests(c: RpcClient, model: Trace): r = c.call("trace", seek="!0") check("seek !0 (0%) goes to the start", r["trace"]["idx"] == 0) r = c.call("trace", seek="!100") - check("seek !100 (100%) goes to the end", - r["trace"]["idx"] == model.length - 1, - f"got {r['trace']['idx']}, want {model.length - 1}") + check( + "seek !100 (100%) goes to the end", + r["trace"]["idx"] == model.length - 1, + f"got {r['trace']['idx']}, want {model.length - 1}", + ) r = c.call("trace", seek="!50") - check("seek !50 (50%) goes to the midpoint", - abs(r["trace"]["idx"] - mid) <= 1, - f"got {r['trace']['idx']}, want ~{mid}") + check( + "seek !50 (50%) goes to the midpoint", + abs(r["trace"]["idx"] - mid) <= 1, + f"got {r['trace']['idx']}, want ~{mid}", + ) # ------------------------------------------------------------------ # # step forward/backward # ------------------------------------------------------------------ # c.call("trace", seek=0) r = c.call("trace", step=1) - check("step=1 advances one timestamp", r["trace"]["idx"] == 1, - f"got {r['trace']['idx']}") + check( + "step=1 advances one timestamp", + r["trace"]["idx"] == 1, + f"got {r['trace']['idx']}", + ) r = c.call("trace", step=1) - check("another step=1 reaches 2", r["trace"]["idx"] == 2, - f"got {r['trace']['idx']}") + check( + "another step=1 reaches 2", r["trace"]["idx"] == 2, f"got {r['trace']['idx']}" + ) r = c.call("trace", step=-1) - check("step=-1 goes backward", r["trace"]["idx"] == 1, - f"got {r['trace']['idx']}") + check("step=-1 goes backward", r["trace"]["idx"] == 1, f"got {r['trace']['idx']}") # Step backward at the start should clamp to 0. c.call("trace", seek=0) r = c.call("trace", step=-1) - check("step=-1 at t=0 stays at 0", r["trace"]["idx"] == 0, - f"got {r['trace']['idx']}") + check( + "step=-1 at t=0 stays at 0", r["trace"]["idx"] == 0, f"got {r['trace']['idx']}" + ) # Multi-step. c.call("trace", seek=0) r = c.call("trace", step=5) - check("step=5 advances five timestamps", r["trace"]["idx"] == 5, - f"got {r['trace']['idx']}") + check( + "step=5 advances five timestamps", + r["trace"]["idx"] == 5, + f"got {r['trace']['idx']}", + ) # ------------------------------------------------------------------ # # step over (follows SP) @@ -212,8 +259,14 @@ def run_trace_tests(c: RpcClient, model: Trace): a = model.register(sp_name, i) b = model.register(sp_name, i + 1) if a and b and b < a: - ret = next((j for j in range(i + 1, model.length) - if (model.register(sp_name, j) or 0) >= a), None) + ret = next( + ( + j + for j in range(i + 1, model.length) + if (model.register(sp_name, j) or 0) >= a + ), + None, + ) if ret and ret > i + 3: call_at = (i, ret) break @@ -221,11 +274,12 @@ def run_trace_tests(c: RpcClient, model: Trace): i, ret = call_at c.call("trace", seek=i) r = c.call("trace", step=1, over=True) - check("step over skips the callee", - r["trace"]["idx"] == ret, - f"from {i}, got {r['trace']['idx']}, expected {ret}") - check("step over goes further than a plain step", - r["trace"]["idx"] > i + 1) + check( + "step over skips the callee", + r["trace"]["idx"] == ret, + f"from {i}, got {r['trace']['idx']}, expected {ret}", + ) + check("step over goes further than a plain step", r["trace"]["idx"] > i + 1) else: check("found a call to step over", False, "none in this short trace") @@ -235,25 +289,28 @@ def run_trace_tests(c: RpcClient, model: Trace): r = c.call("trace", goto="main") tr = r["trace"] check("goto main lands on a timestamp", tr["idx"] >= 0) - check("and the function context says main", - r.get("function", {}).get("name") == "main", - f"function={r.get('function')}") + check( + "and the function context says main", + r.get("function", {}).get("name") == "main", + f"function={r.get('function')}", + ) # goto by hex address. main_ea = r["function"]["ea"] c.call("trace", seek=0) # reset position r = c.call("trace", goto=hex(main_ea)) - check("goto by hex address works", - r["trace"]["idx"] >= 0 and r["function"]["ea"] == main_ea, - f"idx={r['trace']['idx']}, ea={r.get('function', {}).get('ea')}") + check( + "goto by hex address works", + r["trace"]["idx"] >= 0 and r["function"]["ea"] == main_ea, + f"idx={r['trace']['idx']}, ea={r.get('function', {}).get('ea')}", + ) # goto a function that was never executed. try: c.call("trace", goto="0xDEADBEEF") check("goto an unexecuted address raises", False, "no error raised") except RpcError as e: - check("goto an unexecuted address raises", "never executed" in str(e), - str(e)) + check("goto an unexecuted address raises", "never executed" in str(e), str(e)) # ------------------------------------------------------------------ # # changed registers in the response @@ -261,9 +318,11 @@ def run_trace_tests(c: RpcClient, model: Trace): c.call("trace", seek=0) r = c.call("trace", step=1) changed = r["trace"]["changed"] - check("changed is a list of register names", - isinstance(changed, list) and all(isinstance(s, str) for s in changed), - f"{changed}") + check( + "changed is a list of register names", + isinstance(changed, list) and all(isinstance(s, str) for s in changed), + f"{changed}", + ) # The PC always changes on a step (it's a different instruction). check("rip is always in changed", "rip" in changed, f"{changed}") @@ -275,9 +334,13 @@ def run_trace_tests(c: RpcClient, model: Trace): r2 = c.call("trace", seek=min(50, model.length - 1)) ea1 = r1.get("cursor", {}).get("ea") ea2 = r2.get("cursor", {}).get("ea") - check("the cursor ea follows the trace pc", - ea1 is not None and ea2 is not None and (ea1 != ea2 or r1["trace"]["pc"] == r2["trace"]["pc"]), - f"ea1={ea1}, ea2={ea2}") + check( + "the cursor ea follows the trace pc", + ea1 is not None + and ea2 is not None + and (ea1 != ea2 or r1["trace"]["pc"] == r2["trace"]["pc"]), + f"ea1={ea1}, ea2={ea2}", + ) # ------------------------------------------------------------------ # # trace verb without a trace raises cleanly @@ -289,8 +352,10 @@ def run_trace_tests(c: RpcClient, model: Trace): # Actually, if none of seek/goto/step is given, it just settles and # returns the current state — that's fine, it's a status query. r = c.call("trace") - check("trace with no action is a status query", - "trace" in r and r["trace"]["idx"] >= 0) + check( + "trace with no action is a status query", + "trace" in r and r["trace"]["idx"] >= 0, + ) except RpcError: check("trace with no action is a status query", False, "raised an error") @@ -299,14 +364,20 @@ def run_trace_tests(c: RpcClient, model: Trace): # ------------------------------------------------------------------ # r = c.call("trace", seek=3) for key in ("idx", "length", "pc", "changed"): - check(f"trace response has '{key}'", key in r.get("trace", {}), - f"trace={r.get('trace')}") + check( + f"trace response has '{key}'", + key in r.get("trace", {}), + f"trace={r.get('trace')}", + ) # Standard snapshot fields are ALSO present (the trace response is a # superset of a normal snapshot). for key in ("active", "function", "cursor", "status", "ready"): - check(f"trace response also has snapshot key '{key}'", key in r, - f"keys={list(r.keys())}") + check( + f"trace response also has snapshot key '{key}'", + key in r, + f"keys={list(r.keys())}", + ) # ------------------------------------------------------------------ # # edge cases: seek beyond bounds @@ -314,25 +385,30 @@ def run_trace_tests(c: RpcClient, model: Trace): r = c.call("trace", seek=-1) check("seek -1 clamps to 0", r["trace"]["idx"] == 0) r = c.call("trace", seek=model.length + 1000) - check("seek beyond length clamps to the end", - r["trace"]["idx"] == model.length - 1, - f"got {r['trace']['idx']}") + check( + "seek beyond length clamps to the end", + r["trace"]["idx"] == model.length - 1, + f"got {r['trace']['idx']}", + ) # ------------------------------------------------------------------ # # seek with comma-separated numbers (ergonomic) # ------------------------------------------------------------------ # r = c.call("trace", seek="100") - check("seek accepts a string number", - r["trace"]["idx"] == min(100, model.length - 1)) + check( + "seek accepts a string number", r["trace"]["idx"] == min(100, model.length - 1) + ) # ------------------------------------------------------------------ # # pseudocode still works with a trace loaded # ------------------------------------------------------------------ # c.call("trace", goto="main") r = c.call("pseudocode", target="main", lines=5) - check("pseudocode works alongside the trace", - "code" in r and "main" in r.get("code", ""), - f"keys={list(r.keys())}") + check( + "pseudocode works alongside the trace", + "code" in r and "main" in r.get("code", ""), + f"keys={list(r.keys())}", + ) # ------------------------------------------------------------------ # # state includes trace position @@ -341,17 +417,18 @@ def run_trace_tests(c: RpcClient, model: Trace): r = c.call("state") # The state verb doesn't include trace info (that's trace-specific), # but the standard snapshot fields should be consistent. - check("state works with a trace loaded", - r.get("ready") is True and "cursor" in r) + check("state works with a trace loaded", r.get("ready") is True and "cursor" in r) # ------------------------------------------------------------------ # # view_lines works with trail painted # ------------------------------------------------------------------ # c.call("trace", seek=min(40, model.length - 1)) r = c.call("view", lines=10) - check("view returns lines with a trace active", - "lines" in r and len(r["lines"]) > 0, - f"keys={list(r.keys())}") + check( + "view returns lines with a trace active", + "lines" in r and len(r["lines"]) > 0, + f"keys={list(r.keys())}", + ) # ------------------------------------------------------------------ # # navigation works alongside trace: goto a function, trace follows @@ -359,14 +436,18 @@ def run_trace_tests(c: RpcClient, model: Trace): c.call("trace", goto="main") start_idx = c.call("trace")["trace"]["idx"] r = c.call("goto", target="error_at_line") - check("goto still works with trace loaded", - r.get("function", {}).get("name") == "error_at_line") + check( + "goto still works with trace loaded", + r.get("function", {}).get("name") == "error_at_line", + ) # The trace timestamp should NOT change from a regular goto — the trace # position is independent of navigation. r2 = c.call("trace") - check("regular goto does not change the trace position", - r2["trace"]["idx"] == start_idx, - f"was {start_idx}, now {r2['trace']['idx']}") + check( + "regular goto does not change the trace position", + r2["trace"]["idx"] == start_idx, + f"was {start_idx}, now {r2['trace']['idx']}", + ) if __name__ == "__main__": diff --git a/tests/test_trace_ui.py b/tests/test_trace_ui.py index 586ce2e..768d18b 100644 --- a/tests/test_trace_ui.py +++ b/tests/test_trace_ui.py @@ -17,18 +17,21 @@ import sys sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) -from textual.widgets import Input, OptionList, Static # noqa: E402 - from _fixtures import fast_keys, staged # noqa: E402 +from textual.widgets import Input, OptionList, Static # noqa: E402 -fast_keys() # ~85ms -> ~2ms per keypress; see _fixtures.fast_keys +fast_keys() # ~85ms -> ~2ms per keypress; see _fixtures.fast_keys from idatui._sync import settle # noqa: E402 -from idatui.app import (DecompView, IdaTui, ListingView, # noqa: E402 - RegWriteScreen, TraceDock) +from idatui.app import ( # noqa: E402 + DecompView, + IdaTui, + ListingView, + RegWriteScreen, + TraceDock, +) REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) -TRACER = os.path.expanduser( - "~/.pi/agent/skills/tenet-trace/scripts/tenet-trace") +TRACER = os.path.expanduser("~/.pi/agent/skills/tenet-trace/scripts/tenet-trace") PASS = FAIL = 0 @@ -56,8 +59,12 @@ async def wait(pred, pilot, t=240.0): def make_trace(tmp, binary): out = os.path.join(tmp, "t") try: - subprocess.run([TRACER, "-o", out, binary, "hi"], - capture_output=True, timeout=180, check=False) + subprocess.run( + [TRACER, "-o", out, binary, "hi"], + capture_output=True, + timeout=180, + check=False, + ) except (OSError, subprocess.TimeoutExpired): return None log = out + ".0.log" @@ -71,8 +78,9 @@ async def run() -> int: # anything tracked. On echo the seeding saves only ~0.2s (it analyses fast); # it is here so a suite pointed at a bigger target doesn't pay for analysis # on every run. - async with staged(binary, lambda p: IdaTui(open_path=p, keepalive=False), - prefix="idatui-traceui-") as target: + async with staged( + binary, lambda p: IdaTui(open_path=p, keepalive=False), prefix="idatui-traceui-" + ) as target: tmp = os.path.dirname(target) log = make_trace(tmp, target) if not log: @@ -90,22 +98,29 @@ async def run() -> int: # Rebasing: the tracer runs the binary relocated, so without a slide # nothing in the trace matches anything on screen. - check("trace addresses were rebased onto the database", - t.slide != 0 and t.ip(0) < 0x1000000, - f"slide={t.slide:#x} ip0={t.ip(0):#x}") + check( + "trace addresses were rebased onto the database", + t.slide != 0 and t.ip(0) < 0x1000000, + f"slide={t.slide:#x} ip0={t.ip(0):#x}", + ) idx = app._func_index touched = [f.name for f in idx.all_loaded() if t.executions(f.addr)] - check("and now line up with real functions", - len(touched) > 1 and "main" in touched, f"{touched[:6]}") + check( + "and now line up with real functions", + len(touched) > 1 and "main" in touched, + f"{touched[:6]}", + ) dock = app.query_one(TraceDock) check("the dock is docked and visible", dock.display) head = str(dock.query_one("#trace-head", Static).render()) - check("it shows where we are in time", "0" in head and "%" in head, - head[:60]) + check( + "it shows where we are in time", "0" in head and "%" in head, head[:60] + ) regs = str(dock.query_one("#trace-regs", Static).render()) - check("and the register state at that time", "rip" in regs.lower(), - regs[:60]) + check( + "and the register state at that time", "rip" in regs.lower(), regs[:60] + ) # -- stepping --------------------------------------------------- # lst = app.query_one(ListingView) @@ -117,15 +132,19 @@ async def run() -> int: # runs in a worker. Waiting on it and then reading the cursor was a # race that the (slower) IDA Nexus backend loses. Gate on the thing # the check is about. - await settle(app, lambda: app._t == 1 and lst._cursor_ea() == t.ip(1), - timeout=20) + await settle( + app, lambda: app._t == 1 and lst._cursor_ea() == t.ip(1), timeout=20 + ) check("] steps forward one instruction", app._t == 1, f"t={app._t}") - check("the code view follows the trace", - lst._cursor_ea() == t.ip(1), - f"{lst._cursor_ea()} vs {t.ip(1)}") + check( + "the code view follows the trace", + lst._cursor_ea() == t.ip(1), + f"{lst._cursor_ea()} vs {t.ip(1)}", + ) await pilot.press("[") - await settle(app, lambda: app._t == 0 and lst._cursor_ea() == t.ip(0), - timeout=20) + await settle( + app, lambda: app._t == 0 and lst._cursor_ea() == t.ip(0), timeout=20 + ) check("[ steps backward", app._t == 0, f"t={app._t}") await pilot.press("[") # Nothing should happen, so there is no signal to wait FOR: the @@ -144,8 +163,14 @@ async def run() -> int: for i in range(1, min(t.length - 1, 400)): a, b = t.register(sp, i), t.register(sp, i + 1) if a and b and b < a: - ret = next((j for j in range(i + 1, t.length) - if (t.register(sp, j) or 0) >= a), None) + ret = next( + ( + j + for j in range(i + 1, t.length) + if (t.register(sp, j) or 0) >= a + ), + None, + ) if ret and ret > i + 3: call_at = (i, ret) break @@ -157,8 +182,11 @@ async def run() -> int: await wait(lambda: app._t == i, pilot, 20) await pilot.press("}") await wait(lambda: app._t != i, pilot, 30) - check("} steps OVER a call instead of into it", - app._t == ret, f"{i} -> {app._t}, expected {ret}") + check( + "} steps OVER a call instead of into it", + app._t == ret, + f"{i} -> {app._t}, expected {ret}", + ) check("which is further than a plain step", app._t > i + 1) # -- memory at time T -------------------------------------------- # @@ -168,36 +196,56 @@ async def run() -> int: # image would have nothing to show. dock = app.query_one(TraceDock) app._seek(min(60, t.length - 1)) - await wait(lambda: "stack (" in - str(dock.query_one("#trace-stack", Static).render()), - pilot, 5) + await wait( + lambda: ( + "stack (" in str(dock.query_one("#trace-stack", Static).render()) + ), + pilot, + 5, + ) stack = str(dock.query_one("#trace-stack", Static).render()) - check("the dock shows the stack at this timestamp", - "stack (" in stack and len(stack.splitlines()) > 4, stack[:60]) + check( + "the dock shows the stack at this timestamp", + "stack (" in stack and len(stack.splitlines()) > 4, + stack[:60], + ) sp_name = next(r for r in ("rsp", "esp", "sp") if r in t.reg_at) sp = t.register(sp_name, app._t) - check("anchored at the stack pointer", - f"{sp:012x}" in stack, f"sp={sp:#x} / {stack[:80]}") + check( + "anchored at the stack pointer", + f"{sp:012x}" in stack, + f"sp={sp:#x} / {stack[:80]}", + ) # A trace knows what it observed and nothing else. Unseen bytes are # printed as '?', never as zeros — rendering them as zero would # invent facts about memory nobody looked at. data, known = t.memory_raw(sp, 8, app._t) if not all(known): - check("memory the trace never saw is marked unknown", - "?" in stack, stack[:80]) + check( + "memory the trace never saw is marked unknown", + "?" in stack, + stack[:80], + ) else: - check("known stack words are shown as values", - any(c in "0123456789abcdef" for c in stack), stack[:60]) + check( + "known stack words are shown as values", + any(c in "0123456789abcdef" for c in stack), + stack[:60], + ) # Stepping must move the memory view with time. before = stack app._seek(min(80, t.length - 1)) - await wait(lambda: str(dock.query_one("#trace-stack", - Static).render()) != before, - pilot, 5) - check("and it follows as you move through time", - str(dock.query_one("#trace-stack", Static).render()) != before) + await wait( + lambda: str(dock.query_one("#trace-stack", Static).render()) != before, + pilot, + 5, + ) + check( + "and it follows as you move through time", + str(dock.query_one("#trace-stack", Static).render()) != before, + ) # -- trails ------------------------------------------------------ # # Not "every address the trace ever touched": on a loop-heavy @@ -207,17 +255,27 @@ async def run() -> int: await wait(lambda: bool(lst.trail), pilot, 5) trail = lst.trail kinds = {k for k in trail.values()} - check("the listing is painted with an execution trail", - {"now", "past", "future"} <= kinds, f"{sorted(kinds)}") - check("'now' is the instruction we're standing on", - trail.get(t.ip(app._t)) == "now", f"{trail.get(t.ip(app._t))}") - check("the step behind is past, the step ahead is future", - trail.get(t.ip(app._t - 1)) == "past" - and trail.get(t.ip(app._t + 1)) == "future", - f"{trail.get(t.ip(app._t - 1))}, {trail.get(t.ip(app._t + 1))}") - painted = [y for y in range(min(lst.size.height, 30)) - if any(seg.style and seg.style.bgcolor - for seg in lst.render_line(y))] + check( + "the listing is painted with an execution trail", + {"now", "past", "future"} <= kinds, + f"{sorted(kinds)}", + ) + check( + "'now' is the instruction we're standing on", + trail.get(t.ip(app._t)) == "now", + f"{trail.get(t.ip(app._t))}", + ) + check( + "the step behind is past, the step ahead is future", + trail.get(t.ip(app._t - 1)) == "past" + and trail.get(t.ip(app._t + 1)) == "future", + f"{trail.get(t.ip(app._t - 1))}, {trail.get(t.ip(app._t + 1))}", + ) + painted = [ + y + for y in range(min(lst.size.height, 30)) + if any(seg.style and seg.style.bgcolor for seg in lst.render_line(y)) + ] check("and it actually reaches the screen", painted, "no tinted rows") # -- the same trail on PSEUDOCODE -------------------------------- # @@ -234,23 +292,33 @@ async def run() -> int: await wait(lambda: app._t == first + 12, pilot, 20) lst.focus() await pilot.press("tab") - got = await wait(lambda: app.query_one(DecompView).display - and app.query_one(DecompView)._texts, pilot, 120) + got = await wait( + lambda: ( + app.query_one(DecompView).display + and app.query_one(DecompView)._texts + ), + pilot, + 120, + ) dec = app.query_one(DecompView) check("pseudocode is available for the traced function", got) app._seek(first + 12) await wait(lambda: len(dec.trail) > 2, pilot, 5) - check("pseudocode lines are painted with the trail", - len(dec.trail) > 2, f"{len(dec.trail)} lines") + check( + "pseudocode lines are painted with the trail", + len(dec.trail) > 2, + f"{len(dec.trail)} lines", + ) now = [i for i, k in dec.trail.items() if k == "now"] - check("exactly one pseudocode line is 'now'", - len(now) == 1, f"{now}") + check("exactly one pseudocode line is 'now'", len(now) == 1, f"{now}") # The 'now' line must be the one covering the current # instruction, not merely some executed line. covered = app._trail_map[now[0]] if now and app._trail_map else [] - check("and it's the line covering the current instruction", - t.ip(app._t) in covered, - f"pc={t.ip(app._t):#x} line covers {[hex(a) for a in covered][:4]}") + check( + "and it's the line covering the current instruction", + t.ip(app._t) in covered, + f"pc={t.ip(app._t):#x} line covers {[hex(a) for a in covered][:4]}", + ) # Stepping must not throw you out of the view you're reading. # A step navigates to an address, and navigating to an address @@ -260,13 +328,19 @@ async def run() -> int: was = app._t await pilot.press("]") await settle(app, lambda: app._t != was) - check("stepping in pseudocode stays in pseudocode", - app._active == "decomp", f"active={app._active}") + check( + "stepping in pseudocode stays in pseudocode", + app._active == "decomp", + f"active={app._active}", + ) was = app._t await pilot.press("[") await settle(app, lambda: app._t != was) - check("and so does stepping backward", - app._active == "decomp", f"active={app._active}") + check( + "and so does stepping backward", + app._active == "decomp", + f"active={app._active}", + ) # -- split view: a step is a GLOBAL move ------------------------ # # Normal navigation moves one pane and gives the companion a band, @@ -274,8 +348,13 @@ async def run() -> int: # navigation though: both panes show the same instant, so the # listing cursor must sit on the current instruction. app.action_toggle_split() - await wait(lambda: app._split and lst.display - and app.query_one(DecompView).display, pilot, 10) + await wait( + lambda: ( + app._split and lst.display and app.query_one(DecompView).display + ), + pilot, + 10, + ) if not app._split: check("split view toggled on", False) else: @@ -286,14 +365,18 @@ async def run() -> int: # Wait for the cursor to arrive rather than sleeping a flat # 0.5s and hoping. Same question -- does the listing follow # the pc? -- but it costs what it costs instead of 3s. - if await wait(lambda: lst._cursor_ea() == t.ip(app._t), - pilot, 5): + if await wait(lambda: lst._cursor_ea() == t.ip(app._t), pilot, 5): tracked += 1 - check("stepping in split moves the listing cursor to the pc", - tracked == 6, f"{tracked}/6 steps tracked") - check("and the trail follows in both panes", - lst.trail.get(t.ip(app._t)) == "now", - f"{lst.trail.get(t.ip(app._t))}") + check( + "stepping in split moves the listing cursor to the pc", + tracked == 6, + f"{tracked}/6 steps tracked", + ) + check( + "and the trail follows in both panes", + lst.trail.get(t.ip(app._t)) == "now", + f"{lst.trail.get(t.ip(app._t))}", + ) # The pseudocode cursor follows too — but only for instructions # the decompiler actually attributes to a line. About half @@ -313,8 +396,14 @@ async def run() -> int: # on `pc in _trail_line_of` instead would burn the timeout on # every unmapped instruction -- about half of them -- and be # slower than the flat sleep it replaces. - await wait(lambda: lst._cursor_ea() == pc - and app._trail_map_ea == dec.loaded_ea, pilot, 5) + await wait( + lambda: ( + lst._cursor_ea() == pc + and app._trail_map_ea == dec.loaded_ea + ), + pilot, + 5, + ) if app._trail_map_ea == dec.loaded_ea and pc in app._trail_line_of: mapped += 1 # Mapped: the pseudocode cursor is expected, so it's fair @@ -324,9 +413,11 @@ async def run() -> int: await wait(lambda: dec.cursor == line, pilot, 3) if dec.cursor != line: missed += 1 - check("the pseudocode cursor follows every mapped instruction", - mapped > 3 and missed == 0, - f"{mapped} mapped, {missed} not followed") + check( + "the pseudocode cursor follows every mapped instruction", + mapped > 3 and missed == 0, + f"{mapped} mapped, {missed} not followed", + ) # -- a late navigation must not drag the view back --------------- # # Navigations run in workers and finish out of order. The trace's @@ -347,10 +438,12 @@ async def run() -> int: # for the workers to drain rather than for three seconds and a # hope: same question, ~50ms instead of 3s. await settle(app) - check("a stale navigation doesn't drag the cursor away", - lst._cursor_ea() == dbaddr and app._cur.ea == dbaddr, - f"cursor={lst._cursor_ea():#x} cur={app._cur.ea:#x} " - f"want {dbaddr:#x}") + check( + "a stale navigation doesn't drag the cursor away", + lst._cursor_ea() == dbaddr and app._cur.ea == dbaddr, + f"cursor={lst._cursor_ea():#x} cur={app._cur.ea:#x} " + f"want {dbaddr:#x}", + ) # -- seeking, as opposed to stepping ---------------------------- # # "When else did this instruction run?" — the question that makes a @@ -359,8 +452,11 @@ async def run() -> int: stamps = list(t.by_ip[hot]) db = hot + t.slide if len(stamps) < 2 or lst.model is None: - check("found an address executed more than once", False, - f"{len(stamps)} executions") + check( + "found an address executed more than once", + False, + f"{len(stamps)} executions", + ) else: if app._split: app.action_toggle_split() @@ -377,31 +473,48 @@ async def run() -> int: lst.cursor = row lst._scroll_cursor_into_view() await settle(app, lambda: lst._cursor_ea() == db) - check("cursor is on the repeated instruction", - lst._cursor_ea() == db, f"{lst._cursor_ea():#x} vs {db:#x}") + check( + "cursor is on the repeated instruction", + lst._cursor_ea() == db, + f"{lst._cursor_ea():#x} vs {db:#x}", + ) await pilot.press(">") await settle(app, lambda: app._t == stamps[1]) - check("> seeks to the next execution of it", - app._t == stamps[1], f"t={app._t}, expected {stamps[1]}") + check( + "> seeks to the next execution of it", + app._t == stamps[1], + f"t={app._t}, expected {stamps[1]}", + ) status = str(app.query_one("#status", Static).render()) - check("and says which execution this is", - f"2 of {len(stamps)}" in status, status[:80]) + check( + "and says which execution this is", + f"2 of {len(stamps)}" in status, + status[:80], + ) lst.cursor = row await settle(app) await pilot.press("<") await settle(app, lambda: app._t == stamps[0]) - check("< seeks back to the previous one", - app._t == stamps[0], f"t={app._t}, expected {stamps[0]}") + check( + "< seeks back to the previous one", + app._t == stamps[0], + f"t={app._t}, expected {stamps[0]}", + ) # An edge must SAY it's an edge rather than silently doing # nothing, which is indistinguishable from a broken key. lst.cursor = row await settle(app) await pilot.press("<") - await settle(app, lambda: "first" in str( - app.query_one("#status", Static).render())) + await settle( + app, + lambda: "first" in str(app.query_one("#status", Static).render()), + ) status = str(app.query_one("#status", Static).render()) - check("and the first execution says so instead of moving", - app._t == stamps[0] and "first" in status, status[:80]) + check( + "and the first execution says so instead of moving", + app._t == stamps[0] and "first" in status, + status[:80], + ) # -- "which instruction set this register?" ---------------------- # want_t = min(200, t.length - 1) @@ -409,14 +522,24 @@ async def run() -> int: await settle(app, lambda: app._t == want_t) lst.focus() await pilot.press("W") - opened = await wait(lambda: isinstance(app.screen, RegWriteScreen), - pilot, 20) - check("W lists the registers and where each was set", opened, - f"screen={type(app.screen).__name__}") + opened = await wait( + lambda: isinstance(app.screen, RegWriteScreen), pilot, 20 + ) + check( + "W lists the registers and where each was set", + opened, + f"screen={type(app.screen).__name__}", + ) if opened: sc = app.screen - pick = next((k for k, (n, v, l, x) in enumerate(sc._rows) - if l is not None and l != app._t), None) + pick = next( + ( + k + for k, (n, v, l, x) in enumerate(sc._rows) + if l is not None and l != app._t + ), + None, + ) if pick is None: check("a register was set by an earlier instruction", False) await pilot.press("escape") @@ -426,13 +549,18 @@ async def run() -> int: await settle(app) await pilot.press("enter") await wait(lambda: app._t == last, pilot, 30) - check("choosing one seeks to the write that set it", - app._t == last, f"t={app._t}, expected {last}") + check( + "choosing one seeks to the write that set it", + app._t == last, + f"t={app._t}, expected {last}", + ) # The real check: that instruction must actually have # written the register we asked about. - check("and that instruction really wrote it", - name in t.changed(app._t), - f"{name} not in {sorted(t.changed(app._t))}") + check( + "and that instruction really wrote it", + name in t.changed(app._t), + f"{name} not in {sorted(t.changed(app._t))}", + ) print(f"\n{PASS} passed, {FAIL} failed") return 1 if FAIL else 0 diff --git a/tests/test_trace_vs_tenet.py b/tests/test_trace_vs_tenet.py index 30d74b0..c2a0ea2 100644 --- a/tests/test_trace_vs_tenet.py +++ b/tests/test_trace_vs_tenet.py @@ -46,6 +46,7 @@ def _reference(): log.pmsg = lambda *a, **k: None import tenet # noqa: F401 import tenet.util # noqa: F401 + sys.modules["tenet.util.log"] = log from tenet.trace.arch import ArchAMD64 from tenet.trace.reader import TraceReader @@ -57,6 +58,7 @@ def _reference(): # stay in raw trace addresses, which is what we want to compare. def get_instruction_addresses(self): return [0xDEAD0000] + return TraceReader, ArchAMD64, FakeDctx @@ -88,9 +90,11 @@ def compare(path, ref_cls, arch, dctx, samples=200): theirs = TraceReader(path, ArchAMD64(), FakeDctx()) name = os.path.basename(path) - check(f"{name}: same length", - ours.length == theirs.trace.length, - f"{ours.length} vs {theirs.trace.length}") + check( + f"{name}: same length", + ours.length == theirs.trace.length, + f"{ours.length} vs {theirs.trace.length}", + ) n = min(ours.length, theirs.trace.length) if not n: return @@ -99,8 +103,11 @@ def compare(path, ref_cls, arch, dctx, samples=200): idxs = sorted({0, n - 1, n // 2} | {rnd.randrange(n) for _ in range(samples)}) bad = [i for i in idxs if ours.raw_ip(i) != theirs.get_ip(i)] - check(f"{name}: same PC at every sampled timestamp", not bad, - f"first mismatch at {bad[:1]}") + check( + f"{name}: same PC at every sampled timestamp", + not bad, + f"first mismatch at {bad[:1]}", + ) # Register reconstruction is the part that is easy to get subtly wrong: a # delta belongs to the line that CAUSED it, and an off-by-one here silently @@ -121,11 +128,16 @@ def compare(path, ref_cls, arch, dctx, samples=200): continue true = _truth(path, r, i) (ours_wrong if mine != true else ref_wrong).append((i, r, mine, ref, true)) - check(f"{name}: register state matches the trace text everywhere", - not ours_wrong, f"{ours_wrong[:3]}") + check( + f"{name}: register state matches the trace text everywhere", + not ours_wrong, + f"{ours_wrong[:3]}", + ) if ref_wrong: - print(f" (reference disagrees at {len(ref_wrong)} sampled points; " - f"the text backs us, e.g. idx {ref_wrong[0][0]} {ref_wrong[0][1]})") + print( + f" (reference disagrees at {len(ref_wrong)} sampled points; " + f"the text backs us, e.g. idx {ref_wrong[0][0]} {ref_wrong[0][1]})" + ) # Execution queries: what painting is built on. hot = sorted(ours.by_ip, key=lambda a: -len(ours.by_ip[a]))[:5] @@ -135,8 +147,11 @@ def compare(path, ref_cls, arch, dctx, samples=200): ref = list(theirs.get_executions(ea)) if mine != ref: ex_bad.append((hex(ea), len(mine), len(ref))) - check(f"{name}: same execution timestamps for the hottest addresses", - not ex_bad, f"{ex_bad[:3]}") + check( + f"{name}: same execution timestamps for the hottest addresses", + not ex_bad, + f"{ex_bad[:3]}", + ) # Memory STATE at a timestamp — reconstructed from the deltas, which is the # hard part and the whole point of reading memory from a trace. @@ -153,11 +168,14 @@ def compare(path, ref_cls, arch, dctx, samples=200): # own coverage separately and a byte neither has seen is not a # disagreement. for j in range(n): - if known[j] and refb[j:j + 1] and mine[j] != refb[j]: + if known[j] and refb[j : j + 1] and mine[j] != refb[j]: mem_bad.append((i, hex(op.addr + j), mine[j], refb[j])) mem_checked += 1 - check(f"{name}: memory state at a timestamp matches the reference", - not mem_bad, f"{mem_bad[:3]}") + check( + f"{name}: memory state at a timestamp matches the reference", + not mem_bad, + f"{mem_bad[:3]}", + ) # Memory: the bytes an instruction touched, and which way. with_mem = [i for i in idxs if ours.memory_ops(i)][:40] @@ -167,18 +185,24 @@ def compare(path, ref_cls, arch, dctx, samples=200): ref = theirs.get_memory(op.addr, len(op.data), i + 1) if op.write else None if ref is not None and bytes(ref.data) != op.data: mem_bad.append((i, hex(op.addr), op.data.hex(), bytes(ref.data).hex())) - check(f"{name}: written bytes match the reference's memory state", - not mem_bad, f"{mem_bad[:2]}") - print(f" ({ours.length:,} instructions, {len(idxs)} sampled, " - f"{len(with_mem)} with memory)") + check( + f"{name}: written bytes match the reference's memory state", + not mem_bad, + f"{mem_bad[:2]}", + ) + print( + f" ({ours.length:,} instructions, {len(idxs)} sampled, " + f"{len(with_mem)} with memory)" + ) def main(argv): if not os.path.isdir(TENET): print(f" skip: reference not found at {TENET}") return 0 - traces = argv or [p for p in ("/tmp/echotrace.0.log", "/tmp/big.0.log") - if os.path.exists(p)] + traces = argv or [ + p for p in ("/tmp/echotrace.0.log", "/tmp/big.0.log") if os.path.exists(p) + ] if not traces: print(" skip: no traces to compare (pass one, or run tenet-trace first)") return 0 diff --git a/tools/demo.py b/tools/demo.py index fa82a98..5cf5207 100644 --- a/tools/demo.py +++ b/tools/demo.py @@ -24,6 +24,7 @@ Edits (rename/comment) are reverted at the end, so the tour is repeatable and a scratch database is not left renamed. --spawn works on a COPY of the target so the tracked .i64 is never touched at all. """ + from __future__ import annotations import argparse @@ -223,8 +224,10 @@ class Demo: """Undo the demo's edits so the take is repeatable.""" for kind, args in reversed(self.undo): if kind == "rename" and args.get("addr") is not None: - self.do("rename_many", - items=[{"addr": hex(args["addr"]), "name": args["name"]}]) + self.do( + "rename_many", + items=[{"addr": hex(args["addr"]), "name": args["name"]}], + ) elif kind == "comment": self.do("comment", text="") # Re-navigate so the view shows the reverted name: the nav entry caches @@ -252,15 +255,20 @@ SCENES = [ def spawn_pane(target: str) -> tuple[str, str, str]: """Spawn a TUI pane on a COPY of ``target``. Returns (sock, pane, tmpdir).""" import json + tmp = tempfile.mkdtemp(prefix="idatui-demo-") copy = os.path.join(tmp, os.path.basename(target)) shutil.copy2(target, copy) - for suffix in (".i64",): # reuse the analysis if present + for suffix in (".i64",): # reuse the analysis if present if os.path.exists(target + suffix): shutil.copy2(target + suffix, copy + suffix) out = subprocess.run( [sys.executable, "-m", "idatui.pane", "spawn", "--open", copy], - cwd=REPO, capture_output=True, text=True, check=True).stdout + cwd=REPO, + capture_output=True, + text=True, + check=True, + ).stdout row = json.loads(out) return row["sock"], row.get("pane", ""), tmp @@ -280,8 +288,9 @@ def run_here(target: str) -> tuple[subprocess.Popen, str, str]: shutil.copy2(target + ".i64", copy + ".i64") sockdir = os.environ.get("XDG_RUNTIME_DIR") or "/tmp" sock = os.path.join(sockdir, f"idatui-demo-{os.getpid()}.sock") - proc = subprocess.Popen([os.path.join(REPO, "ida-tui"), copy, "--rpc", sock], - cwd=REPO) # stdio inherited on purpose + proc = subprocess.Popen( + [os.path.join(REPO, "ida-tui"), copy, "--rpc", sock], cwd=REPO + ) # stdio inherited on purpose return proc, sock, tmp @@ -291,28 +300,43 @@ def wait_for_socket(proc: subprocess.Popen, sock: str, timeout: float = 600.0) - while time.time() < deadline: if os.path.exists(sock): return True - if proc.poll() is not None: # died before it ever listened + if proc.poll() is not None: # died before it ever listened return False time.sleep(0.1) return False def main(argv=None) -> int: - ap = argparse.ArgumentParser(description=__doc__, - formatter_class=argparse.RawDescriptionHelpFormatter) + ap = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) ap.add_argument("--sock", help="RPC socket of a running TUI (see --rpc)") - ap.add_argument("--spawn", action="store_true", - help="spawn a pane on a scratch copy, then tear it down") - ap.add_argument("--here", "--inline", dest="here", action="store_true", - help="run the TUI in THIS terminal (single-pane recording)") - ap.add_argument("--target", default=DEFAULT_TARGET, - help="binary for --here/--spawn") - ap.add_argument("--speed", type=float, default=1.0, - help="pause multiplier: <1 snappier, >1 slower (default 1.0)") + ap.add_argument( + "--spawn", + action="store_true", + help="spawn a pane on a scratch copy, then tear it down", + ) + ap.add_argument( + "--here", + "--inline", + dest="here", + action="store_true", + help="run the TUI in THIS terminal (single-pane recording)", + ) + ap.add_argument( + "--target", default=DEFAULT_TARGET, help="binary for --here/--spawn" + ) + ap.add_argument( + "--speed", + type=float, + default=1.0, + help="pause multiplier: <1 snappier, >1 slower (default 1.0)", + ) ap.add_argument("--only", help="comma-separated scene names") ap.add_argument("--list", action="store_true", help="list scenes and exit") - ap.add_argument("--no-revert", action="store_true", - help="keep the demo's rename/comment") + ap.add_argument( + "--no-revert", action="store_true", help="keep the demo's rename/comment" + ) ap.add_argument("--quiet", action="store_true", help="no operator narration") args = ap.parse_args(argv) @@ -384,8 +408,11 @@ def main(argv=None) -> int: rc = 1 finally: if args.spawn and sock: - subprocess.run([sys.executable, "-m", "idatui.pane", "stop", - "--sock", sock], cwd=REPO, capture_output=True) + subprocess.run( + [sys.executable, "-m", "idatui.pane", "stop", "--sock", sock], + cwd=REPO, + capture_output=True, + ) if proc is not None: try: proc.wait(timeout=30) @@ -397,10 +424,12 @@ def main(argv=None) -> int: proc.kill() if tmp: shutil.rmtree(tmp, ignore_errors=True) - if args.here and transcript: # the alt screen is gone: safe to print + if args.here and transcript: # the alt screen is gone: safe to print print("\n\033[1m-- ida-tui demo --\033[0m") for line in transcript: - print(f" {line}" if not line.startswith("[") else f"\033[1m{line}\033[0m") + print( + f" {line}" if not line.startswith("[") else f"\033[1m{line}\033[0m" + ) print("done.") return rc diff --git a/tools/make_logo_ans.py b/tools/make_logo_ans.py index 9736e30..41c8b18 100644 --- a/tools/make_logo_ans.py +++ b/tools/make_logo_ans.py @@ -16,6 +16,7 @@ Needs Pillow, so run it with a python that has it (NOT ~/ida-venv): /usr/bin/python3 tools/make_logo_ans.py [--cols 60] [-o logo.ans] """ + from __future__ import annotations import argparse @@ -25,9 +26,9 @@ import sys from PIL import Image REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) -ALPHA_ON = 128 # at/above this a pixel counts as present +ALPHA_ON = 128 # at/above this a pixel counts as present -UPPER, LOWER = "\u2580", "\u2584" # upper half block, lower half block +UPPER, LOWER = "\u2580", "\u2584" # upper half block, lower half block def main() -> int: @@ -35,8 +36,11 @@ def main() -> int: ap.add_argument("--png", default=os.path.join(REPO, "logo.png")) ap.add_argument("-o", "--out", default=os.path.join(REPO, "logo.ans")) ap.add_argument("--cols", type=int, default=60) - ap.add_argument("--cell", default="9x22", - help="terminal cell size WxH in px, for the aspect ratio") + ap.add_argument( + "--cell", + default="9x22", + help="terminal cell size WxH in px, for the aspect ratio", + ) args = ap.parse_args() cw, ch = (int(v) for v in args.cell.lower().split("x")) @@ -62,13 +66,15 @@ def main() -> int: if not t_on and not b_on: sgr, ch_ = "\033[0m", " " elif t_on and b_on: - sgr = (f"\033[38;2;{bot[0]};{bot[1]};{bot[2]}m" - f"\033[48;2;{top[0]};{top[1]};{top[2]}m") + sgr = ( + f"\033[38;2;{bot[0]};{bot[1]};{bot[2]}m" + f"\033[48;2;{top[0]};{top[1]};{top[2]}m" + ) ch_ = LOWER - elif b_on: # only the lower pixel is present + elif b_on: # only the lower pixel is present sgr = f"\033[0m\033[38;2;{bot[0]};{bot[1]};{bot[2]}m" ch_ = LOWER - else: # only the upper pixel is present + else: # only the upper pixel is present sgr = f"\033[0m\033[38;2;{top[0]};{top[1]};{top[2]}m" ch_ = UPPER if sgr != prev: @@ -81,8 +87,10 @@ def main() -> int: text = "\n".join(out) + "\n" with open(args.out, "w", encoding="utf-8") as f: f.write(text) - print(f"{args.png} {w}x{h} -> {args.out} {cols}x{rows} cells " - f"({len(text):,} bytes, cell {cw}x{ch})") + print( + f"{args.png} {w}x{h} -> {args.out} {cols}x{rows} cells " + f"({len(text):,} bytes, cell {cw}x{ch})" + ) return 0 diff --git a/tools/verify_procs.py b/tools/verify_procs.py index 323bf1f..f112cbc 100644 --- a/tools/verify_procs.py +++ b/tools/verify_procs.py @@ -30,6 +30,7 @@ from idatui.formats import PROCESSORS # noqa: E402 def main() -> int: import idapro + idapro.enable_console_messages(False) import ida_auto import ida_ida @@ -55,9 +56,12 @@ def main() -> int: bits = "" if rc == 0: import ida_ida + bits = f" bitness={ida_ida.inf_get_app_bitness()}" ok = rc == 0 and got.lower() == base.lower() - print(f" {'ok ' if ok else 'BAD '} {name:<14} rc={rc} -> {got!r}{bits} {desc}") + print( + f" {'ok ' if ok else 'BAD '} {name:<14} rc={rc} -> {got!r}{bits} {desc}" + ) if rc == 0: idapro.close_database(save=False) if not ok: -- cgit v1.3.1-sl0p