path: root/tests/test_trace_rpc.py

#!/usr/bin/env python3
"""Trace integration over the RPC socket — the path an agent actually takes.

The UI tests (test_trace_ui.py) drive the trace via in-process keystrokes and
the Textual pilot. This one spawns a real tmux pane with ``--trace`` and drives
every trace verb through the RPC socket, validating the JSON responses an agent
would see. It exercises the serialization, the settle/timeout machinery, and
the response shape that a driver depends on.

Requires: tmux, IDA (idalib), and a trace. Records a fresh trace with the QEMU
tracer if built; falls back to /tmp/echotrace.0.log if present; skips with a
message otherwise.

    ~/ida-venv/bin/python tests/test_trace_rpc.py
"""
import json
import os
import shutil
import subprocess
import sys
import tempfile

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from idatui.rpcclient import RpcClient, RpcError  # noqa: E402
from idatui.trace import Trace  # noqa: E402

REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
TRACER = os.path.expanduser(
    "~/.pi/agent/skills/tenet-trace/scripts/tenet-trace")
FALLBACK_TRACE = "/tmp/echotrace.0.log"
BINARY = os.path.join(REPO, "targets", "echo")

PASS = FAIL = 0


def check(name, ok, detail=""):
    global PASS, FAIL
    if ok:
        PASS += 1
        print(f"  ok   [{name}]")
    else:
        FAIL += 1
        print(f"  FAIL [{name}]   {detail}")


def make_trace(tmp, binary):
    """Record a short trace of the echo binary."""
    out = os.path.join(tmp, "t")
    try:
        subprocess.run([TRACER, "-o", out, binary, "hello"],
                       capture_output=True, timeout=180, check=False)
    except (OSError, subprocess.TimeoutExpired):
        return None
    log = out + ".0.log"
    return log if os.path.exists(log) else None


def spawn_pane(target, trace_log, timeout=300):
    """Spawn an idatui pane with --trace and wait for readiness."""
    cmd = [sys.executable, "-m", "idatui.pane", "spawn",
           "--open", target, "--trace", trace_log,
           "--detached", "--size", "60%", "--timeout", str(timeout)]
    r = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout + 30,
                       cwd=REPO)
    if r.returncode != 0:
        print(f"  spawn failed: {r.stderr.strip()}", file=sys.stderr)
        return None
    return json.loads(r.stdout)


def stop_pane(sock, timeout=60):
    cmd = [sys.executable, "-m", "idatui.pane", "stop",
           "--sock", sock, "--timeout", str(timeout)]
    subprocess.run(cmd, capture_output=True, text=True, timeout=timeout + 10,
                   cwd=REPO)


def main() -> int:
    global PASS, FAIL
    if not os.environ.get("TMUX"):
        print("  skip: not inside tmux (test spawns a pane)")
        return 0

    with tempfile.TemporaryDirectory() as tmp:
        # Scratch copy so the suite never writes into the tracked tree.
        target = os.path.join(tmp, "echo")
        shutil.copy2(BINARY, target)

        # Get a trace.
        trace_log = make_trace(tmp, target)
        if trace_log is None and os.path.exists(FALLBACK_TRACE):
            trace_log = FALLBACK_TRACE
        if trace_log is None:
            print(f"  skip: no trace available (tracer at {TRACER}, "
                  f"fallback {FALLBACK_TRACE})")
            return 0

        # Load our own model for ground-truth comparisons.
        model = Trace.load(trace_log)
        check("model loaded for ground truth", model.length > 10,
              f"length={model.length}")

        # Spawn the pane.
        info = spawn_pane(target, trace_log)
        if info is None or not info.get("ready"):
            print(f"  FAIL: pane did not become ready: {info}")
            return 1
        sock = info["sock"]
        check("pane spawned and ready", True)

        try:
            with RpcClient(sock) as c:
                run_trace_tests(c, model)
        finally:
            stop_pane(sock)

    print(f"\n{PASS} passed, {FAIL} failed")
    return 1 if FAIL else 0


