angr 9.2.166__cp310-abi3-manylinux_2_28_x86_64.whl

angr/analyses/forward_analysis/visitors/function_graph.py

@@ -0,0 +1,86 @@
+from __future__ import annotations
+import logging
+
+import networkx
+
+from angr.utils.graph import dfs_back_edges, GraphUtils
+from .graph import GraphVisitor, NodeType
+
+_l = logging.getLogger(__name__)
+
+
+class FunctionGraphVisitor(GraphVisitor):
+    """
+    :param knowledge.Function func:
+    """
+
+    def __init__(self, func, graph=None):
+        super().__init__()
+        self.function = func
+
+        if graph is None:
+            self.graph = self.function.graph
+        else:
+            self.graph = graph
+
+        self.reset()
+
+    def resume_with_new_graph(self, graph: networkx.DiGraph) -> bool:
+        """
+        We can only reasonably reuse existing results if the node index of the already traversed nodes are the same as
+        the ones from the new graph. Otherwise, we always restart.
+
+        :return:    True if we are resuming, False if reset() is called.
+        """
+        # update the graph
+        self.graph = graph
+
+        must_restart = False
+        new_sorted_nodes = list(self.sort_nodes())
+        # check if new sorted_nodes is an extension of the existing sorted_nodes
+        if len(new_sorted_nodes) < len(self._sorted_nodes):
+            must_restart = True
+        else:
+            must_restart = new_sorted_nodes[: len(self._sorted_nodes)] != self._sorted_nodes
+
+        if must_restart:
+            _l.debug("Cannot resume for function %r with the new graph.", self.function)
+            self.reset()
+            return False
+
+        # update related data structures
+        for i, n in enumerate(new_sorted_nodes):
+            if i >= len(self._sorted_nodes):
+                self._node_to_index[n] = i
+        # update worklist and nodes_set
+        for n in new_sorted_nodes[len(self._sorted_nodes) :]:
+            self.revisit_node(n)
+        # update sorted_nodes in the end
+        self._sorted_nodes = new_sorted_nodes
+
+        return True
+
+    def successors(self, node):
+        return list(self.graph.successors(node))
+
+    def predecessors(self, node):
+        return list(self.graph.predecessors(node))
+
+    def sort_nodes(self, nodes=None):
+        sorted_nodes = GraphUtils.quasi_topological_sort_nodes(self.graph)
+
+        if nodes is not None:
+            sorted_nodes = [n for n in sorted_nodes if n in set(nodes)]
+
+        return sorted_nodes
+
+    def back_edges(self) -> list[tuple[NodeType, NodeType]]:
+        start_nodes = [node for node in self.graph if node.addr == self.function.addr]
+        if not start_nodes:
+            start_nodes = [node for node in self.graph if self.graph.in_degree(node) == 0]
+
+        if not start_nodes:
+            raise NotImplementedError
+
+        start_node = start_nodes[0]
+        return list(dfs_back_edges(self.graph, start_node))

angr/analyses/forward_analysis/visitors/graph.py

