path: root/idatui/pane.py

"""Spawn/stop/list idatui TUI panes in tmux, for an agent to drive over RPC.

The agent (running inside a tmux pane) can open a fresh pane with the TUI
running against a binary, wait until it's ready, drive it over the RPC socket,
then close it — all without a human touching the keyboard.

    # open a binary in a new pane, block until analysed + drivable, print JSON
    python -m idatui.pane spawn --open /abs/path/to/bin
    # -> {"sock": "/run/user/1000/idatui-3f2a.sock", "pane": "%7", "ready": true, ...}

    # or attach to an existing session id
    python -m idatui.pane spawn --db 80d83396

    # drive it (see docs/RPC.md / the idatui-rpc skill)
    python -m idatui.rpcclient --sock <sock> pseudocode target=main

    # inventory + teardown
    python -m idatui.pane list
    python -m idatui.pane stop --sock <sock>        # graceful quit + kill pane

Requires: running inside tmux, and the ida-pro-mcp supervisor already up
(./spawn.sh). Uses ~/ida-venv/bin/python for the TUI (needs textual) unless
--python / IDATUI_PYTHON says otherwise.
"""
from __future__ import annotations

import argparse
import json
import os
import secrets
import signal
import socket
import subprocess
import sys
import time
from typing import Any
from urllib.parse import urlparse

from .client import DEFAULT_URL
from .rpcclient import RpcClient, RpcError

REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
DEFAULT_PY = os.environ.get(
    "IDATUI_PYTHON", os.path.expanduser("~/ida-venv/bin/python"))


def _sockdir() -> str:
    return os.environ.get("XDG_RUNTIME_DIR") or "/tmp"


def _registry_path() -> str:
    return os.path.join(_sockdir(), "idatui-panes.json")


def _load_registry() -> list[dict[str, Any]]:
    try:
        with open(_registry_path()) as f:
            return json.load(f)
    except (OSError, ValueError):
        return []


def _save_registry(rows: list[dict[str, Any]]) -> None:
    tmp = _registry_path() + ".tmp"
    with open(tmp, "w") as f:
        json.dump(rows, f)
    os.replace(tmp, _registry_path())


def _pane_alive(pane: str) -> bool:
    out = subprocess.run(["tmux", "list-panes", "-a", "-F", "#{pane_id}"],
                         capture_output=True, text=True)
    return pane in out.stdout.split()


def _tmux(*args: str) -> str:
    return subprocess.run(["tmux", *args], capture_output=True, text=True,
                          check=True).stdout.strip()


# --------------------------------------------------------------------------- #
# idalib worker reaping
#
# ``pane stop`` kills the TUI pane, but the ida-pro-mcp supervisor does not
# reliably reap the ``ida_pro_mcp.idalib_server`` worker it forked for that
# session. Leaked workers accumulate against IDA_MCP_MAX_WORKERS until the next
# ``spawn`` blocks forever waiting for a free slot (the TUI comes up but never
# becomes ready). A worker is only *safe* to reap when no idatui pane is live
# (then every worker is orphaned) — that mirrors the hand workaround
# `pkill -f ida_pro_mcp.idalib_server` and avoids killing an in-use analyser.
# --------------------------------------------------------------------------- #
_WORKER_PATTERN = r"ida_pro_mcp\.idalib_server"


def _worker_pids() -> list[int]:
    """PIDs of the supervisor's per-binary idalib worker processes.

    Matches the worker module invocation only (not the ``idalib-mcp``
    supervisor, whose command line does not contain the module path), and
    never our own PID.
    """
    try:
        out = subprocess.run(["pgrep", "-f", _WORKER_PATTERN],
                             capture_output=True, text=True)
    except OSError:
        return []
    me = os.getpid()
    pids: list[int] = []
    for tok in out.stdout.split():
        try:
            pid = int(tok)
        except ValueError:
            continue
        if pid != me:
            pids.append(pid)
    return pids


def _count_live_panes() -> int:
    return sum(1 for r in _load_registry() if _pane_alive(r.get("pane", "")))


