angr cfgfast, cfgemulated path finder

angr_find_paths.py

#!python3

import angr
import networkx
from pathlib import Path


def main(args):
    
    p = angr.Project(args.binary, load_options={
        'auto_load_libs': False
    })

    print(f"base address is 0x{p.loader.main_object.mapped_base:x}")

    print(f"constructing {args.cfg_strategy} cfg...")
    if args.cfg_strategy == 'fast':
        cfg: angr.analyses.cfg.cfg_fast.CFGFast = p.analyses.CFGFast(show_progressbar=True)
        cfg.normalize()
    elif args.cfg_strategy == 'emulated':
        cfg: angr.analyses.cfg.cfg_emulated.CFGEmulated = p.analyses.CFGEmulated(show_progressbar=True)
    else:
        raise Exception

    print(f"nodes: {len(cfg.graph.nodes())}")
    print(f"edges: {len(cfg.graph.edges())}")

    if args.add_base_to_address:
        source_address = p.loader.main_object.mapped_base + args.source
        target_address = p.loader.main_object.mapped_base + args.target
    else:
        source_address = args.source
        target_address = args.target

    sources = cfg.get_all_nodes(source_address)
    if not sources:
        print(cfg.graph.nodes())
        print(f"can't find source: 0x{source_address:x}")
        return -1
    elif len(sources) > 1:
        print(f"there are more than 1 node about source: 0x{source_address:x}")

    source = sources[-1]

    targets = cfg.get_all_nodes(target_address)
    if not targets:
        print(cfg.graph.nodes())
        print(f"can't find source: 0x{target_address:x}")
        return -1
    elif len(targets) > 1:
        print(f"there are more than 1 node about source: 0x{target_address:x}")

    target = targets[-1]

    print("searching...")
    if args.iterate:
        try:
            for path in networkx.all_simple_paths(cfg.graph, source, target):
                print("*-- found! " + '-' * 30)
                print(path)
                input("press ENTER to find next paths.")
        except KeyboardInterrupt:
            print("keyboard interrupt detected.")
    else:
        paths = list(networkx.all_simple_paths(cfg.graph, source, target))
        print(f"has {len(paths)} paths.")


if __name__ == '__main__':

    hexint = lambda x: int(x, 16)

    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("binary", type=Path)
    parser.add_argument("-b", "--add-base-to-address", action='store_true')
    parser.add_argument("-i", "--iterate", action='store_true')
    parser.add_argument("-s", "--source", type=hexint, required=True)
    parser.add_argument("-t", "--target", type=hexint, required=True)
    parser.add_argument("--cfg-strategy", default='fast', choices=['fast', 'emulated'])
    args = parser.parse_args()
    main(args)