angr 9.2.135__py3-none-win_amd64.whl → 9.2.136__py3-none-win_amd64.whl

angr/knowledge_plugins/functions/function.py

@@ -1,11 +1,14 @@
+# pylint:disable=too-many-boolean-expressions
 from __future__ import annotations
 import os
 import logging
-import networkx
 import itertools
 from collections import defaultdict
 from collections.abc import Iterable
+import contextlib
+from typing import overload
 
+import networkx
 from itanium_demangler import parse
 
 from cle.backends.symbol import Symbol
@@ -13,7 +16,6 @@ from archinfo.arch_arm import get_real_address_if_arm
 import claripy
 
 from angr.knowledge_plugins.cfg.memory_data import MemoryDataSort
-
 from angr.codenode import CodeNode, BlockNode, HookNode, SyscallNode
 from angr.serializable import Serializable
 from angr.errors import AngrValueError, SimEngineError, SimMemoryError
@@ -22,14 +24,12 @@ from angr.procedures.definitions import SimSyscallLibrary
 from angr.protos import function_pb2
 from angr.calling_conventions import DEFAULT_CC, default_cc
 from angr.misc.ux import deprecated
-from .function_parser import FunctionParser
-
-l = logging.getLogger(name=__name__)
-
 from angr.sim_type import SimTypeFunction, parse_defns
 from angr.calling_conventions import SimCC
 from angr.project import Project
-import contextlib
+from .function_parser import FunctionParser
+
+l = logging.getLogger(name=__name__)
 
 
 class Function(Serializable):
@@ -41,7 +41,6 @@ class Function(Serializable):
         "_addr_to_block_node",
         "_argument_registers",
         "_argument_stack_variables",
-        "_block_cache",
         "_block_sizes",
         "_call_sites",
         "_callout_sites",
@@ -158,7 +157,6 @@ class Function(Serializable):
         # function, it may not be removed from _addr_to_block_node. if you want to list
         # all blocks of a function, access .blocks.
         self._block_sizes = {}  # map addresses to block sizes
-        self._block_cache = {}  # a cache of real, hard data Block objects
         self._local_blocks = {}  # a dict of all blocks inside the function
         self._local_block_addrs = set()  # a set of addresses of all blocks inside the function
 
@@ -374,13 +372,6 @@ class Function(Serializable):
         :param byte_string:
         :return:
         """
-        if addr in self._block_cache:
-            b = self._block_cache[addr]
-            if size is None or b.size == size:
-                return b
-            # size seems to be updated. remove this cached entry from the block cache
-            del self._block_cache[addr]
-
         if size is None and addr in self.block_addrs:
             # we know the size
             size = self._block_sizes[addr]
@@ -389,7 +380,6 @@ class Function(Serializable):
         if size is None:
             # update block_size dict
             self._block_sizes[addr] = block.size
-        self._block_cache[addr] = block
         return block
 
     # compatibility
@@ -440,7 +430,7 @@ class Function(Serializable):
 
     @classmethod
     def _get_cmsg(cls):
-        return function_pb2.Function()
+        return function_pb2.Function()  # pylint:disable=no-member
 
     def serialize_to_cmessage(self):
         return FunctionParser.serialize(self)
@@ -613,7 +603,6 @@ class Function(Serializable):
         d["_local_transition_graph"] = None
         d["_project"] = None
         d["_function_manager"] = None
-        d["_block_cache"] = {}
         return d
 
     @property
@@ -642,7 +631,7 @@ class Function(Serializable):
 
     @property
     def size(self):
-        return sum([b.size for b in self.blocks])
+        return sum(b.size for b in self.blocks)
 
     @property
     def binary(self):
@@ -675,7 +664,7 @@ class Function(Serializable):
         dec = self.project.analyses.Decompiler(self, cfg=self._function_manager._kb.cfgs.get_most_accurate())
         return dec.codegen.text
 
-    def add_jumpout_site(self, node):
+    def add_jumpout_site(self, node: CodeNode):
         """
         Add a custom jumpout site.
 