def _reap_orphan_workers(force: bool = False) -> int:
    """Kill leaked idalib workers when it is safe (no live pane) or ``force``.

    Returns the number of workers signalled. Best-effort; never raises.
    """
    if not force and _count_live_panes() > 0:
        return 0
    reaped = 0
    for pid in _worker_pids():
        try:
            os.kill(pid, signal.SIGKILL)
            reaped += 1
        except OSError:
            pass
    return reaped


# --------------------------------------------------------------------------- #
# supervisor (ida-pro-mcp server) — auto-start if down
# --------------------------------------------------------------------------- #
def _server_addr(url: str) -> tuple[str, int]:
    u = urlparse(url)
    return (u.hostname or "127.0.0.1", u.port or 8745)


def _server_up(host: str, port: int, timeout: float = 0.75) -> bool:
    """Is something listening on host:port? (Cheap TCP probe; the readiness poll
    that follows catches a half-up server.)"""
    try:
        with socket.create_connection((host, port), timeout=timeout):
            return True
    except OSError:
        return False


def _ensure_server(host: str, port: int, timeout: float,
                   detached: bool = True) -> dict[str, Any]:
    """Make sure the supervisor is up; start ./spawn.sh in a tmux pane if not.

    Only auto-starts a *local* server (can't launch a remote one). spawn.sh binds
    the port, so starting a duplicate is impossible — the probe guards that.
    ``IDATUI_SERVER_CMD`` overrides the launch command (used by the tests).
    """
    if _server_up(host, port):
        return {"server_started": False, "server_up": True}
    if host not in ("127.0.0.1", "localhost", "::1"):
        return {"server_started": False, "server_up": False,
                "error": f"server at {host}:{port} is down and not local; "
                         "cannot auto-start"}
    cmd_str = os.environ.get("IDATUI_SERVER_CMD", "./spawn.sh")
    cmd = f"cd {_q(REPO)} && exec {cmd_str}"
    split = ["split-window", "-v", "-P", "-F", "#{pane_id}"]
    if detached:
        split += ["-d"]
    anchor = os.environ.get("TMUX_PANE")
    if anchor:
        split += ["-t", anchor]
    split.append(cmd)
    pane = _tmux(*split)
    deadline = time.time() + timeout
    while time.time() < deadline:
        if not _pane_alive(pane):
            return {"server_started": True, "server_up": False, "server_pane": pane,
                    "error": "supervisor pane exited during startup (check it)"}
        if _server_up(host, port):
            return {"server_started": True, "server_up": True, "server_pane": pane}
        time.sleep(0.5)
    return {"server_started": True, "server_up": False, "server_pane": pane,
            "error": "supervisor did not come up in time"}


