Merge branch 'master' into loader_simplify

Author: yhql

examples/ledger_ctf2/ripped.py

@@ -29,7 +29,7 @@ def srand(em):
 # We'd like to trace everything we can
 # - memory accesses
 # - modified registers
-e = rainbow_x64(trace_config=TraceConfig(mem_value=HammingWeight(), register=HammingWeight()), allow_stubs=True)
+e = rainbow_x64(trace_config=TraceConfig(mem_value=HammingWeight(), register=HammingWeight()))
 e.stubbed_functions = {
     "time": time,
     "clock_gettime": clock_gettime,

examples/ledger_ctf2/ripped2.py

@@ -21,7 +21,7 @@ def rand(em):
     em["rax"] = 7
 
 
-e = rainbow_x64(trace_config=TraceConfig(mem_value=Identity()), allow_stubs=True)
+e = rainbow_x64(trace_config=TraceConfig(mem_value=Identity()))
 e.load("ctf2", typ=".elf", except_missing_libs=False)
 e.setup()
 

rainbow/generics/arm.py

@@ -52,11 +52,12 @@ def reset_stack(self):
     def return_force(self):
         self["pc"] = self["lr"]
 
-    def _block_hook(self, uci, address, size, user_data):
-        if self.thumb_bit == 0:
-            # switch disassembler to ARM mode
-            self.disasm.mode = cs.CS_MODE_ARM
-        else:
-            self.disasm.mode = cs.CS_MODE_THUMB
-
-        super()._block_hook(uci, address | self.thumb_bit, size, user_data)
+    def disassemble_single(self, addr: int, size: int):
+        # Switch disassembler to ARM or Thumb mode
+        self.disasm.mode = cs.CS_MODE_ARM if self.thumb_bit == 0 else cs.CS_MODE_THUMB
+        return super().disassemble_single(addr, size)
+
+    def disassemble_single_detailed(self, addr: int, size: int) -> cs.CsInsn:
+        # Switch disassembler to ARM or Thumb mode
+        self.disasm.mode = cs.CS_MODE_ARM if self.thumb_bit == 0 else cs.CS_MODE_THUMB
+        return super().disassemble_single_detailed(addr, size)

rainbow/generics/cortexm.py

@@ -50,6 +50,13 @@ def __init__(self, *args, **kwargs):
 
         self.emu.hook_add(uc.UC_HOOK_INTR, HookWeakMethod(self.intr_hook))
 
+        self.emu.hook_add(
+            uc.UC_HOOK_BLOCK,
+            HookWeakMethod(self.endmem_hook),
+            begin=0xfffffff0,
+            end=0xffffffff,
+        )
+
     def reset_stack(self):
         self.emu.reg_write(uc.arm_const.UC_ARM_REG_SP, self.STACK_ADDR)
 
@@ -81,7 +88,7 @@ def start(self, begin, *args, **kwargs):
     def return_force(self):
         self["pc"] = self["lr"]
 
-    def _block_hook(self, uci, address, size, user_data):
+    def endmem_hook(self, _uci, address, _size, _user_data):
         if (address & 0xfffffff0) == 0xfffffff0:
             is_psp = (address >> 2) & 1
             is_unpriv = (address >> 3) & 1
@@ -96,8 +103,3 @@ def _block_hook(self, uci, address, size, user_data):
 
             self['ipsr'] = 0
             self['sp'] = sp + 32
-            return
-
-        # In ARM Cortex-M, all code is in thumb mode
-        # So all function addresses are odd
-        super()._block_hook(uci, address | 1, size, user_data)

rainbow/rainbow.py

@@ -44,17 +44,22 @@ class Print(Flag):
 class TraceConfig:
     """Tracing configuration."""
 
-    def __init__(self,
-                 mem_address: Optional[LeakageModel] = None,
-                 mem_value: Optional[LeakageModel] = None,
-                 register: Optional[LeakageModel] = None,
-                 instruction: bool = False,
-                 ignored_registers: Optional[Set[str]] = None):
+    def __init__(
+        self,
+        mem_address: Optional[LeakageModel] = None,
+        mem_value: Optional[LeakageModel] = None,
+        register: Optional[LeakageModel] = None,
+        instruction: bool = False,
+        ignored_registers: Optional[Set[str]] = None,
+        only_when_registers: bool = True,
+    ):
         self.mem_address = mem_address
         self.mem_value = mem_value
         self.register = register
         self.instructions = instruction
         self.ignored_registers = ignored_registers
+        # When true, only output register leakage when at least one register is written
+        self.only_when_registers = only_when_registers
 
 
 class Rainbow(abc.ABC):
@@ -83,18 +88,21 @@ class Rainbow(abc.ABC):
     ENDIANNESS: str
     PC: int
 
-    last_regs: Optional[List[str]]
+    last_regs: List[str]
     last_reg_values: Optional[Dict[str, int]]
     last_address: Optional[int]
     last_value: Optional[int]
     trace: List[Any]
 
-    block_hook: Optional[int]
     mem_hook: Optional[int]
     code_hook: Optional[int]
 
-    def __init__(self, print_config: Print = Print(0), trace_config: TraceConfig = TraceConfig(),
-                 allow_breakpoints: bool = False, allow_stubs: bool = False):
+    def __init__(
+        self,
+        print_config: Print = Print(0),
+        trace_config: TraceConfig = TraceConfig(),
+        allow_breakpoints: bool = False,
+    ):
         self.breakpoints = set()
         self.functions = {}
         self.function_names = {}
@@ -104,7 +112,6 @@ def __init__(self, print_config: Print = Print(0), trace_config: TraceConfig = T
         self.print_config = print_config
         self.trace_config = trace_config
         self.allow_breakpoints = allow_breakpoints
-        self.allow_stubs = allow_stubs
 
         # Leak storage
         self.last_reg_values = {}
@@ -311,12 +318,17 @@ def start_and_fault(self, fault_model, fault_index: int, begin: int, end: int, *
         return pc_fault
 
     def setup(self):
-        """Add base hooks to the engine."""
-        # We need the block hook only if we are
-        # printing functions or need to handle stubs.
-        # if self.print_config & Print.Functions or self.allow_stubs:
-        self.block_hook = self.emu.hook_add(uc.UC_HOOK_BLOCK,
-                                            HookWeakMethod(self._block_hook))
+        """Setup engine hooks."""
+        # Hook functions calls for printing
+        if self.print_config & Print.Functions:
+            for addr, name in self.function_names.items():
+                self.emu.hook_add(
+                    uc.UC_HOOK_BLOCK,
+                    self._print_function_hook,
+                    begin=addr,
+                    end=addr,
+                    user_data=name,
+                )
 
         # We need the mem hook only if we are
         # printing memory or tracing memory values or addresses.
@@ -394,91 +406,72 @@ def disassemble_single_detailed(self, addr: int, size: int) -> cs.CsInsn:
         insn = self.emu.mem_read(addr, 2 * size)
         return self._disassemble_cache(self.disasm.disasm, bytes(insn), addr)
 
+    def _get_addrs(self, name_or_addr):
+        if isinstance(name_or_addr, str):
+            # Stub all function addresses matching this name
+            addrs = [a for a, n in self.function_names.items() if n == name_or_addr]
+            if not addrs:
+                raise IndexError(f"'{name_or_addr}' could not be found.")
+            return addrs
+        elif isinstance(name_or_addr, int):
+            # Name is only one address
+            return [name_or_addr]
+        raise TypeError("name_or_addr should be function name or address")
+
+    def _stub_hook(self, _uci, _address, _size, userdata):
+        """Call user stub set up with hook_prolog/hook_bypass."""
+        fn, bypass = userdata
+        if fn is not None:
+            fn(self)
+        if bypass:
+            # Make the function return early
+            self.return_force()
+
     def hook_prolog(self, name, fn):
         """
         Add a call to function 'fn' when 'name' is called during execution.
         After executing 'fn, execution resumes into 'name'.
         """
-        if not self.allow_stubs:
-            raise ValueError("Cannot use stubs, allow_stubs is False.")
-
-        def to_hook(x):
-            if fn is not None:
-                fn(x)
-            return False
-
-        if isinstance(name, str):
-            # Stub all function addresses matching this name
-            addrs = [a for a, n in self.function_names.items() if n == name]
-            if not addrs:
-                raise IndexError(f"'{name}' could not be found.")
-            for addr in addrs:
-                self.stubbed_functions[addr] = to_hook
-        elif isinstance(name, int):
-            # Name is an address
-            self.stubbed_functions[name] = to_hook
-        else:
-            raise TypeError("name should be function name or address")
+        for addr in self._get_addrs(name):
+            self.stubbed_functions[addr] = self.emu.hook_add(
+                uc.UC_HOOK_BLOCK,
+                HookWeakMethod(self._stub_hook),
+                begin=addr,
+                end=addr,
+                user_data=(fn, False),
+            )
 
     def hook_bypass(self, name, fn=None):
         """
         Add a call to function 'fn' when 'name' is called during execution.
         After executing 'fn', execution returns to the caller.
         """
-        if not self.allow_stubs:
-            raise ValueError("Cannot use stubs, allow_stubs is False.")
-
-        def to_hook(x):
-            if fn is not None:
-                fn(x)
-            return True
-
-        if isinstance(name, str):
-            # Stub all function addresses matching this name
-            addrs = [a for a, n in self.function_names.items() if n == name]
-            if not addrs:
-                raise IndexError(f"'{name}' could not be found.")
-            for addr in addrs:
-                self.stubbed_functions[addr] = to_hook
-        elif isinstance(name, int):
-            # Name is an address
-            self.stubbed_functions[name] = to_hook
-        else:
-            raise TypeError("name should be function name or address")
+        for addr in self._get_addrs(name):
+            self.stubbed_functions[addr] = self.emu.hook_add(
+                uc.UC_HOOK_BLOCK,
+                HookWeakMethod(self._stub_hook),
+                begin=addr,
+                end=addr,
+                user_data=(fn, True),
+            )
 
     def remove_hook(self, name):
         """Remove the hook."""
-        if not self.allow_stubs:
-            raise ValueError("Cannot use stubs, allow_stubs is False.")
-        del self.stubbed_functions[name]
+        for addr in self._get_addrs(name):
+            if addr in self.stubbed_functions:
+                self.emu.hook_del(self.stubbed_functions[addr])
+                del self.stubbed_functions[addr]
 
     def remove_hooks(self):
-        """Remove the hooked functions."""
-        if not self.allow_stubs:
-            raise ValueError("Cannot use stubs, allow_stubs is False.")
-        self.stubbed_functions = {}
+        """Remove all hooked functions."""
+        for addr, hook in self.stubbed_functions.items():
+            self.emu.hook_del(hook)
+            del self.stubbed_functions[addr]
 
-    def _block_hook(self, _uci, address: int, _size, _):
-        """
-        Hook called on every jump to a basic block that checks if a known
-        address+function is redefined in the user's python script and if so,
-        calls that instead.
-        """
-        # Print function calls
-        if address in self.function_names and (self.allow_stubs or self.print_config & Print.Functions):
-            # Handle the function call printing
-            f = self.function_names[address]
-            if self.print_config & Print.Functions:
-                print(f"{color('MAGENTA', f)}(...) @ 0x{address:x}")
-
-        # If stub is enabled and set at this address, run it
-        if self.allow_stubs:
-            stub_func = self.stubbed_functions.get(address)
-            if stub_func is not None:
-                r = stub_func(self)
-                if r:
-                    # If stub returns True, then make the function return early
-                    self.return_force()
+    @staticmethod
+    def _print_function_hook(_uci, address: int, _size, name: str):
+        """Print function call."""
+        print(f"{color('MAGENTA', name)}(...) @ 0x{address:x}")
 
     def _mem_hook(self, uci, access, address, size, value, _):
         # Get the value
@@ -566,7 +559,7 @@ def _code_hook(self, uci, address, size, _):
             #  - last_reg_values are register values as they were before the previous instruction
             #
             # So we need to go over last_regs, get their prev values from last_reg_values and get their current values.
-            if self.last_regs:
+            if self.last_regs or not self.trace_config.only_when_registers:
                 reg_values = {r: uci.reg_read(self.REGS[r]) for r in self.last_regs}
                 leak = sum(
                     self.trace_config.register(reg_values[r], self.last_reg_values.get(r, 0)) for r in self.last_regs)
@@ -580,19 +573,18 @@ def _code_hook(self, uci, address, size, _):
             if ins is None:
                 ins = self.disassemble_single_detailed(address, size)
                 _, regs_written = ins.regs_access()
-                if regs_written:
-                    regs = list(filter(lambda r: r not in self.IGNORED_REGS and (
-                                not self.trace_config.ignored_registers or r not in self.trace_config.ignored_registers),
-                                       map(ins.reg_name, regs_written)))  # type: ignore
-                else:
-                    regs = None
+                regs = list(filter(lambda r: r not in self.IGNORED_REGS and (
+                            not self.trace_config.ignored_registers or r not in self.trace_config.ignored_registers),
+                                    map(ins.reg_name, regs_written)))  # type: ignore
 
         if self.trace_config.instructions:
-            if ins is None:
-                ins = self.disassemble_single_detailed(address, size)
+            if ins is not None:
+                address, mnemonic, op_str = ins.address, ins.mnemonic, ins.op_str
+            else:
+                address, _size, mnemonic, op_str = self.disassemble_single(address, size)
             if event is None:
                 event = {"type": "code"}
-            event["instruction"] = f"{ins.address:8X} {ins.mnemonic:<6}  {ins.op_str}"
+            event["instruction"] = f"{address:8X} {mnemonic:<6}  {op_str}"
         if event is not None:
             self.trace.append(event)
         self.last_regs = regs

tests/test_hook.py

@@ -3,7 +3,7 @@
 
 
 def test_hook_bypass_ctf2():
-    emu = rainbow_x64(allow_stubs=True)
+    emu = rainbow_x64()
     emu.load("examples/ledger_ctf2/ctf2")
     emu.setup()
 
@@ -15,22 +15,22 @@ def strtol(e):
 
 
 def test_hook_bypass_ctf2_empty():
-    emu = rainbow_x64(allow_stubs=True)
+    emu = rainbow_x64()
     emu.load("examples/ledger_ctf2/ctf2")
     emu.setup()
     emu.hook_bypass("strtol")
     emu.start(0xCA9, 0xDCE)
 
 
 def test_hook_bypass_missing_name():
-    emu = rainbow_x64(allow_stubs=True)
+    emu = rainbow_x64()
     emu.load("examples/ledger_ctf2/ctf2")
     with pytest.raises(IndexError):
         emu.hook_bypass("strtol_blabla")
 
 
 def test_hook_prolog_missing_name():
-    emu = rainbow_x64(allow_stubs=True)
+    emu = rainbow_x64()
     emu.load("examples/ledger_ctf2/ctf2")
 
     def strtol(e):
@@ -41,7 +41,7 @@ def strtol(e):
 
 
 def test_remove_hooks():
-    emu = rainbow_x64(allow_stubs=True)
+    emu = rainbow_x64()
     emu.load("examples/ledger_ctf2/ctf2")
     emu.setup()