@@ -683,11 +672,11 @@ class Function(Serializable):
         :return:        None
         """
 
-        self._register_nodes(True, node)
+        node = self._register_node(True, node)
         self._jumpout_sites.add(node)
         self._add_endpoint(node, "transition")
 
-    def add_retout_site(self, node):
+    def add_retout_site(self, node: CodeNode):
         """
         Add a custom retout site.
 
@@ -703,7 +692,7 @@ class Function(Serializable):
         :return:     None
         """
 
-        self._register_nodes(True, node)
+        node = self._register_node(True, node)
         self._retout_sites.add(node)
         self._add_endpoint(node, "return")
 
@@ -785,7 +774,6 @@ class Function(Serializable):
         return None
 
     def _clear_transition_graph(self):
-        self._block_cache = {}
         self._block_sizes = {}
         self._addr_to_block_node = {}
         self._local_blocks = {}
@@ -812,11 +800,13 @@ class Function(Serializable):
 
         # it's confirmed. register the node if needed
         if "outside" not in data or data["outside"] is False:
-            self._register_nodes(True, dst)
+            dst = self._register_node(True, dst)
 
         self.transition_graph[src][dst]["confirmed"] = True
 
-    def _transit_to(self, from_node, to_node, outside=False, ins_addr=None, stmt_idx=None, is_exception=False):
+    def _transit_to(
+        self, from_node: CodeNode, to_node, outside=False, ins_addr=None, stmt_idx=None, is_exception=False
+    ):
         """
         Registers an edge between basic blocks in this function's transition graph.
         Arguments are CodeNode objects.
@@ -830,16 +820,15 @@ class Function(Serializable):
         """
 
         if outside:
-            self._register_nodes(True, from_node)
+            from_node = self._register_node(True, from_node)
             if to_node is not None:
-                self._register_nodes(False, to_node)
+                to_node = self._register_node(False, to_node)
 
             self._jumpout_sites.add(from_node)
         else:
+            from_node = self._register_node(True, from_node)
             if to_node is not None:
-                self._register_nodes(True, from_node, to_node)
-            else:
-                self._register_nodes(True, from_node)
+                to_node = self._register_node(True, to_node)
 
         type_ = "transition" if not is_exception else "exception"
         if to_node is not None:
@@ -871,19 +860,22 @@ class Function(Serializable):
         :type  ins_addr:    int or None
         """
 
-        self._register_nodes(True, from_node)
+        from_node = self._register_node(True, from_node)
 
         if to_func.is_syscall:
             self.transition_graph.add_edge(from_node, to_func, type="syscall", stmt_idx=stmt_idx, ins_addr=ins_addr)
         else:
             self.transition_graph.add_edge(from_node, to_func, type="call", stmt_idx=stmt_idx, ins_addr=ins_addr)
             if ret_node is not None:
+                ret_node = self._register_node(return_to_outside is False, ret_node)
                 self._fakeret_to(from_node, ret_node, to_outside=return_to_outside)
 
         self._local_transition_graph = None
 
     def _fakeret_to(self, from_node, to_node, confirmed=None, to_outside=False):
-        self._register_nodes(True, from_node)
+        from_node = self._register_node(True, from_node)
+        if confirmed:
+            to_node = self._register_node(not to_outside, to_node)
 
         if confirmed is None:
             self.transition_graph.add_edge(from_node, to_node, type="fake_return", outside=to_outside)
@@ -891,8 +883,6 @@ class Function(Serializable):
             self.transition_graph.add_edge(
                 from_node, to_node, type="fake_return", confirmed=confirmed, outside=to_outside
             )
-            if confirmed:
-                self._register_nodes(not to_outside, to_node)
 
         self._local_transition_graph = None
 
