angr 9.2.192__cp311-cp311-macosx_10_12_x86_64.whl

angr/analyses/proximity_graph.py

@@ -0,0 +1,444 @@
+from __future__ import annotations
+from typing import TYPE_CHECKING
+import logging
+
+import networkx
+
+from angr import ailment
+from angr.ailment.block_walker import AILBlockViewer
+from angr.codenode import BlockNode
+from angr.sim_variable import SimMemoryVariable
+from angr.knowledge_plugins.functions import Function
+from .analysis import Analysis, AnalysesHub
+
+if TYPE_CHECKING:
+    from angr.analyses.decompiler.decompiler import Decompiler
+    from angr.knowledge_plugins.cfg import CFGModel
+    from angr.knowledge_plugins.xrefs import XRefManager
+
+_l = logging.getLogger(name=__name__)
+
+
+class ProxiNodeTypes:
+    """
+    Node Type Enums
+    """
+
+    Empty = 0
+    String = 1
+    Function = 2
+    FunctionCall = 3
+    Integer = 4
+    Unknown = 5
+    Variable = 6
+
+
+class BaseProxiNode:
+    """
+    Base class for all nodes in a proximity graph.
+    """
+
+    def __init__(self, type_: int, ref_at: set[int] | None = None):
+        self.type_ = type_
+        self.ref_at = ref_at
+
+    def __eq__(self, other):
+        return isinstance(other, BaseProxiNode) and other.type_ == self.type_ and self.ref_at == other.ref_at
+
+    def __hash__(self):
+        return hash((BaseProxiNode, self.type_))
+
+
+class FunctionProxiNode(BaseProxiNode):
+    """
+    Proximity node showing current and expanded function calls in graph.
+    """
+
+    def __init__(self, func, ref_at: set[int] | None = None):
+        super().__init__(ProxiNodeTypes.Function, ref_at=ref_at)
+        self.func = func
+
+    def __eq__(self, other):
+        return isinstance(other, FunctionProxiNode) and other.type_ == self.type_ and self.func == other.func
+
+    def __hash__(self):
+        return hash((FunctionProxiNode, self.func))
+
+
+class VariableProxiNode(BaseProxiNode):
+    """
+    Variable arg node
+    """
+
+    def __init__(self, addr, name, ref_at: set[int] | None = None):
+        super().__init__(ProxiNodeTypes.Variable, ref_at=ref_at)
+        self.addr = addr
+        self.name = name
+
+    def __eq__(self, other):
+        return isinstance(other, VariableProxiNode) and other.type_ == self.type_ and self.addr == other.addr
+
+    def __hash__(self):
+        return hash((VariableProxiNode, self.addr))
+
+
+class StringProxiNode(BaseProxiNode):
+    """
+    String arg node
+    """
+
+    def __init__(self, addr, content, ref_at: set[int] | None = None):
+        super().__init__(ProxiNodeTypes.String, ref_at=ref_at)
+        self.addr = addr
+        self.content = content
+
+    def __eq__(self, other):
+        return isinstance(other, StringProxiNode) and other.type_ == self.type_ and self.addr == other.addr
+
+    def __hash__(self):
+        return hash((StringProxiNode, self.addr))
+
+
+class CallProxiNode(BaseProxiNode):
+    """
+    Call node
+    """
+
+    def __init__(self, callee, ref_at: set[int] | None = None, args: tuple[BaseProxiNode] | None = None):
+        super().__init__(ProxiNodeTypes.FunctionCall, ref_at=ref_at)
+        self.callee = callee
+        self.args = args
+
+    def __eq__(self, other):
+        return (
+            isinstance(other, CallProxiNode)
+            and other.type_ == self.type_
+            and self.callee == other.callee
+            and self.args == other.args
+            and self.ref_at == other.ref_at
+        )
+
+    def __hash__(self):
+        return hash((CallProxiNode, self.callee, self.args))
+
+
+class IntegerProxiNode(BaseProxiNode):
+    """
+    Int arg node
+    """
+
+    def __init__(self, value: int, ref_at: set[int] | None = None):
+        super().__init__(ProxiNodeTypes.Integer, ref_at=ref_at)
+        self.value = value
+
+    def __eq__(self, other):
+        return isinstance(other, IntegerProxiNode) and self.type_ == other.type_ and self.value == other.value
+
+    def __hash__(self):
+        return hash((IntegerProxiNode, self.value))
+
+
+class UnknownProxiNode(BaseProxiNode):
+    """
+    Unknown arg node
+    """
+
+    def __init__(self, dummy_value: str):
+        super().__init__(ProxiNodeTypes.Unknown)
+        self.dummy_value = dummy_value
+
+    def __eq__(self, other):
+        return (
+            isinstance(other, UnknownProxiNode) and self.type_ == other.type_ and self.dummy_value == other.dummy_value
+        )
+
+    def __hash__(self):
+        return hash((UnknownProxiNode, self.dummy_value))
+
+
+class ProximityGraphAnalysis(Analysis):
+    """
+    Generate a proximity graph.
+    """
+
+    def __init__(
+        self,
+        func: Function,
+        cfg_model: CFGModel,
+        xrefs: XRefManager,
+        decompilation: Decompiler | None = None,
+        expand_funcs: set[int] | None = None,
+    ):
+        self._function = func
+        self._cfg_model = cfg_model
+        self._xrefs = xrefs
+        self._decompilation = decompilation
+        self._expand_funcs = expand_funcs.copy() if expand_funcs else None
+
+        self.graph: networkx.DiGraph | None = None
+        self.handled_stmts = []
+
+        self._work()
+
+    def _condense_blank_nodes(self, graph: networkx.DiGraph) -> None:
+        nodes = list(graph.nodes)
+        blank_nodes: list[BaseProxiNode] = []
+
+        for node in nodes:
+            if isinstance(node, BaseProxiNode) and node.type_ == ProxiNodeTypes.Empty:
+                blank_nodes.append(node)
+            else:
+                if blank_nodes:
+                    self._merge_nodes(graph, blank_nodes)
+                    blank_nodes = []
+
+        if blank_nodes:
+            self._merge_nodes(graph, blank_nodes)
+
+    def _merge_nodes(self, graph: networkx.DiGraph, nodes: list[BaseProxiNode]) -> None:
+        for node in nodes:
+            predecessors = set(graph.predecessors(node))
+            successors = set(graph.successors(node))
+
+            for pred in predecessors:
+                for succ in successors:
+                    edge_data = graph.get_edge_data(pred, node) or {}
+                    graph.add_edge(pred, succ, **edge_data)
+
+            graph.remove_node(node)
+
+    def _work(self):
+        self.graph = networkx.DiGraph()
+
+        # initial function
+        func_proxi_node = FunctionProxiNode(self._function)
+
+        if not self._decompilation:
+            to_expand = self._process_function(self._function, self.graph, func_proxi_node=func_proxi_node)
+        else:
+            to_expand = self._process_decompilation(
+                self.graph, decompilation=self._decompilation, func_proxi_node=func_proxi_node
+            )
+
+        for func_node in to_expand:
+            if self._expand_funcs:
+                self._expand_funcs.discard(func_node.func.addr)
+
+            subgraph = networkx.DiGraph()
+            dec = self._decompilation.project.analyses.Decompiler(func_node.func, cfg=self._decompilation._cfg)
+            if not dec:
+                sub_expand = self._process_function(func_node.func, subgraph, func_proxi_node=func_node)
+            else:
+                sub_expand = self._process_decompilation(subgraph, decompilation=dec, func_proxi_node=func_node)
+
+            to_expand.extend(sub_expand)
+
+            # merge subgraph into the original graph
+            self.graph.add_nodes_from(subgraph.nodes())
+            self.graph.add_edges_from(subgraph.edges())
+
+        # condense blank nodes after the graph has been constructed
+        self._condense_blank_nodes(self.graph)
+
+    def _endnode_connector(self, func: Function, subgraph: networkx.DiGraph):
+        """
+        Properly connect expanded function call's to proximity graph.
+        """
+        current_node = None
+        successors = []
+        # Get successor node of the current function node
+        for node in self.graph.nodes():
+            if isinstance(node, FunctionProxiNode) and node.func == func:
+                current_node = node
+                successors = list(self.graph.succ[node])
+                # Remove edge FunctionProxiNode->successors
+                for succ in successors:
+                    self.graph.remove_edge(node, succ)
+                break
+
+        if successors:
+            # add edges subgraph_end_nodes->successor
+            end_nodes = [n for n in subgraph.nodes() if subgraph.in_degree(n) >= 1 and subgraph.out_degree(n) == 0]
+
+            # handle subgraphs that are empty
+            if not end_nodes:
+                end_nodes.append(current_node)
+            for end_node in end_nodes:
+                for succ in successors:
+                    subgraph.add_edge(end_node, succ)
+
+    def _process_function(
+        self, func: Function, graph: networkx.DiGraph, func_proxi_node: FunctionProxiNode | None = None
+    ) -> list[FunctionProxiNode]:
+        to_expand: list[FunctionProxiNode] = []
+        found_blocks: dict[BlockNode, BaseProxiNode] = {}
+
+        # function calls
+        for n_ in func.nodes:
+            if isinstance(n_, Function):
+                func_node = n_
+                for block, _, data in func.transition_graph.in_edges(func_node, data=True):
+                    if "ins_addr" in data:
+                        if (
+                            self._expand_funcs and func_node.addr in self._expand_funcs
+                        ):  # pylint:disable=unsupported-membership-test
+                            node = FunctionProxiNode(func_node, ref_at={data["ins_addr"]})
+                            to_expand.append(node)
+                        else:
+                            node = CallProxiNode(func_node, ref_at={data["ins_addr"]})
+                        found_blocks[block] = node
+
+        # subgraph check - do before in case of recursion
+        subgraph = self.graph != graph
+
+        for edge in func.graph.edges:
+            nodes = ()
+            for block in edge:
+                if block in found_blocks:
+                    nodes += (found_blocks[block],)
+                else:
+                    nodes += (BaseProxiNode(ProxiNodeTypes.Empty, {block.addr}),)
+            graph.add_edge(*nodes)
+
+        # Append FunctionProxiNode before Graph
+        root_node = [n for n, d in graph.in_degree() if d == 0]
+        if root_node:
+            graph.add_edge(func_proxi_node, root_node[0])
+
+        # Draw edge from subgraph endnodes to current node's successor
+        if subgraph:
+            self._endnode_connector(func_proxi_node.func, graph)
+
+        return to_expand
+
+    # TODO look into harnessing ailblock_walker to find Strings, Constants, or Function Calls
+    def _arg_handler(self, arg, args, string_refs):
+        if isinstance(arg, ailment.Expr.Const):
+            # is it a reference to a string?
+            if arg.value in self._cfg_model.memory_data:
+                md = self._cfg_model.memory_data[arg.value]
+                if md.sort == "string":
+                    # Yes!
+                    args.append(StringProxiNode(arg.value, md.content))
+                    string_refs.add(arg.value)
+            else:
+                # not a string. present it as a constant integer
+                args.append(IntegerProxiNode(arg.value, None))
+        elif isinstance(arg, ailment.expression.Load):
+            if arg.variable is not None:
+                args.append(VariableProxiNode(arg.variable.addr, arg.variable.name))
+            elif arg.addr.variable is not None:
+                if isinstance(arg.addr.variable, SimMemoryVariable):
+                    args.append(VariableProxiNode(arg.addr.variable.addr, arg.addr.variable.name))
+                else:
+                    args.append(VariableProxiNode(None, arg.addr.variable.name))
+            else:
+                args.append(UnknownProxiNode("l_"))
+        elif isinstance(arg, ailment.expression.StackBaseOffset):
+            if arg.variable is not None:
+                args.append(VariableProxiNode(arg.variable, arg.variable.name))
+            else:
+                args.append(UnknownProxiNode("s_"))
+        else:
+            args.append(UnknownProxiNode("_"))
+
+    def _process_decompilation(
+        self, graph: networkx.DiGraph, decompilation: Decompiler, func_proxi_node: FunctionProxiNode | None = None
+    ) -> list[FunctionProxiNode]:
+        to_expand: list[FunctionProxiNode] = []
+
+        # dedup
+        string_refs: set[int] = set()
+
+        # Walk the clinic structure to dump string references and function calls
+        ail_graph = decompilation.clinic.cc_graph
+
+        def _handle_Call(
+            stmt_idx: int, stmt: ailment.Stmt.Call, block: ailment.Block | None  # pylint:disable=unused-argument
+        ):  # pylint:disable=unused-argument
+            if isinstance(stmt.target, ailment.Expr.Const) and self.kb.functions.contains_addr(stmt.target.value):
+                func_node = self.kb.functions[stmt.target.value]
+                ref_at = {stmt.tags["ins_addr"]}
+
+                # extract arguments
+                args = []
+                if stmt.args:
+                    for arg in stmt.args:
+                        self._arg_handler(arg, args, string_refs)
+
+                if (
+                    self._expand_funcs and func_node.addr in self._expand_funcs
+                ):  # pylint:disable=unsupported-membership-test
+                    new_node = FunctionProxiNode(func_node, ref_at=ref_at)
+                    if new_node not in to_expand:
+                        to_expand.append(new_node)
+                else:
+                    new_node = CallProxiNode(func_node, ref_at=ref_at, args=tuple(args) if args is not None else None)
+
+                # stmt has been properly handled, add the proxi node to the list
+                self.handled_stmts.append(new_node)
+
+        # This should have the same functionality as the previous handler
+        def _handle_CallExpr(
+            expr_idx: int,
+            expr: ailment.Stmt.Call,
+            stmt_idx: int,  # pylint:disable=unused-argument
+            stmt: ailment.Stmt.Statement,  # pylint:disable=unused-argument
+            block: ailment.Block | None,
+        ):  # pylint:disable=unused-argument
+            _handle_Call(stmt_idx, expr, block)
+
+        # subgraph check - do before in case of recursion
+        subgraph = self.graph != graph
+
+        # Keep all default handlers, but overwrite necessary ones:
+        bw = AILBlockViewer()
+        bw.stmt_handlers[ailment.Stmt.Call] = _handle_Call
+        bw.expr_handlers[ailment.Stmt.Call] = _handle_CallExpr
+
+        # Custom Graph walker, go through AIL edges
+        for ail_edge in ail_graph.edges:
+            new_edge = ()
+            # Handle both blocks in edge
+            for block in ail_edge:
+                bw.walk(block)
+
+                # Check to see if our custom _handle's already dealt with the nodes
+                if self.handled_stmts:
+                    # Add each handled stmt to the graph in a linear fashion []->[]->[]
+                    # Linear is fine since stmts come from the same block and are handled in order of appearance
+                    for idx, current in enumerate(self.handled_stmts):
+                        # Skip Head of List to avoid out of bounds error
+                        if idx > 0:
+                            graph.add_edge(self.handled_stmts[idx - 1], current)
+                    # If first block in edge, LAST handled node -> next node
+                    # Else (second block in edge), prev node -> FIRST handled node
+                    proxi_node = self.handled_stmts[-1] if block == ail_edge[0] else self.handled_stmts[0]
+
+                    # Clear the handled stmts for next block
+                    self.handled_stmts = []
+
+                # Block stmts weren't handled, make a simple empty proxi node
+                else:
+                    proxi_node = BaseProxiNode(ProxiNodeTypes.Empty, {block.addr})
+
+                # Create edge from block_1 (proxi_node_1) -> block_2 (proxi_node_2)
+                # both blocks may have multiple proxi nodes, but these are already added to the graph
+                new_edge += (proxi_node,)
+
+            # Add the new edge for this pair of blocks
+            graph.add_edge(*new_edge)
+
+        # Append FunctionProxiNode before Graph
+        root_node = [n for n, d in graph.in_degree() if d == 0]
+        if root_node:
+            graph.add_edge(func_proxi_node, root_node[0])
+
+        # Draw edge from subgraph endnodes to current node's successor
+        if subgraph:
+            self._endnode_connector(func_proxi_node.func, graph)
+
+        return to_expand
+
+
+AnalysesHub.register_default("Proximity", ProximityGraphAnalysis)

