angr 9.2.117__py3-none-manylinux2014_x86_64.whl → 9.2.118__py3-none-manylinux2014_x86_64.whl

angr/state_plugins/solver.py

@@ -1,3 +1,4 @@
+from __future__ import annotations
 import functools
 import time
 import logging
@@ -5,7 +6,6 @@ import os
 from typing import TypeVar, overload
 
 import claripy
-from claripy import backend_manager
 
 from angr import sim_options as o
 from angr.errors import SimValueError, SimUnsatError, SimSolverModeError, SimSolverOptionError
@@ -45,10 +45,10 @@ def timed_function(f):
                     location = "bbl {:#x}, stmt {} (inst {})".format(
                         s.scratch.bbl_addr,
                         s.scratch.stmt_idx,
-                        ("%s" % s.scratch.ins_addr if s.scratch.ins_addr is None else "%#x" % s.scratch.ins_addr),
+                        (f"{s.scratch.ins_addr}" if s.scratch.ins_addr is None else f"{s.scratch.ins_addr:#x}"),
                     )
                 elif s.scratch.sim_procedure is not None:
-                    location = "sim_procedure %s" % s.scratch.sim_procedure
+                    location = f"sim_procedure {s.scratch.sim_procedure}"
                 else:
                     location = "unknown"
             except Exception:  # pylint:disable=broad-except
@@ -61,8 +61,7 @@ def timed_function(f):
             return r
 
         return timing_guy
-    else:
-        return f
+    return f
 
 
 # pylint:disable=global-variable-undefined
@@ -110,24 +109,18 @@ def error_converter(f):
 def _concrete_bool(e):
     if isinstance(e, bool):
         return e
-    elif isinstance(e, claripy.ast.Base) and e.op == "BoolV":
+    if isinstance(e, claripy.ast.Base) and e.op == "BoolV" or isinstance(e, SimActionObject) and e.op == "BoolV":
         return e.args[0]
-    elif isinstance(e, SimActionObject) and e.op == "BoolV":
-        return e.args[0]
-    else:
-        return None
+    return None
 
 
 def _concrete_value(e):
     # shortcuts for speed improvement
     if isinstance(e, (int, float, bool)):
         return e
-    elif isinstance(e, claripy.ast.Base) and e.op in claripy.operations.leaf_operations_concrete:
-        return e.args[0]
-    elif isinstance(e, SimActionObject) and e.op in claripy.operations.leaf_operations_concrete:
+    if isinstance(e, claripy.ast.Base | SimActionObject) and e.is_leaf() and not e.symbolic:
         return e.args[0]
-    else:
-        return None
+    return None
 
 
 def concrete_path_bool(f):
@@ -136,8 +129,7 @@ def concrete_path_bool(f):
         v = _concrete_bool(args[0])
         if v is None:
             return f(self, *args, **kwargs)
-        else:
-            return v
+        return v
 
     return concrete_shortcut_bool
 
@@ -148,8 +140,7 @@ def concrete_path_not_bool(f):
         v = _concrete_bool(args[0])
         if v is None:
             return f(self, *args, **kwargs)
-        else:
-            return not v
+        return not v
 
     return concrete_shortcut_not_bool
 
@@ -160,8 +151,7 @@ def concrete_path_scalar(f):
         v = _concrete_value(args[0])
         if v is None:
             return f(self, *args, **kwargs)
-        else:
-            return v
+        return v
 
     return concrete_shortcut_scalar
 
@@ -172,8 +162,7 @@ def concrete_path_tuple(f):
         v = _concrete_value(args[0])
         if v is None:
             return f(self, *args, **kwargs)
-        else:
-            return (v,)
+        return (v,)
 
     return concrete_shortcut_tuple
 
@@ -184,8 +173,7 @@ def concrete_path_list(f):
         v = _concrete_value(args[0])
         if v is None:
             return f(self, *args, **kwargs)
-        else:
-            return [v]
+        return [v]
 
     return concrete_shortcut_list
 
@@ -270,10 +258,10 @@ class SimSolver(SimStatePlugin):
             self.eternal_tracked_variables[key] = v
         else:
             self.temporal_tracked_variables = dict(self.temporal_tracked_variables)
