Disassembly Control Flow Graph

cfg.py

# usage:
# otool -tv program | python3 graph.py | dot -Tsvg -o program.svg; open program.svg
# (objdump will need a slightly different parser)

import sys

def split(items, before=None, after=None):
    result = []
    for item in items:
        if before and before(item):
            yield result
            result = []
        result.append(item)
        if after and after(item):
            yield result
            result = []

def mk_block(instrs):
    return {'addr': instrs[0][0], 'instrs': instrs, 'succ': []}

lines = [line.strip().split() for line in sys.stdin]
lines = [l for l in lines if len(l) >= 2 and 'section' not in line]
lines = [[int(a, 16), *r] for a, *r in lines]

targets = {int(i[2], 16) for i in lines if 'j' in i[1]}

blocks = [mk_block(b) for b in split(
    lines,
    before=lambda i: i[0] in targets,
    after=lambda i: 'j' in i[1] or 'ret' in i[1],
) if b]

block_map = {b['addr']: b for b in blocks}

instr_names = {l[1] for l in lines}

for i, block in enumerate(blocks):
    tail_instr = block['instrs'][-1]
    tail_op = tail_instr[1]

    if tail_op.startswith('j'):
        block['succ'].append((int(tail_instr[2], 16), 'jmp'))
    if tail_op == 'jmp' or tail_op.startswith('ret'):
        continue
    if i + 1 < len(blocks):
        block['succ'].append((blocks[i + 1]['addr'], 'next'))

print('digraph G {')
print('node [shape=box, fontname="fira code"]')

for block in blocks:
    instrs = ['{:12} {}\\l'.format(a, ' '.join('{:8}'.format(x) for x in b))
              for _, a, *b in block['instrs']]
    code = ''.join(instrs)
    print('{} [label="{}"]'.format(block['addr'], code))

for block in blocks:
    for dest, reason in block['succ']:
        print('{} -> {} [taillabel="{}"]'.format(block['addr'], dest, reason))

print('}')