Merge pull request #152 from angr/feat/concrete_setter

Feat/concrete setter

Author: Kyle-Kyle

angrop/chain_builder/mem_writer.py

@@ -3,6 +3,7 @@
 from collections import defaultdict, Counter
 from functools import cmp_to_key
 
+import claripy
 import networkx as nx
 from angr.errors import SimUnsatError
 
@@ -15,12 +16,76 @@
 
 l = logging.getLogger(__name__)
 
+class ConcreteRegSetter(Builder):
+    """
+    a chain builder that aims to find best gadgets that can set registers to
+    concrete values
+    """
+    def __init__(self, chain_builder):
+        super().__init__(chain_builder)
+        self._concrete_setting_gadgets = None
+        self._concrete_setting_cache = defaultdict(list)
+
+    def bootstrap(self):
+        self._concrete_setting_gadgets = self.filter_gadgets(self.chain_builder.gadgets)
+        for g in self._concrete_setting_gadgets:
+            reg = list(g.concrete_regs)[0]
+            self._concrete_setting_cache[reg].append(g)
+
+    def _effect_tuple(self, g):
+        assert len(g.concrete_regs) == 1
+        reg = list(g.concrete_regs)[0]
+        return (reg, g.concrete_regs[reg])
+
+    def _comparison_tuple(self, g):
+        assert len(g.concrete_regs) == 1
+        reg = list(g.concrete_regs)[0]
+        return (len(g.changed_regs-{reg}), g.stack_change, g.num_sym_mem_access,
+                   g.isn_count, int(g.has_conditional_branch is True))
+
+    def filter_gadgets(self, gadgets):
+        gadgets = [g for g in gadgets if len(g.concrete_regs) == 1 and not g.num_sym_mem_access]
+        return self._filter_gadgets(gadgets)
+
+class ConcreteRegChanger(Builder):
+    """
+    a chain builder that aims to craft register values using concrete value
+    effect in gadgets
+    """
+    def __init__(self, chain_builder):
+        super().__init__(chain_builder)
+        self._concrete_reg_changing_gadgets = None
+        self._concrete_reg_changing_cache = defaultdict(list)
+
+    def bootstrap(self):
+        self._concrete_reg_changing_gadgets = self.filter_gadgets(self.chain_builder.gadgets)
+        for g in self._concrete_reg_changing_gadgets:
+            reg = list(g.concrete_reg_changes)[0]
+            self._concrete_reg_changing_cache[reg].append(g)
+
+    def _effect_tuple(self, g):
+        reg = list(g.concrete_reg_changes.keys())[0]
+        init_ast, final_ast = g.concrete_reg_changes[reg]
+        val = claripy.algorithm.replace(expr=final_ast,
+                                        old=init_ast,
+                                        new=claripy.BVV(0, self.project.arch.bits))
+        op = final_ast.op
+        if op in ('ZeroExt', 'SignExt'):
+            op = final_ast.args[1].op
+        return (reg, op, val.concrete_value)
+
+    def _comparison_tuple(self, g):
+        reg = list(g.concrete_reg_changes)[0]
+        return (len(g.changed_regs-{reg}), g.stack_change, g.num_sym_mem_access,
+                   g.isn_count, int(g.has_conditional_branch is True))
+
+    def filter_gadgets(self, gadgets):
+        gadgets = [g for g in gadgets if len(g.concrete_reg_changes) == 1 and not g.num_sym_mem_access]
+        return self._filter_gadgets(gadgets)
+
 class RegSetter(Builder):
     """
-    a chain builder that aims to set registers using different algorithms
-    1. algo1: graph-search, fast, not reliable
-    2. algo2: pop-only bfs search, fast, reliable, can generate chains to bypass bad-bytes
-    3. algo3: riscy-rop inspired backward search, slow, can utilize gadgets containing conditional branches
+    a chain builder that aims to set registers using the graph search algorithm
     """
 
     #### Inits ####
@@ -32,8 +97,12 @@ def __init__(self, chain_builder):
         # Estimate of how difficult it is to set each register.
         # all self-contained and not symbolic access
         self._reg_setting_dict: dict[str, list] = defaultdict(list)
