diff options
Diffstat (limited to 'idatui/worker_client.py')
| -rw-r--r-- | idatui/worker_client.py | 234 |
1 files changed, 234 insertions, 0 deletions
diff --git a/idatui/worker_client.py b/idatui/worker_client.py new file mode 100644 index 0000000..79a6db9 --- /dev/null +++ b/idatui/worker_client.py @@ -0,0 +1,234 @@ +"""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 .errors import IDAToolError, IDAConnectionError, Session +from .worker import recv as _recv +from .worker import send as _send + +_WORKER_PY = os.path.join(os.path.dirname(os.path.abspath(__file__)), "worker.py") +_worker_python_cache: str | None = None + + +def _find_worker_python() -> str: + """A python that can import ``ida_pro_mcp`` (and thus idalib) — NOT necessarily + the TUI's python. On a typical box the TUI runs under a venv that has textual + + idalib but not ida_pro_mcp, while the system python has idalib + + ida_pro_mcp. Override with IDATUI_WORKER_PYTHON.""" + global _worker_python_cache + if _worker_python_cache: + return _worker_python_cache + override = os.environ.get("IDATUI_WORKER_PYTHON") + candidates = [override] if override else [] + candidates += ["/usr/bin/python", "/usr/bin/python3", sys.executable] + for py in candidates: + if not py or not os.path.exists(py): + continue + try: + r = subprocess.run([py, "-c", "import ida_pro_mcp"], + capture_output=True, timeout=30) + if r.returncode == 0: + _worker_python_cache = py + return py + except Exception: # noqa: BLE001 + continue + return sys.executable # last resort; the worker will report the real error + + +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, load_args: str = "") -> None: + self._bin = os.path.abspath(os.path.expanduser(binary_path)) + self._load_args = load_args or "" # IDA switches for a headerless blob + self._python = python or _find_worker_python() + tag = f"{os.getpid()}-{uuid.uuid4().hex[:8]}" + self._sock_path = f"/tmp/idatui-worker-{tag}.sock" + self._log_path = f"/tmp/idatui-worker-{tag}.log" + 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: + # run worker.py as a SCRIPT (not -m idatui.worker) so we don't + # import the textual-dependent idatui package __init__ under the + # IDA python, which usually has no textual. + argv = [self._python, _WORKER_PY, self._sock_path, self._bin] + if self._load_args: + argv.append(self._load_args) + self._proc = subprocess.Popen( + argv, + stdout=open(self._log_path, "wb"), + stderr=subprocess.STDOUT, stdin=subprocess.DEVNULL, + ) + 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 (code {self._proc.returncode}): " + f"{self._log_tail()} [full log: {self._log_path}]") + 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") + + @property + def pid(self) -> int | None: + """The worker process id (for memory accounting), or None if not spawned.""" + return self._proc.pid if self._proc is not None else None + + def close(self, grace: float = 20.0) -> None: + """Shut the worker down cleanly. + + After ``__shutdown__`` the worker still has to ``close_database()``, which + re-packs the ``.i64`` and removes the unpacked ``.id0/.id1/...`` scratch. + Signalling it before that finishes is what leaves databases wedged, so + wait out the grace period first and only escalate if it really is stuck. + """ + 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.wait(timeout=grace) # let it close the DB properly + except Exception: # noqa: BLE001 -- TimeoutExpired: it's stuck + 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() + + def _log_tail(self, n: int = 400) -> str: + """Last meaningful line(s) of the worker log (skip IDA's licence banner), + so a startup crash surfaces the real cause instead of just 'code 1'.""" + try: + with open(self._log_path, encoding="utf-8", errors="replace") as f: + lines = [ln.strip() for ln in f if ln.strip()] + except OSError: + return "(no worker log)" + # the worker prints a clean 'WORKER-FATAL: ...' line on a startup crash + for ln in reversed(lines): + if ln.startswith("WORKER-FATAL:"): + return ln[len("WORKER-FATAL:"):].strip()[-n:] + skip = ("thank you", "licensed to", "[mcp]", "ida ", "hex-rays") + meaningful = [ln for ln in lines + if not any(s in ln.lower() for s in skip)] + return " | ".join((meaningful or lines)[-3:])[-n:] + + # context manager parity with IDAClient + def __enter__(self) -> "WorkerClient": + return self.connect() + + def __exit__(self, *exc) -> None: + self.close() |