def run_trace_tests(c: RpcClient, model: Trace):
    """All RPC-level trace tests, given a connected client and the model."""

    # ------------------------------------------------------------------ #
    # Basic: ping should show the trace is loaded
    # ------------------------------------------------------------------ #
    p = c.call("ping")
    check("ping reports ready", p.get("ready") is True)

    # ------------------------------------------------------------------ #
    # trace verb: seek to timestamp 0
    # ------------------------------------------------------------------ #
    r = c.call("trace", seek=0)
    check("trace seek=0 returns a snapshot", "active" in r and "trace" in r)
    tr = r["trace"]
    check("response includes trace metadata",
          "idx" in tr and "length" in tr and "pc" in tr and "changed" in tr,
          f"keys={list(tr.keys())}")
    check("idx is 0 after seeking to 0", tr["idx"] == 0)
    check("length matches the model", tr["length"] == model.length,
          f"{tr['length']} vs {model.length}")
    check("pc at 0 is a hex string", isinstance(tr["pc"], str)
          and tr["pc"].startswith("0x"), tr["pc"])

    # ------------------------------------------------------------------ #
    # seek to a mid-trace timestamp
    # ------------------------------------------------------------------ #
    mid = model.length // 2
    r = c.call("trace", seek=mid)
    tr = r["trace"]
    check("seek to midpoint lands correctly", tr["idx"] == mid,
          f"got {tr['idx']}, want {mid}")
    # The model's IP at this point, rebased — the RPC should agree.
    # We can't compare directly because the model isn't rebased yet, but
    # the pc should be a small address (database-space, not ASLR'd).
    pc = int(tr["pc"], 16)
    check("pc is in database space (not ASLR'd)",
          pc < 0x100000, f"pc={tr['pc']}")

    # ------------------------------------------------------------------ #
    # seek by percentage (Tenet shell syntax)
    # ------------------------------------------------------------------ #
    r = c.call("trace", seek="!0")
    check("seek !0 (0%) goes to the start", r["trace"]["idx"] == 0)
    r = c.call("trace", seek="!100")
    check("seek !100 (100%) goes to the end",
          r["trace"]["idx"] == model.length - 1,
          f"got {r['trace']['idx']}, want {model.length - 1}")
    r = c.call("trace", seek="!50")
    check("seek !50 (50%) goes to the midpoint",
          abs(r["trace"]["idx"] - mid) <= 1,
          f"got {r['trace']['idx']}, want ~{mid}")

    # ------------------------------------------------------------------ #
    # step forward/backward
    # ------------------------------------------------------------------ #
    c.call("trace", seek=0)
    r = c.call("trace", step=1)
    check("step=1 advances one timestamp", r["trace"]["idx"] == 1,
          f"got {r['trace']['idx']}")
    r = c.call("trace", step=1)
    check("another step=1 reaches 2", r["trace"]["idx"] == 2,
          f"got {r['trace']['idx']}")
    r = c.call("trace", step=-1)
    check("step=-1 goes backward", r["trace"]["idx"] == 1,
          f"got {r['trace']['idx']}")

    # Step backward at the start should clamp to 0.
    c.call("trace", seek=0)
    r = c.call("trace", step=-1)
    check("step=-1 at t=0 stays at 0", r["trace"]["idx"] == 0,
          f"got {r['trace']['idx']}")

    # Multi-step.
    c.call("trace", seek=0)
    r = c.call("trace", step=5)
    check("step=5 advances five timestamps", r["trace"]["idx"] == 5,
          f"got {r['trace']['idx']}")

    # ------------------------------------------------------------------ #
    # step over (follows SP)
    # ------------------------------------------------------------------ #
    # Find a call in the model: SP drops and later recovers.
    sp_name = "rsp" if "rsp" in model.reg_at else "esp"
    call_at = None
    for i in range(1, min(model.length - 1, 400)):
        a = model.register(sp_name, i)
        b = model.register(sp_name, i + 1)
        if a and b and b < a:
            ret = next((j for j in range(i + 1, model.length)
                        if (model.register(sp_name, j) or 0) >= a), None)
            if ret and ret > i + 3:
                call_at = (i, ret)
                break
    if call_at is not None:
        i, ret = call_at
        c.call("trace", seek=i)
        r = c.call("trace", step=1, over=True)
        check("step over skips the callee",
              r["trace"]["idx"] == ret,
              f"from {i}, got {r['trace']['idx']}, expected {ret}")
        check("step over goes further than a plain step",
              r["trace"]["idx"] > i + 1)
    else:
        check("found a call to step over", False, "none in this short trace")

    # ------------------------------------------------------------------ #
    # goto: first execution of a function by name
    # ------------------------------------------------------------------ #
    r = c.call("trace", goto="main")
    tr = r["trace"]
    check("goto main lands on a timestamp", tr["idx"] >= 0)
    check("and the function context says main",
          r.get("function", {}).get("name") == "main",
          f"function={r.get('function')}")

    # goto by hex address.
    main_ea = r["function"]["ea"]
    c.call("trace", seek=0)  # reset position
    r = c.call("trace", goto=hex(main_ea))
    check("goto by hex address works",
          r["trace"]["idx"] >= 0 and r["function"]["ea"] == main_ea,
          f"idx={r['trace']['idx']}, ea={r.get('function', {}).get('ea')}")

    # goto a function that was never executed.
    try:
        c.call("trace", goto="0xDEADBEEF")
        check("goto an unexecuted address raises", False, "no error raised")
    except RpcError as e:
        check("goto an unexecuted address raises", "never executed" in str(e),
              str(e))

    # ------------------------------------------------------------------ #
    # changed registers in the response
    # ------------------------------------------------------------------ #
    c.call("trace", seek=0)
    r = c.call("trace", step=1)
    changed = r["trace"]["changed"]
    check("changed is a list of register names",
          isinstance(changed, list) and all(isinstance(s, str) for s in changed),
          f"{changed}")
    # The PC always changes on a step (it's a different instruction).
    check("rip is always in changed", "rip" in changed, f"{changed}")

    # ------------------------------------------------------------------ #
    # the cursor tracks the trace: seek moves the code view
    # ------------------------------------------------------------------ #
    # Seek to two different timestamps and verify the cursor ea changes.
    r1 = c.call("trace", seek=10)
    r2 = c.call("trace", seek=min(50, model.length - 1))
    ea1 = r1.get("cursor", {}).get("ea")
    ea2 = r2.get("cursor", {}).get("ea")
    check("the cursor ea follows the trace pc",
          ea1 is not None and ea2 is not None and (ea1 != ea2 or r1["trace"]["pc"] == r2["trace"]["pc"]),
          f"ea1={ea1}, ea2={ea2}")

    # ------------------------------------------------------------------ #
    # trace verb without a trace raises cleanly
    # ------------------------------------------------------------------ #
    # (Can't test this on THIS pane since it has a trace. But we can test
    # the error path by verifying the error message shape for bad params.)
    try:
        c.call("trace")  # no seek/goto/step — should still return a snapshot
        # Actually, if none of seek/goto/step is given, it just settles and
        # returns the current state — that's fine, it's a status query.
        r = c.call("trace")
        check("trace with no action is a status query",
              "trace" in r and r["trace"]["idx"] >= 0)
    except RpcError:
        check("trace with no action is a status query", False, "raised an error")

    # ------------------------------------------------------------------ #
    # snapshot shape: the trace key is present on every trace response
    # ------------------------------------------------------------------ #
    r = c.call("trace", seek=3)
    for key in ("idx", "length", "pc", "changed"):
        check(f"trace response has '{key}'", key in r.get("trace", {}),
              f"trace={r.get('trace')}")

    # Standard snapshot fields are ALSO present (the trace response is a
    # superset of a normal snapshot).
    for key in ("active", "function", "cursor", "status", "ready"):
        check(f"trace response also has snapshot key '{key}'", key in r,
              f"keys={list(r.keys())}")

    # ------------------------------------------------------------------ #
    # edge cases: seek beyond bounds
    # ------------------------------------------------------------------ #
    r = c.call("trace", seek=-1)
    check("seek -1 clamps to 0", r["trace"]["idx"] == 0)
    r = c.call("trace", seek=model.length + 1000)
    check("seek beyond length clamps to the end",
          r["trace"]["idx"] == model.length - 1,
          f"got {r['trace']['idx']}")

    # ------------------------------------------------------------------ #
    # seek with comma-separated numbers (ergonomic)
    # ------------------------------------------------------------------ #
    r = c.call("trace", seek="100")
    check("seek accepts a string number",
          r["trace"]["idx"] == min(100, model.length - 1))

    # ------------------------------------------------------------------ #
    # pseudocode still works with a trace loaded
    # ------------------------------------------------------------------ #
    c.call("trace", goto="main")
    r = c.call("pseudocode", target="main", lines=5)
    check("pseudocode works alongside the trace",
          "code" in r and "main" in r.get("code", ""),
          f"keys={list(r.keys())}")

    # ------------------------------------------------------------------ #
    # state includes trace position
    # ------------------------------------------------------------------ #
    c.call("trace", seek=7)
    r = c.call("state")
    # The state verb doesn't include trace info (that's trace-specific),
    # but the standard snapshot fields should be consistent.
    check("state works with a trace loaded",
          r.get("ready") is True and "cursor" in r)

    # ------------------------------------------------------------------ #
    # view_lines works with trail painted
    # ------------------------------------------------------------------ #
    c.call("trace", seek=min(40, model.length - 1))
    r = c.call("view", lines=10)
    check("view returns lines with a trace active",
          "lines" in r and len(r["lines"]) > 0,
          f"keys={list(r.keys())}")

    # ------------------------------------------------------------------ #
    # navigation works alongside trace: goto a function, trace follows
    # ------------------------------------------------------------------ #
    c.call("trace", goto="main")
    start_idx = c.call("trace")["trace"]["idx"]
    r = c.call("goto", target="error_at_line")
    check("goto still works with trace loaded",
          r.get("function", {}).get("name") == "error_at_line")
    # The trace timestamp should NOT change from a regular goto — the trace
    # position is independent of navigation.
    r2 = c.call("trace")
    check("regular goto does not change the trace position",
          r2["trace"]["idx"] == start_idx,
          f"was {start_idx}, now {r2['trace']['idx']}")


if __name__ == "__main__":
    raise SystemExit(main())