Python angr.BP_BEFORE Examples

def enter_prologue_callback(self, state):
        """
        this is a bp on the first instruction of the prologue gadget
        we add a bp over call instruction to handle future indirect call
        TODO: BUG here: we did not consider the indirect jump
        :param state:
        :return:
        """
        self.reach_current_prologue_entry = True
        print(colorama.Fore.RED + 'enter prologue gadget' + colorama.Style.RESET_ALL)
        if not self.is_dfs_search_routine:
            state.inspect.remove_breakpoint("call", self.first_fork_site_bp)
            print('[+] removed the call bp at the first fork site..')
        #
        self.bp_enforce_prologue_to_copy_to_user = state.inspect.b("call", when=angr.BP_BEFORE
                                                                   , action=self.enforce_prologue_to_copy_to_user)
        print('[+] enforced a bp on call for disclosure')
        #import IPython; IPython.embed()
        return 

def prep_tracer(state):
        if state.has_plugin("zen_plugin"):
            zen_plugin = state.get_plugin("zen_plugin")
        else:
            zen_plugin = ZenPlugin()

        state.register_plugin("zen_plugin", zen_plugin)
        state.inspect.b(
            'reg_write',
            angr.BP_BEFORE,
            action=zen_register_write
        )
        state.inspect.b(
            'mem_write',
            angr.BP_BEFORE,
            action=zen_memory_write
        )

        # setup the byte dict
        byte_dict = zen_plugin.byte_dict
        for i, b in enumerate(state.cgc.flag_bytes):
            var = list(b.variables)[0]
            byte_dict[var] = {i}

        state.preconstrainer.preconstraints.extend(zen_plugin.preconstraints) 

def __init__(self):
        #self.cfg = self.project.analyses.CFGAccurate(keep_state=True, enable_symbolic_back_traversal=True, normalize=True)
        self.cfg = self.project.analyses.CFGFast(collect_data_references=True, normalize=True)
        self.accesses = {}
        self.ty_backend = TypeBackend()
        self.global_struct = SimStructAbstract(label='global')
        #self.identer = Identifier(self.project, self.cfg)
        #self.ident_result = list(self.identer.run())

        self.struct_mapping = {}
        self.struct_base = 0x40000000
        self.pass_results = []
        self.function_initial_regs = {}
        self.function_return_vals = {}

        self.syscall_mapping = self.project._simos.syscall_table

        self.initial_state = self.project.factory.blank_state(remove_options={o.SIMPLIFY_MEMORY_WRITES, o.SIMPLIFY_REGISTER_WRITES}, add_options={o.UNSUPPORTED_BYPASS_ZERO_DEFAULT, o.BYPASS_UNSUPPORTED_IROP, o.BYPASS_UNSUPPORTED_IRCCALL, o.AVOID_MULTIVALUED_READS, o.AVOID_MULTIVALUED_WRITES})
        self.initial_state.inspect.b('mem_read', when=angr.BP_AFTER, action=self._memory_access)
        self.initial_state.inspect.b('mem_write', when=angr.BP_AFTER, action=self._memory_access)
        self.initial_state.inspect.b('exit', when=angr.BP_BEFORE, action=self._exit_taken)

        for func in self.real_functions(self.cfg):
            l.info("Working on %s", func.name)
            self._init_analysis(func) 

def arm(self, state):
        """
        Setup hooks and breakpoints to perform Spectre gadget vulnerability detection.
        Also set up concretization to ensure addresses are always made to be OOB when possible.
        """
        state.inspect.b('mem_read',  when=angr.BP_AFTER, condition=_tainted_read, action=detected_spectre_read)
        state.inspect.b('mem_write', when=angr.BP_AFTER, condition=_tainted_write, action=detected_spectre_write)
        state.inspect.b('exit', when=angr.BP_BEFORE, condition=_tainted_branch, action=detected_spectre_branch)

        state.memory.read_strategies.insert(0, OOBStrategy())
        state.memory.write_strategies.insert(0, OOBStrategy())
        state.inspect.b('address_concretization', when=angr.BP_AFTER, condition=concretization_succeeded, action=log_concretization)

        state.options.add(angr.options.SYMBOL_FILL_UNCONSTRAINED_MEMORY)
        state.options.add(angr.options.SYMBOL_FILL_UNCONSTRAINED_REGISTERS)
        state.options.add(angr.options.SYMBOLIC_INITIAL_VALUES)
        state.options.add(angr.options.SPECIAL_MEMORY_FILL)
        state._special_memory_filler = oob_memory_fill
        state.options.add(angr.options.SYMBOLIC_WRITE_ADDRESSES)

        self._armed = True 

