diff options
Diffstat (limited to 'experiments/profile_remote.py')
| -rw-r--r-- | experiments/profile_remote.py | 111 |
1 files changed, 44 insertions, 67 deletions
diff --git a/experiments/profile_remote.py b/experiments/profile_remote.py index ca673db..fbf17dc 100644 --- a/experiments/profile_remote.py +++ b/experiments/profile_remote.py @@ -1,94 +1,71 @@ -"""Profile an operation INSIDE the database process. +"""Profile persistent ida-tui operations inside the IDA process. -`bench_ops.py` says how long an operation takes; this says where that time -goes. The snippet ships cProfile into the Code Mode sandbox, runs the real -remote-library function there in a loop, and returns the stats as text -- so -the split between IDA's own calls and OUR python in `remote_tools.py` is -visible, which no client-side timer can see. +The profiler itself is a typed ``RemoteModule`` function in ``remote_tools.py``; +this file contains no generated Python source or knowledge of remote module names. - PYTHONPATH=. ~/ida-venv/bin/python experiments/profile_remote.py [BINARY] - PYTHONPATH=. ~/ida-venv/bin/python experiments/profile_remote.py --op decompile +Usage:: -Read the `tottime` column: time in that function excluding subcalls. IDA -builtins (generate_disasm_line, get_flags, next_head...) are the floor; a -python frame from ida_tui_remote near the top is ours, and ours is fixable. + PYTHONPATH=. python experiments/profile_remote.py [BINARY] + PYTHONPATH=. python experiments/profile_remote.py --op decompile """ + from __future__ import annotations import argparse import os -import sys - -from idatui.codemode_client import CodeModeClient, _REMOTE_MODULE, _script -# Runs in the database process. `a` is the bound argument dict. -PROFILE = ''' -import cProfile, pstats, io, sys -_m = sys.modules.get(%(mod)r) -if _m is None: - result = {"error": "remote lib not installed yet"} -else: - call = a["call"] - reps = int(a["reps"]) - ns = {"_m": _m, "a": a} - src = "for _ in range(%%d):\\n _m.%%s" %% (reps, call) - code = compile(src, "<profile>", "exec") - pr = cProfile.Profile() - pr.enable() - exec(code, ns) - pr.disable() - buf = io.StringIO() - st = pstats.Stats(pr, stream=buf).sort_stats("tottime") - st.print_stats(int(a["lines"])) - result = {"stats": buf.getvalue(), "total": st.total_tt, "reps": reps} -''' % {"mod": _REMOTE_MODULE} +from idatui import remote_ops +from idatui.nexus_client import NexusClient CALLS = { - # One full listing page, exactly as the background grower asks for it. - "heads": 'heads(addr=a["addr"], count=500, annotate=True)', - "heads_plain": 'heads(addr=a["addr"], count=500, annotate=False)', - "heads_skeleton": 'heads(addr=a["addr"], count=500, annotate=True, text=False)', - "decompile": 'decompile(a["addr"])', - "disasm": 'disasm(a["addr"], 500)', + "heads": ("heads", {"count": 500, "annotate": True}), + "heads_plain": ("heads", {"count": 500, "annotate": False}), + "heads_skeleton": ( + "heads", + {"count": 500, "annotate": True, "text": False}, + ), + "decompile": ("decompile", {}), } def main() -> int: - ap = argparse.ArgumentParser() - ap.add_argument("binary", nargs="?", default="targets/bash") - ap.add_argument("--op", default="heads", choices=sorted(CALLS)) - ap.add_argument("--reps", type=int, default=20) - ap.add_argument("--lines", type=int, default=18) - ap.add_argument("--addr", default=None, help="default: the .text start") - args = ap.parse_args() - - client = CodeModeClient(os.path.abspath(args.binary)) - client.connect() + parser = argparse.ArgumentParser() + parser.add_argument("binary", nargs="?", default="targets/bash") + parser.add_argument("--op", default="heads", choices=sorted(CALLS)) + parser.add_argument("--reps", type=int, default=20) + parser.add_argument("--addr", default=None, help="default: the .text start") + args = parser.parse_args() + client = NexusClient(os.path.abspath(args.binary)).connect() addr = args.addr if addr is None: - regions = client.invoke("file_regions") + regions = client.call(remote_ops.file_regions) rows = regions.get("regions") or regions.get("result") or [] - text = next((r for r in rows if ".text" in str(r.get("name", ""))), None) + text = next( + (row for row in rows if ".text" in str(row.get("name", ""))), + None, + ) addr = (text or rows[0])["start"] if rows else "0x0" - print(f"# {os.path.basename(args.binary)} op={args.op} addr={addr} reps={args.reps}") - # Prime: the remote lib installs lazily, and its lru_caches must be warm or - # the profile measures cache misses that the real workload never pays. - client.invoke("heads", addr=addr, count=500, annotate=True) + operation, call_args = CALLS[args.op] + call_args = {"addr": addr, **call_args} + print( + f"# {os.path.basename(args.binary)} op={args.op} " + f"addr={addr} reps={args.reps}" + ) - # _script binds the args as JSON and adds the pack epilogue, exactly as a - # real operation is shipped -- so this measures the same path, not a - # special one. - out = client._unpack(client.execute_python(_script( - {"call": CALLS[args.op], "addr": addr, - "reps": args.reps, "lines": args.lines}, PROFILE), timeout=600)) - if "error" in out: - print("FAILED:", out["error"]) - return 1 + # Install the persistent tool module and warm its caches before profiling. + client.call(remote_ops.heads, addr=addr, count=500, annotate=True) + out = client.call( + remote_ops.profile_remote, + operation=operation, + args=call_args, + reps=args.reps, + ) per = out["total"] / out["reps"] * 1000 - print(f"# {out['total']*1000:.0f}ms total, {per:.1f}ms per call\n") + print(f"# {out['total'] * 1000:.0f}ms total, {per:.1f}ms per call\n") print(out["stats"]) + client.close() return 0 |