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
|
#!/usr/bin/env python3
"""Annotate a TA in a live idatui session from the decompiler's own switch output.
`fw/ta_symbols.py` recovers entry points statically, but its jump-table walk misses
dispatchers reached by a `b.w` tail call and two-tier switches. The decompiler already
resolves both, so this reads `case N: ... = sub_XXXX(...)` straight out of the
pseudocode and renames each handler `<prefix>_cmdNN`.
Run against a pane already opened on the TA (socket auto-resolved when only one is live).
usage: annotate_ta.py --prefix gcp --dispatch gcp_dispatch [--note "..."]
"""
import argparse
import json
import os
import re
import subprocess
import sys
DRIVE = [sys.executable, '-m', 'idatui.drive']
CWD = os.path.expanduser('~/dev/ida-tui-maybe')
def drive(*args, timeout=180):
r = subprocess.run(DRIVE + list(args), cwd=CWD, capture_output=True,
text=True, timeout=timeout)
return (r.stdout or '') + (r.stderr or '')
def parse_ifchain(pc):
"""Some TAs (CSU, OC, WSPE) dispatch with an if/else-if chain rather than a
switch, which has no `case N:` lines at all. Match `== N` guards followed within
a few lines by a call."""
out = {}
lines = pc.splitlines()
for n, line in enumerate(lines):
m = re.search(r'==\s*(0x[0-9A-Fa-f]+|\d+)\s*\)', line)
if not m or not re.search(r'\bif\b', line):
continue
try:
cid = int(m.group(1), 0)
except ValueError:
continue
if cid > 0x2000:
continue
for k in range(n + 1, min(n + 4, len(lines))):
c = re.search(r'=\s*([A-Za-z_]\w*)\s*\(', lines[k])
if c and c.group(1) not in ('if', 'while', 'for'):
out.setdefault(cid, c.group(1))
break
return out
def parse_cases(pc):
"""case <id>: ... <something> = sub_XXXX( -> {id: handler}"""
out, pending = {}, []
for line in pc.splitlines():
m = re.search(r'case\s+(0x[0-9A-Fa-f]+|\d+)[uU]?:', line)
if m:
pending.append(int(m.group(1), 0))
continue
m2 = (re.search(r'=\s*([A-Za-z_]\w*)\s*\(', line)
or re.search(r'\)\s*([A-Za-z_]\w*)\s*\)\s*\(', line))
if m2 and pending:
for cid in pending:
out.setdefault(cid, m2.group(1))
pending = []
continue
m3 = re.search(r'return\s+([A-Za-z_]\w*)\s*\(', line)
if m3 is None and 'return' in line:
m3 = re.search(r'\)\s*([A-Za-z_]\w*)\s*\)\s*\(', line)
if m3 and pending:
for cid in pending:
out.setdefault(cid, m3.group(1))
pending = []
if re.match(r'\s*(default|LABEL_|\})', line):
pending = []
return out
def main():
ap = argparse.ArgumentParser()
ap.add_argument('--prefix', required=True)
ap.add_argument('--dispatch', required=True)
ap.add_argument('--note')
ap.add_argument('--out')
a = ap.parse_args()
def promote_locs(text):
"""Handlers are often reached as `loc_XXXX)(args)` because IDA never made a
function there; the decompiler then shows a cast-through-label instead of a
callee, and no rename is possible. Promote them to functions first."""
locs = set(re.findall(r'loc_([0-9A-Fa-f]+)\s*\)\s*\(', text))
for L in locs:
drive('raw', 'define', 'kind=func', 'target=0x%s' % L)
return len(locs)
pc = drive('pc', a.dispatch)
n_promoted = promote_locs(pc)
if n_promoted:
pc = drive('pc', a.dispatch)
cases = parse_cases(pc) or parse_ifchain(pc)
disp = a.dispatch
# Wrapper-following: several TAs' invoke entry only validates paramTypes and then
# tail-calls (b.w) or calls the real dispatcher. Follow one level when the entry
# itself shows no cases.
if not cases:
inner = re.findall(r'(?:return|=)\s*([A-Za-z_]\w*)\s*\(', pc)
inner = [x for x in dict.fromkeys(inner)
if x not in ('if', 'while', 'for', 'return', 'sizeof')]
for cand in inner:
pc2 = drive('pc', cand)
if promote_locs(pc2):
pc2 = drive('pc', cand)
c2 = parse_cases(pc2) or parse_ifchain(pc2)
if c2:
if re.fullmatch(r'sub_[0-9A-Fa-f]+', cand):
import tempfile as _tf
with _tf.NamedTemporaryFile('w', suffix='.json', delete=False) as fh:
json.dump([dict(addr=int(cand.split('_')[1], 16),
name='%s_dispatch' % a.prefix)], fh)
dp = fh.name
drive('syms', dp)
os.unlink(dp)
cand = '%s_dispatch' % a.prefix
cases, pc, disp = c2, pc2, cand
break
# a handler shared by several ids gets the lowest id in its name
byfn = {}
for cid, fn in sorted(cases.items()):
byfn.setdefault(fn, cid)
# Rename by ADDRESS via rename_many, not by token: `drive mv` resolves the old
# name against the *current view*, so it silently fails for anything off-screen
# (8 of 14 handlers on the first GCP run).
# Resolve EVERY handler to an address so the decompiler-derived map is
# authoritative. The earlier static pass guessed names from a jump-table walk and
# got some wrong (it labelled OPM's precheck `opm_cmd00`); those must be corrected,
# not skipped just because the function already has a name.
name2addr = {}
for ln in drive('names', a.prefix, '400').splitlines():
mm = re.match(r'\s*(0x[0-9A-Fa-f]+)\s+(\S+)', ln)
if mm:
name2addr[mm.group(2)] = int(mm.group(1), 16)
syms, already = [], []
for fn, cid in byfn.items():
if re.fullmatch(r'sub_[0-9A-Fa-f]+', fn):
addr = int(fn.split('_')[1], 16)
elif fn in name2addr:
addr = name2addr[fn]
already.append(fn)
else:
continue
syms.append(dict(addr=addr, name='%s_cmd%02X' % (a.prefix, cid)))
import tempfile
with tempfile.NamedTemporaryFile('w', suffix='.json', delete=False) as fh:
json.dump(syms, fh)
sp = fh.name
res = drive('syms', sp)
os.unlink(sp)
m = re.search(r'(\d+)/(\d+) renamed', res)
ok = int(m.group(1)) if m else 0
if n_promoted:
print(' promoted %d loc_ labels to functions' % n_promoted)
print('%s: %d commands, %d distinct handlers, %d renamed%s'
% (a.prefix, len(cases), len(byfn), max(ok, 0),
(', %d already named' % len(already)) if already else ''))
print(' ids: %s' % ', '.join(hex(c) for c in sorted(cases)))
if a.note:
drive('note', disp, a.note)
drive('save')
if a.out:
json.dump({hex(k): v for k, v in sorted(cases.items())},
open(a.out, 'w'), indent=1)
return 0
if __name__ == '__main__':
raise SystemExit(main())
|
