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
|
#!/usr/bin/env python3
"""Differential gate for the fast decompile_function_safe (run by checks.sh).
`idatui/worker.py` rebinds ida-pro-mcp's `decompile_function_safe` to our own
loop, which skips two of the three SWIG allocations per line and memoises the
per-line `dstr()` by ctree obj_id. That is a pure speed change and the text it
returns is what the pseudocode pane shows, markers and all -- so it has to be
byte-identical, not merely similar.
This runs BOTH implementations against the same cfunc for every function of a
real binary and compares the strings, with `include_addresses` both ways (the
marker path is the whole point, and the no-marker path must not regress either).
~/ida-venv/bin/python .auto/check_decomp.py [targets/echo] [max_funcs]
"""
from __future__ import annotations
import os
import shutil
import sys
HERE = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.dirname(HERE)
sys.path.insert(0, ROOT)
sys.path.insert(0, HERE)
from bench import stage # noqa: E402
def _add_mcp_path() -> None:
"""ida-pro-mcp is installed for the interpreter the WORKER runs, which is
not necessarily the one running this check."""
import glob
for pat in ("/home/user/.local/lib/python3.*/site-packages",
os.path.expanduser("~/.local/lib/python3.*/site-packages")):
for d in glob.glob(pat):
if os.path.isdir(os.path.join(d, "ida_pro_mcp")) and d not in sys.path:
sys.path.append(d)
def main() -> int:
target = sys.argv[1] if len(sys.argv) > 1 else "targets/echo"
limit = int(sys.argv[2]) if len(sys.argv) > 2 else 400
d, path = stage(os.path.join(ROOT, target))
os.environ["IDA_MCP_TOOL_TIMEOUT_SEC"] = "0"
_add_mcp_path()
try:
import idapro
idapro.open_database(path, run_auto_analysis=True)
try:
import ida_funcs
import ida_hexrays
import idautils
from ida_pro_mcp.ida_mcp import utils
original = utils.decompile_function_safe
sys.path.insert(0, os.path.join(ROOT, "server"))
import patch_server
# exec only our function out of the decoded BODY: the rest of it
# needs api_types' namespace (@tool, @idasync, ...).
body = patch_server.BODY
i = body.find("def _idatui_decompile_function_safe(")
j = body.find("\n_idatui_strings_cache", i)
assert i > 0 and j > i, "could not slice the function out of BODY"
g = {}
exec(body[i:j], g)
fast = g["_idatui_decompile_function_safe"]
ida_hexrays.init_hexrays_plugin()
fails: list[str] = []
n = ok = 0
for ea in idautils.Functions():
f = ida_funcs.get_func(ea)
if not f:
continue
n += 1
if n > limit:
break
for markers in (True, False):
a, ea_err = original(f.start_ea, include_addresses=markers)
b, eb_err = fast(f.start_ea, include_addresses=markers)
if a != b or (ea_err is None) != (eb_err is None):
fails.append(f"{f.start_ea:#x} (markers={markers})")
if len(fails) <= 3:
la = (a or "").splitlines()
lb = (b or "").splitlines()
i = next((k for k, (x, y) in enumerate(zip(la, lb))
if x != y), None)
if i is None:
print(f" FAIL {f.start_ea:#x}: {len(la)} lines "
f"vs {len(lb)}, errs {ea_err!r}/{eb_err!r}")
else:
print(f" FAIL {f.start_ea:#x} line {i}:")
print(f" original: {la[i][:100]!r}")
print(f" fast : {lb[i][:100]!r}")
break
else:
ok += 1
print(f"decompile text: {ok}/{min(n, limit)} functions of {target} "
f"byte-identical to ida-pro-mcp's own loop, "
f"{len(fails)} problems")
return 1 if fails else 0
finally:
idapro.close_database(save=False)
finally:
shutil.rmtree(d, ignore_errors=True)
if __name__ == "__main__":
sys.exit(main())
|