-            ctrkey = key + (None,)
+            ctrkey = (*key, None)
             ctrval = self.temporal_tracked_variables.get(ctrkey, 0) + 1
             self.temporal_tracked_variables[ctrkey] = ctrval
-            tempkey = key + (ctrval,)
+            tempkey = (*key, ctrval)
             self.temporal_tracked_variables[tempkey] = v
 
     def describe_variables(self, v):
@@ -302,14 +290,7 @@ class SimSolver(SimStatePlugin):
         track = o.CONSTRAINT_TRACKING_IN_SOLVER in self.state.options
         approximate_first = o.APPROXIMATE_FIRST in self.state.options
 
-        if o.STRINGS_ANALYSIS in self.state.options:
-            if o.COMPOSITE_SOLVER in self.state.options:
-                self._stored_solver = claripy.SolverComposite(
-                    template_solver_string=claripy.SolverCompositeChild(
-                        backend=backend_manager.backends.z3, track=track
-                    )
-                )
-        elif o.ABSTRACT_SOLVER in self.state.options:
+        if o.ABSTRACT_SOLVER in self.state.options:
             self._stored_solver = claripy.SolverVSA()
         elif o.SYMBOLIC in self.state.options and o.REPLACEMENT_SOLVER in self.state.options:
             self._stored_solver = claripy.SolverReplacement(auto_replace=False)
@@ -317,9 +298,11 @@ class SimSolver(SimStatePlugin):
             self._stored_solver = claripy.SolverCacheless(track=track)
         elif o.SYMBOLIC in self.state.options and o.COMPOSITE_SOLVER in self.state.options:
             self._stored_solver = claripy.SolverComposite(track=track)
-        elif o.SYMBOLIC in self.state.options and any(opt in self.state.options for opt in o.approximation):
-            self._stored_solver = claripy.SolverHybrid(track=track, approximate_first=approximate_first)
-        elif o.HYBRID_SOLVER in self.state.options:
+        elif (
+            o.SYMBOLIC in self.state.options
+            and any(opt in self.state.options for opt in o.approximation)
+            or o.HYBRID_SOLVER in self.state.options
+        ):
             self._stored_solver = claripy.SolverHybrid(track=track, approximate_first=approximate_first)
         elif o.SYMBOLIC in self.state.options:
             self._stored_solver = claripy.Solver(track=track)
