aboutsummaryrefslogtreecommitdiffstats
path: root/tests/test_trace_ui.py
blob: a85c0890e9d57e7183227cc1a3efacf469f3ed9a (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
#!/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 (DecompView, IdaTui, ListingView,  # noqa: E402
                        TraceDock)

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)

            # -- trails ------------------------------------------------------ #
            # Not "every address the trace ever touched": on a loop-heavy
            # program that's almost everything and says nothing. The last/next
            # few dozen steps say how you got here and where you're going.
            app._seek(min(40, t.length - 1))
            await pilot.pause(0.6)
            trail = lst.trail
            kinds = {k for k in trail.values()}
            check("the listing is painted with an execution trail",
                  {"now", "past", "future"} <= kinds, f"{sorted(kinds)}")
            check("'now' is the instruction we're standing on",
                  trail.get(t.ip(app._t)) == "now", f"{trail.get(t.ip(app._t))}")
            check("the step behind is past, the step ahead is future",
                  trail.get(t.ip(app._t - 1)) == "past"
                  and trail.get(t.ip(app._t + 1)) == "future",
                  f"{trail.get(t.ip(app._t - 1))}, {trail.get(t.ip(app._t + 1))}")
            painted = [y for y in range(min(lst.size.height, 30))
                       if any(seg.style and seg.style.bgcolor
                              for seg in lst.render_line(y))]
            check("and it actually reaches the screen", painted, "no tinted rows")

            # -- the same trail on PSEUDOCODE -------------------------------- #
            # The point of doing this in our app rather than using Tenet: a
            # trace's addresses are instructions, but decomp_map (built for the
            # split view) says which instructions each pseudocode line covers,
            # so the trail lands on C.
            main_ea = app.program.resolve("main")
            first = t.first_execution(main_ea)
            if first is None:
                check("main was executed in the trace", False)
            else:
                app._seek(first + 12)
                await wait(lambda: app._t == first + 12, pilot, 20)
                lst.focus()
                await pilot.press("tab")
                got = await wait(lambda: app.query_one(DecompView).display
                                 and app.query_one(DecompView)._texts, pilot, 120)
                dec = app.query_one(DecompView)
                check("pseudocode is available for the traced function", got)
                app._seek(first + 12)
                await pilot.pause(1.0)
                check("pseudocode lines are painted with the trail",
                      len(dec.trail) > 2, f"{len(dec.trail)} lines")
                now = [i for i, k in dec.trail.items() if k == "now"]
                check("exactly one pseudocode line is 'now'",
                      len(now) == 1, f"{now}")
                # The 'now' line must be the one covering the current
                # instruction, not merely some executed line.
                covered = app._trail_map[now[0]] if now and app._trail_map else []
                check("and it's the line covering the current instruction",
                      t.ip(app._t) in covered,
                      f"pc={t.ip(app._t):#x} line covers {[hex(a) for a in covered][:4]}")

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


if __name__ == "__main__":
    raise SystemExit(asyncio.run(run()))