@@ -0,0 +1,242 @@
+from __future__ import annotations
+from typing import TypeVar, Generic
+from collections.abc import Collection, Iterator
+from collections import defaultdict
+
+from angr.utils.algo import binary_insert
+
+NodeType = TypeVar("NodeType")
+
+
+class GraphVisitor(Generic[NodeType]):
+    """
+    A graph visitor takes a node in the graph and returns its successors. Typically, it visits a control flow graph,
+    and returns successors of a CFGNode each time. This is the base class of all graph visitors.
+    """
+
+    __slots__ = (
+        "_back_edges_by_dst",
+        "_back_edges_by_src",
+        "_node_to_index",
+        "_nodes_set",
+        "_pending_nodes",
+        "_reached_fixedpoint",
+        "_sorted_nodes",
+        "_worklist",
+    )
+
+    def __init__(self):
+        self._sorted_nodes: list[NodeType] = []  # a list of sorted nodes. do not change until we get a new graph
+        self._worklist: list[NodeType] = []  # a list of nodes that the analysis should work on and finally exhaust
+        self._nodes_set: set[NodeType] = set()
+        self._node_to_index: dict[NodeType, int] = {}
+        self._reached_fixedpoint: set[NodeType] = set()
+        self._back_edges_by_src: dict[NodeType, set[NodeType]] | None = None
+        self._back_edges_by_dst: dict[NodeType, set[NodeType]] | None = None
+
+        self._pending_nodes: dict[NodeType, set[NodeType]] = defaultdict(set)
+
+    #
+    # Interfaces
+    #
+
+    def successors(self, node: NodeType) -> list[NodeType]:
+        """
+        Get successors of a node. The node should be in the graph.
+
+        :param node: The node to work with.
+        :return:     A list of successors.
+        :rtype:      list
+        """
+
+        raise NotImplementedError
+
+    def predecessors(self, node: NodeType) -> list[NodeType]:
+        """
+        Get predecessors of a node. The node should be in the graph.
+
+        :param node: The node to work with.
+        :return:     A list of predecessors.
+        """
+
+        raise NotImplementedError
+
+    def sort_nodes(self, nodes: Collection[NodeType] | None = None) -> list[NodeType]:
+        """
+        Get a list of all nodes sorted in an optimal traversal order.
+
+        :param iterable nodes: A collection of nodes to sort. If none, all nodes in the graph will be used to sort.
+        :return:               A list of sorted nodes.
+        """
+
+        raise NotImplementedError
+
+    def back_edges(self) -> list[tuple[NodeType, NodeType]]:
+        """
+        Get a list of back edges. This function is optional. If not overridden, the traverser cannot achieve an optimal
+        graph traversal order.
+
+        :return:                A list of back edges (source -> destination).
+        """
+        raise NotImplementedError
+
+    #
+    # Public methods
+    #
+
+    def nodes(self) -> Iterator[NodeType]:
+        """
+        Return an iterator of nodes following an optimal traversal order.
+
+        :return:
+        """
+
+        return iter(self.sort_nodes())
+
+    # Traversal
+
+    def reset(self):
+        """
+        Reset the internal node traversal state. Must be called prior to visiting future nodes.
+
+        :return: None
+        """
+
+        self._sorted_nodes.clear()
+        self._worklist.clear()
+        self._nodes_set.clear()
+        self._node_to_index.clear()
+        self._reached_fixedpoint.clear()
+
+        self._sorted_nodes = list(self.sort_nodes())
+        for i, n in enumerate(self._sorted_nodes):
+            self._node_to_index[n] = i
+            binary_insert(self._worklist, n, lambda elem: self._node_to_index[elem])
+            self._nodes_set.add(n)
+
+        self._populate_back_edges()
+
+    def next_node(self) -> NodeType | None:
+        """
+        Get the next node to visit.
+
+        :return: A node in the graph.
+        """
+
+        if not self._worklist:
+            return None
+
+        node = None
+        for idx in range(len(self._worklist)):  # pylint:disable=consider-using-enumerate
+            node_ = self._worklist[idx]
+            if node_ in self._pending_nodes:
+                if not self._pending_nodes[node_]:
+                    # this pending node is cleared - take it
+                    node = node_
+                    del self._pending_nodes[node_]
+                    del self._worklist[idx]
+                    break
+                # try the next node
+                continue
+
+            node = node_
+            del self._worklist[idx]
+            break
+
+        if node is None:
+            # all nodes are pending... we will just pick the first one
+            node = self._worklist.pop(0)
+
+        self._nodes_set.discard(node)
+
+        #  check if this node should be added to pending
+        if self._back_edges_by_dst and node in self._back_edges_by_dst:
+            for back_edge_src in self._back_edges_by_dst[node]:
+                self._pending_nodes[node].add(back_edge_src)
+
+        # check if this node is being pended on by any other node
+        if self._back_edges_by_src and node in self._back_edges_by_src:
+            for back_edge_dst in self._back_edges_by_src[node]:
+                self._pending_nodes[back_edge_dst].discard(node)
+
+        return node
+
+    def all_successors(self, node: NodeType, skip_reached_fixedpoint=False) -> set[NodeType]:
+        """
+        Returns all successors to the specific node.
+
+        :param node: A node in the graph.
+        :return:     A set of nodes that are all successors to the given node.
+        :rtype:      set
+        """
+
+        successors = set()
+
+        stack = [node]
+        while stack:
+            n = stack.pop()
+            successors.add(n)
+            stack.extend(
+                succ
+                for succ in self.successors(n)
+                if succ not in successors and (not skip_reached_fixedpoint or succ not in self._reached_fixedpoint)
+            )
+
+        return successors
+
+    def revisit_successors(self, node: NodeType, include_self=True) -> None:
+        """
+        Revisit a node in the future. As a result, the successors to this node will be revisited as well.
+
+        :param node: The node to revisit in the future.
+        :return:     None
+        """
+
+        successors = self.successors(node)  # , skip_reached_fixedpoint=True)
+
+        if include_self and node not in self._nodes_set:
+            binary_insert(self._worklist, node, lambda elem: self._node_to_index[elem])
+            self._nodes_set.add(node)
+
+        for succ in successors:
+            if succ not in self._nodes_set:
+                binary_insert(self._worklist, succ, lambda elem: self._node_to_index[elem])
+                self._nodes_set.add(succ)
+
+    def revisit_node(self, node: NodeType) -> None:
+        """
+        Revisit a node in the future. Do not include its successors immediately.
+
+        :param node:    The node to revisit in the future.
+        :return:        None
+        """
+
+        if node not in self._nodes_set:
+            binary_insert(self._worklist, node, lambda elem: self._node_to_index[elem])
+            self._nodes_set.add(node)
+
+    def reached_fixedpoint(self, node: NodeType) -> None:
+        """
+        Mark a node as reached fixed-point. This node as well as all its successors will not be visited in the future.
+
+        :param node: The node to mark as reached fixed-point.
+        :return:     None
+        """
+
+        self._reached_fixedpoint.add(node)
+
+    #
+    # Private methods
+    #
+
+    def _populate_back_edges(self):
+        try:
+            back_edges = self.back_edges()
+        except NotImplementedError:
+            return
+
+        self._back_edges_by_src = defaultdict(set)
+        self._back_edges_by_dst = defaultdict(set)
+        for src, dst in back_edges:
+            self._back_edges_by_src[src].add(dst)
+            self._back_edges_by_dst[dst].add(src)