@@ -332,7 +315,16 @@ class SimSolver(SimStatePlugin):
     # Get unconstrained stuff
     #
     def Unconstrained(
-        self, name, bits, uninitialized=True, inspect=True, events=True, key=None, eternal=False, **kwargs
+        self,
+        name,
+        bits,
+        uninitialized=True,
+        inspect=True,
+        events=True,
+        key=None,
+        eternal=False,
+        uc_alloc_depth=None,
+        **kwargs,
     ):
         """
         Creates an unconstrained symbol or a default concrete value (0), based on the state options.
@@ -345,7 +337,7 @@ class SimSolver(SimStatePlugin):
         :param events:          Set to False to avoid generating a SimEvent for the occasion
         :param key:             Set this to a tuple of increasingly specific identifiers (for example,
                                 ``('mem', 0xffbeff00)`` or ``('file', 4, 0x20)`` to cause it to be tracked, i.e.
-                                accessable through ``solver.get_variables``.
+                                accessible through ``solver.get_variables``.
         :param eternal:         Set to True in conjunction with setting a key to cause all states with the same
                                 ancestry to retrieve the same symbol when trying to create the value. If False, a
                                 counter will be appended to the key.
@@ -359,33 +351,22 @@ class SimSolver(SimStatePlugin):
                 r = claripy.TSI(bits=bits, name=name, uninitialized=uninitialized, **kwargs)
             else:
                 l.debug("Creating new unconstrained BV named %s", name)
-                if o.UNDER_CONSTRAINED_SYMEXEC in self.state.options:
-                    r = self.BVS(
-                        name,
-                        bits,
-                        uninitialized=uninitialized,
-                        key=key,
-                        eternal=eternal,
-                        inspect=inspect,
-                        events=events,
-                        **kwargs,
-                    )
-                else:
-                    r = self.BVS(
-                        name,
-                        bits,
-                        uninitialized=uninitialized,
-                        key=key,
-                        eternal=eternal,
-                        inspect=inspect,
-                        events=events,
-                        **kwargs,
-                    )
+                r = self.BVS(
+                    name,
+                    bits,
+                    uninitialized=uninitialized,
+                    key=key,
+                    eternal=eternal,
+                    inspect=inspect,
+                    events=events,
+                    **kwargs,
+                )
+                if uc_alloc_depth is not None:
+                    self.state.uc_manager.set_alloc_depth(r, uc_alloc_depth)
 
             return r
-        else:
-            # Return a default value, aka. 0
-            return claripy.BVV(0, bits)
+        # Return a default value, aka. 0
+        return claripy.BVV(0, bits)
 
     def BVS(
         self,
@@ -416,7 +397,7 @@ class SimSolver(SimStatePlugin):
         :param explicit_name:   Set to True to prevent an identifier from appended to the name to ensure uniqueness.
         :param key:             Set this to a tuple of increasingly specific identifiers (for example,
                                 ``('mem', 0xffbeff00)`` or ``('file', 4, 0x20)`` to cause it to be tracked, i.e.
-                                accessable through ``solver.get_variables``.
+                                accessible through ``solver.get_variables``.
         :param eternal:         Set to True in conjunction with setting a key to cause all states with the same
                                 ancestry to retrieve the same symbol when trying to create the value. If False, a
                                 counter will be appended to the key.
@@ -438,7 +419,7 @@ class SimSolver(SimStatePlugin):
                 or uninitialized != r.args[4]
                 or bool(explicit_name) ^ (r.args[0] == name)
             ):
-                l.warning("Variable %s being retrieved with differnt settings than it was tracked with", name)
+                l.warning("Variable %s being retrieved with different settings than it was tracked with", name)
         else:
             r = claripy.BVS(
                 name,
@@ -496,8 +477,7 @@ class SimSolver(SimStatePlugin):
     def widen(self, others):
         c = claripy.BVS("random_widen_condition", 32)
         merge_conditions = [[c == i] for i in range(len(others) + 1)]
-        merging_occurred = self.merge(others, merge_conditions)
-        return merging_occurred
+        return self.merge(others, merge_conditions)
 
     #
     # Frontend passthroughs
@@ -520,19 +500,17 @@ class SimSolver(SimStatePlugin):
     def _adjust_constraint(self, c):
         if self.state._global_condition is None:
             return c
-        elif c is None:  # this should never happen
+        if c is None:  # this should never happen
             l.critical("PLEASE REPORT THIS MESSAGE, AND WHAT YOU WERE DOING, TO YAN")
             return self.state._global_condition
-        else:
-            return claripy.Or(claripy.Not(self.state._global_condition), c)
+        return claripy.Or(claripy.Not(self.state._global_condition), c)
 
     def _adjust_constraint_list(self, constraints):
         if self.state._global_condition is None:
             return constraints
         if len(constraints) == 0:
             return constraints.__class__((self.state._global_condition,))
-        else:
-            return constraints.__class__((self._adjust_constraint(claripy.And(*constraints)),))
+        return constraints.__class__((self._adjust_constraint(claripy.And(*constraints)),))
 
     @timed_function
     @ast_stripping_decorator
@@ -767,7 +745,7 @@ class SimSolver(SimStatePlugin):
         if type(solution) is bool:
             if cast_to is bytes:
                 return bytes([int(solution)])
-            elif cast_to is int:
+            if cast_to is int:
                 return int(solution)
         elif type(solution) is float:
             solution = _concrete_value(claripy.FPV(solution, claripy.fp.FSort.from_size(len(e))).raw_to_bv())
@@ -1034,10 +1012,9 @@ class SimSolver(SimStatePlugin):
         if len(r) == 1:
             self.add(e == r[0])
             return True
-        elif len(r) == 0:
+        if len(r) == 0:
             raise SimValueError("unsatness during uniqueness check(ness)")
-        else:
-            return False
+        return False
 
     def symbolic(self, e):  # pylint:disable=R0201
         """