# --------------------------------------------------------------------------- #
# spawn
# --------------------------------------------------------------------------- #
def spawn(args) -> int:
    if not os.environ.get("TMUX"):
        print("error: not inside tmux (spawn creates a tmux pane)", file=sys.stderr)
        return 2
    if not args.open and not args.db:
        print("error: pass --open <binary> or --db <session>", file=sys.stderr)
        return 2

    sock = args.sock or os.path.join(_sockdir(), f"idatui-{secrets.token_hex(3)}.sock")
    target = os.path.abspath(os.path.expanduser(args.open)) if args.open else args.db
    if args.open and not os.path.exists(target):
        print(f"error: no such binary: {target}", file=sys.stderr)
        return 2

    # Reap workers leaked by previously-stopped/crashed panes so we don't spawn
    # into a full IDA_MCP_MAX_WORKERS (which makes the new TUI hang forever,
    # never reaching ready). No-op while any pane is live.
    reaped = _reap_orphan_workers()
    if reaped:
        print(f"reaped {reaped} orphaned idalib worker(s) before spawn",
              file=sys.stderr)

    # make sure the ida-pro-mcp supervisor is up (auto-start it if not)
    srv: dict[str, Any] = {"server_started": False, "server_up": True}
    if not args.no_ensure_server:
        host, port = _server_addr(args.url or DEFAULT_URL)
        srv = _ensure_server(host, port, args.server_timeout)
        if srv.get("server_started"):
            print(f"supervisor was down — started it ({srv.get('server_pane')})",
                  file=sys.stderr)
        if not srv.get("server_up"):
            print(json.dumps({"ready": False, **srv}), file=sys.stderr)
            return 3

    # the command the pane runs: become the TUI so kill-pane kills it cleanly
    inner = [args.python, "-m", "idatui.tui", "--rpc", sock]
    if args.open:
        inner += ["--open", target]
    else:
        inner += ["--db", target]
    if args.url:
        inner += ["--url", args.url]
    cmd = f"cd {REPO!r} && exec " + " ".join(_q(a) for a in inner)

    split = ["split-window", "-v" if args.vertical else "-h",
             "-P", "-F", "#{pane_id}"]
    if args.size:
        split += ["-l", str(args.size)]
    if args.detached:
        split += ["-d"]
    anchor = os.environ.get("TMUX_PANE")
    if anchor:
        split += ["-t", anchor]
    split.append(cmd)
    pane = _tmux(*split)

    row = {"sock": sock, "pane": pane, "target": target,
           "kind": "open" if args.open else "db", "started": time.time(),
           "server_started": srv.get("server_started", False)}
    if srv.get("server_pane"):
        row["server_pane"] = srv["server_pane"]
    reg = [r for r in _load_registry() if r.get("sock") != sock]
    reg.append(row)
    _save_registry(reg)

    ready = _wait_ready(sock, args.timeout, pane)
    row.update(ready)
    print(json.dumps(row))
    return 0 if ready.get("ready") else 1


def _q(s: str) -> str:
    import shlex
    return shlex.quote(s)


def _wait_ready(sock: str, timeout: float, pane: str,
                stuck_after: float = 45.0) -> dict[str, Any]:
    """Poll the socket + ping until the TUI reports ready (or timeout).

    Emits a one-time hint to stderr if it's still not ready after ``stuck_after``
    seconds, so a wedged idalib worker / full worker pool surfaces a diagnostic
    instead of an unexplained silent hang.
    """
    start = time.time()
    deadline = start + timeout
    warned = False
    last: dict[str, Any] = {"ready": False}
    while time.time() < deadline:
        if not _pane_alive(pane):
            return {"ready": False, "error": "pane exited during startup"}
        if os.path.exists(sock):
            try:
                with RpcClient(sock) as c:
                    last = c.call("ping")
                if last.get("ready"):
                    return last
            except (OSError, RpcError, ConnectionError):
                pass
        if not warned and (time.time() - start) > stuck_after:
            warned = True
            why = ("RPC socket not created yet" if not os.path.exists(sock)
                   else "TUI up but analysis not ready")
            print(f"still waiting ({int(time.time() - start)}s): {why}. If this "
                  f"hangs, the idalib worker may be stuck or IDA_MCP_MAX_WORKERS "
                  f"is full — try `python -m idatui.pane reap`.", file=sys.stderr)
        time.sleep(0.4)
    last = dict(last)
    last["ready"] = False
    last.setdefault("error", "timed out waiting for the TUI to become ready")
    return last


# --------------------------------------------------------------------------- #
# stop / list
# --------------------------------------------------------------------------- #
def stop(args) -> int:
    reg = _load_registry()
    rows = [r for r in reg
            if (args.sock and r.get("sock") == args.sock)
            or (args.pane and r.get("pane") == args.pane)]
    if not rows and args.sock:  # allow stopping an untracked socket
        rows = [{"sock": args.sock, "pane": args.pane}]
    if not rows:
        print("error: no matching pane (need --sock or --pane)", file=sys.stderr)
        return 2
    for r in rows:
        sock, pane = r.get("sock"), r.get("pane")
        if sock and os.path.exists(sock):
            try:  # ask it to quit gracefully first
                with RpcClient(sock) as c:
                    c.call("quit")
                time.sleep(0.4)
            except (OSError, RpcError, ConnectionError):
                pass
        if pane and _pane_alive(pane):
            subprocess.run(["tmux", "kill-pane", "-t", pane],
                           capture_output=True)
        if sock:
            try:
                os.unlink(sock)
            except OSError:
                pass
    _save_registry([r for r in reg if r not in rows])
    # Reap the workers those panes leaked (safe: only fires once no pane is live).
    reaped = _reap_orphan_workers()
    out = {"stopped": [r.get("sock") or r.get("pane") for r in rows]}
    if reaped:
        out["reaped_workers"] = reaped
    print(json.dumps(out))
    return 0


