1from .. import engines
2from ..errors import SimError, AngrError, AngrExplorationTechniqueError
3
4
5def condition_to_lambda(condition, default=False):
6    """
7    Translates an integer, set, list or function into a lambda that checks if state's current basic block matches
8    some condition.
9
10    :param condition:   An integer, set, list or lambda to convert to a lambda.
11    :param default:     The default return value of the lambda (in case condition is None). Default: false.
12
13    :returns:           A tuple of two items: a lambda that takes a state and returns the set of addresses that it
14                        matched from the condition, and a set that contains the normalized set of addresses to stop
15                        at, or None if no addresses were provided statically.
16    """
17    if condition is None:
18        condition_function = lambda state: default
19        static_addrs = set()
20
21    elif isinstance(condition, int):
22        return condition_to_lambda((condition,))
23
24    elif isinstance(condition, (tuple, set, list)):
25        static_addrs = set(condition)
26        def condition_function(state):
27            if state.addr in static_addrs:
28                # returning {state.addr} instead of True to properly handle find/avoid conflicts
29                return {state.addr}
30
31            if not isinstance(state.project.factory.default_engine, engines.vex.VEXLifter):
32                return False
33
34            try:
35                # If the address is not in the set (which could mean it is
36                # not at the top of a block), check directly in the blocks
37                # (Blocks are repeatedly created for every check, but with
38                # the IRSB cache in angr lifter it should be OK.)
39                return static_addrs.intersection(set(state.block().instruction_addrs))
40            except (AngrError, SimError):
41                return False
42
43    elif hasattr(condition, '__call__'):
44        condition_function = condition
45        static_addrs = None
46    else:
47        raise AngrExplorationTechniqueError("ExplorationTechnique is unable to convert given type (%s) to a callable condition function." % condition.__class__)
48
49    return condition_function, static_addrs
50