@@ -1056,12 +1033,10 @@ class SimSolver(SimStatePlugin):
         if self.state.mode == "static":
             if type(e) in (int, bytes, float, bool):
                 return True
-            else:
-                return e.cardinality <= 1
+            return e.cardinality <= 1
 
-        else:
-            # All symbolic expressions are not single-valued
-            return not self.symbolic(e)
+        # All symbolic expressions are not single-valued
+        return not self.symbolic(e)
 
     def simplify(self, e=None):
         """
@@ -1070,16 +1045,13 @@ class SimSolver(SimStatePlugin):
         """
         if e is None:
             return self._solver.simplify()
-        elif isinstance(e, (int, float, bool)):
-            return e
-        elif isinstance(e, claripy.ast.Base) and e.op in claripy.operations.leaf_operations_concrete:
+        if (
+            isinstance(e, (int, float, bool))
+            or (isinstance(e, claripy.ast.Base | SimActionObject) and e.is_leaf() and not e.symbolic)
+            or (not isinstance(e, claripy.ast.Base | SimActionObject))
+        ):
             return e
-        elif isinstance(e, SimActionObject) and e.op in claripy.operations.leaf_operations_concrete:
-            return e.ast
-        elif not isinstance(e, (SimActionObject, claripy.ast.Base)):
-            return e
-        else:
-            return self._claripy_simplify(e)
+        return self._claripy_simplify(e)
 
     @timed_function
     @ast_stripping_decorator

angr/state_plugins/symbolizer.py

@@ -1,3 +1,4 @@
+from __future__ import annotations
 import logging
 import claripy
 import struct
@@ -181,16 +182,15 @@ class SimSymbolizer(SimStatePlugin):  # pylint:disable=abstract-method
     def _resymbolize_int(self, be, le=0, base=0, offset=0, skip=()):
         if base + offset in skip:
             return None
-        elif self._min_addr <= be < self._max_addr and self._should_symbolize(be):
+        if self._min_addr <= be < self._max_addr and self._should_symbolize(be):
             s = self._preconstrain(be)
             l.debug("Replacing %#x (at %#x, endness BE) with %s!", be, base + offset, s)
             return s
-        elif self._min_addr <= le < self._max_addr and self._should_symbolize(le):
+        if self._min_addr <= le < self._max_addr and self._should_symbolize(le):
             s = self._preconstrain(le).reversed
             l.debug("Replacing %#x (at %#x, endness LE) with %s!", le, base + offset, s)
             return s
-        else:
-            return None
+        return None
 
     def _resymbolize_data(self, data, prefix=b"", base=0, skip=()):
         ws = self.state.arch.bytes
@@ -222,10 +222,9 @@ class SimSymbolizer(SimStatePlugin):  # pylint:disable=abstract-method
             values_squashed.append(data[last_idx * ws :])
         values_squashed.append(suffix)
 
-        new_data = claripy.Concat(*values_squashed)
+        return claripy.Concat(*values_squashed)
         # assert len(new_data)/8 == len(data) + len(prefix)
         # assert self.state.solver.eval_one(new_data) == self.state.solver.eval_one(claripy.BVV(data))
-        return new_data
 
     def _resymbolize_region(self, storage: PagedMemoryMixin, addr, length):
         assert type(addr) is int

angr/state_plugins/trace_additions.py

