diff options
| -rw-r--r-- | docs/PROJECTS.md | 17 | ||||
| -rw-r--r-- | idatui/app.py | 243 | ||||
| -rw-r--r-- | idatui/launch.py | 63 | ||||
| -rw-r--r-- | tests/test_project_ui.py | 165 |
4 files changed, 466 insertions, 22 deletions
diff --git a/docs/PROJECTS.md b/docs/PROJECTS.md index d1e2806..1f39c05 100644 --- a/docs/PROJECTS.md +++ b/docs/PROJECTS.md @@ -108,10 +108,19 @@ basename and must be unique (it names the staged file). ## Phases -**Phase 1 — project model + switching.** Project file, staging dir, -`WorkerPool` (lazy spawn, budget eviction, save-on-evict, clean shutdown), -`BinaryState` save/restore, switcher palette, active binary in the status bar, -`--project` CLI. One active binary; no cross-binary search yet. +**Phase 1 — project model + switching. DONE.** Project file + staging +(`idatui/project.py`), `WorkerPool` with budget eviction / save-on-evict / +clean shutdown (`idatui/pool.py`), `BinaryState` snapshot+restore and the switch +itself, the `Ctrl+O` switcher palette, the active binary in the status line, and +`--project` (which creates the project when given binaries). One active binary; +no cross-binary search yet. + +Project mode is **additive**: with no `--project` the app is byte-for-byte the +single-binary tool it was, which is what keeps the 167-check pilot honest. +Switching reuses the `_after_reconnect` shape — swap client+program, rebuild the +index, reopen the entry. A binary whose worker is still resident restores +instantly (its `Program` and index are still in memory); an evicted one comes +back with a fresh worker but keeps its nav history, since that is just addresses. **Phase 2 — index cache + project-wide search.** Per-binary index (functions, strings) persisted after first open, keyed by source size+mtime. Scope toggle in diff --git a/idatui/app.py b/idatui/app.py index 2576db5..a436eb4 100644 --- a/idatui/app.py +++ b/idatui/app.py @@ -20,7 +20,7 @@ import asyncio import os import re import subprocess -from dataclasses import dataclass +from dataclasses import dataclass, field from rich.segment import Segment from rich.style import Style @@ -93,6 +93,24 @@ _ADDR_MARK_STRIP_RE = re.compile(r"\s*/\*\s*0x[0-9A-Fa-f]+\s*\*/") @dataclass +class BinaryState: + """Everything that makes one project binary's session resumable across a + switch. Addresses outlive the worker, so nav history survives eviction; the + Program/index only survive while that worker is still resident.""" + + label: str + program: object | None = None + func_index: object | None = None + nav: list = field(default_factory=list) + cur: object | None = None + pref: str = "listing" + active: str = "listing" + split: bool = False + filter_term: str = "" + dirty: bool = False + + +@dataclass class NavEntry: ea: int name: str @@ -2213,6 +2231,93 @@ class StringsPalette(ModalScreen): self.dismiss(None) +class ProjectPalette(ModalScreen): + """The project's binaries; Enter switches to one. Shows which are resident + (a live worker, so switching is instant) vs cold (needs an open).""" + + BINDINGS = [ + Binding("escape", "close", "Close"), + Binding("down,ctrl+n", "cursor_down", show=False), + Binding("up,ctrl+p", "cursor_up", show=False), + ] + + def __init__(self, entries: list[dict]) -> None: + super().__init__() + self._entries = entries + self._results: list[dict] = [] + + def compose(self) -> ComposeResult: + with Vertical(id="pal-box"): + yield Static(" binaries", id="pal-title") + 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: + self._apply("") + self.query_one("#pal-input", Input).focus() + + def on_input_changed(self, event: Input.Changed) -> None: + event.stop() + self._apply(event.value.strip()) + + def on_input_submitted(self, event: Input.Submitted) -> None: + event.stop() + self.action_choose() + + 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()] + 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(f"{e['label']:<22}", _S_LABEL) + if e["resident"]: + mb = e.get("memory_mb") or 0 + label.append(f"resident {mb:>4}MB ", _S_MNEM) + elif e["analysed"]: + label.append("analysed ", _S_DIM) + else: + label.append("not opened ", _S_DIM) + if e["pinned"]: + label.append("pin ", _S_ADDR) + label.append(e["source"], _S_DIM) + opts.append(Option(label)) + ol.add_options(opts) + if rows: + ol.highlighted = 0 + self.query_one("#pal-title", Static).update( + f" binaries: {len(rows)} of {len(self._entries)}") + + def action_cursor_down(self) -> None: + ol = self.query_one(OptionList) + if ol.option_count: + ol.highlighted = min((ol.highlighted or 0) + 1, ol.option_count - 1) + + def action_cursor_up(self) -> None: + ol = self.query_one(OptionList) + if ol.option_count: + ol.highlighted = max((ol.highlighted or 0) - 1, 0) + + def action_choose(self) -> None: + i = self.query_one(OptionList).highlighted + if i is not None and 0 <= i < len(self._results): + self.dismiss(self._results[i]["label"]) + + def on_option_list_option_selected(self, event: OptionList.OptionSelected) -> None: + if 0 <= event.option_index < len(self._results): + self.dismiss(self._results[event.option_index]["label"]) + + def action_close(self) -> None: + self.dismiss(None) + + # --------------------------------------------------------------------------- # # Confirmation dialog # --------------------------------------------------------------------------- # @@ -2571,6 +2676,8 @@ class IdaCommands(Provider): ("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), ("Follow symbol under cursor", "jump to the referenced symbol (Enter)", lambda: va("follow")), ("Show xrefs to symbol", "cross-references to the cursor symbol (x)", @@ -2713,6 +2820,7 @@ class IdaTui(App): Binding("backslash", "hex", "Hex"), Binding("s", "toggle_split", "Split", show=False), Binding("quotation_mark,shift+f12", "strings", "Strings", show=False), + Binding("ctrl+o", "switch_binary", "Binaries", show=False), Binding("g", "goto", "Goto"), Binding("slash", "filter", "Filter", show=False), Binding("ctrl+b", "toggle_functions", "Names", show=False), @@ -2721,9 +2829,22 @@ class IdaTui(App): Binding("escape", "back", "Back"), ] - def __init__(self, open_path: str, keepalive: bool = True, - rpc_path: str | None = None, ttl: int = 1800) -> None: + def __init__(self, open_path: str | None = None, keepalive: bool = True, + rpc_path: str | None = None, ttl: int = 1800, + project=None) -> 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._states: dict[str, BinaryState] = {} + self._pending_restore = None # entry to reopen after a switch + if project is not None: + from .pool import WorkerPool + self._pool = WorkerPool(project, ttl=ttl) + self._binary = project.refs[0].label + open_path = project.refs[0].staged self._open_path = open_path self._ttl = ttl self._do_keepalive = keepalive @@ -2802,7 +2923,11 @@ class IdaTui(App): gi.display = False gi.can_focus = False yield gi - yield Static("connecting…", id="status") + # markup=False: the status is plain text full of [listing]/[split]/[label] + # markers and symbol names that may contain brackets. With Textual markup + # on, a single-word marker parses as a style tag and is silently eaten — + # which is why [listing] and [pseudocode] never actually rendered. + yield Static("connecting\u2026", id="status", markup=False) yield Footer() def on_mount(self) -> None: @@ -2848,6 +2973,8 @@ class IdaTui(App): # -- status helper ----------------------------------------------------- # def _status(self, text: str) -> None: + if self._binary: # project mode: always say which binary you're in + text = f"[{self._binary}] {text}" try: self.query_one("#status", Static).update(text) except Exception: # noqa: BLE001 -- status bar transiently unavailable @@ -2969,6 +3096,14 @@ class IdaTui(App): """Our idalib-worker path: spawn the worker (it opens + analyzes the binary in its own process) and connect. Returns the client, or None.""" from .worker_client import WorkerClient + if self._pool is not None: # project mode: the pool owns the workers + 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)) + self._binary = label + self._pool.set_active(label) + self._open_path = self._project.by_label(label).staged + return client if not self._open_path: self.app.call_from_thread( self._status, "the worker backend needs a binary path") @@ -3018,6 +3153,13 @@ class IdaTui(App): main() if it exists; else open the symbol picker so you're never staring at a blank pane. Runs once (guarded by ``_cur``) and never steals focus from a user who already navigated.""" + entry = self._pending_restore + if entry is not None: # switched back to a binary we'd already explored + self._pending_restore = None + self._did_auto_land = True + self._dismiss_loading() + self._open_entry(entry, push=False) + return if self._cur is not None or self._did_auto_land or self._func_index is None: return self._did_auto_land = True @@ -3165,6 +3307,95 @@ class IdaTui(App): f = self._func_index.by_addr(addr) if self._func_index else None self._open_function(addr, f.name if f else hex(addr)) + # -- projects: switching between the binaries of one target ------------ # + def action_switch_binary(self) -> None: + """Ctrl+O: pick another binary from the project (Ghidra-style).""" + if self._pool is None: + self._status("not a project — open one with --project") + return + if self._prompt_active(): + return + 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: + self._switch_binary(label) + + def _switch_binary(self, label: str) -> None: + # Snapshot what we're leaving so coming back restores the view, then let + # the pool hand us a worker (spawning + evicting as the budget dictates). + 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, + pref=self._pref, 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) + + @work(thread=True, exclusive=True, group="switch-binary") + 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)) + except Exception as e: # noqa: BLE001 + self.app.call_from_thread(self._switch_failed, label, str(e)) + return + st = self._states.get(label) + # The Program (and its caches) only survive while that worker does; a + # binary that was evicted comes back with a fresh one. Either way the nav + # 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) + program = st.program if reuse else Program(client) + 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 + self.program = program + self._binary = label + self._pool.set_active(label) + self._open_path = self._project.by_label(label).staged + self._pref = st.pref if st else "listing" + self._active = st.active if st else "listing" + self._split = st.split if st else False + self._filter_term = st.filter_term if st else "" + self._dirty = st.dirty if st else False + self._nav = list(st.nav) if st else [] + self._decomp_return = None + self._split_eamap, self._split_ea2line, self._split_range = [], {}, None + self.query_one(DecompView).loaded_ea = None # belongs to the old binary + if reuse and st.func_index is not None: + self._cur = st.cur + self._func_index = st.func_index + self._apply_filter(self._filter_term) # repopulate the names table + self._dismiss_loading() + if self._cur is not None: + self._open_entry(self._cur, push=False) + else: + self._did_auto_land = False + self._auto_land() + return + # Cold (first visit, or the worker was evicted): rebuild the index, then + # land back where we were via _pending_restore. + self._cur = None + self._func_index = None + self._pending_restore = st.cur if st else None + # Landing is per-binary: without this the app-wide "already landed" flag + # from the first binary would stop the new one landing at all (and leave + # the switch overlay up forever). + self._did_auto_land = False + self._status("loading functions\u2026") + self._load_functions() + + def _switch_failed(self, label: str, why: str) -> None: + self._dismiss_loading() + self._status(f"could not open {label}: {why}") + def action_strings(self) -> None: """'\"' / Shift+F12: browse every string in the binary (filterable); Enter jumps to it in the unified listing.""" @@ -4934,5 +5165,7 @@ class IdaTui(App): self._ka.stop() if self.program is not None: self.program.close() - if self.client is not None: + if self._pool is not None: + self._pool.close_all() # saves each DB, then closes cleanly + elif self.client is not None: self.client.close() diff --git a/idatui/launch.py b/idatui/launch.py index e2845e7..0ae63a1 100644 --- a/idatui/launch.py +++ b/idatui/launch.py @@ -45,7 +45,12 @@ def main(argv: list[str] | None = None) -> int: p = argparse.ArgumentParser( prog="ida-tui", description="Open a binary in the IDA TUI (private idalib worker).") - p.add_argument("binary", help="binary to open and analyze") + 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="worker idle-TTL seconds (default 1800)") p.add_argument("--no-keepalive", action="store_true", @@ -54,17 +59,49 @@ def main(argv: list[str] | None = None) -> int: help="listen for RPC on this unix socket (puppeteer the TUI)") args = p.parse_args(argv) - binary = os.path.abspath(os.path.expanduser(args.binary)) - if not os.path.isfile(binary): - _log(f"no such file: {binary}") - return 2 - if not os.access(os.path.dirname(binary), os.W_OK): - _log(f"directory not writable (IDA writes a .i64 there): " - f"{os.path.dirname(binary)}") - return 2 - swept = _sweep_locks(binary) # a crashed worker can leave the DB wedged - if swept: - _log(f"cleared {swept} stale lock file(s) from a crashed worker") + project = None + 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): + project = Project.load(ppath) + for b in args.binary: # extend an existing project + project.add(b) + if args.binary: + project.save() + elif args.binary: + project = Project.create(ppath, args.binary) + _log(f"created project {ppath} with {len(project.refs)} binaries") + else: + _log(f"no such project: {ppath} (pass binaries to create it)") + return 2 + except ProjectError as e: + _log(str(e)) + return 2 + # Everything IDA writes lives in the project's sidecar, so the source + # tree is never touched and never needs to be writable. + try: + project.stage_all(progress=lambda m: _log(m)) + except ProjectError as e: + _log(str(e)) + return 2 + else: + if len(args.binary) != 1: + _log("give exactly one binary, or use --project for several") + return 2 + binary = os.path.abspath(os.path.expanduser(args.binary[0])) + if not os.path.isfile(binary): + _log(f"no such file: {binary}") + return 2 + if not os.access(os.path.dirname(binary), os.W_OK): + _log(f"directory not writable (IDA writes a .i64 there): " + f"{os.path.dirname(binary)}") + return 2 + swept = _sweep_locks(binary) # a crashed worker can leave the DB wedged + if swept: + _log(f"cleared {swept} stale lock file(s) from a crashed worker") # Hand off to the TUI (imported late so --help works without textual). It # spawns the worker behind its loading overlay while auto-analysis runs. @@ -75,7 +112,7 @@ def main(argv: list[str] | None = None) -> int: return 1 rpc_path = os.path.abspath(os.path.expanduser(args.rpc)) if args.rpc else None IdaTui(open_path=binary, keepalive=not args.no_keepalive, - rpc_path=rpc_path, ttl=args.ttl).run() + rpc_path=rpc_path, ttl=args.ttl, project=project).run() return 0 diff --git a/tests/test_project_ui.py b/tests/test_project_ui.py new file mode 100644 index 0000000..025107d --- /dev/null +++ b/tests/test_project_ui.py @@ -0,0 +1,165 @@ +#!/usr/bin/env python3 +"""End-to-end pilot for project mode: two binaries, switching between them. + +Needs idalib (it spawns real workers, one per binary) and textual: + + ~/ida-venv/bin/python tests/test_project_ui.py [bin1 bin2] + +Defaults to targets/echo + targets/cat. The binaries are copied into a temp +source dir first, so the "source tree stays pristine" promise is checkable. +""" +import asyncio +import os +import shutil +import sys +import tempfile + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +from idatui._sync import wait_for # noqa: E402 +from idatui.app import IdaTui, ProjectPalette # noqa: E402 +from idatui.project import Project # noqa: E402 +from textual.widgets import Input, Static # noqa: E402 + +PASS = FAIL = 0 + + +def check(name, cond, detail=""): + global PASS, FAIL + if cond: + PASS += 1 + print(f" ok {name}") + else: + FAIL += 1 + print(f" FAIL {name} {detail}") + + +async def run(bins): + repo = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + with tempfile.TemporaryDirectory() as tmp: + src = os.path.join(tmp, "src") + os.makedirs(src) + srcs = [] + for b in bins: + dst = os.path.join(src, os.path.basename(b)) + shutil.copy2(b, dst) + srcs.append(dst) + proj = Project.create(os.path.join(tmp, "proj.json"), srcs, name="proj") + proj.stage_all() + first, second = (r.label for r in proj.refs) + + 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) + + # -- boots on the project's first binary ----------------------- # + ok = await settle(lambda: app.program is not None + and app._func_index is not None + and app._func_index.complete) + 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]) + + # -- 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__}") + 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)) + + # -- switch to the second binary -------------------------------- # + pal.query_one(Input).value = second + await pilot.pause(0.2) + 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) + 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}") + 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__}") + if not reopened: + return + app.screen.query_one(Input).value = first + await pilot.pause(0.2) + await pilot.press("enter") + back = await settle(lambda: app._binary == first + and app._func_index is not None + and app._func_index.complete, 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") + return + app.screen.query_one(Input).value = second + await pilot.pause(0.2) + 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}") + + # -- 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}") + 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}") + + print(f"\n{PASS} passed, {FAIL} failed") + return 1 if FAIL else 0 + + +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")] + for b in bins: + if not os.path.isfile(b): + print(f"no such binary: {b}") + return 2 + return asyncio.run(run(bins)) + + +if __name__ == "__main__": + raise SystemExit(main(sys.argv[1:])) |