building a miniaturized control flow graph

cfg.sh

#!/bin/sh
r2 -qc "e anal.bb.maxsize=2048;s sym.$2;af;agfj" "$1" | python3 control-flow-graph.py

control-flow-graph.py

# building a miniaturized control flow graph with Radare (https://rada.re).
#
# usage: r2 -qc "s sym.$function;af;agfj" "$binary" | python3 control-flow-graph.py
#                              .------.  .-----------.
#          .-------------------|------|--|-.         |
#     .----|-.  .--------------|------|--|-|---------|----.
#.----|----|-|--|--------------|------|--|-|---------|    |
#*--->*--->* '->*--->*-.->*--->*-.->*-+->*-'->*-.->*-'->* '->*--->*
#                      '--'      |    '---------|--|---------|----'
#                                '--------------|--|---------'
#                                               '--'

from json import load
from sys import stdin

d = load(stdin)

blocks = d[0]['blocks']

# list of blocks and the blocks they link to
l1 = \
[ ( block['offset']
  , [ block[key]
      for key in ['jump', 'fail']
      if key in block
    ]
  )
  for block in blocks
]

# mapping from block offset to list offset
m = \
{ value: index
  for (index, (value, _)) in enumerate(l1)
}

# simplified list of blocks and the blocks they link to
# by list offset
l2 = \
[ [m[x] for x in xs]
  for (_, xs) in l1
]

# whether or not this block links to the next
next_link = \
[ i + 1 in xs
  for (i, xs) in enumerate(l2)
]

# list of blocks (zero or one) after the next that this block links to
forward = \
[ list(filter(lambda x: x > i + 1, xs))
  for (i, xs) in enumerate(l2)
]

# list of blocks (zero or one) prior to the next that this block links to
backward = \
[ list(filter(lambda x: x < i + 1, xs))
  for (i, xs) in enumerate(l2)
]


def pairs(xs):
    return [
        (i, ys[0])
        for (i, ys) in enumerate(xs)
        if ys
    ]


def bar(xs):
    x = 1
    while x in xs:
        x = x + 1
    return x


def forward_bars(xs, pair):
    return xs + [
        ( i, y, bar([
            z for (_, l, z) in xs
            if l > i
        ]) )
        for (i, y) in [pair]
    ]


def backward_bars(xs, pair):
    return xs + [
        ( i, y, bar([
            z for (_, l, z) in xs
            if l <= i
        ]) )
        for (i, y) in [pair]
    ]


from functools import reduce

forward = pairs(forward)
forward_bars = reduce(forward_bars, forward, [])

backward = list(reversed(pairs(backward)))
backward_bars = reduce(backward_bars, backward, [])

ymin = max([0] + [z for (_, _, z) in backward_bars])
ymax = max([0] + [z for (_, _, z) in forward_bars])

width = 1 + 5 * (len(next_link) - 1)
height = sum([ymin, 1, ymax])


first = lambda t: t[0]
second = lambda t: t[1]
third = lambda t: t[2]

cmp = lambda x, y: [-1, 0, 1][int(x == y) + 2 * int(x > y)]

backward_ends = lambda x: x // 5 + 1 in map(second, backward_bars)
contains_next_link = lambda x: next_link[x // 5]
forward_ends = lambda x: x // 5 + 1 in map(second, forward_bars)
any_terminal = lambda x: any([
    backward_ends(x), 
    contains_next_link(x),
    forward_ends(x),
])


def render(i):
    (x, y) = (i % width, i // width)
    y = (height - 1 - ymin) - y

    if y < 0:
        return [
            "|",
            "'",
            "-" if [
                () for (a, b, c) in backward_bars
                if 5 * a > x > 5 * b - 3 and c == -y
            ] else " ",
        ][1 + cmp(-y, max([0] + [
            c for (a, b, c) in backward_bars
            if 5 * a == x or 5 * b - 3 == x
        ]))]
    if y == 0:
        return [
            "*",
            "-" if contains_next_link(x) else " ",
            [" ", ".", "-", ".", "'", "|", "'", "+"][
                backward_ends(x) + 2 * contains_next_link(x) + 4 * forward_ends(x)
            ],
            "-" if any_terminal(x) else " ",
            ">" if any_terminal(x) else " ",
        ][x % 5]
    if y > 0:
        return [
            "|",
            ".",
            "-" if [
                () for (a, b, c) in forward_bars
                if 5 * a < x < 5 * b - 3 and c == y
            ] else " ",
        ][1 + cmp(y, max([0] + [
            c for (a, b, c) in forward_bars
            if 5 * a == x or 5 * b - 3 == x
        ]))]


output = [
    render(i)
    for i in range(width * height)
]

[
    print(''.join(
        output[width * y:width * (y+1)]
    ))
    for y in range(height)
]