Computes cyclomatic complexity of all functions in a binary

cyclomatic.py

import bap
import networkx as nx

def build_cfg(sub):
    G = nx.DiGraph()
    entry = sub.blks[0].id.number
    G.add_node(entry)
    for blk in sub.blks:
        for jmp in blk.jmps:
            if jmp.constr == 'Goto' and jmp.target.constr == 'Direct':
                G.add_edge(blk.id.number, jmp.target.arg.number)
    return {'entry': entry, 'graph': G}


def connect_graph(cfg):
    G = cfg['graph']
    for node in G.nodes:
        degree = len(list(G.successors(node)))
        if not degree:
            G.add_edge(node, cfg['entry'])


def complexity(G):
    edges = len(G.edges())
    nodes = len(G.nodes())
    parts = nx.components.number_strongly_connected_components(G)
    return edges - nodes + parts


def main():
    proj = bap.run('echo')
    for sub in proj.program.subs:
        cfg = build_cfg(sub)
        connect_graph(cfg)
        print('{}:{}'.format(sub.name, complexity(cfg['graph'])))


if __name__ == '__main__':
    main()