angr/analyses/forward_analysis/visitors/loop.py

@@ -0,0 +1,29 @@
+from __future__ import annotations
+from angr.utils.graph import GraphUtils
+from .graph import GraphVisitor
+
+
+class LoopVisitor(GraphVisitor):
+    """
+    :param angr.analyses.loopfinder.Loop loop: The loop to visit.
+    """
+
+    def __init__(self, loop):
+        super().__init__()
+        self.loop = loop
+
+        self.reset()
+
+    def successors(self, node):
+        return self.loop.graph.successors(node)
+
+    def predecessors(self, node):
+        return self.loop.graph.predecessors(node)
+
+    def sort_nodes(self, nodes=None):
+        sorted_nodes = GraphUtils.quasi_topological_sort_nodes(self.loop.graph)
+
+        if nodes is not None:
+            sorted_nodes = [n for n in sorted_nodes if n in set(nodes)]
+
+        return sorted_nodes

angr/analyses/forward_analysis/visitors/single_node_graph.py

@@ -0,0 +1,38 @@
+from __future__ import annotations
+from .graph import GraphVisitor
+
+
+class SingleNodeGraphVisitor(GraphVisitor):
+    """
+    :param node: The single node that should be in the graph.
+    """
+
+    __slots__ = (
+        "node",
+        "node_returned",
+    )
+
+    def __init__(self, node):
+        super().__init__()
+        self.node = node
+        self.node_returned = False
+
+    def reset(self):
+        self.node_returned = False
+
+    def next_node(self):
+        if self.node_returned:
+            return None
+        self.node_returned = True
+        return self.node
+
+    def successors(self, node):
+        return []
+
+    def predecessors(self, node):
+        return []
+
+    def sort_nodes(self, nodes=None):
+        if nodes:
+            return nodes
+        return [self.node]

