angr: custom hook lookup

angr-stub-demo.py

#!/usr/bin/env python
"""
angr demo: stub out all functions that are not in the main object's .text
section. Because we are just stubbing out all functions, simulation may not be
very meaningful. This could be extended though, for example to provide some
mock data.

Example target.c:

/* gcc -o target target.c */
#include <stdio.h>

int internal_func(int x)
{
    return x * 2;
}

int main(int argc, char const *argv[])
{
    int x = 10;

    if (scanf("%d", &x) == 1) {
        x = internal_func(x);
    }

    printf("New value: %d", x);
    return x;
}


"""

import angr
import logging

log = logging.getLogger(__name__)
log.setLevel(logging.INFO)


class Stub(angr.SimProcedure):
    IS_FUNCTION = True
    RETURNS = True

    def run(self):
        try:
            name = self.project.kb.labels.get(self.addr)
        except KeyError:
            name = "unk"
        log.info("Hello from Stub @ %#x (%s)!", self.addr, name)
        return self.state.solver.BVS("stub", 64)


class CustomExecutionEngine(angr.engines.UberEngine):
    """
    A simple engine mixin that overrides HooksMixin::_lookup_hook to add a
    default hook on *everything* outside of the main object's .text region.
    """

    def process_successors(self, successors, **kwargs):
        log.debug(
            "Processing successors of %s, rax = %s", self.state, self.state.regs.rax
        )
        super().process_successors(successors, **kwargs)

    def _lookup_hook(self, state, procedure):
        hook = super()._lookup_hook(state, procedure)
        if hook:
            return hook
        if self._should_hook(state):
            return Stub()
        return None

    def _should_hook(self, state):
        ip = state.solver.eval(state.regs._ip)
        t = self.project.loader.main_object.sections_map[".text"]
        return not (t.min_addr <= ip < t.max_addr)


p = angr.Project("target", auto_load_libs=False, engine=CustomExecutionEngine)
main = p.kb.labels.lookup("main")
log.info("Calling main @ %#x", main)
s = p.factory.call_state(
    main,
    ret_addr=p.simos.return_deadend,
    add_options={"SYMBOL_FILL_UNCONSTRAINED_MEMORY", "SYMBOL_FILL_UNCONSTRAINED_REGISTERS"}
)
simgr = p.factory.simgr(s)
simgr.run()

log.info("Simulation Manager: %s", simgr)
for state in simgr.deadended:
    log.info("Return Value: %s", state.regs.rax)