def list_panes(args) -> int:
    reg = _load_registry()
    alive = []
    for r in reg:
        r = dict(r)
        r["pane_alive"] = _pane_alive(r.get("pane", ""))
        r["sock_up"] = bool(r.get("sock") and os.path.exists(r["sock"]))
        if args.prune and not r["pane_alive"]:
            if r.get("sock") and os.path.exists(r["sock"]):
                try:
                    os.unlink(r["sock"])
                except OSError:
                    pass
            continue
        alive.append(r)
    if args.prune:
        _save_registry(alive)
        reaped = _reap_orphan_workers()
        if reaped:
            print(f"reaped {reaped} orphaned idalib worker(s)", file=sys.stderr)
    print(json.dumps(alive, indent=2))
    return 0


def reap(args) -> int:
    """Kill leaked idalib workers (safe when no pane is live; --force overrides)."""
    live = _count_live_panes()
    n = _reap_orphan_workers(force=args.force)
    print(json.dumps({"reaped_workers": n, "live_panes": live, "forced": args.force}))
    if n == 0 and not args.force and live > 0:
        print(f"note: {live} live pane(s) — not reaping in-use workers; pass "
              f"--force to reap anyway", file=sys.stderr)
    return 0


def main(argv: list[str]) -> int:
    p = argparse.ArgumentParser(prog="idatui.pane",
                                description="spawn/manage idatui TUI panes in tmux")
    sub = p.add_subparsers(dest="cmd", required=True)

    sp = sub.add_parser("spawn", help="open a TUI pane and wait until ready")
    sp.add_argument("--open", metavar="PATH", help="binary to open (dir must be writable)")
    sp.add_argument("--db", metavar="SESSION", help="attach to an existing session id")
    sp.add_argument("--sock", help="RPC socket path (default: auto in $XDG_RUNTIME_DIR)")
    sp.add_argument("--python", default=DEFAULT_PY, help=f"python for the TUI ({DEFAULT_PY})")
    sp.add_argument("--url", help="MCP server URL (default: idatui's default)")
    sp.add_argument("--no-ensure-server", action="store_true",
                    help="don't auto-start ./spawn.sh if the supervisor is down")
    sp.add_argument("--server-timeout", type=float, default=90.0,
                    help="seconds to wait for an auto-started supervisor")
    sp.add_argument("--vertical", action="store_true", help="split vertically (stacked)")
    sp.add_argument("--size", help="new pane size (tmux -l value, e.g. 60%% or 120)")
    sp.add_argument("--detached", action="store_true", help="don't focus the new pane")
    sp.add_argument("--timeout", type=float, default=300.0,
                    help="seconds to wait for readiness (fresh --open analysis is slow)")
    sp.set_defaults(fn=spawn)

    st = sub.add_parser("stop", help="graceful quit + kill the pane")
    st.add_argument("--sock")
    st.add_argument("--pane")
    st.set_defaults(fn=stop)

    ls = sub.add_parser("list", help="list tracked panes")
    ls.add_argument("--prune", action="store_true", help="drop dead panes (and their sockets)")
    ls.set_defaults(fn=list_panes)

    rp = sub.add_parser("reap", help="kill leaked idalib workers (frees worker slots)")
    rp.add_argument("--force", action="store_true",
                    help="reap even while panes are live (may kill an in-use analyser)")
    rp.set_defaults(fn=reap)

    args = p.parse_args(argv)
    return args.fn(args)


if __name__ == "__main__":
    raise SystemExit(main(sys.argv[1:]))