angr directed symbolic execution

angr_directed_symexec.py

#!python3

from angr                                       import Project
from angr.sim_state                             import SimState
from angr.sim_manager                           import SimulationManager
from angr.exploration_techniques.director       import Director, ExecuteAddressGoal
from angr.exploration_techniques.veritesting    import Veritesting


hits = {}


def goal_satisfied_callback(goal: ExecuteAddressGoal, state: SimState, simgr: SimulationManager):

    global hits

    print(f"hit: 0x{goal.addr:x}, {state}")

    if goal.addr not in hits:
        hits[hex(goal.addr)] = {
            'hit': 0
        }

    hits[hex(goal.addr)]['hit'] += 1


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

    director = Director(goal_satisfied_callback=goal_satisfied_callback)
    for goal in args.goal:
        if args.add_mapped_base_to_goal:
            goal += project.loader.main_object.mapped_base
        director.add_goal(ExecuteAddressGoal(goal))

    if args.blank_state:
        addr = args.blank_state
        state_initializer = project.factory.blank_state
    elif args.call_state:
        addr = args.call_state
        state_initializer = project.factory.call_state
    elif args.entry_state:
        addr = args.entry_state
        state_initializer = project.factory.entry_state
    else:
        addr = None
        state_initializer = project.factory.entry_state

    if addr and args.add_mapped_base_to_goal:
        addr += project.loader.main_object.mapped_base

    initial_state = state_initializer(addr=addr)

    simgr = project.factory.simgr(initial_state)
    simgr.use_technique(director)
    
    try:
        if args.interactive:
            raise NotImplementedError
        else:
            out = simgr.run()
    except KeyboardInterrupt:
        print("keyboard interrupted.")

    print(f"done: {out}")
    print(__import__("json").dumps(hits))


if __name__ == '__main__':

    hexint = lambda x: int(x, 16)

    import argparse
    from pathlib import Path

    parser = argparse.ArgumentParser()

    parser.add_argument('binary', type=Path)

    state_group = parser.add_mutually_exclusive_group()
    state_group.add_argument("--blank-state", type=hexint)
    state_group.add_argument("--call-state", type=hexint)
    state_group.add_argument("--entry-state", type=hexint)

    parser.add_argument('-g', '--goal', nargs='+', required=True, type=hexint)
    parser.add_argument('-b', '--add-mapped-base-to-goal', action='store_true')
    parser.add_argument('-i', '--interactive', action='store_true')

    args = parser.parse_args()
    main(args)