+        self._concrete_reg_setter = ConcreteRegSetter(chain_builder)
+        self._concrete_reg_changer = ConcreteRegChanger(chain_builder)
 
     def bootstrap(self):
+        self._concrete_reg_setter.bootstrap()
+        self._concrete_reg_changer.bootstrap()
         self._reg_setting_gadgets = self.filter_gadgets(self.chain_builder.gadgets)
 
         # update reg_setting_dict
@@ -603,7 +672,7 @@ def _handle_hard_regs(self, gadgets, registers, preserve_regs) -> list[RopGadget
             if hard_chains:
                 hard_chain = hard_chains[0]
             else:
-                hard_chain = self._find_add_chain(gadgets, reg, val)
+                hard_chain = self._find_add_chain(reg, val)
             if hard_chain:
                 self.hard_chain_cache[key] = hard_chain # we cache the result even if it fails
 
@@ -627,16 +696,22 @@ def _find_concrete_chains(gadgets, registers):
                     chains.append([g])
         return chains
 
-    def _find_add_chain(self, gadgets, reg, val) -> list[RopGadget|RopBlock]:
+    def _find_add_chain(self, reg, val) -> list[RopGadget|RopBlock]:
         """
         find one chain to set one single register to a specific value using concrete values only through add/dec
         """
         val = rop_utils.cast_rop_value(val, self.project)
-        concrete_setter_gadgets = [ x for x in gadgets if reg in x.concrete_regs ]
-        delta_gadgets = [ x for x in gadgets if len(x.reg_dependencies) == 1 and reg in x.reg_dependencies\
-                            and len(x.reg_dependencies[reg]) == 1 and reg in x.reg_dependencies[reg]]
+        arch_bits = self.project.arch.bits
+        concrete_setter_gadgets = self._concrete_reg_setter._concrete_setting_cache[reg]
+        delta_gadgets = self._concrete_reg_changer._concrete_reg_changing_cache[reg]
         for g1 in concrete_setter_gadgets:
             for g2 in delta_gadgets:
+                init_ast, final_ast = g2.concrete_reg_changes[reg]
+                ast = claripy.algorithm.replace(expr=final_ast,
+                                                old=init_ast,
+                                                new=claripy.BVV(g1.concrete_regs[reg], arch_bits))
+                if ast.concrete_value != val.concreted:
+                    continue
                 try:
                     chain = self._build_reg_setting_chain([g1, g2], {reg: val})
                     state = chain.exec()
@@ -653,13 +728,9 @@ def _find_add_chain(self, gadgets, reg, val) -> list[RopGadget|RopBlock]:
 
     def _effect_tuple(self, g):
         v1 = tuple(sorted(g.popped_regs))
-        v2 = tuple(sorted(g.concrete_regs.items()))
-        v3 = []
-        for x,y in g.reg_dependencies.items():
-            v3.append((x, tuple(sorted(y))))
-        v3 = tuple(sorted(v3))
-        v4 = g.transit_type
-        return (v1, v2, v3, v4)
+        v2 = tuple(sorted(g.reg_moves))
+        v3 = g.transit_type
+        return (v1, v2, v3)
 
     def _comparison_tuple(self, g):
         return (len(g.changed_regs-g.popped_regs), g.stack_change, g.num_sym_mem_access,

angrop/chain_builder/reg_setter.py

@@ -151,6 +151,6 @@ def sigreturn(self, **registers):
             chain.add_value(word)
 
         # save sigreturn frame information for pretty printing
-        chain._sigreturn_frames.append((frame, frame_start_offset))
+        chain._sigreturn_frame = (frame, frame_start_offset)
 
         return chain

angrop/chain_builder/sigreturn.py

@@ -596,25 +596,36 @@ def _check_reg_changes(self, final_state, _, gadget):
             else:
                 gadget.changed_regs.add(reg)
 
-    def _check_reg_change_dependencies(self, symbolic_state, symbolic_p, gadget):
+    def _check_reg_change_dependencies(self, init_state, final_state, gadget):
         """
         Checks which registers affect register changes
         :param symbolic_state: the input state for testing
         :param symbolic_p: the stepped path, symbolic_state is an ancestor of it.
         :param gadget: the gadget to store the reg change dependencies in
         """
+        arch_bits = self.project.arch.bits
         for reg in gadget.changed_regs:
             # skip popped regs
             if reg in gadget.popped_regs:
                 continue
             # check its dependencies and controllers
-            dependencies = self._get_reg_dependencies(symbolic_p, reg)
+            dependencies = self._get_reg_dependencies(final_state, reg)
             if len(dependencies) != 0:
                 gadget.reg_dependencies[reg] = set(dependencies)
-                controllers = self._get_reg_controllers(symbolic_state, symbolic_p, reg, dependencies)
+                controllers = self._get_reg_controllers(init_state, final_state, reg, dependencies)
                 if controllers:
                     gadget.reg_controllers[reg] = set(controllers)
 
+            # record simple register changing effects
+            init_reg = init_state.registers.load(reg)
+            final_reg = final_state.registers.load(reg)
+            if init_reg is final_reg:
+                continue
+            ast = claripy.algorithm.replace(expr=final_reg, old=init_reg, new=claripy.BVV(0, arch_bits))
+            if ast.symbolic:
+                continue
+            gadget.concrete_reg_changes[reg] = (init_reg, final_reg)
+
     def _check_pop_equal_set(self, gadget, final_state):
         """
         identify the situation where the final registers are dependent on each other

angrop/gadget_finder/gadget_analyzer.py

@@ -40,8 +40,8 @@ def __init__(self, project, builder, state=None, badbytes=None):
         self._pivoted = False
         self._init_sp = None
 
-        # sigreturn frame information: list of (frame_object, start_offset_in_values)
-        self._sigreturn_frames = []
+        # sigreturn frame information: (frame_object, start_offset_in_values)
+        self._sigreturn_frame = None
 
     def __add__(self, other):
         # need to add the values from the other's stack and the constraints to the result state
@@ -66,9 +66,10 @@ def __add__(self, other):
 
         # merge sigreturn frames: adjust offsets from other
         word_count_before_merge = len(self._values) - (1 if idx is not None else 0)
-        for frame, offset in other._sigreturn_frames:
+        if other._sigreturn_frame:
+            frame, offset = other._sigreturn_frame
             adjusted_offset = word_count_before_merge + offset
-            result._sigreturn_frames.append((frame, adjusted_offset))
+            result._sigreturn_frame = (frame, adjusted_offset)
 
         # FIXME: cannot handle cases where a rop_block is used twice and have different constraints
         # because right now symbolic values go with rop_blocks
@@ -229,7 +230,7 @@ def copy(self):
         cp.payload_len = self.payload_len
         cp._blank_state = self._blank_state.copy()
         cp.badbytes = self.badbytes
-        cp._sigreturn_frames = list(self._sigreturn_frames)
+        cp._sigreturn_frame = self._sigreturn_frame
 
         cp._pivoted = self._pivoted
         cp._init_sp = self._init_sp
@@ -406,29 +407,25 @@ def dstr(self):
         prefix_len = bs*2+2
         prefix = " "*prefix_len
 
-        sigreturn_map = {} # start_offset -> frame and end offset
-        for frame, start_offset in self._sigreturn_frames:
-            # iterate through frame registers to build a map
-            frame_words = frame.to_words()
-            sigreturn_map[start_offset] = (frame, start_offset + len(frame_words))
-        idx = 0
-        while idx < len(self._values):
-            v = self._values[idx]
+        # a chain may end with a sigreturn frame, in that case, we print up to the frame
+        # and then print the frame differently
+        top = len(self._values) if self._sigreturn_frame is None else self._sigreturn_frame[1]
+        for i in range(top):
+            v = self._values[i]
             if v.symbolic:
                 res += prefix + f"  {v.ast}\n"
-                idx += 1
                 continue
             for g in self._gadgets:
                 if g.addr == v.concreted:
                     fmt = f"%#0{prefix_len}x"
                     res += fmt % g.addr + f": {g.dstr()}\n"
                     break
             else:
-                if idx in sigreturn_map:
-                    sigframe, idx = sigreturn_map[idx]
-                    res += sigframe.dstr(prefix=prefix)
-                    continue
-            idx += 1
+                res += prefix + f"  {v.concreted:#x}\n"
+
+        # if there is a sigreturn frame, print it
+        if self._sigreturn_frame:
+            res += self._sigreturn_frame[0].dstr(prefix=prefix)
         return res
 
     def pp(self):

angrop/rop_chain.py

@@ -130,6 +130,7 @@ def __init__(self):
         # Stores the stack variables that each register depends on.
         # Used to check for cases where two registers are popped from the same location.
         self.concrete_regs = {}
+        self.concrete_reg_changes = {}
         self.reg_dependencies = {}  # like rax might depend on rbx, rcx
         self.reg_controllers = {}  # like rax might be able to be controlled by rbx (for any value of rcx)
         self.reg_pops = set()
@@ -231,6 +232,7 @@ def copy_effect(self, cp):
         cp.stack_change = self.stack_change
         cp.changed_regs = set(self.changed_regs)
         cp.concrete_regs = dict(self.concrete_regs)
+        cp.concrete_reg_changes = dict(self.concrete_reg_changes)
         cp.reg_dependencies = dict(self.reg_dependencies)
         cp.reg_controllers = dict(self.reg_controllers)
         cp.reg_pops = set(self.reg_pops)

angrop/rop_effect.py

@@ -1,8 +1,10 @@
 import os
+import time
 import logging
 
 import angr
 import angrop  # pylint: disable=unused-import
+from angrop.errors import RopException
 
 l = logging.getLogger(__name__)
 
@@ -87,9 +89,10 @@ def test_badbyte_transform():
     rop = proj.analyses.ROP()
 
     if os.path.exists(cache_path):
-        rop.load_gadgets(cache_path, optimize=False)
+        rop.load_gadgets(cache_path)
     else:
         rop.find_gadgets()
+        rop.save_gadgets(cache_path)
 
     rop.set_badbytes([0x00, 0x0A])
     chain = rop.write_to_mem(0xdeadbeef, b"\x00", fill_byte=b"A")
@@ -106,24 +109,42 @@ def test_badbyte_multibyte():
     rop = proj.analyses.ROP()
 
     if os.path.exists(cache_path):
-        rop.load_gadgets(cache_path, optimize=False)
+        rop.load_gadgets(cache_path)
     else:
         rop.find_gadgets()
+        rop.save_gadgets(cache_path)
 
     rop.set_badbytes([0x00, 0x0A])
     target = b"\x00\x00\x00\x42"
     chain = rop.write_to_mem(0xdeadbf00, target, fill_byte=b"\xff")
     state = chain.exec()
-    endian = "little" if proj.arch.memory_endness == "Iend_LE" else "big"
-    loaded = state.memory.load(
-        0xdeadbf00, len(target), endness=proj.arch.memory_endness
-    ).concrete_value
-    assert loaded == int.from_bytes(target, endian)
+    data = state.solver.eval(state.memory.load(0xdeadbf00, len(target)), cast_to=bytes)
+    assert data == target
 
     payload = chain.payload_str()
     assert b"\x00" not in payload
     assert b"\x0A" not in payload
 
+def test_hard_regs_loop():
+    cache_path = os.path.join(data_dir, "bronze_ropchain")
+    proj = angr.Project(os.path.join(tests_dir, "i386", "bronze_ropchain"), auto_load_libs=False)
+    rop = proj.analyses.ROP()
+
+    if os.path.exists(cache_path):
+        rop.load_gadgets(cache_path)
+    else:
+        rop.find_gadgets()
+        rop.save_gadgets(cache_path)
+
+    rop.set_badbytes([0x00, 0x0A])
+    start = time.time()
+    try:
+        rop.set_regs(eax=0x4200009a, edx=0xdeadbefc)
+    except RopException:
+        pass
+
+    # if should take less than 0.1s, but let's be lenient here
+    assert time.time() - start < 2
 
 def run_all():
     functions = globals()