angr 9.2.126__tar.gz → 9.2.128__tar.gz

{angr-9.2.126/angr.egg-info → angr-9.2.128}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: angr
-Version: 9.2.126
+Version: 9.2.128
 Summary: A multi-architecture binary analysis toolkit, with the ability to perform dynamic symbolic execution and various static analyses on binaries
 Home-page: https://github.com/angr/angr
 License: BSD-2-Clause
@@ -16,13 +16,13 @@ Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: CppHeaderParser
 Requires-Dist: GitPython
-Requires-Dist: ailment==9.2.126
-Requires-Dist: archinfo==9.2.126
+Requires-Dist: ailment==9.2.128
+Requires-Dist: archinfo==9.2.128
 Requires-Dist: cachetools
 Requires-Dist: capstone==5.0.3
 Requires-Dist: cffi>=1.14.0
-Requires-Dist: claripy==9.2.126
-Requires-Dist: cle==9.2.126
+Requires-Dist: claripy==9.2.128
+Requires-Dist: cle==9.2.128
 Requires-Dist: itanium-demangler
 Requires-Dist: mulpyplexer
 Requires-Dist: nampa
@@ -31,7 +31,7 @@ Requires-Dist: protobuf>=5.28.2
 Requires-Dist: psutil
 Requires-Dist: pycparser>=2.18
 Requires-Dist: pyformlang
-Requires-Dist: pyvex==9.2.126
+Requires-Dist: pyvex==9.2.128
 Requires-Dist: rich>=13.1.0
 Requires-Dist: sortedcontainers
 Requires-Dist: sympy

{angr-9.2.126 → angr-9.2.128}/angr/__init__.py

@@ -2,7 +2,7 @@
 # pylint: disable=wrong-import-position
 from __future__ import annotations
 
-__version__ = "9.2.126"
+__version__ = "9.2.128"
 
 if bytes is str:
     raise Exception(

{angr-9.2.126 → angr-9.2.128}/angr/analyses/analysis.py

@@ -9,6 +9,7 @@ from typing import TYPE_CHECKING, TypeVar, Generic, cast
 from collections.abc import Callable
 from types import NoneType
 from itertools import chain
+from traceback import format_exception
 
 import logging
 import time
@@ -75,6 +76,11 @@ class AnalysisLogEntry:
 
         self.message = message
 
+    def format(self) -> str:
+        if self.exc_traceback is None:
+            return self.message
+        return "\n".join((*format_exception(self.exc_type, self.exc_value, self.exc_traceback), "", self.message))
+
     def __getstate__(self):
         return (
             str(self.__dict__.get("exc_type")),
@@ -281,8 +287,8 @@ class Analysis:
     kb: KnowledgeBase
     _fail_fast: bool
     _name: str
-    errors = []
-    named_errors = defaultdict(list)
+    errors: list[AnalysisLogEntry] = []
+    named_errors: defaultdict[str, list[AnalysisLogEntry]] = defaultdict(list)
     _progress_callback = None
     _show_progressbar = False
     _progressbar = None

{angr-9.2.126 → angr-9.2.128}/angr/analyses/cfg/cfg_fast.py

@@ -3309,7 +3309,18 @@ class CFGFast(ForwardAnalysis[CFGNode, CFGNode, CFGJob, int], CFGBase):  # pylin
                         # this is not a no-op block. Determine where nop instructions terminate.
                         insns = block.capstone.insns
                         if insns:
-                            nop_length = self._get_nop_length(insns)
+                            if (
+                                self.project.simos is not None
+                                and self.project.simos.name == "Win32"
+                                and insns[0].mnemonic == "mov"
+                            ):
+                                op0, op1 = insns[0].operands
+                                if op0.type == 1 and op1.type == 1 and op0.reg == op1.reg:
+                                    # hot-patch points on Windows DLLs
+                                    # https://devblogs.microsoft.com/oldnewthing/20110921-00/?p=9583
+                                    nop_length = None
+                            else:
+                                nop_length = self._get_nop_length(insns)
 
                     if nop_length is None or nop_length <= 0:
                         continue

{angr-9.2.126 → angr-9.2.128}/angr/analyses/decompiler/clinic.py

@@ -113,6 +113,8 @@ class Clinic(Analysis):
         vvar_id_start: int = 0,
         optimization_scratch: dict[str, Any] | None = None,
         desired_variables: set[str] | None = None,
+        force_loop_single_exit: bool = True,
+        complete_successors: bool = False,
     ):
         if not func.normalized and mode == ClinicMode.DECOMPILE:
             raise ValueError("Decompilation must work on normalized function graphs.")
@@ -157,6 +159,8 @@ class Clinic(Analysis):
         self._inlined_counts = {} if inlined_counts is None else inlined_counts
         self._inlining_parents = inlining_parents or ()
         self._desired_variables = desired_variables
+        self._force_loop_single_exit = force_loop_single_exit
+        self._complete_successors = complete_successors
 
         self._register_save_areas_removed: bool = False
 
@@ -333,6 +337,7 @@ class Clinic(Analysis):
             optimization_passes=[StackCanarySimplifier],
             sp_shift=self._max_stack_depth,
             vvar_id_start=self.vvar_id_start,
+            fail_fast=self._fail_fast,
         )
         self.vvar_id_start = callee_clinic.vvar_id_start + 1
         self._max_stack_depth = callee_clinic._max_stack_depth