angr/analyses/purity/__init__.py

@@ -0,0 +1,15 @@
+from __future__ import annotations
+
+__all__ = [
+    "AILPurityAnalysis",
+    "AILPurityDataSource",
+    "AILPurityDataUsage",
+    "AILPurityResultType",
+]
+
+from .analysis import AILPurityAnalysis
+from .engine import (
+    ResultType as AILPurityResultType,
+    DataSource as AILPurityDataSource,
+    DataUsage as AILPurityDataUsage,
+)

angr/analyses/purity/analysis.py

@@ -0,0 +1,78 @@
+from __future__ import annotations
+from collections.abc import Callable
+
+from angr import ailment
+from angr.analyses import Analysis
+from angr.analyses.decompiler.clinic import Clinic
+from angr.analyses.forward_analysis.forward_analysis import ForwardAnalysisForClinic
+from angr.errors import AngrAnalysisError
+from angr.knowledge_plugins.functions.function import Function
+from .engine import PurityEngineAIL, ResultType, StateType
+
+
+class AILPurityAnalysis(Analysis, ForwardAnalysisForClinic[StateType]):
+    """
+    Determined how data sources are used in a given function.
+    """
+
+    def __init__(
+        self,
+        subject: str | int | Function | Clinic,
+        memo: dict[Function, ResultType | None] | None = None,
+        lifter: Callable[[Function], Clinic] | None = None,
+    ):
+        if memo is None:
+            memo = {}
+        self.memo = memo
+        self._lifter = lifter or self._lift_default
+        if isinstance(subject, (str, int)):
+            subject = self.kb.functions[subject]
+        if isinstance(subject, Function):
+            subject = self._lifter(subject)
+        self.engine = PurityEngineAIL(self.project, subject, self._recurse)
+        self.result = ResultType()
+        ForwardAnalysisForClinic.__init__(self, allow_merging=True, clinic=subject)
+
+        memo[subject.function] = None
+        self._analyze()
+        memo[subject.function] = self.result
+
+    def _lift_default(self, function: Function) -> Clinic:
+        return self.project.analyses.Clinic(function)
+
+    def _recurse(self, function: Function) -> ResultType:
+        if function not in self.memo:
+            self.project.analyses[AILPurityAnalysis].prep(kb=self.kb)(function, self.memo, self._lifter)
+        r = self.memo[function]
+        if r is None:
+            raise AngrAnalysisError("Cannot do purity for recursive call tree")
+        return r
+
+    def _initial_abstract_state(self, node: ailment.Block) -> StateType:
+        return self.engine.initial_state(node)
+
+    def _step_node(self, node, state: StateType) -> tuple[bool | None, StateType]:
+        state = state.copy()
+        state.addr = (node.addr, node.idx)
+        result = self.engine.process(state, block=node)
+        self.result.update(result)
+        return None, state
+
+    def _handle_phi(self, node, succ, state, dst_vvar, src_vvar):
+        state.vars[dst_vvar.varid] |= state.vars[src_vvar.varid]
+
+    def _merge_states(self, node, *states) -> tuple[StateType, bool]:
+        changed = False
+        result = states[0]
+        for state in states[1:]:
+            for varid, val in state.vars.items():
+                oval = result.vars[varid]
+                if val - oval:
+                    changed = True
+                    oval |= val
+
+        assert self.engine.clinic.graph is not None
+        return result, changed
+
+    def _compare_states(self, node, old_state, new_state) -> bool:
+        return old_state.vars == new_state.vars