angr/analyses/identifier/__init__.py

@@ -0,0 +1,5 @@
+from __future__ import annotations
+
+from .identify import Identifier
+
+__all__ = ("Identifier",)

angr/analyses/identifier/custom_callable.py

@@ -0,0 +1,137 @@
+from __future__ import annotations
+import logging
+
+from angr.errors import AngrCallableError, AngrCallableMultistateError
+from angr.calling_conventions import default_cc
+
+
+l = logging.getLogger(name=__name__)
+
+
+class IdentifierCallable:
+    """
+    Callable is a representation of a function in the binary that can be
+    interacted with like a native python function.
+
+    If you set perform_merge=True (the default), the result will be returned to you, and
+    you can get the result state with callable.result_state.
+
+    Otherwise, you can get the resulting path group (immutable) at callable.result_path_group.
+    """
+
+    def __init__(
+        self,
+        project,
+        addr,
+        concrete_only=False,
+        perform_merge=False,
+        base_state=None,
+        toc=None,
+        cc=None,
+        max_steps=None,
+    ):
+        """
+        :param project:         The project to operate on
+        :param addr:            The address of the function to use
+
+        The following parameters are optional:
+
+        :param concrete_only:   Throw an exception if the execution splits into multiple paths
+        :param perform_merge:   Merge all result states into one at the end (only relevant if concrete_only=False)
+        :param base_state:      The state from which to do these runs
+        :param toc:             The address of the table of contents for ppc64
+        :param cc:              The SimCC to use for a calling convention
+        """
+
+        self._project = project
+        self._addr = addr
+        self._concrete_only = concrete_only
+        self._perform_merge = perform_merge
+        self._base_state = base_state
+        self._toc = toc
+        self._caller = None
+        self._cc = (
+            cc
+            if cc is not None
+            else default_cc(project.arch.name, platform=project.simos.name if project.simos is not None else None)(
+                project.arch
+            )
+        )
+        self._deadend_addr = project.simos.return_deadend
+        self._max_steps = max_steps
+
+        self.result_path_group = None
+        self.result_state = None
+
+    def set_base_state(self, state):
+        """
+        Swap out the state you'd like to use to perform the call
+        :param state: The state to use to perform the call
+        """
+        self._base_state = state
+
+    def __call__(self, *args):
+        prototype = self._cc.guess_prototype(args)
+        self.perform_call(*args, prototype=prototype)
+        if self.result_state is not None:
+            loc = self._cc.return_val(prototype.returnty)
+            return self.result_state.solver.simplify(
+                loc.get_value(self.result_state, stack_base=self.result_state.regs.sp - self._cc.STACKARG_SP_DIFF)
+            )
+        return None
+
+    def get_base_state(self, *args):
+        prototype = self._cc.guess_prototype(args)
+        self._base_state.ip = self._addr
+        return self._project.factory.call_state(
+            self._addr,
+            *args,
+            prototype=prototype,
+            cc=self._cc,
+            base_state=self._base_state,
+            ret_addr=self._deadend_addr,
+            toc=self._toc,
+        )
+
+    def perform_call(self, *args, prototype=None):
+        if prototype is None:
+            prototype = self._cc.guess_prototype(args)
+        self._base_state.ip = self._addr
+        state = self._project.factory.call_state(
+            self._addr,
+            *args,
+            cc=self._cc,
+            prototype=prototype,
+            base_state=self._base_state,
+            ret_addr=self._deadend_addr,
+            toc=self._toc,
+        )
+
+        def step_func(pg):
+            pg2 = pg.prune()
+            if len(pg2.active) > 1:
+                raise AngrCallableMultistateError("Execution split on symbolic condition!")
+            return pg2
+
+        caller = self._project.factory.simulation_manager(state)
+        for _ in range(self._max_steps):
+            if len(caller.active) == 0:
+                break
+            if caller.active[0].history.block_count > 100000:
+                l.debug("super long path %s", caller.active[0])
+                raise AngrCallableError("Super long path")
+            caller.step(step_func=step_func if self._concrete_only else None)
+        if len(caller.active) > 0:
+            raise AngrCallableError("didn't make it to the end of the function")
+
+        caller.unstash(from_stash="deadended")
+        caller.prune(filter_func=lambda pt: pt.addr == self._deadend_addr)
+
+        if len(caller.active) == 0:
+            raise AngrCallableError("No paths returned from function")
+
+        self.result_path_group = caller.copy()
+
+        if self._perform_merge:
+            caller.merge()
+        self.result_state = caller.active[0]

