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
|
#!/usr/bin/env python3
"""Dump real basic-block CFGs to JSON, so graph-layout work can be done offline.
Run with a python that has ``idapro`` (i.e. the worker python, /usr/bin/python3):
/usr/bin/python3 experiments/cfg_dump.py targets/echo -o /tmp/cfg-echo.json
Copies the binary to a scratch dir first (never touches the tracked .i64), opens
it with idalib, and writes one record per function:
{"name","ea","blocks":[{"id","start","end","lines":[str],"succs":[[id,kind]]}]}
``kind`` is "fall" (falls through to the next address), "jump" (unconditional or
taken branch) or "switch" (n-way). That is exactly the input a graph view needs;
everything after this point (layering, ordering, routing, rendering) is pure
python and needs no IDA at all.
"""
from __future__ import annotations
import argparse
import json
import os
import shutil
import sys
import tempfile
def dump(path: str, want: list[str], max_blocks: int) -> list[dict]:
import idapro
scratch = tempfile.mkdtemp(prefix="cfgdump-")
local = os.path.join(scratch, os.path.basename(path))
shutil.copy2(path, local)
if idapro.open_database(local, run_auto_analysis=True) != 0:
raise SystemExit(f"failed to open {local}")
import ida_funcs
import ida_gdl
import ida_lines
import ida_bytes
import idautils
import idaapi
out = []
try:
for fea in idautils.Functions():
fn = ida_funcs.get_func(fea)
if not fn:
continue
name = ida_funcs.get_func_name(fea)
if want and not any(w in name for w in want):
continue
fc = ida_gdl.FlowChart(fn, flags=ida_gdl.FC_PREDS)
blocks = []
index = {}
for i, bb in enumerate(fc):
index[bb.start_ea] = i
for bb in fc:
lines = []
ea = bb.start_ea
while ea < bb.end_ea and ea != idaapi.BADADDR:
txt = ida_lines.tag_remove(
ida_lines.generate_disasm_line(ea, 0) or "")
lines.append(txt.rstrip())
nxt = ida_bytes.next_head(ea, bb.end_ea)
if nxt <= ea:
break
ea = nxt
succs = []
sl = list(bb.succs())
for s in sl:
if s.start_ea not in index:
continue
if len(sl) > 2:
kind = "switch"
elif s.start_ea == bb.end_ea:
kind = "fall"
else:
kind = "jump"
succs.append([index[s.start_ea], kind])
blocks.append({
"id": index[bb.start_ea],
"start": bb.start_ea,
"end": bb.end_ea,
"lines": lines,
"succs": succs,
})
if max_blocks and len(blocks) > max_blocks:
continue
out.append({"name": name, "ea": fea, "blocks": blocks})
finally:
idapro.close_database(save=False)
shutil.rmtree(scratch, ignore_errors=True)
return out
def main() -> int:
ap = argparse.ArgumentParser()
ap.add_argument("binary")
ap.add_argument("-o", "--out", default="/tmp/cfg.json")
ap.add_argument("-f", "--func", action="append", default=[],
help="only functions whose name contains this (repeatable)")
ap.add_argument("--max-blocks", type=int, default=0,
help="skip functions with more blocks than this")
args = ap.parse_args()
recs = dump(os.path.abspath(args.binary), args.func, args.max_blocks)
recs.sort(key=lambda r: len(r["blocks"]), reverse=True)
with open(args.out, "w") as f:
json.dump(recs, f)
tot = sum(len(r["blocks"]) for r in recs)
print(f"{len(recs)} functions, {tot} blocks -> {args.out}")
for r in recs[:12]:
print(f" {len(r['blocks']):4d} blocks {r['name']}")
return 0
if __name__ == "__main__":
sys.exit(main())
|