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
|
#!/usr/bin/env python3
# SPDX-License-Identifier: BSD-3-Clause
#
# annotate_corpus.py - turn the raw decompiled corpus into an annotated,
# human-navigable corpus.
#
# For every function in the decompiled buckets it rewrites the header comment
# with:
# * the recovered name (from a names.tsv map) or the Ghidra FUN_ name;
# * file virtual address and size;
# * direct callers and callees (from callgraph.tsv), labelled with names;
# * a "seen in" marker for functions referenced by the docs (optional).
# It also substitutes recovered names for FUN_<addr> references inside the
# function bodies, so a renamed function is readable everywhere.
#
# Inputs (produced by the toolchain):
# <outdir>/buckets/bucket_*.c raw decompiled C
# <outdir>/decompiled_index.tsv addr file_addr size name file line status
# <indexdir>/callgraph.tsv caller callee caller_file callee_file
# <names.tsv> file_vaddr<TAB>name
#
# Usage:
# annotate_corpus.py <outdir> <indexdir> <names.tsv> <annotated_outdir>
#
# The output directory is gitignored like the raw corpus; this tool is tracked
# and reusable. Addresses are file virtual addresses in all inputs; the corpus
# uses Ghidra program addresses, so the tool translates via the index's
# addr/file_addr columns (no hard-coded image base).
import os
import re
import sys
from collections import defaultdict
def load_index(path):
rows = {}
with open(path) as f:
header = f.readline()
for line in f:
p = line.rstrip("\n").split("\t")
if len(p) < 6:
continue
prog = int(p[0], 16)
rows[prog] = dict(file=int(p[1], 16), size=int(p[2]),
name=p[3], file_name=p[4], line=int(p[5]))
return rows
def load_callgraph(path):
callees = defaultdict(list)
callers = defaultdict(list)
with open(path) as f:
f.readline()
for line in f:
p = line.rstrip("\n").split("\t")
if len(p) < 4:
continue
cf, kf = int(p[2], 16), int(p[3], 16)
callees[cf].append(kf)
callers[kf].append(cf)
return callees, callers
def load_names(path):
names = {}
with open(path) as f:
for line in f:
line = line.strip()
if not line or line.startswith("#"):
continue
p = line.split("\t")
if len(p) >= 2:
try:
names[int(p[0], 16)] = p[1]
except ValueError:
pass
return names
def label(file_addr, names, index):
name = names.get(file_addr)
if name:
return f"{name} (libcuda+0x{file_addr:x})"
return f"libcuda+0x{file_addr:x}"
def main():
if len(sys.argv) < 5:
print(__doc__)
sys.exit(2)
outdir, indexdir, namemap, out = sys.argv[1:5]
index = load_index(os.path.join(outdir, "decompiled_index.tsv"))
callees, callers = load_callgraph(os.path.join(indexdir, "callgraph.tsv"))
names = load_names(namemap)
# program address -> recovered name, for body substitution
prog_names = {p: names[r["file"]] for p, r in index.items() if r["file"] in names}
os.makedirs(os.path.join(out, "buckets"), exist_ok=True)
n = 0
for bucket in sorted(os.listdir(os.path.join(outdir, "buckets"))):
if not bucket.endswith(".c"):
continue
src = os.path.join(outdir, "buckets", bucket)
dst = os.path.join(out, "buckets", bucket)
with open(src) as f:
text = f.read()
def annotate(m):
nonlocal n
prog = int(m.group("addr"), 16)
file_addr = int(m.group("file"), 16)
size = int(m.group("size"))
rec = index.get(prog, {})
name = names.get(file_addr, rec.get("name", f"FUN_{prog:08x}"))
cl = sorted(set(callees.get(file_addr, [])))
cr = sorted(set(callers.get(file_addr, [])))
lines = [
"/* ------------------------------------------------------------------",
f" * {name}",
f" * file_addr=libcuda+0x{file_addr:x} size={size}",
]
if cr:
shown = [label(a, names, index) for a in cr[:8]]
more = f" (+{len(cr)-8} more)" if len(cr) > 8 else ""
lines.append(" * called by: " + "; ".join(shown) + more)
if cl:
shown = [label(a, names, index) for a in cl[:8]]
more = f" (+{len(cl)-8} more)" if len(cl) > 8 else ""
lines.append(" * calls: " + "; ".join(shown) + more)
lines.append(" *----------------------------------------------------------------- */")
n += 1
return "\n" + "\n".join(lines) + "\n"
text = re.sub(
r"\n/\* -+\n \* (?P<name>[^\n]*)\n \* addr=(?P<addr>0x[0-9a-f]+) "
r"file_addr=libcuda\+(?P<file>0x[0-9a-f]+) size=(?P<size>\d+)\n"
r" \*-+ \*/",
annotate, text)
if prog_names:
def sub(m):
a = int(m.group(1), 16)
return prog_names.get(a, m.group(0))
text = re.sub(r"\bFUN_([0-9a-f]{8})\b", sub, text)
with open(dst, "w") as f:
f.write(text)
print(f"[annotate_corpus] annotated {n} functions -> {out}/buckets")
print(f"[annotate_corpus] names applied: {len(names)}")
if __name__ == "__main__":
main()
|