@@ -910,45 +900,56 @@ class Function(Serializable):
         self._local_transition_graph = None
 
     def _update_local_blocks(self, node: CodeNode):
-        self._local_blocks[node.addr] = node
-        self._local_block_addrs.add(node.addr)
+        if node.addr not in self._local_blocks or self._local_blocks[node.addr] != node:
+            self._local_blocks[node.addr] = node
+            self._local_block_addrs.add(node.addr)
 
     def _update_addr_to_block_cache(self, node: BlockNode):
         if node.addr not in self._addr_to_block_node:
             self._addr_to_block_node[node.addr] = node
 
-    def _register_nodes(self, is_local, *nodes):
-        if not isinstance(is_local, bool):
-            raise AngrValueError('_register_nodes(): the "is_local" parameter must be a bool')
-
-        for node in nodes:
-            if node.addr not in self:
-                # only add each node once
-                self.transition_graph.add_node(node)
-
-            if not isinstance(node, CodeNode):
-                continue
-            node._graph = self.transition_graph
-            if self._block_sizes.get(node.addr, 0) == 0:
-                self._block_sizes[node.addr] = node.size
-            if node.addr == self.addr and (self.startpoint is None or not self.startpoint.is_hook):
-                self.startpoint = node
-            if is_local and node.addr not in self._local_blocks:
-                self._update_local_blocks(node)
-            # add BlockNodes to the addr_to_block_node cache if not already there
-            if isinstance(node, BlockNode):
-                self._update_addr_to_block_cache(node)
-                # else:
-                #    # checks that we don't have multiple block nodes at a single address
-                #    assert node == self._addr_to_block_node[node.addr]
-
-    def _add_return_site(self, return_site):
+    @overload
+    def _register_node(self, is_local: bool, node: CodeNode) -> CodeNode: ...
+
+    @overload
+    def _register_node(self, is_local: bool, node: Function) -> Function: ...
+
+    def _register_node(self, is_local: bool, node: CodeNode | Function) -> CodeNode | Function:
+        # if the node already exists and is the same, we reuse the existing node
+        if is_local and self._local_blocks.get(node.addr, None) == node:
+            return self._local_blocks[node.addr]
+
+        if node.addr not in self and node not in self.transition_graph:
+            # only add each node to the graph once
+            self.transition_graph.add_node(node)
+
+        if not isinstance(node, CodeNode):
+            # function and other things bail here
+            return node
+
+        # this is either a new node or a different node at the same address
+        node._graph = self.transition_graph
+        if self._block_sizes.get(node.addr, 0) == 0:
+            self._block_sizes[node.addr] = node.size
+        if node.addr == self.addr and (self.startpoint is None or not self.startpoint.is_hook):
+            self.startpoint = node
+        if is_local and node.addr not in self._local_blocks:
+            self._update_local_blocks(node)
+        # add BlockNodes to the addr_to_block_node cache if not already there
+        if isinstance(node, BlockNode):
+            self._update_addr_to_block_cache(node)
+            # else:
+            #    # checks that we don't have multiple block nodes at a single address
+            #    assert node == self._addr_to_block_node[node.addr]
+        return node
+
+    def _add_return_site(self, return_site: CodeNode):
         """
         Registers a basic block as a site for control flow to return from this function.
 
-        :param CodeNode return_site:     The block node that ends with a return.
+        :param return_site:     The block node that ends with a return.
         """
-        self._register_nodes(True, return_site)
+        return_site = self._register_node(True, return_site)
 
         self._ret_sites.add(return_site)
         # A return site must be an endpoint of the function - you cannot continue execution of the current function
@@ -1285,8 +1286,8 @@ class Function(Serializable):
         """
         Draw the graph and save it to a PNG file.
         """
