diff options
Diffstat (limited to '')
| -rw-r--r-- | tests/test_trace_rpc.py | 259 |
1 files changed, 170 insertions, 89 deletions
diff --git a/tests/test_trace_rpc.py b/tests/test_trace_rpc.py index 015ce8b..9b0f2e0 100644 --- a/tests/test_trace_rpc.py +++ b/tests/test_trace_rpc.py @@ -30,8 +30,7 @@ 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") +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") @@ -52,8 +51,12 @@ 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) + subprocess.run( + [TRACER, "-o", out, binary, "hello"], + capture_output=True, + timeout=180, + check=False, + ) except (OSError, subprocess.TimeoutExpired): return None log = out + ".0.log" @@ -62,11 +65,24 @@ def make_trace(tmp, binary): 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) + 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 @@ -74,10 +90,17 @@ def spawn_pane(target, trace_log, timeout=300): 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) + 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: @@ -96,14 +119,17 @@ def main() -> int: 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})") + 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}") + check( + "model loaded for ground truth", model.length > 10, f"length={model.length}" + ) # Spawn the pane. info = spawn_pane(target, trace_log) @@ -138,14 +164,22 @@ def run_trace_tests(c: RpcClient, model: Trace): 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( + "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"]) + 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 @@ -153,14 +187,16 @@ def run_trace_tests(c: RpcClient, model: Trace): 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}") + 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']}") + check("pc is in database space (not ASLR'd)", pc < 0x100000, f"pc={tr['pc']}") # ------------------------------------------------------------------ # # seek by percentage (Tenet shell syntax) @@ -168,39 +204,50 @@ def run_trace_tests(c: RpcClient, model: Trace): 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}") + 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}") + 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']}") + 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']}") + 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']}") + 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']}") + 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']}") + check( + "step=5 advances five timestamps", + r["trace"]["idx"] == 5, + f"got {r['trace']['idx']}", + ) # ------------------------------------------------------------------ # # step over (follows SP) @@ -212,8 +259,14 @@ def run_trace_tests(c: RpcClient, model: Trace): 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) + 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 @@ -221,11 +274,12 @@ def run_trace_tests(c: RpcClient, model: Trace): 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) + 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") @@ -235,25 +289,28 @@ def run_trace_tests(c: RpcClient, model: Trace): 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')}") + 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')}") + 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)) + check("goto an unexecuted address raises", "never executed" in str(e), str(e)) # ------------------------------------------------------------------ # # changed registers in the response @@ -261,9 +318,11 @@ def run_trace_tests(c: RpcClient, model: Trace): 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}") + 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}") @@ -275,9 +334,13 @@ def run_trace_tests(c: RpcClient, model: Trace): 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}") + 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 @@ -289,8 +352,10 @@ def run_trace_tests(c: RpcClient, model: Trace): # 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) + 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") @@ -299,14 +364,20 @@ def run_trace_tests(c: RpcClient, model: Trace): # ------------------------------------------------------------------ # 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')}") + 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())}") + check( + f"trace response also has snapshot key '{key}'", + key in r, + f"keys={list(r.keys())}", + ) # ------------------------------------------------------------------ # # edge cases: seek beyond bounds @@ -314,25 +385,30 @@ def run_trace_tests(c: RpcClient, model: Trace): 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']}") + 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)) + 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())}") + check( + "pseudocode works alongside the trace", + "code" in r and "main" in r.get("code", ""), + f"keys={list(r.keys())}", + ) # ------------------------------------------------------------------ # # state includes trace position @@ -341,17 +417,18 @@ def run_trace_tests(c: RpcClient, model: Trace): 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) + 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())}") + 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 @@ -359,14 +436,18 @@ def run_trace_tests(c: RpcClient, model: Trace): 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") + 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']}") + 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__": |