@@ -12,6 +12,8 @@ This can happen in challenge response if all of the values in the flag page are
 before being printed.
 """
 
+from __future__ import annotations
+
 import logging
 import string
 
@@ -52,10 +54,9 @@ class FormatInfoStrToInt(FormatInfo):
         self.input_base = None
 
     def copy(self):
-        out = FormatInfoStrToInt(
+        return FormatInfoStrToInt(
             self.addr, self.func_name, self.str_arg_num, self.base, self.base_arg, self.allows_negative
         )
-        return out
 
     def compute(self, state):
         self.input_val = angr.calling_conventions.SimCCCdecl(state.arch).arg(state, self.str_arg_num)
@@ -92,10 +93,9 @@ class FormatInfoIntToStr(FormatInfo):
         self.str_dst_addr = None
 
     def copy(self):
-        out = FormatInfoIntToStr(
+        return FormatInfoIntToStr(
             self.addr, self.func_name, self.int_arg_num, self.str_dst_num, self.base, self.base_arg
         )
-        return out
 
     def compute(self, state):
         self.input_val = angr.calling_conventions.SimCCCdecl(state.arch).arg(state, self.int_arg_num)
@@ -120,8 +120,7 @@ class FormatInfoDontConstrain(FormatInfo):
         self.check_symbolic_arg = check_symbolic_arg
 
     def copy(self):
-        out = FormatInfoDontConstrain(self.addr, self.func_name, self.check_symbolic_arg)
-        return out
+        return FormatInfoDontConstrain(self.addr, self.func_name, self.check_symbolic_arg)
 
     def compute(self, state):
         pass
@@ -253,7 +252,7 @@ def end_info_hook(state):
     state.solver._solver.add_replacement(new_var, result, invalidate_cache=False)
     # dont add this constraint to preconstraints or we lose real constraints
     # chall_resp_plugin.tracer.preconstraints.append(constraint)
-    chall_resp_plugin.state.preconstrainer.variable_map[list(new_var.variables)[0]] = constraint
+    chall_resp_plugin.state.preconstrainer.variable_map[next(iter(new_var.variables))] = constraint
     chall_resp_plugin.pop_from_backup()
 
 
@@ -285,7 +284,7 @@ def syscall_hook(state):
         stdin_min_stdout_reads = state.get_plugin("chall_resp_info").stdin_min_stdout_reads
         stdout_pos = state.solver.eval(state.posix.fd[1].write_pos)
         stdin_pos = state.solver.eval(state.posix.fd[0].read_pos)
-        for i in range(0, stdin_pos):
+        for i in range(stdin_pos):
             if i not in stdin_min_stdout_reads:
                 stdin_min_stdout_reads[i] = stdout_pos
 
@@ -456,7 +455,7 @@ class ChallRespInfo(angr.state_plugins.SimStatePlugin):
     def get_same_length_constraints(self):
         constraints = []
         for str_var, int_var in self.str_to_int_pairs:
-            int_var_name = list(int_var.variables)[0]
+            int_var_name = next(iter(int_var.variables))
             base = int(int_var_name.split("_")[1], 10)
             original_len = str_var.size() // 8
             abs_max = (1 << int_var.size()) - 1
@@ -490,18 +489,14 @@ class ChallRespInfo(angr.state_plugins.SimStatePlugin):
             for _, int_var in chall_resp_plugin.str_to_int_pairs:
                 vars_to_solve.append(int_var)
 
-            if require_same_length:
-                extra_constraints = chall_resp_plugin.get_same_length_constraints()
-            else:
-                extra_constraints = []
+            extra_constraints = chall_resp_plugin.get_same_length_constraints() if require_same_length else []
 
             solns = state.solver._solver.batch_eval(vars_to_solve, 1, extra_constraints=extra_constraints)
             if len(solns) == 0:
                 if require_same_length:
                     l.warning("could not satisfy with same length, falling back to different lengths")
                     return ChallRespInfo.atoi_dumps(state, require_same_length=False)
-                else:
-                    return state.posix.dumps(0)
+                return state.posix.dumps(0)
             solns = solns[0]
 
             # now make the real stdin
@@ -509,7 +504,7 @@ class ChallRespInfo(angr.state_plugins.SimStatePlugin):
 
             stdin_replacements = []
             for soln, (_, int_var) in zip(solns[1:], chall_resp_plugin.str_to_int_pairs):
-                int_var_name = list(int_var.variables)[0]
+                int_var_name = next(iter(int_var.variables))
                 indices = chall_resp_plugin.get_stdin_indices(int_var_name)
                 if len(indices) == 0:
                     continue
@@ -526,9 +521,8 @@ class ChallRespInfo(angr.state_plugins.SimStatePlugin):
 
             # filter for same start with value 0
             for i in list(stdin_replacements):
-                if any(ii[0] == i[0] and ii[2] != i[2] for ii in stdin_replacements):
-                    if int(i[2]) == 0:
-                        stdin_replacements.remove(i)
+                if any(ii[0] == i[0] and ii[2] != i[2] for ii in stdin_replacements) and int(i[2]) == 0:
+                    stdin_replacements.remove(i)
 
             # now do the replacing
             offset = 0
@@ -569,7 +563,7 @@ def zen_hook(state, expr):
     if state.has_plugin("chall_resp_info") and state.get_plugin("chall_resp_info").pending_info is not None:
         return None
 
-    if expr.op not in claripy.operations.leaf_operations and expr.op != "Concat":
+    if expr.is_leaf() and expr.op != "Concat":
         # if there is more than one symbolic argument we replace it and preconstrain it
         flag_args = ZenPlugin.get_flag_rand_args(expr)
         if len(flag_args) > 1:
@@ -593,17 +587,17 @@ def zen_hook(state, expr):
                     con = replacement == expr
                     state.add_constraints(con)
                     contained_bytes = zen_plugin.get_flag_bytes(expr)
-                    zen_plugin.byte_dict[list(replacement.variables)[0]] = contained_bytes
+                    zen_plugin.byte_dict[next(iter(replacement.variables))] = contained_bytes
                     zen_plugin.zen_constraints.append(con)
                     # saves a ton of memory to do this here rather than later
                     zen_plugin.zen_constraints.append(state.solver.simplify(con))
                 else:
                     # otherwise don't add the constraint, just replace
                     depth = 0
-                    zen_plugin.byte_dict[list(replacement.variables)[0]] = set()
+                    zen_plugin.byte_dict[next(iter(replacement.variables))] = set()
 
                 # save and replace
-                var = list(replacement.variables)[0]
+                var = next(iter(replacement.variables))
                 zen_plugin.depths[var] = depth
                 constraint = replacement == concrete_val
                 zen_plugin.state.preconstrainer.preconstraints.append(constraint)
@@ -663,18 +657,17 @@ class ZenPlugin(angr.state_plugins.SimStatePlugin):
         symbolic_args = tuple(a for a in expr.args if isinstance(a, claripy.ast.Base) and a.symbolic)
         flag_args = []
         for a in symbolic_args:
-            if any(v.startswith("cgc-flag") or v.startswith("random") for v in a.variables):
+            if any(v.startswith(("cgc-flag", "random")) for v in a.variables):
                 flag_args.append(a)
         return flag_args
 
     def get_expr_depth(self, expr):
         flag_args = self.get_flag_rand_args(expr)
         flag_arg_vars = set.union(*[set(v.variables) for v in flag_args])
-        flag_arg_vars = {v for v in flag_arg_vars if v.startswith("cgc-flag") or v.startswith("random")}
+        flag_arg_vars = {v for v in flag_arg_vars if v.startswith(("cgc-flag", "random"))}
         if len(flag_arg_vars) == 0:
             return 0
-        depth = max(self.depths.get(v, 0) for v in flag_arg_vars) + 1
-        return depth
+        return max(self.depths.get(v, 0) for v in flag_arg_vars) + 1
 
     @angr.state_plugins.SimStatePlugin.memo
     def copy(self, memo):  # pylint: disable=unused-argument
@@ -713,7 +706,7 @@ class ZenPlugin(angr.state_plugins.SimStatePlugin):
         for con in constraints:
             if (
                 con.cache_key in zen_cache_keys
-                or not all(v.startswith("cgc-flag") or v.startswith("random") for v in con.variables)
+                or not all(v.startswith(("cgc-flag", "random")) for v in con.variables)
                 or len(con.variables) == 0
             ):
                 new_cons.append(con)
@@ -724,16 +717,13 @@ class ZenPlugin(angr.state_plugins.SimStatePlugin):
         try:
             state.memory.permissions(state.solver.eval(buf))
         except angr.SimMemoryError:
-            l.warning("detected possible arbitary transmit to fd %d", fd)
+            l.warning("detected possible arbitrary transmit to fd %d", fd)
             if fd == 0 or fd == 1:
                 self.controlled_transmits.append((state.copy(), buf))
 
     @staticmethod
     def prep_tracer(state):
-        if state.has_plugin("zen_plugin"):
-            zen_plugin = state.get_plugin("zen_plugin")
-        else:
-            zen_plugin = ZenPlugin()
+        zen_plugin = state.get_plugin("zen_plugin") if state.has_plugin("zen_plugin") else ZenPlugin()
 
         state.register_plugin("zen_plugin", zen_plugin)
         state.inspect.b("reg_write", angr.BP_BEFORE, action=zen_register_write)
@@ -742,7 +732,7 @@ class ZenPlugin(angr.state_plugins.SimStatePlugin):
         # setup the byte dict
         byte_dict = zen_plugin.byte_dict
         for i, b in enumerate(state.cgc.flag_bytes):
-            var = list(b.variables)[0]
+            var = next(iter(b.variables))
             byte_dict[var] = {i}
 
         state.preconstrainer.preconstraints.extend(zen_plugin.preconstraints)

angr/state_plugins/uc_manager.py

@@ -1,3 +1,4 @@
+from __future__ import annotations
 import logging
 
 import claripy
@@ -9,6 +10,8 @@ l = logging.getLogger(name=__name__)
 
 
 class SimUCManager(SimStatePlugin):
+    _uc_alloc_depth: dict[claripy.ast.Base, int]
+
     def __init__(self, man=None):
         SimStatePlugin.__init__(self)
 
@@ -32,6 +35,8 @@ class SimUCManager(SimStatePlugin):
         # The maximum allocation depth
         self._max_alloc_depth = 20
 
+        self._uc_alloc_depth = {}
+
     def assign(self, dst_addr_ast):
         """
         Assign a new region for under-constrained symbolic execution.