-        import matplotlib.pyplot as pyplot  # pylint: disable=import-error
-        from networkx.drawing.nx_agraph import graphviz_layout  # pylint: disable=import-error
+        import matplotlib.pyplot as pyplot  # pylint: disable=import-error,import-outside-toplevel
+        from networkx.drawing.nx_agraph import graphviz_layout  # pylint: disable=import-error,import-outside-toplevel
 
         tmp_graph = networkx.classes.digraph.DiGraph()
         for from_block, to_block in self.transition_graph.edges():
@@ -1416,11 +1417,8 @@ class Function(Serializable):
                     self._local_blocks[n.addr] = new_node
 
                 # update block_cache and block_sizes
-                if (n.addr in self._block_cache and self._block_cache[n.addr].size != new_node.size) or (
-                    n.addr in self._block_sizes and self._block_sizes[n.addr] != new_node.size
-                ):
+                if n.addr in self._block_sizes and self._block_sizes[n.addr] != new_node.size:
                     # the cache needs updating
-                    self._block_cache.pop(n.addr, None)
                     self._block_sizes[n.addr] = new_node.size
 
                 for p, _, data in original_predecessors:
@@ -1680,7 +1678,6 @@ class Function(Serializable):
         func.startpoint = self.startpoint
         func._addr_to_block_node = self._addr_to_block_node.copy()
         func._block_sizes = self._block_sizes.copy()
-        func._block_cache = self._block_cache.copy()
         func._local_blocks = self._local_blocks.copy()
         func._local_block_addrs = self._local_block_addrs.copy()
         func.info = self.info.copy()

angr/knowledge_plugins/functions/function_manager.py

@@ -6,6 +6,7 @@ from collections.abc import Generator
 import logging
 import collections.abc
 import re
+import weakref
 from sortedcontainers import SortedDict
 
 import networkx
@@ -31,7 +32,7 @@ class FunctionDict(SortedDict):
     """
 
     def __init__(self, backref, *args, **kwargs):
-        self._backref = backref
+        self._backref = weakref.proxy(backref) if backref is not None else None
         self._key_types = kwargs.pop("key_types", int)
         super().__init__(*args, **kwargs)
 
@@ -39,7 +40,7 @@ class FunctionDict(SortedDict):
         try:
             return super().__getitem__(addr)
         except KeyError as ex:
-            if not isinstance(addr, self._key_types):
+            if isinstance(addr, bool) or not isinstance(addr, self._key_types):
                 raise TypeError(f"FunctionDict only supports {self._key_types} as key type") from ex
 
             if isinstance(addr, SootMethodDescriptor):
@@ -148,7 +149,7 @@ class FunctionManager(KnowledgeBasePlugin, collections.abc.Mapping):
         dst_func = self._function_map[function_addr]
         if syscall in (True, False):
             dst_func.is_syscall = syscall
-        dst_func._register_nodes(True, node)
+        dst_func._register_node(True, node)
         self.block_map[node.addr] = node
 
     def _add_call_to(
@@ -160,7 +161,7 @@ class FunctionManager(KnowledgeBasePlugin, collections.abc.Mapping):
         syscall=None,
         stmt_idx=None,
         ins_addr=None,
-        return_to_outside=False,
+        return_to_outside: bool = False,
     ):
         """
         Add a call to a function.
@@ -172,7 +173,7 @@ class FunctionManager(KnowledgeBasePlugin, collections.abc.Mapping):
         :param bool syscall:        If this is a call to a syscall or not.
         :param int stmt_idx:        ID of the statement where this call happens.
         :param int ins_addr:        Address of the instruction where this call happens.
