Implement symbolic comparison and match traces via Miasm

Co-authored-by: Theofilos Augoustis <theofilos.augoustis@gmail.com>
Co-authored-by: Nicola Crivellin <nicola.crivellin98@gmail.com>

1 files changed, 72 insertions, 46 deletions

diff --git a/lldb_target.py b/lldb_target.py
index f587b37..b96f66b 100644
--- a/lldb_target.py
+++ b/lldb_target.py
@@ -1,6 +1,6 @@
 import lldb
 
-from arch import x86
+from arch import supported_architectures, x86
 from snapshot import ProgramState
 
 class MemoryMap:
@@ -53,15 +53,7 @@ class LLDBConcreteTarget:
             raise RuntimeError(f'[In LLDBConcreteTarget.__init__]: Failed to'
                                f' launch process.')
 
-    def set_breakpoint(self, addr):
-        command = f'b -a {addr} -s {self.module.GetFileSpec().GetFilename()}'
-        result = lldb.SBCommandReturnObject()
-        self.interpreter.HandleCommand(command, result)
-
-    def remove_breakpoint(self, addr):
-        command = f'breakpoint delete {addr}'
-        result = lldb.SBCommandReturnObject()
-        self.interpreter.HandleCommand(command, result)
+        self.archname = self.target.GetPlatform().GetTriple().split('-')[0]
 
     def is_exited(self):
         """Signals whether the concrete process has exited.
@@ -84,11 +76,48 @@ class LLDBConcreteTarget:
         thread: lldb.SBThread = self.process.GetThreadAtIndex(0)
         thread.StepInstruction(False)
 
+    def record_snapshot(self) -> ProgramState:
+        """Record the concrete target's state in a ProgramState object."""
+        # Determine current arch
+        if self.archname not in supported_architectures:
+            print(f'[ERROR] LLDBConcreteTarget: Recording snapshots is not'
+                  f' supported for architecture {self.archname}!')
+            raise NotImplementedError()
+        arch = supported_architectures[self.archname]
+
+        state = ProgramState(arch)
+
+        # Query and store register state
+        for regname in arch.regnames:
+            try:
+                conc_val = self.read_register(regname)
+                state.set(regname, conc_val)
+            except KeyError:
+                pass
+            except ConcreteRegisterError:
+                # Special rule for flags on X86
+                if arch.archname == x86.archname:
+                    rflags = x86.decompose_rflags(self.read_register('rflags'))
+                    if regname in rflags:
+                        state.set(regname, rflags[regname])
+
+        # Query and store memory state
+        for mapping in self.get_mappings():
+            assert(mapping.end_address > mapping.start_address)
+            size = mapping.end_address - mapping.start_address
+            try:
+                data = self.read_memory(mapping.start_address, size)
+                state.write_memory(mapping.start_address, data)
+            except ConcreteMemoryError:
+                pass
+
+        return state
+
     def _get_register(self, regname: str) -> lldb.SBValue:
         """Find a register by name.
 
-        :raise SimConcreteRegisterError: If no register with the specified name
-                                         can be found.
+        :raise ConcreteRegisterError: If no register with the specified name
+                                      can be found.
         """
         frame = self.process.GetThreadAtIndex(0).GetFrameAtIndex(0)
         reg = frame.FindRegister(regname)
@@ -99,6 +128,12 @@ class LLDBConcreteTarget:
         return reg
 
     def read_register(self, regname: str) -> int:
+        """Read the value of a register.
+
+        :raise ConcreteRegisterError: If `regname` is not a valid register name
+                                      or the target is otherwise unable to read
+                                      the register's value.
+        """
         reg = self._get_register(regname)
         val = reg.GetValue()
         if val is None:
@@ -109,6 +144,12 @@ class LLDBConcreteTarget:
         return int(val, 16)
 
     def write_register(self, regname: str, value: int):
+        """Read the value of a register.
+
+        :raise ConcreteRegisterError: If `regname` is not a valid register name
+                                      or the target is otherwise unable to set
+                                      the register's value.
+        """
         reg = self._get_register(regname)
         error = lldb.SBError()
         reg.SetValueFromCString(hex(value), error)
@@ -118,19 +159,27 @@ class LLDBConcreteTarget:
                 f' {regname} to value {hex(value)}!')
 
     def read_memory(self, addr, size):
+        """Read bytes from memory.
+
+        :raise ConcreteMemoryError: If unable to read `size` bytes from `addr`.
+        """
         err = lldb.SBError()
         content = self.process.ReadMemory(addr, size, err)
         if not err.success:
             raise ConcreteMemoryError(f'Error when reading {size} bytes at'
-                                         f' address {hex(addr)}: {err}')
+                                      f' address {hex(addr)}: {err}')
         return content
 
-    def write_memory(self, addr, value):
+    def write_memory(self, addr, value: bytes):
+        """Write bytes to memory.
+
+        :raise ConcreteMemoryError: If unable to write at `addr`.
+        """
         err = lldb.SBError()
         res = self.process.WriteMemory(addr, value, err)
         if not err.success or res != len(value):
             raise ConcreteMemoryError(f'Error when writing to address'
-                                         f' {hex(addr)}: {err}')
+                                      f' {hex(addr)}: {err}')
 
     def get_mappings(self) -> list[MemoryMap]:
         mmap = []
@@ -151,35 +200,12 @@ class LLDBConcreteTarget:
                                   perms))
         return mmap
 
-def record_snapshot(target: LLDBConcreteTarget) -> ProgramState:
-    """Record a concrete target's state in a ProgramState object.
+    def set_breakpoint(self, addr):
+        command = f'b -a {addr} -s {self.module.GetFileSpec().GetFilename()}'
+        result = lldb.SBCommandReturnObject()
+        self.interpreter.HandleCommand(command, result)
 
-    :param target: The target from which to query state. Currently assumes an
-                   X86 target.
-    """
-    state = ProgramState(x86.ArchX86())
-
-    # Query and store register state
-    rflags = x86.decompose_rflags(target.read_register('rflags'))
-    for regname in x86.regnames:
-        try:
-            conc_val = target.read_register(regname)
-            state.set(regname, conc_val)
-        except KeyError:
-            pass
-        except ConcreteRegisterError:
-            if regname in rflags:
-                state.set(regname, rflags[regname])
-
-    # Query and store memory state
-    for mapping in target.get_mappings():
-        assert(mapping.end_address > mapping.start_address)
-        size = mapping.end_address - mapping.start_address
-        try:
-            data = target.read_memory(mapping.start_address, size)
-            state.write_memory(mapping.start_address, data)
-        except ConcreteMemoryError:
-            # Unable to read memory from mapping
-            pass
-
-    return state
+    def remove_breakpoint(self, addr):
+        command = f'breakpoint delete {addr}'
+        result = lldb.SBCommandReturnObject()
+        self.interpreter.HandleCommand(command, result)