@@ -40,17 +45,17 @@ class SimUCManager(SimStatePlugin):
         :return: as ast of memory address that points to a new region
         """
 
-        if dst_addr_ast.uc_alloc_depth > self._max_alloc_depth:
+        dst_uc_alloc_depth = self._uc_alloc_depth[dst_addr_ast]
+        if dst_uc_alloc_depth > self._max_alloc_depth:
             raise SimUCManagerAllocationError(
-                "Current allocation depth %d is greater than the cap (%d)"
-                % (dst_addr_ast.uc_alloc_depth, self._max_alloc_depth)
+                "Current allocation depth %d is greater than the cap (%d)" % (dst_uc_alloc_depth, self._max_alloc_depth)
             )
 
         abs_addr = self._region_base + self._pos
         ptr = claripy.BVV(abs_addr, self.state.arch.bits)
         self._pos += self._region_size
 
-        self._alloc_depth_map[(abs_addr - self._region_base) // self._region_size] = dst_addr_ast.uc_alloc_depth
+        self._alloc_depth_map[(abs_addr - self._region_base) // self._region_size] = dst_uc_alloc_depth
 
         l.debug("Assigned new memory region %s", ptr)
         return ptr
@@ -59,13 +64,15 @@ class SimUCManager(SimStatePlugin):
     def copy(self, memo):  # pylint: disable=unused-argument
         return SimUCManager(man=self)
 
-    def get_alloc_depth(self, addr):
-        block_pos = (addr - self._region_base) // self._region_size
+    def get_alloc_depth(self, addr: int | claripy.ast.Base) -> int | None:
+        if isinstance(addr, claripy.ast.Base):
+            return self._uc_alloc_depth.get(addr, None)
 
-        if block_pos not in self._alloc_depth_map:
-            return None
+        block_pos = (addr - self._region_base) // self._region_size
+        return self._alloc_depth_map.get(block_pos, None)
 
-        return self._alloc_depth_map[block_pos]
+    def set_alloc_depth(self, addr: claripy.ast.Base, depth: int):
+        self._uc_alloc_depth[addr] = depth
 
     def is_bounded(self, ast):
         """