Add development scripts for comparing Arancini dumps

1 files changed, 214 insertions, 0 deletions

diff --git a/run.py b/run.py
new file mode 100755
index 0000000..f1f1060
--- /dev/null
+++ b/run.py
@@ -0,0 +1,214 @@
+#! /bin/python3
+import os
+import re
+import sys
+import lldb
+import shutil
+import argparse
+
+from utils import print_separator
+
+verbose = False
+
+regnames = ['PC',
+            'RAX',
+            'RBX',
+            'RCX',
+            'RDX',
+            'RSI',
+            'RDI',
+            'RBP',
+            'RSP',
+            'R8',
+            'R9',
+            'R10',
+            'R11',
+            'R12',
+            'R13',
+            'R14',
+            'R15',
+            'RFLAGS']
+
+class DebuggerCallback:
+    def __init__(self, ostream=sys.stdout, skiplist: set = {}):
+        self.stream = ostream
+        self.regex = re.compile('(' + '|'.join(regnames) + ')$')
+        self.skiplist = skiplist
+
+    @staticmethod
+    def parse_flags(flag_reg: int):
+        flags = {'ZF': 0,
+                 'CF': 0,
+                 'OF': 0,
+                 'SF': 0,
+                 'PF': 0,
+                 'DF': 0}
+
+        # CF (Carry flag) Bit 0
+        # PF (Parity flag) Bit 2
+        # ZF (Zero flag) Bit 6
+        # SF (Sign flag) Bit 7
+        # TF (Trap flag) Bit 8
+        # IF (Interrupt enable flag) Bit 9
+        # DF (Direction flag) Bit 10
+        # OF (Overflow flag) Bit 11
+        flags['CF'] = int(0 != flag_reg & 1)
+        flags['ZF'] = int(0 != flag_reg & (1 << 6))
+        flags['OF'] = int(0 != flag_reg & (1 << 11))
+        flags['SF'] = int(0 != flag_reg & (1 << 7))
+        flags['DF'] = int(0 != flag_reg & (1 << 10))
+        flags['PF'] = int(0 != flag_reg & (1 << 1))
+        return flags
+
+
+    def print_regs(self, frame):
+        for reg in frame.GetRegisters():
+            for sub_reg in reg:
+                match = self.regex.match(sub_reg.GetName().upper())
+                if match and match.group() == 'RFLAGS':
+                    flags = DebuggerCallback.parse_flags(int(sub_reg.GetValue(),
+                                                             base=16))
+                    for flag in flags:
+                        print(f'flag {flag}:\t{hex(flags[flag])}',
+                                                     file=self.stream)
+                elif match:
+                    print(f"{sub_reg.GetName().upper()}:\t\t {hex(int(sub_reg.GetValue(), base=16))}",
+                          file=self.stream)
+
+    def print_stack(self, frame, element_count: int):
+        first = True
+        for i in range(element_count):
+            addr = frame.GetSP() + i * frame.GetThread().GetProcess().GetAddressByteSize()
+            error = lldb.SBError()
+            stack_value = int(frame.GetThread().GetProcess().ReadPointerFromMemory(addr, error))
+            if error.Success() and not first:
+                print(f'{hex(stack_value)}', file=self.stream)
+            elif error.Success():
+                print(f'{hex(stack_value)}\t\t<- rsp', file=self.stream)
+            else:
+                print(f"Error reading memory at address 0x{addr:x}",
+                      file=self.stream)
+            first=False
+
+    def __call__(self, frame):
+        pc = frame.GetPC()
+
+        # Skip this PC
+        if pc in self.skiplist:
+            self.skiplist.discard(pc)
+            return False
+
+        print_separator('=', stream=self.stream, count=20)
+        print(f'INVOKE PC={hex(pc)}', file=self.stream)
+        print_separator('=', stream=self.stream, count=20)
+
+        print("Register values:", file=self.stream)
+        self.print_regs(frame)
+        print_separator(stream=self.stream)
+
+        print("STACK:", file=self.stream)
+        self.print_stack(frame, 20)
+
+        return True  # Continue execution
+
+class Debugger:
+    def __init__(self, program):
+        self.debugger = lldb.SBDebugger.Create()
+        self.debugger.SetAsync(False)
+        self.target = self.debugger.CreateTargetWithFileAndArch(program,
+                                                                lldb.LLDB_ARCH_DEFAULT)
+        self.module = self.target.FindModule(self.target.GetExecutable())
+        self.interpreter = self.debugger.GetCommandInterpreter()
+
+    def set_breakpoint_by_addr(self, address: int):
+        command = f"b -a {address} -s {self.module.GetFileSpec().GetFilename()}"
+        result = lldb.SBCommandReturnObject()
+        self.interpreter.HandleCommand(command, result)
+
+        if verbose:
+            print(f'Set breakpoint at address {hex(address)}')
+
+    def get_breakpoints_count(self):
+        return self.target.GetNumBreakpoints()
+
+    def execute(self, callback: callable):
+        error = lldb.SBError()
+        listener = self.debugger.GetListener()
+        process = self.target.Launch(listener, None, None, None, None, None, None, 0,
+                                     True, error)
+
+        # Check if the process has launched successfully
+        if process.IsValid():
+            print(f'Launched process: {process}')
+        else:
+            print('Failed to launch process', file=sys.stderr)
+
+        while True:
+            state = process.GetState()
+            if state == lldb.eStateStopped:
+                 for thread in process:
+                    callback(thread.GetFrameAtIndex(0))
+                 process.Continue()
+            if state == lldb.eStateExited:
+                break
+
+        self.debugger.Terminate()
+
+        print(f'Process state: {process.GetState()}')
+        print('Program output:')
+        print(process.GetSTDOUT(1024))
+        print(process.GetSTDERR(1024))
+
+class ListWriter:
+    def __init__(self):
+        self.data = []
+
+    def write(self, s):
+        self.data.append(s)
+
+    def __str__(self):
+        return "".join(self.data)
+
+class Runner:
+    def __init__(self, dbt_log: list, oracle_program: str):
+        self.log = dbt_log
+        self.program = oracle_program
+        self.debugger = Debugger(self.program)
+        self.writer = ListWriter()
+
+    @staticmethod
+    def get_addresses(lines: list):
+        addresses = []
+
+        backlist = []
+        backlist_regex = re.compile(r'^\s\s\d*:')
+
+        skiplist = set()
+        for l in lines:
+            if l.startswith('INVOKE'):
+                addresses.append(int(l.split('=')[1].strip(), base=16))
+
+                if addresses[-1] in backlist:
+                    skiplist.add(addresses[-1])
+                backlist = []
+
+            if backlist_regex.match(l):
+                backlist.append(int(l.split()[0].split(':')[0], base=16))
+
+        return set(addresses), skiplist
+
+    def run(self):
+        # Get all addresses to stop at
+        addresses, skiplist = Runner.get_addresses(self.log)
+
+        # Set breakpoints
+        for address in addresses:
+            self.debugger.set_breakpoint_by_addr(address)
+
+        # Sanity check
+        assert(self.debugger.get_breakpoints_count() == len(addresses))
+
+        self.debugger.execute(DebuggerCallback(self.writer, skiplist))
+
+        return self.writer.data
+