diff options
| -rw-r--r-- | experiments/worker_smoke.py | 58 | ||||
| -rw-r--r-- | idatui/worker.py | 160 | ||||
| -rw-r--r-- | idatui/worker_client.py | 165 |
3 files changed, 383 insertions, 0 deletions
diff --git a/experiments/worker_smoke.py b/experiments/worker_smoke.py new file mode 100644 index 0000000..b215a57 --- /dev/null +++ b/experiments/worker_smoke.py @@ -0,0 +1,58 @@ +"""Runnable read-path smoke: drives the REAL domain.Program through WorkerClient +(our idalib worker over a unix socket). Run when idalib can spawn: + ~/ida-venv/bin/python experiments/worker_smoke.py +""" +import os, sys, shutil, time +REPO=os.path.expanduser("~/dev/ida-tui-maybe"); sys.path.insert(0, REPO); os.chdir(REPO) +# fresh copy so the worker's idalib doesn't fight any running server +src=f"{REPO}/targets/echo"; tmp="/tmp/echo_worker" +shutil.copy(src, tmp) +for e in ".i64 .id0 .id1 .id2 .nam .til".split(): + try: os.remove(tmp+e) + except OSError: pass + +from idatui.worker_client import WorkerClient +from idatui.domain import Program + +print("spawning worker + opening echo…", flush=True) +t=time.time() +cl=WorkerClient(tmp) +cl.connect(progress=lambda m: None) +print(f" worker ready in {time.time()-t:.2f}s session={cl.resolve_db()}", flush=True) +prog=Program(cl) + +# --- drive the REAL domain layer through the worker (read path) --- +main=prog.resolve("main") +print("resolve('main') =", hex(main), flush=True) + +idx=prog.functions(); idx.load_all() +print("functions() ->", len(idx), "funcs", flush=True) + +fn=prog.function_of(main) +print("function_of(main) ->", fn.name, hex(fn.addr), "size", fn.size, flush=True) + +b=prog.read_bytes(main, 16) +print("read_bytes(main,16) ->", b.hex(), flush=True) + +lm=prog.listing(main) +for _ in range(3): lm.load_next_page() +rows=[lm.get(i) for i in range(min(6,len(lm)))] +print("listing() first rows:", flush=True) +for h in rows: + if h: print(" ", hex(h.ea), h.kind, repr(h.text[:44]), flush=True) + +d=prog.decompile(main) +print("decompile(main) -> failed?", d.failed, "lines:", len((d.code or '').splitlines()), flush=True) + +regs=prog.file_regions() +print("file_regions ->", len(regs), "segments", flush=True) + +# xrefs to a called function +callee=next((f.addr for f in idx.all_loaded() if f.name.startswith("sub_")), None) +if callee: + xr=prog.xrefs_to(callee) + print("xrefs_to(", hex(callee), ") ->", len(xr), "refs", flush=True) + +ok = (fn.name=="main" and len(idx)>100 and b and not d.failed and len(regs)>0) +print("VERDICT:", "OK — domain.Program runs unchanged on the worker" if ok else "FAIL", flush=True) +cl.close() diff --git a/idatui/worker.py b/idatui/worker.py new file mode 100644 index 0000000..d5fab20 --- /dev/null +++ b/idatui/worker.py @@ -0,0 +1,160 @@ +"""idatui's own idalib worker — the replacement for the ida-pro-mcp supervisor. + +Opens ONE database in-process (on the main thread, as idalib requires) and +serves ida-pro-mcp's *tool functions* over a unix socket with length-prefixed +pickle. Same tool implementations as the MCP path (we call +``MCP_SERVER.tools.methods[name](**args)`` directly), so return shapes are +byte-identical — but with ~50us/call instead of the HTTP path's ~5ms, and no +supervisor / HTTP / JSON / 50KB-truncation machinery. + + python -m idatui.worker <sock_path> <binary_path> + +The socket only appears once the database is open + analyzed, so a client can +poll ``connect()`` to know when the worker is ready. Requests are served +serially on the main thread (idalib is single-threaded; every tool runs inline +through its own execute_sync, which is a no-op on the main thread). + +Protocol (both directions length-prefixed: 4-byte big-endian len + pickle): + request = (tool_name: str, kwargs: dict) + response = (ok: bool, result_or_error) + tool_name == "__shutdown__" ends the worker. +""" +from __future__ import annotations + +import os +import pickle +import socket +import struct +import sys +import uuid + + +# --------------------------------------------------------------------------- # +# framing +# --------------------------------------------------------------------------- # +def _recvn(sock: socket.socket, n: int) -> bytes | None: + buf = bytearray() + while len(buf) < n: + chunk = sock.recv(n - len(buf)) + if not chunk: + return None + buf += chunk + return bytes(buf) + + +def send(sock: socket.socket, obj) -> None: + data = pickle.dumps(obj, protocol=pickle.HIGHEST_PROTOCOL) + sock.sendall(struct.pack(">I", len(data)) + data) + + +def recv(sock: socket.socket): + hdr = _recvn(sock, 4) + if hdr is None: + return None + (n,) = struct.unpack(">I", hdr) + body = _recvn(sock, n) + return None if body is None else pickle.loads(body) + + +# --------------------------------------------------------------------------- # +# worker +# --------------------------------------------------------------------------- # +def _open_and_register(binpath: str): + """Open the DB (main thread) then import ida-pro-mcp so every @tool registers + against this live database. Returns (tools_dict, module_name, save_fn).""" + import idapro + idapro.enable_console_messages(False) + rc = idapro.open_database(binpath, run_auto_analysis=True) + if rc: + raise RuntimeError(f"open_database({binpath!r}) failed rc={rc}") + + # importing the package registers all api_*/patched tools against MCP_SERVER + from ida_pro_mcp.ida_mcp import MCP_SERVER # noqa: WPS433 + + import ida_nalt + module = os.path.basename(ida_nalt.get_root_filename() or binpath) + + def save(): + import idc + try: + idc.save_database(idc.get_idb_path(), 0) + except Exception: # noqa: BLE001 + import ida_loader, ida_pro # noqa: WPS433 + ida_loader.save_database(idc.get_idb_path(), 0) + + return MCP_SERVER.tools.methods, module, save + + +def serve(sockpath: str, binpath: str) -> None: + tools, module, save = _open_and_register(binpath) + sid = uuid.uuid4().hex[:8] + + def dispatch(name: str, args: dict): + args = dict(args) + args.pop("database", None) # single-DB worker: no session routing + # session-management shims (were the supervisor's job): + if name in ("idb_open",): + return {"success": True, + "session": {"session_id": sid, "module": module, + "input_path": binpath}} + if name in ("idb_save", "save"): + save() + return {"success": True} + if name in ("server_health", "ping", "health", "state"): + return {"module": module, "ok": True, "session_id": sid} + if name in ("idb_list",): + return {"sessions": [{"session_id": sid, "module": module, + "input_path": binpath}]} + fn = tools.get(name) + if fn is None: + raise KeyError(f"unknown tool: {name!r}") + return fn(**args) + + try: + os.unlink(sockpath) + except OSError: + pass + srv = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + srv.bind(sockpath) + srv.listen(8) + try: + while True: + conn, _ = srv.accept() + try: + while True: + req = recv(conn) + if req is None: + break + name, args = req + if name == "__shutdown__": + return + try: + send(conn, (True, dispatch(name, args))) + except Exception as e: # noqa: BLE001 -- report, keep serving + send(conn, (False, f"{type(e).__name__}: {e}")) + except (ConnectionError, OSError): + pass + finally: + conn.close() + finally: + try: + import idapro + idapro.close_database(save=False) + except Exception: # noqa: BLE001 + pass + try: + os.unlink(sockpath) + except OSError: + pass + + +def main(argv=None) -> None: + argv = argv if argv is not None else sys.argv[1:] + if len(argv) < 2: + sys.stderr.write("usage: python -m idatui.worker <sock> <binary>\n") + raise SystemExit(2) + serve(argv[0], argv[1]) + + +if __name__ == "__main__": + main() diff --git a/idatui/worker_client.py b/idatui/worker_client.py new file mode 100644 index 0000000..49da8fa --- /dev/null +++ b/idatui/worker_client.py @@ -0,0 +1,165 @@ +"""WorkerClient — a drop-in replacement for ``IDAClient`` backed by our own +idalib worker (``idatui.worker``) over a unix socket instead of ida-pro-mcp's +HTTP/JSON transport. + +It exposes exactly the surface the app/domain use on the client +(``call``/``call_envelope``/``connect``/``set_db``/``resolve_db``/ +``list_sessions``/``health``/``keepalive``/``close``) and returns byte-identical +payloads (the worker calls the same tool functions), so ``domain.py`` and the +app are unchanged — you just construct a WorkerClient instead of an IDAClient. + +Concurrency: the app fires calls from several worker threads over one client; +the worker is single-threaded, so calls are serialized under a lock (the worker +processes one tool at a time anyway — and at ~50us/call that's free). +""" +from __future__ import annotations + +import os +import socket +import subprocess +import sys +import threading +import time +import uuid +from typing import Any + +from .client import IDAToolError, IDAConnectionError, Session +from .worker import recv as _recv +from .worker import send as _send + + +class _NoopKeepAlive: + """The worker is ours and never idles out, so keepalive is a no-op.""" + + def __init__(self) -> None: + self.beats = self.failures = 0 + + def start(self): + return self + + def stop(self) -> None: + pass + + +class WorkerClient: + def __init__(self, binary_path: str, *, ttl: int = 0, + python: str | None = None) -> None: + self._bin = os.path.abspath(os.path.expanduser(binary_path)) + self._python = python or sys.executable + self._sock_path = f"/tmp/idatui-worker-{os.getpid()}-{uuid.uuid4().hex[:8]}.sock" + self._proc: subprocess.Popen | None = None + self._sock: socket.socket | None = None + self._sid = uuid.uuid4().hex[:8] + self._lock = threading.Lock() # serialize socket use + self._spawn_lock = threading.Lock() + + # -- lifecycle --------------------------------------------------------- # + def connect(self, timeout: float = 1800.0, progress=None) -> "WorkerClient": + """Spawn the worker (opens + analyzes the DB) and connect once ready.""" + with self._spawn_lock: + if self._sock is not None: + return self + if self._proc is None or self._proc.poll() is not None: + self._proc = subprocess.Popen( + [self._python, "-m", "idatui.worker", + self._sock_path, self._bin], + cwd=os.path.dirname(os.path.dirname(os.path.abspath(__file__))), + ) + deadline = time.time() + timeout + t0 = time.time() + while time.time() < deadline: + try: + s = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + s.connect(self._sock_path) + self._sock = s + return self + except OSError: + if self._proc.poll() is not None: + raise IDAConnectionError( + f"worker exited during startup (code " + f"{self._proc.returncode})") + if progress: + progress(f"auto-analyzing {os.path.basename(self._bin)}… " + f"({int(time.time() - t0)}s)") + time.sleep(0.2) + raise IDAConnectionError("worker did not become ready in time") + + def close(self) -> None: + with self._lock: + s = self._sock + self._sock = None + if s is not None: + try: + _send(s, ("__shutdown__", {})) + except Exception: # noqa: BLE001 + pass + try: + s.close() + except Exception: # noqa: BLE001 + pass + if self._proc is not None: + try: + self._proc.terminate() + self._proc.wait(timeout=5) + except Exception: # noqa: BLE001 + try: + self._proc.kill() + except Exception: # noqa: BLE001 + pass + + # -- the call surface -------------------------------------------------- # + def call(self, tool: str, *, timeout: float | None = None, **args) -> Any: + if self._sock is None: + self.connect() + with self._lock: + s = self._sock + if s is None: + raise IDAConnectionError("worker connection is closed") + try: + _send(s, (tool, args)) + reply = _recv(s) + except (OSError, ConnectionError) as e: + self._sock = None + raise IDAConnectionError(f"worker transport failed: {e}") from e + if reply is None: + self._sock = None + raise IDAConnectionError("worker closed the connection") + ok, payload = reply + if not ok: + raise IDAToolError(tool, str(payload)) + return payload + + def call_envelope(self, tool: str, *, timeout: float | None = None, + **args) -> dict: + # domain.decompile() reads result.structuredContent — mirror that shape. + return {"result": {"structuredContent": self.call(tool, timeout=timeout, + **args)}} + + # -- session shims (single-DB worker) --------------------------------- # + def set_db(self, db: str | None) -> None: + if db: + self._sid = db + + def resolve_db(self) -> str: + return self._sid + + def list_sessions(self) -> list[Session]: + return [Session(session_id=self._sid, + filename=os.path.basename(self._bin), + input_path=self._bin, is_active=True)] + + def health(self) -> dict: + try: + return self.call("server_health") + except IDAToolError: + return {"module": os.path.basename(self._bin), "ok": True} + + def keepalive(self, interval: float = 120.0) -> _NoopKeepAlive: + return _NoopKeepAlive() + + # context manager parity with IDAClient + def __enter__(self) -> "WorkerClient": + return self.connect() + + def __exit__(self, *exc) -> None: + self.close() |