@@ -787,7 +792,7 @@ class Clinic(Analysis):
 
             # case 2: the callee is a SimProcedure
             if target_func.is_simprocedure:
-                cc = self.project.analyses.CallingConvention(target_func)
+                cc = self.project.analyses.CallingConvention(target_func, fail_fast=self._fail_fast)
                 if cc.cc is not None and cc.prototype is not None:
                     target_func.calling_convention = cc.cc
                     target_func.prototype = cc.prototype
@@ -795,7 +800,7 @@ class Clinic(Analysis):
 
             # case 3: the callee is a PLT function
             if target_func.is_plt:
-                cc = self.project.analyses.CallingConvention(target_func)
+                cc = self.project.analyses.CallingConvention(target_func, fail_fast=self._fail_fast)
                 if cc.cc is not None and cc.prototype is not None:
                     target_func.calling_convention = cc.cc
                     target_func.prototype = cc.prototype
@@ -834,6 +839,7 @@ class Clinic(Analysis):
                     callsite_block_addr=callsite.addr,
                     callsite_insn_addr=callsite_ins_addr,
                     func_graph=func_graph,
+                    fail_fast=self._fail_fast,
                 )
 
                 if cc.cc is not None and cc.prototype is not None:
@@ -864,6 +870,7 @@ class Clinic(Analysis):
         # finally, recover the calling convention of the current function
         if self.function.prototype is None or self.function.calling_convention is None:
             self.project.analyses.CompleteCallingConventions(
+                fail_fast=self._fail_fast,
                 recover_variables=True,
                 prioritize_func_addrs=[self.function.addr],
                 skip_other_funcs=True,
@@ -896,6 +903,7 @@ class Clinic(Analysis):
         spt = self.project.analyses.StackPointerTracker(
             self.function,
             regs,
+            fail_fast=self._fail_fast,
             track_memory=self._sp_tracker_track_memory,
             cross_insn_opt=False,
             initial_reg_values=initial_reg_values,
@@ -1043,6 +1051,7 @@ class Clinic(Analysis):
                                 reg_name=self.project.arch.translate_register_name(
                                     ret_reg_offset, size=self.project.arch.bits
                                 ),
+                                **target.tags,
                             )
                             call_stmt = ailment.Stmt.Call(
                                 None,
@@ -1129,6 +1138,7 @@ class Clinic(Analysis):
         simp = self.project.analyses.AILBlockSimplifier(
             ail_block,
             self.function.addr,
+            fail_fast=self._fail_fast,
             remove_dead_memdefs=remove_dead_memdefs,
             stack_pointer_tracker=stack_pointer_tracker,
             peephole_optimizations=self.peephole_optimizations,
@@ -1200,6 +1210,7 @@ class Clinic(Analysis):
 
         simp = self.project.analyses.AILSimplifier(
             self.function,
+            fail_fast=self._fail_fast,
             func_graph=ail_graph,
             remove_dead_memdefs=remove_dead_memdefs,
             unify_variables=unify_variables,
@@ -1258,6 +1269,8 @@ class Clinic(Analysis):
                 vvar_id_start=self.vvar_id_start,
                 entry_node_addr=self.entry_node_addr,
                 scratch=self.optimization_scratch,
+                force_loop_single_exit=self._force_loop_single_exit,
+                complete_successors=self._complete_successors,
                 **kwargs,
             )
             if a.out_graph:
@@ -1341,6 +1354,7 @@ class Clinic(Analysis):
         ssailification = self.project.analyses.Ssailification(
             self.function,
             ail_graph,
+            fail_fast=self._fail_fast,
             entry=next(iter(bb for bb in ail_graph if (bb.addr, bb.idx) == self.entry_node_addr)),
             ail_manager=self._ail_manager,
             ssa_stackvars=False,
@@ -1354,6 +1368,7 @@ class Clinic(Analysis):
         ssailification = self.project.analyses.Ssailification(
             self.function,
             ail_graph,
+            fail_fast=self._fail_fast,
             entry=next(iter(bb for bb in ail_graph if (bb.addr, bb.idx) == self.entry_node_addr)),
             ail_manager=self._ail_manager,
             ssa_tmps=True,
@@ -1368,6 +1383,7 @@ class Clinic(Analysis):
         dephication = self.project.analyses.GraphDephicationVVarMapping(
             self.function,
             ail_graph,
+            fail_fast=self._fail_fast,
             entry=next(iter(bb for bb in ail_graph if (bb.addr, bb.idx) == self.entry_node_addr)),
             vvar_id_start=self.vvar_id_start,
         )
@@ -1420,6 +1436,7 @@ class Clinic(Analysis):
         rd = self.project.analyses.SReachingDefinitions(
             subject=self.function,
             func_graph=ail_graph,
+            fail_fast=self._fail_fast,
             # use_callee_saved_regs_at_return=not self._register_save_areas_removed,  FIXME
         )
 
@@ -1430,6 +1447,7 @@ class Clinic(Analysis):
         def _handler(block):
             csm = self.project.analyses.AILCallSiteMaker(
                 block,
+                fail_fast=self._fail_fast,
                 reaching_definitions=rd,
                 stack_pointer_tracker=stack_pointer_tracker,
                 ail_manager=self._ail_manager,
@@ -1443,6 +1461,7 @@ class Clinic(Analysis):
                 simp = self.project.analyses.AILBlockSimplifier(
                     ail_block,
                     self.function.addr,
+                    fail_fast=self._fail_fast,
                     stack_pointer_tracker=stack_pointer_tracker,
                     peephole_optimizations=self.peephole_optimizations,
                 )
@@ -1526,6 +1545,7 @@ class Clinic(Analysis):
         tmp_kb.functions = self.kb.functions
         vr = self.project.analyses.VariableRecoveryFast(
             self.function,  # pylint:disable=unused-variable
+            fail_fast=self._fail_fast,
             func_graph=ail_graph,
             kb=tmp_kb,
             track_sp=False,
@@ -1558,6 +1578,7 @@ class Clinic(Analysis):
                 vr.type_constraints,
                 vr.func_typevar,
                 kb=tmp_kb,
+                fail_fast=self._fail_fast,
                 var_mapping=vr.var_to_typevars,
                 must_struct=must_struct,
                 ground_truth=groundtruth,

{angr-9.2.126 → angr-9.2.128}/angr/analyses/decompiler/condition_processor.py

@@ -1113,17 +1113,15 @@ class ConditionProcessor:
         r1_with: claripy.ast.Bool,
     ) -> claripy.ast.Bool:
         if ast.op == "And":
-            return ast.make_like(
-                "And", (ConditionProcessor._replace_term_in_ast(arg, r0, r0_with, r1, r1_with) for arg in ast.args)
+            return claripy.And(
+                *(ConditionProcessor._replace_term_in_ast(arg, r0, r0_with, r1, r1_with) for arg in ast.args)
             )
         if ast.op == "Or":
-            return ast.make_like(
-                "Or", (ConditionProcessor._replace_term_in_ast(arg, r0, r0_with, r1, r1_with) for arg in ast.args)
+            return claripy.Or(
+                *(ConditionProcessor._replace_term_in_ast(arg, r0, r0_with, r1, r1_with) for arg in ast.args)
             )
         if ast.op == "Not":
-            return ast.make_like(
-                "Not", (ConditionProcessor._replace_term_in_ast(ast.args[0], r0, r0_with, r1, r1_with),)
-            )
+            return claripy.Not(ConditionProcessor._replace_term_in_ast(ast.args[0], r0, r0_with, r1, r1_with))
         if ast is r0:
             return r0_with
         if ast is r1:

{angr-9.2.126 → angr-9.2.128}/angr/analyses/decompiler/decompilation_cache.py

@@ -23,6 +23,7 @@ class DecompilationCache:
         "clinic",
         "ite_exprs",
         "binop_operators",
+        "errors",
     )
 
     def __init__(self, addr):
@@ -35,7 +36,10 @@ class DecompilationCache:
         self.clinic: Clinic | None = None
         self.ite_exprs: set[tuple[int, Any]] | None = None
         self.binop_operators: dict[OpDescriptor, str] | None = None
+        self.errors: list[str] = []
 
     @property
     def local_types(self):
+        if self.clinic is None or self.clinic.variable_kb is None:
+            return None
         return self.clinic.variable_kb.variables[self.addr].types

{angr-9.2.126 → angr-9.2.128}/angr/analyses/decompiler/decompiler.py

@@ -139,7 +139,22 @@ class Decompiler(Analysis):
         self.expr_collapse_depth = expr_collapse_depth
 
         if decompile:
-            self._decompile()
+            with self._resilience():
+                self._decompile()
+            if self.errors:
+                if (self.func.addr, self._flavor) not in self.kb.decompilations:
+                    self.kb.decompilations[(self.func.addr, self._flavor)] = DecompilationCache(self.func.addr)
+                for error in self.errors:
+                    self.kb.decompilations[(self.func.addr, self._flavor)].errors.append(error.format())
+                with self._resilience():
+                    l.info("Decompilation failed for %s. Switching to basic preset and trying again.")
+                    if preset != DECOMPILATION_PRESETS["basic"]:
+                        self._optimization_passes = DECOMPILATION_PRESETS["basic"].get_optimization_passes(
+                            self.project.arch, self.project.simos.name
+                        )
+                        self._decompile()
+                        for error in self.errors:
+                            self.kb.decompilations[(self.func.addr, self._flavor)].errors.append(error.format())
 
     def _can_use_decompilation_cache(self, cache: DecompilationCache) -> bool:
         a, b = self._cache_parameters, cache.parameters
@@ -155,7 +170,7 @@ class Decompiler(Analysis):
 
         if self._cache_parameters is not None:
             try:
-                cache = self.kb.decompilations[self.func.addr]
+                cache = self.kb.decompilations[(self.func.addr, self._flavor)]
                 if not self._can_use_decompilation_cache(cache):
                     cache = None
             except KeyError:
@@ -220,6 +235,7 @@ class Decompiler(Analysis):
             clinic = self.project.analyses.Clinic(
                 self.func,
                 kb=self.kb,
+                fail_fast=self._fail_fast,
                 variable_kb=variable_kb,
                 reset_variable_names=reset_variable_names,
                 optimization_passes=self._optimization_passes,
@@ -233,6 +249,8 @@ class Decompiler(Analysis):
                 inline_functions=self._inline_functions,
                 desired_variables=self._desired_variables,
                 optimization_scratch=self._optimization_scratch,
+                force_loop_single_exit=self._force_loop_single_exit,
+                complete_successors=self._complete_successors,
                 **self.options_to_params(self.options_by_class["clinic"]),
             )
         else:
@@ -293,7 +311,7 @@ class Decompiler(Analysis):
             self._update_progress(75.0, text="Structuring code")
 
             # structure it
-            rs = self.project.analyses[RecursiveStructurer].prep(kb=self.kb)(
+            rs = self.project.analyses[RecursiveStructurer].prep(kb=self.kb, fail_fast=self._fail_fast)(
                 ri.region,
                 cond_proc=cond_proc,
                 func=self.func,
@@ -306,6 +324,7 @@ class Decompiler(Analysis):
                 self.func,
                 rs.result,
                 kb=self.kb,
+                fail_fast=self._fail_fast,
                 variable_kb=clinic.variable_kb,
                 **self.options_to_params(self.options_by_class["region_simplifier"]),
             )
@@ -330,6 +349,7 @@ class Decompiler(Analysis):
                 flavor=self._flavor,
                 func_args=clinic.arg_list,
                 kb=self.kb,
+                fail_fast=self._fail_fast,
                 variable_kb=clinic.variable_kb,
                 expr_comments=old_codegen.expr_comments if old_codegen is not None else None,
                 stmt_comments=old_codegen.stmt_comments if old_codegen is not None else None,
@@ -347,10 +367,10 @@ class Decompiler(Analysis):
         self.cache.codegen = codegen
         self.cache.clinic = self.clinic
 
-        self.kb.decompilations[self.func.addr] = self.cache
+        self.kb.decompilations[(self.func.addr, self._flavor)] = self.cache
 
     def _recover_regions(self, graph: networkx.DiGraph, condition_processor, update_graph: bool = True):
-        return self.project.analyses[RegionIdentifier].prep(kb=self.kb)(
+        return self.project.analyses[RegionIdentifier].prep(kb=self.kb, fail_fast=self._fail_fast)(
             self.func,
             graph=graph,
             cond_proc=condition_processor,
@@ -403,6 +423,8 @@ class Decompiler(Analysis):
                 reaching_definitions=reaching_definitions,
                 entry_node_addr=self.clinic.entry_node_addr,
                 scratch=self._optimization_scratch,
+                force_loop_single_exit=self._force_loop_single_exit,
+                complete_successors=self._complete_successors,
                 **kwargs,
             )
 
@@ -463,6 +485,8 @@ class Decompiler(Analysis):
                 vvar_id_start=self.vvar_id_start,
                 entry_node_addr=self.clinic.entry_node_addr,
                 scratch=self._optimization_scratch,
+                force_loop_single_exit=self._force_loop_single_exit,
+                complete_successors=self._complete_successors,
                 **kwargs,
             )
 
@@ -546,6 +570,7 @@ class Decompiler(Analysis):
                 type_constraints,
                 func_typevar,
                 kb=var_kb,
+                fail_fast=self._fail_fast,
                 var_mapping=var_to_typevar,
                 must_struct=must_struct,
                 ground_truth=groundtruth,
@@ -613,11 +638,15 @@ class Decompiler(Analysis):
         )
 
     def _transform_graph_from_ssa(self, ail_graph: networkx.DiGraph) -> networkx.DiGraph:
-        dephication = self.project.analyses.GraphDephication(self.func, ail_graph, rewrite=True)
+        dephication = self.project.analyses.GraphDephication(
+            self.func, ail_graph, rewrite=True, kb=self.kb, fail_fast=self._fail_fast
+        )
         return dephication.output
 
     def _transform_seqnode_from_ssa(self, seq_node: SequenceNode) -> SequenceNode:
-        dephication = self.project.analyses.SeqNodeDephication(self.func, seq_node, rewrite=True)
+        dephication = self.project.analyses.SeqNodeDephication(
+            self.func, seq_node, rewrite=True, kb=self.kb, fail_fast=self._fail_fast
+        )
         return dephication.output
 
     @staticmethod

{angr-9.2.126 → angr-9.2.128}/angr/analyses/decompiler/dephication/graph_vvar_mapping.py

@@ -290,8 +290,7 @@ class GraphDephicationVVarMapping(Analysis):  # pylint:disable=abstract-method
         for stmt in dst_block.statements:
             if isinstance(stmt, Label):
                 continue
-            r, _ = is_phi_assignment(stmt)
-            if r:
+            if is_phi_assignment(stmt):
                 for src_, vvar in stmt.src.src_and_vvars:
                     if src_ == src and vvar is not None and vvar.varid == vvar_id:
                         return True

{angr-9.2.126 → angr-9.2.128}/angr/analyses/decompiler/graph_region.py

@@ -382,12 +382,9 @@ class GraphRegion:
                         if src in graph:
                             graph.add_edge(src, dst)
                 else:
-                    # it may happen that the dst node does not exist in sub_graph
-                    # fallback
-                    l.info("Node dst is not found in sub_graph. Enter the fall back logic.")
-                    for src in sub_graph.nodes:
-                        if sub_graph.out_degree[src] == 0:
-                            graph.add_edge(src, dst)
+                    # it may happen that the dst node no longer exists in sub_graph or its successors
+                    # this is because we have deemed that the dst node is no longer a valid successor for sub_graph
+                    pass
 
         graph.add_nodes_from(sub_graph_nodes)
         graph.add_edges_from(sub_graph_edges)

angr-9.2.128/angr/analyses/decompiler/label_collector.py

@@ -0,0 +1,32 @@
+# pylint:disable=unused-argument
+from __future__ import annotations
+from collections import defaultdict
+
+import ailment
+
+from .sequence_walker import SequenceWalker
+
+
+class LabelCollector:
+    """
+    Collect all labels.
+    """
+
+    def __init__(self, node):
+        self.root = node
+        self.labels: defaultdict[str, list[tuple[int, int | None]]] = defaultdict(list)
+
+        handlers = {
+            ailment.Block: self._handle_Block,
+        }
+        self._walker = SequenceWalker(handlers=handlers)
+        self._walker.walk(self.root)
+
+    #
+    # Handlers
+    #
+
+    def _handle_Block(self, block: ailment.Block, **kwargs):
+        for stmt in block.statements:
+            if isinstance(stmt, ailment.Stmt.Label):
+                self.labels[stmt.name].append((block.addr, block.idx))

{angr-9.2.126 → angr-9.2.128}/angr/analyses/decompiler/optimization_passes/__init__.py

@@ -32,6 +32,7 @@ from .inlined_string_transformation_simplifier import InlinedStringTransformatio
 from .const_prop_reverter import ConstPropOptReverter
 from .call_stmt_rewriter import CallStatementRewriter
 from .duplication_reverter import DuplicationReverter
+from .switch_reused_entry_rewriter import SwitchReusedEntryRewriter
 
 if TYPE_CHECKING:
     from angr.analyses.decompiler.presets import DecompilationPreset
@@ -55,6 +56,7 @@ ALL_OPTIMIZATION_PASSES = [
     ReturnDuplicatorHigh,
     DeadblockRemover,
     SwitchDefaultCaseDuplicator,
+    SwitchReusedEntryRewriter,
     ConstPropOptReverter,
     DuplicationReverter,
     LoweredSwitchSimplifier,
@@ -129,6 +131,7 @@ __all__ = (
     "CrossJumpReverter",
     "CodeMotionOptimization",
     "SwitchDefaultCaseDuplicator",
+    "SwitchReusedEntryRewriter",
     "DeadblockRemover",
     "InlinedStringTransformationSimplifier",
     "ConstPropOptReverter",

{angr-9.2.126 → angr-9.2.128}/angr/analyses/decompiler/optimization_passes/optimization_pass.py

@@ -118,6 +118,8 @@ class OptimizationPass(BaseOptimizationPass):
         vvar_id_start=None,
         entry_node_addr=None,
         scratch: dict[str, Any] | None = None,
+        force_loop_single_exit: bool = True,
+        complete_successors: bool = False,
         **kwargs,
     ):
         super().__init__(func)
@@ -134,6 +136,8 @@ class OptimizationPass(BaseOptimizationPass):
         self.entry_node_addr: tuple[int, int | None] = (
             entry_node_addr if entry_node_addr is not None else (func.addr, None)
         )
+        self._force_loop_single_exit = force_loop_single_exit
+        self._complete_successors = complete_successors
 
         # output
         self.out_graph: networkx.DiGraph | None = None
@@ -255,9 +259,8 @@ class OptimizationPass(BaseOptimizationPass):
             graph=graph,
             cond_proc=condition_processor or ConditionProcessor(self.project.arch),
             update_graph=update_graph,
-            # TODO: find a way to pass Phoenix/DREAM options here (see decompiler.py for correct use)
-            force_loop_single_exit=True,
-            complete_successors=False,
+            force_loop_single_exit=self._force_loop_single_exit,
+            complete_successors=self._complete_successors,
             entry_node_addr=self.entry_node_addr,
         )
 

{angr-9.2.126 → angr-9.2.128}/angr/analyses/decompiler/optimization_passes/switch_default_case_duplicator.py

@@ -1,10 +1,15 @@
 # pylint:disable=too-many-boolean-expressions
 from __future__ import annotations
 from itertools import count
+from collections import defaultdict
 import logging
 
 import networkx
 
+from ailment.block import Block
+from ailment.statement import Jump
+from ailment.expression import Const
+
 from angr.knowledge_plugins.cfg import IndirectJumpType
 from .optimization_pass import OptimizationPass, OptimizationPassStage
 
@@ -29,17 +34,19 @@ class SwitchDefaultCaseDuplicator(OptimizationPass):
 
     ARCHES = None
     PLATFORMS = None
-    STAGE = OptimizationPassStage.BEFORE_REGION_IDENTIFICATION
+    STAGE = OptimizationPassStage.AFTER_AIL_GRAPH_CREATION
     NAME = "Duplicate default-case nodes to undo default-case node reuse caused by compiler code deduplication"
     DESCRIPTION = __doc__.strip()
 
     def __init__(self, func, **kwargs):
         super().__init__(func, **kwargs)
 
-        self.node_idx = count(start=0)
+        self.node_idx = count(start=self._scratch.get("node_idx", 0))
 
         self.analyze()
 
+        self._scratch["node_idx"] = next(self.node_idx)
+
     def _check(self):
         jumptables = self.kb.cfgs.get_most_accurate().jump_tables
         switch_jump_block_addrs = {
@@ -77,8 +84,39 @@ class SwitchDefaultCaseDuplicator(OptimizationPass):
         out_graph = None
         duplicated_default_addrs: set[int] = set()
 
+        default_addr_count = defaultdict(int)
+        goto_rewritten_default_addrs = set()
+        for _, _, default_addr in default_case_node_addrs:
+            default_addr_count[default_addr] += 1
+        for default_addr, cnt in default_addr_count.items():
+            if cnt > 1:
+                # rewrite all of them into gotos
+                default_node = self._get_block(default_addr)
+                for switch_head_addr in sorted((sa for sa, _, da in default_case_node_addrs if da == default_addr)):
+                    switch_head_node = self._get_block(switch_head_addr)
+                    goto_stmt = Jump(
+                        None,
+                        Const(None, None, default_addr, self.project.arch.bits, ins_addr=default_addr),
+                        target_idx=None,  # I'm assuming the ID of the default node is None here
+                        ins_addr=default_addr,
+                    )
+                    goto_node = Block(
+                        default_addr,
+                        0,
+                        statements=[goto_stmt],
+                        idx=next(self.node_idx),
+                    )
+
+                    if out_graph is None:
+                        out_graph = self._graph
+                    out_graph.remove_edge(switch_head_node, default_node)
+                    out_graph.add_edge(switch_head_node, goto_node)
+                    out_graph.add_edge(goto_node, default_node)
+
+                goto_rewritten_default_addrs.add(default_addr)
+
         for switch_head_addr, jump_node_addr, default_addr in default_case_node_addrs:
-            if default_addr in duplicated_default_addrs:
+            if default_addr in duplicated_default_addrs or default_addr in goto_rewritten_default_addrs:
                 continue
 
             default_case_node = self._func.get_node(default_addr)