def instrument_forking_gadget(self, state, forking_gadget):
        """
        :param state:  current execution state
        :param bloom_gadget: bloom gadget
        :param forking_gadget: forking gadget
        :param bloom_state:  old state serilized when step() found a good bloom_state
        :return: None
        """
        fork_entry, first_fork_site, second_fork_site = self.get_forking_gadget_entry_and_sites(forking_gadget)
        # the problem here is that when rax is symbolic,
        # mov rax, qword ptr [rax + 0x80] will concretize this shit

        # init the state plugin to keep track of fork site
        state.register_plugin('osokplugin', angr.state_plugins.OsokPlugin(False, False, False))
        self.reach_current_bloom_site = True
        state.osokplugin.reach_bloom_site = True  # we have already reached the bloom site

        # add instrumentation via breakpoint
        # TODO handle indirect jump such as call rdx

        #state.inspect.b('mem_read', when=angr.BP_BEFORE, action=self.track_reads)
        #state.inspect.b('mem_write', when=angr.BP_BEFORE, action=self.track_writes)

        state.inspect.b('instruction', when=angr.BP_BEFORE, instruction=fork_entry, action=self.enter_fork_callback)
        # we want to first verify if we can reach both of the forking site,
        # so we should zero out both of the call stubs
        if not self.b.is_hooked(first_fork_site):
            instSize = self.getInstructionLengthByAddr(first_fork_site)
            self.b.hook(first_fork_site, self.reach_first_fork_site_callback, instSize)
        if not self.b.is_hooked(second_fork_site):
            instSize = self.getInstructionLengthByAddr(second_fork_site)
            self.b.hook(second_fork_site, self.reach_second_fork_site_callback, instSize)
        return 

def instrument_bloom(self, state, bloom_entry, bloom_site, \
                                             bloom_gadget, symbolic_regs=['rdi']):
        #state.inspect.b('instruction', when=angr.BP_BEFORE\
        #state.inspect.b('call', when=angr.BP_BEFORE\
                #, instruction = bloom_site\
                #, action=self.call_check_bloom_regs)
                #, action=self.check_bloom_regs)
        state.inspect.b('mem_read', when=angr.BP_BEFORE, action=self.track_reads)
        state.inspect.b('mem_write', when=angr.BP_BEFORE, action=self.track_writes)
        return 

def prep_tracer(state, format_infos=None):
        format_infos = [] if format_infos is None else format_infos
        state.inspect.b(
            'exit',
            angr.BP_BEFORE,
            action=exit_hook
        )
        state.inspect.b(
            'syscall',
            angr.BP_AFTER,
            action=syscall_hook
        )
        state.inspect.b(
            'constraints',
            angr.BP_BEFORE,
            action=constraint_hook
        )

        if state.has_plugin("chall_resp_info"):
            chall_resp_plugin = state.get_plugin("chall_resp_info")
        else:
            chall_resp_plugin = ChallRespInfo()
        for f in format_infos:
            chall_resp_plugin.format_infos[f.addr] = f

        state.register_plugin("chall_resp_info", chall_resp_plugin)

        for addr in chall_resp_plugin.format_infos:
            state.project.hook(addr, generic_info_hook, length=0)


# THE ZEN HOOK 

def test_inspect_exit():
    class counts: #pylint:disable=no-init
        exit_before = 0
        exit_after = 0

    def handle_exit_before(state):
        counts.exit_before += 1
        exit_target = state.inspect.exit_target
        nose.tools.assert_equal(state.solver.eval(exit_target), 0x3f8)
        # change exit target
        state.inspect.exit_target = 0x41414141
        nose.tools.assert_equal(state.inspect.exit_jumpkind, "Ijk_Boring")
        nose.tools.assert_true(state.inspect.exit_guard.is_true())

    def handle_exit_after(state): #pylint:disable=unused-argument
        counts.exit_after += 1

    s = SimState(arch="AMD64", mode="symbolic")
    irsb = pyvex.IRSB(b"\x90\x90\x90\x90\xeb\x0a", mem_addr=1000, arch=archinfo.ArchAMD64())

    # break on exit
    s.inspect.b('exit', BP_BEFORE, action=handle_exit_before)
    s.inspect.b('exit', BP_AFTER, action=handle_exit_after)

    # step it
    succ = HeavyVEXMixin(None).process(s, irsb=irsb).flat_successors

    # check
    nose.tools.assert_equal( succ[0].solver.eval(succ[0].ip), 0x41414141)
    nose.tools.assert_equal(counts.exit_before, 1)
    nose.tools.assert_equal(counts.exit_after, 1) 

