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
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
|
#!/usr/bin/env python3
"""Propagate function names across builds through the CALL GRAPH, not through byte similarity.
Byte and mnemonic porting names a function only when its own body still resembles the
reference. Everything it cannot reach stays sub_*, and for ASD that is thousands of functions.
But a function's position in the call graph is evidence independent of its body: if reference
function A is matched to target function A', and both make the same number of calls, then the
Nth callee of A' is the Nth callee of A.
That is the rule here, with three deliberate restrictions, because positional alignment is
exactly the kind of heuristic that quietly produces wrong names:
* EQUAL CALL COUNT. If A and A' do not make the same number of calls, the alignment is not
trustworthy and the anchor is skipped entirely rather than aligned partially.
* UNIQUE CANDIDATE. A target function is named only if every anchor that reaches it agrees
on the name. Any disagreement and it is left alone.
* SIZE AGREEMENT. The candidate and its reference must be within a factor of `--size-ratio`
of each other. A name is an identity claim, and a function half the size of the one whose
name it is being given is probably not that function.
Names are applied to the target database, so re-running fw/ta_re.py afterwards picks them up
and re-emits the artifacts. Rounds are iterated: a function named in round N becomes an anchor
in round N+1.
Control: --dry-run reports what would be applied without touching anything, and the
ANCHOR-MISMATCH rate is printed, because a high mismatch rate means the builds have diverged
too far for positional alignment to mean anything.
usage:
./ta_cgname.py --type asd [--rounds 3] [--dry-run]
"""
from __future__ import annotations
import argparse
import glob
import json
import os
import re
import sys
from pathlib import Path
HERE = Path(__file__).resolve().parent
NAMEDIR_REL = "cgnames"
REPO = HERE.parent
sys.path.insert(0, str(Path.home() / "dev" / "ida-tui-maybe"))
sys.path.insert(0, str(HERE))
CALL = re.compile(r"\b([A-Za-z_]\w*)\s*\(")
KEYWORDS = {"if", "while", "for", "switch", "return", "sizeof", "__fastcall", "do",
"else", "case", "break", "continue", "goto"}
SUBNAME = re.compile(r"^sub_[0-9A-Fa-f]+$")
def func_calls(client, remote_ops, addr):
r = client.call(remote_ops.decompile, addr=addr, include_addresses=False) or {}
code = r.get("code") or ""
code = "\n".join(l for l in code.splitlines() if not l.lstrip().startswith("//"))
body = code[code.find("{"):] if "{" in code else ""
return [m.group(1) for m in CALL.finditer(body) if m.group(1) not in KEYWORDS]
def list_funcs(client, remote_ops):
out, off = [], 0
while True:
r = client.call(remote_ops.list_funcs, queries=[{"count": 500, "offset": off}])
blk = r["result"][0]
out.extend(blk["data"])
off = blk["next_offset"]
if off >= blk["total"]:
break
return out
def main() -> int:
ap = argparse.ArgumentParser()
ap.add_argument("--type", required=True)
ap.add_argument("--rounds", type=int, default=3)
ap.add_argument("--size-ratio", type=float, default=2.0)
ap.add_argument("--dry-run", action="store_true")
ap.add_argument("--ref", help="use this build as the reference instead of the registered "
"one (substring match); lets a well-named CORPUS build seed "
"its weaker siblings, which the handover reference cannot "
"always reach")
a = ap.parse_args()
from idatui import remote_ops
from idatui.nexus_client import NexusClient
import ta_re as T
if a.ref:
cands = [p for p in sorted(glob.glob(str(T.ELF_DIR / f"{a.type}_*.elf")))
if a.ref in os.path.basename(p)]
if not cands:
raise SystemExit("no build matching --ref %r" % a.ref)
refelf = cands[0]
elif a.type in T.HANDOVER_REF:
refelf = str(REPO / "psp_out" / "handover" / "elf" / f"{a.type}2.elf")
else:
ver = T.LIBDUMP_REF[a.type]
refelf = sorted(glob.glob(str(T.ELF_DIR / f"{a.type}_{ver}_*.elf")))[0]
targets = [p for p in sorted(glob.glob(str(T.ELF_DIR / f"{a.type}_*.elf")))
if p != refelf]
if not targets:
print("no target builds for %s" % a.type)
return 0
# Names are persisted to fw/cgnames/<ta>.json as well as written into the database.
# Writing only into the .i64 has already cost this work twice: three databases have come
# back corrupted, and rebuilding one discards every name derived here (GCP fell from 83%
# to 74% that way). The JSON is the durable record; fw/ta_re.py re-applies it.
namedir = HERE / NAMEDIR_REL
namedir.mkdir(exist_ok=True)
mapfile = namedir / f"{a.type}.json"
persisted = json.loads(mapfile.read_text()) if mapfile.exists() else {}
rc = NexusClient(refelf)
rc.connect()
try:
rf = list_funcs(rc, remote_ops)
rsize = {f["name"]: f["size"] for f in rf}
raddr = {f["name"]: f["addr"] for f in rf
if not SUBNAME.match(f["name"]) and not f["name"].startswith("j_")}
rcalls: dict[str, list] = {}
print("%s: reference %s, %d named anchors available"
% (a.type, os.path.basename(refelf), len(raddr)), flush=True)
total_applied = 0
for tp in targets:
key = os.path.basename(tp).replace(".elf", "")
tc = NexusClient(tp)
try:
tc.connect()
applied_here = 0
for rnd in range(a.rounds):
tf = list_funcs(tc, remote_ops)
tname = {f["addr"]: f["name"] for f in tf}
tsize = {f["addr"]: f["size"] for f in tf}
byname = {f["name"]: f["addr"] for f in tf}
anchors = [n for n in raddr if n in byname]
cand: dict[str, set] = {}
eq = mismatch = 0
for n in anchors:
if n not in rcalls:
try:
rcalls[n] = func_calls(rc, remote_ops, raddr[n])
except Exception:
rcalls[n] = []
rl = rcalls[n]
if not rl:
continue
try:
tl = func_calls(tc, remote_ops, byname[n])
except Exception:
continue
if len(tl) != len(rl):
mismatch += 1
continue
eq += 1
for x, y in zip(rl, tl):
if SUBNAME.match(y) and not SUBNAME.match(x) \
and not x.startswith("j_"):
cand.setdefault(y, set()).add(x)
# unique candidate + size agreement
renames = []
for tgtnm, names in cand.items():
if len(names) != 1:
continue
new = next(iter(names))
ad_existing = byname.get(new)
if ad_existing is not None:
# Already named here -- most likely by an earlier run of this very
# tool. Record it in the durable map anyway, otherwise a module that
# was processed before the map existed can never be recovered from
# source and stays hostage to its .i64.
persisted.setdefault(key, {})[ad_existing] = new
continue
ad = byname.get(tgtnm)
if ad is None:
continue
rs, ts = rsize.get(new, 0), tsize.get(ad, 0)
if rs and ts:
hi, lo = max(rs, ts), min(rs, ts)
if lo == 0 or hi / lo > a.size_ratio:
continue
renames.append({"addr": ad, "name": new})
if not renames:
print(" %-30s round %d: anchors eq=%d mismatch=%d -> nothing"
% (os.path.basename(tp)[:30], rnd + 1, eq, mismatch), flush=True)
break
if a.dry_run:
print(" %-30s round %d: anchors eq=%d mismatch=%d -> WOULD apply %d"
% (os.path.basename(tp)[:30], rnd + 1, eq, mismatch,
len(renames)), flush=True)
break
rr = tc.call(remote_ops.rename,
batch={"func": renames, "allow_overwrite": False})
# response shape is {"func":[...], "summary":{"ok":N,"failed":M}}
ok = int((rr.get("summary") or {}).get("ok", 0))
applied_here += ok
for r_ in renames:
persisted.setdefault(key, {})[r_["addr"]] = r_["name"]
print(" %-30s round %d: anchors eq=%d mismatch=%d -> applied %d"
% (os.path.basename(tp)[:30], rnd + 1, eq, mismatch, ok), flush=True)
if ok == 0:
break
if applied_here and not a.dry_run:
try:
tc.save_database()
except Exception:
pass
total_applied += applied_here
finally:
try:
tc.close()
except Exception:
pass
if not a.dry_run:
mapfile.write_text(json.dumps(persisted, indent=1, sort_keys=True))
print("persisted %d build maps to %s"
% (len(persisted), mapfile.relative_to(REPO)))
print("%s: %d names applied across %d builds" % (a.type, total_applied, len(targets)))
finally:
try:
rc.close()
except Exception:
pass
return 0
if __name__ == "__main__":
sys.exit(main())
|
