diff options
| -rw-r--r-- | idatui/app.py | 21 | ||||
| -rw-r--r-- | idatui/rpc.py | 273 | ||||
| -rw-r--r-- | idatui/tui.py | 5 | ||||
| -rw-r--r-- | tests/rpc_smoke.py | 121 |
4 files changed, 418 insertions, 2 deletions
diff --git a/idatui/app.py b/idatui/app.py index 11d54aa..ebf928c 100644 --- a/idatui/app.py +++ b/idatui/app.py @@ -16,6 +16,7 @@ Design notes: from __future__ import annotations +import asyncio import os import re import subprocess @@ -1731,12 +1732,14 @@ class IdaTui(App): ] def __init__(self, url: str, db: str | None, keepalive: bool = True, - open_path: str | None = None) -> None: + open_path: str | None = None, rpc_path: str | None = None) -> None: super().__init__() self._url = url self._db = db self._open_path = open_path self._do_keepalive = keepalive + self._rpc_path = rpc_path + self._rpc = None self.client: IDAClient | None = None self.program: Program | None = None self._ka = None @@ -1803,6 +1806,22 @@ class IdaTui(App): # (pseudocode) so app bindings work before anything is opened. self.query_one(DecompView).focus() self._connect() + if self._rpc_path: + self._start_rpc() + + def _start_rpc(self) -> None: + from .rpc import RpcServer + self._rpc = RpcServer(self, self._rpc_path) + + async def _serve() -> None: + await self._rpc.start() + self._status(f"rpc: listening on {self._rpc_path}") + + asyncio.get_running_loop().create_task(_serve()) + + async def on_unmount(self) -> None: + if self._rpc is not None: + await self._rpc.stop() # -- status helper ----------------------------------------------------- # def _status(self, text: str) -> None: diff --git a/idatui/rpc.py b/idatui/rpc.py new file mode 100644 index 0000000..e613a91 --- /dev/null +++ b/idatui/rpc.py @@ -0,0 +1,273 @@ +"""In-process RPC server: puppeteer the live TUI over a unix socket. + +The TUI renders normally in its terminal (what a viewer/livestream sees); this +listener runs on the *same* asyncio loop, so handlers touch the UI directly with +no thread marshalling. Every injected key produces the identical on-screen effect +a human keyboard would — that's the whole point (drive it "as if a user"). + +Protocol: newline-delimited JSON, one request/response per line. + -> {"id": 1, "method": "keys", "params": {"keys": ["g","m","a","i","n","enter"]}} + <- {"id": 1, "result": {...}} (or {"id":1,"error":{"message":..}}) + +Anyone who can r/w the socket can drive the app (no auth by design; the socket is +mode 0600, local only). + +Method tiers: + raw keys, text — inject keystrokes (max fidelity) + introspect state, view, screen, functions + (semantic verbs — open/goto/rename/... — layer on top in a later pass.) +""" +from __future__ import annotations + +import asyncio +import io +import json +import os +from typing import Any + +from rich.console import Console + +from ._sync import settle +from .app import DecompView, DisasmView, HexView + +PROTO_VERSION = 1 + + +# --------------------------------------------------------------------------- # +# State introspection +# --------------------------------------------------------------------------- # +def _active_widget(app): + """The currently *shown* code widget (mirrors app._active).""" + if app._active == "hex": + return app.query_one(HexView) + if app._active == "disasm": + return app.query_one(DisasmView) + return app.query_one(DecompView) + + +_MODALS = ("XrefsScreen", "SymbolPalette", "StructEditor", "ConfirmScreen") + + +def _modal_snapshot(app) -> dict[str, Any] | None: + """Describe the top modal screen, if any, enough to drive it.""" + scr = app.screen + name = type(scr).__name__ + if name not in _MODALS: + return None # the base app Screen is not a modal + info: dict[str, Any] = {"kind": name} + items = getattr(scr, "_items", 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)} + for it in items[:64] + ] + return info + + +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} + # disasm / decomp share the ColumnCursor surface + word = None + try: + word = w.word_under_cursor() + except Exception: # noqa: BLE001 + pass + ea = None + try: + 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)} + + +def snapshot(app) -> dict[str, Any]: + """A structured view of everything an agent needs to decide the next move.""" + cur = app._cur + w = _active_widget(app) + st = "" + try: + st = str(app.query_one("#status").render()) + except Exception: # noqa: BLE001 + pass + return { + "active": app._active, + "pref": app._pref, + "function": ({"ea": cur.ea, "name": cur.name} if cur else None), + "cursor": _cursor_info(app, w), + "status": st, + "filter": app._filter_term, + "nav_depth": len(app._nav), + "dirty": bool(app._dirty), + "modal": _modal_snapshot(app), + "loaded": app.program is not None and app._func_index is not None, + } + + +def view_lines(app, lines: int | None = None) -> dict[str, Any]: + """The visible text of the active code pane, cursor-marked. (disasm/decomp; + 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)} + top = round(w.scroll_offset.y) + height = w.size.height or 40 + n = min(lines or height, max(w.total - top, 0)) + out = [] + 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} + + +def screen_text(app) -> dict[str, Any]: + """Plain-text render of the whole screen — exactly what the viewer sees.""" + 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) + render = app.screen._compositor.render_update( + full=True, screen_stack=app._background_screens, simplify=False) + console.print(render) + return {"width": width, "height": height, "text": console.export_text(styles=False)} + + +def functions(app, flt: str | None = None, limit: int = 50) -> list[dict[str, Any]]: + idx = app._func_index + if idx is None: + return [] + flt = (flt or "").lower() + out = [] + for f in idx.all_loaded(): + if flt and flt not in f.name.lower(): + continue + out.append({"ea": f.addr, "name": f.name, "size": f.size}) + if len(out) >= limit: + break + return out + + +# --------------------------------------------------------------------------- # +# Keystroke injection +# --------------------------------------------------------------------------- # +def _text_to_keys(text: str, delay_ms: int = 0) -> list[str]: + """Each printable char becomes a key; interleave wait:<ms> for the typed-out + aesthetic (App._press_keys understands 'wait:NN').""" + keys: list[str] = [] + for i, ch in enumerate(text): + if delay_ms and i: + keys.append(f"wait:{delay_ms}") + keys.append(ch) + return keys + + +# --------------------------------------------------------------------------- # +# Server +# --------------------------------------------------------------------------- # +class RpcServer: + def __init__(self, app, path: str): + self.app = app + self.path = path + self._server: asyncio.AbstractServer | None = None + + async def start(self) -> None: + if os.path.exists(self.path): + try: + os.unlink(self.path) + except OSError: + pass + self._server = await asyncio.start_unix_server(self._on_client, path=self.path) + try: + os.chmod(self.path, 0o600) + except OSError: + pass + + async def stop(self) -> None: + if self._server is not None: + self._server.close() + try: + await self._server.wait_closed() + except Exception: # noqa: BLE001 + pass + try: + if os.path.exists(self.path): + os.unlink(self.path) + except OSError: + pass + + async def _on_client(self, reader: asyncio.StreamReader, + writer: asyncio.StreamWriter) -> None: + try: + while not reader.at_eof(): + line = await reader.readline() + if not line: + break + resp = await self._handle_line(line) + writer.write(json.dumps(resp).encode() + b"\n") + await writer.drain() + except (ConnectionResetError, BrokenPipeError): + pass + finally: + try: + writer.close() + except Exception: # noqa: BLE001 + pass + + async def _handle_line(self, line: bytes) -> dict[str, Any]: + rid = None + try: + req = json.loads(line.decode()) + rid = req.get("id") + method = req.get("method") + params = req.get("params") or {} + result = await self._dispatch(method, params) + return {"id": rid, "result": result} + except Exception as e: # noqa: BLE001 — report, never kill the connection + return {"id": rid, "error": {"message": f"{type(e).__name__}: {e}"}} + + async def _dispatch(self, method: str, params: dict[str, Any]) -> Any: + app = self.app + if method in (None, "ping"): + return {"ok": True, "proto": PROTO_VERSION, "loaded": app.program is not None} + + if method == "keys": + keys = params.get("keys") or [] + if not isinstance(keys, list): + raise ValueError("keys must be a list") + await app._press_keys([str(k) for k in keys]) + if params.get("settle", True): + await settle(app, timeout=float(params.get("timeout", 20.0))) + return snapshot(app) + + if method == "text": + text = str(params.get("text", "")) + keys = _text_to_keys(text, int(params.get("delay_ms", 0))) + await app._press_keys(keys) + if params.get("settle", True): + await settle(app, timeout=float(params.get("timeout", 20.0))) + return snapshot(app) + + if method == "state": + return snapshot(app) + if method == "view": + return view_lines(app, params.get("lines")) + if method == "screen": + return screen_text(app) + if method == "functions": + return functions(app, params.get("filter"), int(params.get("limit", 50))) + + raise ValueError(f"unknown method: {method!r}") diff --git a/idatui/tui.py b/idatui/tui.py index e787fe6..b399ffe 100644 --- a/idatui/tui.py +++ b/idatui/tui.py @@ -32,9 +32,12 @@ def main(argv: list[str] | None = None) -> int: help="open (or reopen) a binary and drive it (dir must be writable)") p.add_argument("--no-keepalive", action="store_true", help="Do not bump idle-TTL / run the keepalive heartbeat") + p.add_argument("--rpc", metavar="PATH", + help="listen for RPC on this unix socket path (puppeteer the TUI)") args = p.parse_args(argv) + rpc_path = os.path.abspath(os.path.expanduser(args.rpc)) if args.rpc else None IdaTui(url=args.url, db=args.db, open_path=args.open, - keepalive=not args.no_keepalive).run() + keepalive=not args.no_keepalive, rpc_path=rpc_path).run() return 0 diff --git a/tests/rpc_smoke.py b/tests/rpc_smoke.py new file mode 100644 index 0000000..92fc801 --- /dev/null +++ b/tests/rpc_smoke.py @@ -0,0 +1,121 @@ +#!/usr/bin/env python3 +"""End-to-end smoke test for the RPC server, in-process over a real unix socket. + + ~/ida-venv/bin/python tests/rpc_smoke.py --db <session_id> + +Boots the TUI headless with --rpc on a temp socket, connects a JSONL client over +that socket (same loop), and exercises the raw + introspection primitives. +""" +import asyncio +import json +import os +import sys +import tempfile + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +from idatui.app import IdaTui # noqa: E402 +from idatui.client import DEFAULT_URL # noqa: E402 +from idatui._sync import wait_for # noqa: E402 +from textual.widgets import DataTable # 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}") + + +class Conn: + """Tiny JSONL client over a unix socket.""" + def __init__(self, r, w): + self.r, self.w = r, w + self._id = 0 + + async def call(self, method, **params): + self._id += 1 + self.w.write(json.dumps({"id": self._id, "method": method, + "params": params}).encode() + b"\n") + await self.w.drain() + line = await self.r.readline() + return json.loads(line.decode()) + + +async def run(db): + sock = os.path.join(tempfile.gettempdir(), f"idatui-rpc-{os.getpid()}.sock") + url = os.environ.get("IDA_MCP_URL", DEFAULT_URL) + app = IdaTui(url=url, db=db, keepalive=False, rpc_path=sock) + async with app.run_test(size=(140, 44)) as pilot: + # boot: functions loaded + socket up + await wait_for(lambda: app.query_one("#func-table", DataTable).row_count > 0, + pilot.pause, 60) + await wait_for(lambda: os.path.exists(sock), pilot.pause, 20) + if app._func_index is not None and not app._func_index.complete: + app._func_index.load_all() + + r, w = await asyncio.open_unix_connection(sock) + c = Conn(r, w) + + resp = await c.call("ping") + check("ping returns proto", resp.get("result", {}).get("proto") == 1, str(resp)) + + resp = await c.call("functions", limit=5) + funcs = resp.get("result", []) + check("functions lists entries", isinstance(funcs, list) and len(funcs) > 0, + str(resp)[:120]) + + resp = await c.call("state") + st = resp.get("result", {}) + check("state has an active view", st.get("active") in ("decomp", "disasm", "hex"), + str(st)[:120]) + + # drive it like a user: goto a function by name via raw keys + target = next((f for f in funcs if f["name"] == "main"), funcs[0]) + keys = ["g"] + list(target["name"]) + ["enter"] + resp = await c.call("keys", keys=keys) + st = resp.get("result", {}) + check("keys(goto) navigates to the function", + st.get("function", {}).get("name") == target["name"], + f"got {st.get('function')}") + + resp = await c.call("view", lines=6) + v = resp.get("result", {}) + check("view returns visible lines with a cursor", + isinstance(v.get("lines"), list) and any(l.get("cur") for l in v["lines"]), + str(v)[:120]) + + resp = await c.call("screen") + scr = resp.get("result", {}) + check("screen returns a text grid", + isinstance(scr.get("text"), str) and target["name"] in scr["text"], + f"len={len(scr.get('text','')) if scr else 0}") + + # raw text primitive into the goto input, then escape out + from textual.widgets import Input + await c.call("keys", keys=["g"]) + await c.call("text", text="sub_", settle=False) + val = app.query_one("#goto", Input).value + check("text primitive types into the focused input", val == "sub_", f"val={val!r}") + await c.call("keys", keys=["escape"]) + + w.close() + print(f"\n{PASS} passed, {FAIL} failed") + return 1 if FAIL else 0 + + +def main(argv): + db = None + it = iter(argv) + for a in it: + if a == "--db": + db = next(it) + return asyncio.run(run(db)) + + +if __name__ == "__main__": + raise SystemExit(main(sys.argv[1:])) |