def test_inspect_syscall():
    class counts: #pylint:disable=no-init
        exit_before = 0
        exit_after = 0

    def handle_syscall_before(state):
        counts.exit_before += 1
        syscall_name = state.inspect.syscall_name
        nose.tools.assert_equal(syscall_name, "close")

    def handle_syscall_after(state):
        counts.exit_after += 1
        syscall_name = state.inspect.syscall_name
        nose.tools.assert_equal(syscall_name, "close")

    s = SimState(arch="AMD64", mode="symbolic")
    # set up to call so syscall close
    s.regs.rax = 3
    s.regs.rdi = 2

    # break on syscall
    s.inspect.b('syscall', BP_BEFORE, action=handle_syscall_before)
    s.inspect.b('syscall', BP_AFTER, action=handle_syscall_after)

    # step it
    proc = SIM_PROCEDURES['posix']['close'](is_syscall=True)
    ProcedureEngine(None).process(s, procedure=proc, ret_to=s.ip)

    # check counts
    nose.tools.assert_equal(counts.exit_before, 1)
    nose.tools.assert_equal(counts.exit_after, 1) 

def test_inspect_engine_process():
    p = angr.Project(os.path.join(os.path.dirname(os.path.realpath(__file__)), '..', '..', 'binaries', 'tests', 'x86_64', 'fauxware'))
    constraints = []
    def check_first_symbolic_fork(state):
        succs = state.inspect.sim_successors.successors
        succ_addr = [hex(s.addr) for s in succs]
        nose.tools.assert_equal(len(succ_addr), 2)
        nose.tools.assert_in('0x400692L', succ_addr)
        nose.tools.assert_in('0x400699L', succ_addr)
        print('Fork after:', hex(state.addr))
        print('Successors:', succ_addr)

    def check_second_symbolic_fork(state):
        succs = state.inspect.sim_successors.successors
        succ_addr = [hex(s.addr) for s in succs]
        nose.tools.assert_equal(len(succ_addr), 2)
        nose.tools.assert_in('0x4006dfL', succ_addr)
        nose.tools.assert_in('0x4006e6L', succ_addr)
        print('Fork after:', hex(state.addr))
        print('Successors:', succ_addr)

    def first_symbolic_fork(state):
        return hex(state.addr) == '0x40068eL' \
           and isinstance(state.inspect.sim_engine, HeavyVEXMixin)
        # TODO: I think this latter check is meaningless with the eleventh hour refactor

    def second_symbolic_fork(state):
        return hex(state.addr) == '0x4006dbL' \
           and isinstance(state.inspect.sim_engine, HeavyVEXMixin)

    def check_state(state):
        nose.tools.assert_in(hex(state.inspect.sim_successors.addr), ('0x40068eL', '0x4006dbL'))

    state = p.factory.entry_state(addr=p.loader.find_symbol('main').rebased_addr)
    pg = p.factory.simulation_manager(state)
    state.inspect.b('engine_process',
                    when=BP_BEFORE,
                    action=check_state,
                    condition=first_symbolic_fork)
    state.inspect.b('engine_process',
                    when=BP_AFTER,
                    action=check_first_symbolic_fork,
                    condition=first_symbolic_fork)
    pg.run()

    state = p.factory.entry_state(addr=p.loader.find_symbol('main').rebased_addr)
    pg = p.factory.simulation_manager(state)
    state.inspect.b('engine_process',
                    when=BP_BEFORE,
                    action=check_state,
                    condition=second_symbolic_fork)
    state.inspect.b('engine_process',
                    when=BP_AFTER,
                    action=check_second_symbolic_fork,
                    condition=second_symbolic_fork)
    pg.run()