path: root/idatui/worker_client.py

"""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()
        #: Set by close(). A dropped socket is respawned on the next call (the
        #: worker segfaulted and we want it back); a CLOSED one must not be.
        #: Teardown and binary-switch both close while @work threads are still
        #: in flight, so without this, quitting during a decompile spawned a
        #: fresh idalib worker that re-opened the database nobody was looking
        #: at any more -- a stray process holding the .i64 we just released.
        self._closed = False

    # -- 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:
            self._closed = False    # an explicit reconnect revives this client
            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()
            # Poll fast at first, then back off. A flat 0.2s cost every caller a
            # fifth of a second even when the worker was ready in milliseconds
            # (a small binary, or a seeded .i64), which is most of the time in
            # the tests and noticeable on a re-open.
            #
            # Backing off all the way to 0.2s was too eager: a seeded database
            # is ready at ~250ms, by which point the delay has grown to 134ms, so
            # every open waited ~350ms whatever the binary -- the same number for
            # a 47KB `echo` and a 1.2MB `bash`, which is what gives a polling
            # artefact away. Cap the backoff at 25ms instead: the overshoot on a
            # fast open is bounded by that, and 40 probes a second is nothing
            # next to an auto-analysis that runs for minutes.
            #
            # Do NOT be tempted to hold the 5ms rate instead. This poll runs on a
            # background thread while the UI thread is drawing, and 200 wakeups a
            # second through a cold analysis cost enough GIL time to delay the
            # app's own startup -- it left the loading overlay up long enough for
            # project mode's first keypress to land on it.
            delay = 0.005
            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(delay)
                    delay = min(delay * 1.6, 0.025)
            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
            self._closed = True
            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._closed:
            raise IDAConnectionError(
                f"{tool}: this worker was closed (call connect() to revive it)")
        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()