-        :param bool return_to_outside:  True if the return of the call is considered going to outside of the current
+        :param return_to_outside:  True if the return of the call is considered going to outside of the current
                                         function.
         :return:                    None
         """
@@ -306,7 +307,7 @@ class FunctionManager(KnowledgeBasePlugin, collections.abc.Mapping):
         try:
             _ = self[item]
             return True
-        except KeyError:
+        except (KeyError, TypeError):
             return False
 
     def __getitem__(self, k) -> Function:
@@ -406,7 +407,7 @@ class FunctionManager(KnowledgeBasePlugin, collections.abc.Mapping):
 
         try:
             prev_addr = self._function_map.floor_addr(addr)
-            return self._function_map[prev_addr]
+            return self._function_map.get(prev_addr)
 
         except KeyError:
             return None

angr/knowledge_plugins/functions/function_parser.py

@@ -246,7 +246,7 @@ class FunctionParser:
         # add leftover blocks
         for block in blocks.values():
             if block not in added_nodes:
-                obj._register_nodes(True, block)
+                obj._register_node(True, block)
 
         return obj
 

angr/knowledge_plugins/functions/soot_function.py

@@ -108,24 +108,24 @@ class SootFunction(Function):
         # The Shimple CFG is already normalized.
         pass
 
-    def _register_nodes(self, is_local, *nodes):
-        if not isinstance(is_local, bool):
-            raise AngrValueError('_register_nodes(): the "is_local" parameter must be a bool')
+    def _register_node(self, is_local: bool, node):
+        if is_local and self._local_blocks.get(node.addr) == node:
+            return self._local_blocks[node.addr]
 
-        for node in nodes:
+        if node not in self.transition_graph:
             self.transition_graph.add_node(node)
-            node._graph = self.transition_graph
-            if node.addr not in self or self._block_sizes[node.addr] == 0:
-                self._block_sizes[node.addr] = node.size
-            if node.addr == self.addr.addr and (self.startpoint is None or not self.startpoint.is_hook):
-                self.startpoint = node
-            if is_local:
-                self._local_blocks[node.addr] = node
-                self._local_block_addrs.add(node.addr)
-            # add BlockNodes to the addr_to_block_node cache if not already there
-            if isinstance(node, BlockNode) and node.addr not in self._addr_to_block_node:
-                self._addr_to_block_node[node.addr] = node
+        node._graph = self.transition_graph
+        if node.addr not in self or self._block_sizes[node.addr] == 0:
+            self._block_sizes[node.addr] = node.size
+        if node.addr == self.addr.addr and (self.startpoint is None or not self.startpoint.is_hook):
+            self.startpoint = node
+        if is_local:
+            self._local_blocks[node.addr] = node
+            self._local_block_addrs.add(node.addr)
+        # add BlockNodes to the addr_to_block_node cache if not already there
+        if isinstance(node, BlockNode) and node.addr not in self._addr_to_block_node:
+            self._addr_to_block_node[node.addr] = node
+        return node
 
 
 from angr.codenode import BlockNode
-from angr.errors import AngrValueError

angr/knowledge_plugins/propagations/propagation_model.py

@@ -6,7 +6,7 @@ import claripy
 import ailment
 from angr.serializable import Serializable
 from angr.knowledge_plugins.functions.function import Function
-from .states import PropagatorVEXState, PropagatorAILState, PropagatorState
+from .states import PropagatorVEXState, PropagatorState
 
 
 class PropagationModel(Serializable):
@@ -65,10 +65,9 @@ class PropagationModel(Serializable):
         preds = [self.states[pnode.addr] for pnode in self._function.graph.predecessors(node)]
         if not preds:
             if isinstance(node, ailment.Block):
-                state = PropagatorAILState.initial_state(self._function.project, func_addr=self._function.addr)
-            else:
-                state = PropagatorVEXState.initial_state(self._function.project, func_addr=self._function.addr)
-                state.store_register(state.arch.ip_offset, state.arch.bytes, claripy.BVV(block_addr, state.arch.bits))
+                raise NotImplementedError
+            state = PropagatorVEXState.initial_state(self._function.project, func_addr=self._function.addr)
+            state.store_register(state.arch.ip_offset, state.arch.bytes, claripy.BVV(block_addr, state.arch.bits))
         else:
             state, _ = preds[0].merge(*preds[1:])
         return state