diff options
Diffstat (limited to 'tests')
| -rw-r--r-- | tests/test_trace_ui.py | 148 |
1 files changed, 148 insertions, 0 deletions
diff --git a/tests/test_trace_ui.py b/tests/test_trace_ui.py new file mode 100644 index 0000000..a30362f --- /dev/null +++ b/tests/test_trace_ui.py @@ -0,0 +1,148 @@ +#!/usr/bin/env python3 +"""The trace dock: registers, timeline, and stepping through time. + +Needs IDA and a trace. Generates its own trace with the QEMU tracer if the +tracer is built; skips with a message otherwise, since neither the emulator nor +the trace is part of this repo. +""" +import asyncio +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 textual.widgets import Static # noqa: E402 + +from idatui.app import IdaTui, ListingView, TraceDock # 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") +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}") + + +async def wait(pred, pilot, t=240.0): + for _ in range(int(t / 0.05)): + await pilot.pause(0.05) + try: + if pred(): + return True + except Exception: # noqa: BLE001 + pass + return False + + +def make_trace(tmp, binary): + out = os.path.join(tmp, "t") + try: + subprocess.run([TRACER, "-o", out, binary, "hi"], + 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 + + +async def run() -> int: + binary = os.path.join(REPO, "targets", "echo") + with tempfile.TemporaryDirectory() as tmp: + # Scratch copy: opening a binary writes a database beside it, and the + # suite must not edit anything tracked. + target = os.path.join(tmp, "echo") + shutil.copy2(binary, target) + log = make_trace(tmp, target) + if not log: + print(f" skip: could not record a trace (tracer at {TRACER})") + return 0 + + app = IdaTui(open_path=target, keepalive=False, trace_path=log) + async with app.run_test(size=(160, 44)) as pilot: + ok = await wait(lambda: app._trace is not None, pilot, 300) + check("the trace loads alongside the binary", ok) + if not ok: + return 1 + t = app._trace + check("it has instructions", t.length > 10, f"{t.length}") + + # Rebasing: the tracer runs the binary relocated, so without a slide + # nothing in the trace matches anything on screen. + check("trace addresses were rebased onto the database", + t.slide != 0 and t.ip(0) < 0x1000000, + f"slide={t.slide:#x} ip0={t.ip(0):#x}") + idx = app._func_index + touched = [f.name for f in idx.all_loaded() if t.executions(f.addr)] + check("and now line up with real functions", + len(touched) > 1 and "main" in touched, f"{touched[:6]}") + + dock = app.query_one(TraceDock) + check("the dock is docked and visible", dock.display) + head = str(dock.query_one("#trace-head", Static).render()) + check("it shows where we are in time", "0" in head and "%" in head, + head[:60]) + regs = str(dock.query_one("#trace-regs", Static).render()) + check("and the register state at that time", "rip" in regs.lower(), + regs[:60]) + + # -- stepping --------------------------------------------------- # + lst = app.query_one(ListingView) + lst.focus() + await pilot.pause(0.4) + await pilot.press("]") + await wait(lambda: app._t == 1, pilot, 20) + check("] steps forward one instruction", app._t == 1, f"t={app._t}") + check("the code view follows the trace", + lst._cursor_ea() == t.ip(1), + f"{lst._cursor_ea()} vs {t.ip(1)}") + await pilot.press("[") + await wait(lambda: app._t == 0, pilot, 20) + check("[ steps backward", app._t == 0, f"t={app._t}") + await pilot.press("[") + await pilot.pause(0.4) + check("and stops at the start of the trace", app._t == 0) + + # -- step over -------------------------------------------------- # + # A call pushes, so the callee runs with SP below where we started; + # stepping until SP comes back up lands after the return. Find a + # real call in this trace rather than assuming one is at a fixed + # place. + sp = "rsp" if "rsp" in t.reg_at else "esp" + call_at = None + for i in range(1, min(t.length - 1, 400)): + a, b = t.register(sp, i), t.register(sp, i + 1) + if a and b and b < a: + ret = next((j for j in range(i + 1, t.length) + if (t.register(sp, j) or 0) >= a), None) + if ret and ret > i + 3: + call_at = (i, ret) + break + if call_at is None: + check("found a call to step over", False, "none in this trace") + else: + i, ret = call_at + app._seek(i) + await wait(lambda: app._t == i, pilot, 20) + await pilot.press("}") + await wait(lambda: app._t != i, pilot, 30) + check("} steps OVER a call instead of into it", + app._t == ret, f"{i} -> {app._t}, expected {ret}") + check("which is further than a plain step", app._t > i + 1) + + print(f"\n{PASS} passed, {FAIL} failed") + return 1 if FAIL else 0 + + +if __name__ == "__main__": + raise SystemExit(asyncio.run(run())) |