diff options
Diffstat (limited to 'idatui/pane.py')
| -rw-r--r-- | idatui/pane.py | 58 |
1 files changed, 49 insertions, 9 deletions
diff --git a/idatui/pane.py b/idatui/pane.py index e34b822..37a5f0c 100644 --- a/idatui/pane.py +++ b/idatui/pane.py @@ -131,15 +131,20 @@ def spawn(args) -> int: if not os.environ.get("TMUX"): print("error: not inside tmux (spawn creates a tmux pane)", file=sys.stderr) return 2 - if not args.open: - print("error: pass --open <binary>", file=sys.stderr) + if not args.open and not getattr(args, "project", None): + print("error: pass --open <binary> or --project <file>", file=sys.stderr) return 2 sock = args.sock or os.path.join(_sockdir(), f"idatui-{secrets.token_hex(3)}.sock") - target = os.path.abspath(os.path.expanduser(args.open)) - if not os.path.exists(target): + 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) return 2 + if project is not None and target is None and not os.path.exists(project): + print(f"error: no such project: {project}", file=sys.stderr) + return 2 # Reap workers leaked by previously-stopped/crashed panes so we don't spawn # into a full IDA_MCP_MAX_WORKERS (which makes the new TUI hang forever, @@ -151,7 +156,17 @@ def spawn(args) -> int: # the command the pane runs: the launcher spawns a private idalib worker for # this binary and becomes the TUI, so kill-pane tears the whole thing down. - inner = [args.python, "-m", "idatui.launch", target, "--rpc", sock] + if project is not None: + # launch takes: --project FILE [binaries...]; extra binaries are added to + # the project (and a missing project file is created from them). + inner = [args.python, "-m", "idatui.launch", "--project", project] + if target is not None: + inner.append(target) + inner += ["--rpc", sock] + else: + inner = [args.python, "-m", "idatui.launch", target, "--rpc", sock] + if getattr(args, "trace", None): + inner += ["--trace", os.path.abspath(os.path.expanduser(args.trace))] cmd = f"cd {REPO!r} && exec " + " ".join(_q(a) for a in inner) split = ["split-window", "-v" if args.vertical else "-h", @@ -166,8 +181,8 @@ def spawn(args) -> int: split.append(cmd) pane = _tmux(*split) - row = {"sock": sock, "pane": pane, "target": target, - "kind": "open", "started": time.time()} + row = {"sock": sock, "pane": pane, "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) @@ -233,18 +248,29 @@ def stop(args) -> int: if not rows: print("error: no matching pane (need --sock or --pane)", file=sys.stderr) return 2 + killed: list[str] = [] for r in rows: sock, pane = r.get("sock"), r.get("pane") + quit_ok = False if sock and os.path.exists(sock): try: # ask it to quit gracefully first with RpcClient(sock) as c: c.call("quit") - time.sleep(0.4) + quit_ok = True except (OSError, RpcError, ConnectionError): pass + # Wait for the pane to actually go away. Quitting runs App.on_unmount, + # which writes every dirty database; a 90 MB .i64 takes tens of seconds. + # Killing the pane on a fixed short sleep truncated that save and + # silently destroyed the session's work, so block on the real signal. + if pane and quit_ok: + deadline = time.monotonic() + float(args.timeout) + while time.monotonic() < deadline and _pane_alive(pane): + time.sleep(0.25) if pane and _pane_alive(pane): subprocess.run(["tmux", "kill-pane", "-t", pane], capture_output=True) + killed.append(pane) if sock: try: os.unlink(sock) @@ -256,6 +282,12 @@ def stop(args) -> int: out = {"stopped": [r.get("sock") or r.get("pane") for r in rows]} if reaped: out["reaped_workers"] = reaped + if killed: + # 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") print(json.dumps(out)) return 0 @@ -301,8 +333,13 @@ def main(argv: list[str]) -> int: 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", required=True, + 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("--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)") @@ -315,6 +352,9 @@ def main(argv: list[str]) -> int: 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.set_defaults(fn=stop) ls = sub.add_parser("list", help="list tracked panes") |