angr/analyses/identifier/errors.py

@@ -0,0 +1,10 @@
+from __future__ import annotations
+from angr.errors import AngrAnalysisError
+
+
+class IdentifierException(AngrAnalysisError):
+    pass
+
+
+class FunctionNotInitialized(AngrAnalysisError):
+    pass

angr/analyses/identifier/func.py

@@ -0,0 +1,60 @@
+from __future__ import annotations
+
+
+class TestData:
+    def __init__(
+        self,
+        input_args,
+        expected_output_args,
+        expected_return_val,
+        max_steps,
+        preloaded_stdin=None,
+        expected_stdout=None,
+    ):
+        assert isinstance(input_args, (list, tuple))
+        assert isinstance(expected_output_args, (list, tuple))
+        if preloaded_stdin is None:
+            preloaded_stdin = ""
+        if expected_stdout is None:
+            expected_stdout = ""
+        for i, v in enumerate(input_args):
+            if type(v) is str:
+                input_args[i] = v.encode("latin-1")  # py3k transition hack :/
+
+        self.input_args = input_args
+        self.expected_output_args = expected_output_args
+        self.expected_return_val = expected_return_val
+        self.preloaded_stdin = preloaded_stdin
+        self.expected_stdout = expected_stdout
+        self.max_steps = max_steps
+
+
+class Func:
+    def __init__(self):
+        pass
+
+    def get_name(self):
+        raise NotImplementedError
+
+    def num_args(self):
+        raise NotImplementedError
+
+    def gen_input_output_pair(self):
+        raise NotImplementedError
+
+    def var_args(self):  # pylint disable=no-self-use
+        return False
+
+    def can_call_other_funcs(self):  # pylint disable=no-self-use
+        return True
+
+    def pre_test(self, func, runner):  # pylint disable=no-self-use,unused-argument
+        """
+        custom tests run before, return False if it for sure is not the function
+        Use tests here to pick which version of the function
+        :arg func: the cfg function it will be compared against
+        :arg runner: a runner to run the tests with
+        :return: True if we should continue testing
+        """
+
+        return True

angr/analyses/identifier/functions/__init__.py

@@ -0,0 +1,37 @@
+from __future__ import annotations
+
+import importlib
+import logging
+import os
+
+from angr.analyses.identifier.func import Func
+
+l = logging.getLogger(name=__name__)
+
+
+# Import all classes under the current directory, and group them based on
+# lib names.
+Functions = {}
+path = os.path.dirname(os.path.abspath(__file__))
+skip_dirs = ["__init__.py"]
+skip_procs = ["__init__", "func"]
+
+for proc_file_name in os.listdir(path):
+    if not proc_file_name.endswith(".py"):
+        continue
+    proc_module_name = proc_file_name[:-3]
+    if proc_file_name.startswith("skip"):
+        continue
+    if proc_module_name in skip_procs:
+        continue
+
+    try:
+        proc_module = importlib.import_module(f".{proc_module_name}", "angr.analyses.identifier.functions")
+    except ImportError:
+        l.warning("Unable to import procedure %s", proc_module_name)
+        continue
+
+    for attr_name in dir(proc_module):
+        attr = getattr(proc_module, attr_name)
+        if isinstance(attr, type) and issubclass(attr, Func) and attr_name != "Func":
+            Functions[attr_name] = attr