aboutsummaryrefslogtreecommitdiffstats
path: root/idatui/worker_client.py
diff options
context:
space:
mode:
Diffstat (limited to 'idatui/worker_client.py')
-rw-r--r--idatui/worker_client.py234
1 files changed, 0 insertions, 234 deletions
diff --git a/idatui/worker_client.py b/idatui/worker_client.py
deleted file mode 100644
index 79a6db9..0000000
--- a/idatui/worker_client.py
+++ /dev/null
@@ -1,234 +0,0 @@
-"""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()
generated by cgit v1.3.1-sl0p (git 2.54.0) at 2026-08-08 22:49:09 +0000 — sl0p fork