add some basic trace analysis to help auto-resolve ASLR from traces

plugins/tenet/context.py

@@ -294,7 +294,8 @@ def interactive_next_execution(self):
         Handle UI actions for seeking to the next execution of the selected address.
         """
         address = disassembler[self].get_current_address()
-        result = self.reader.seek_to_next(address, BreakpointType.EXEC)
+        rebased_address = self.reader.analysis.rebase_pointer(address)
+        result = self.reader.seek_to_next(rebased_address, BreakpointType.EXEC)
 
         # TODO: blink screen? make failure more visible...
         if not result:
@@ -305,7 +306,8 @@ def interactive_prev_execution(self):
         Handle UI actions for seeking to the previous execution of the selected address.
         """
         address = disassembler[self].get_current_address()
-        result = self.reader.seek_to_prev(address, BreakpointType.EXEC)
+        rebased_address = self.reader.analysis.rebase_pointer(address)
+        result = self.reader.seek_to_prev(rebased_address, BreakpointType.EXEC)
 
         # TODO: blink screen? make failure more visible...
         if not result:
@@ -316,7 +318,8 @@ def interactive_first_execution(self):
         Handle UI actions for seeking to the first execution of the selected address.
         """
         address = disassembler[self].get_current_address()
-        result = self.reader.seek_to_first(address, BreakpointType.EXEC)
+        rebased_address = self.reader.analysis.rebase_pointer(address)
+        result = self.reader.seek_to_first(rebased_address, BreakpointType.EXEC)
 
         # TODO: blink screen? make failure more visible...
         if not result:
@@ -327,7 +330,8 @@ def interactive_final_execution(self):
         Handle UI actions for seeking to the final execution of the selected address.
         """
         address = disassembler[self].get_current_address()
-        result = self.reader.seek_to_final(address, BreakpointType.EXEC)
+        rebased_address = self.reader.analysis.rebase_pointer(address)
+        result = self.reader.seek_to_final(rebased_address, BreakpointType.EXEC)
 
         # TODO: blink screen? make failure more visible...
         if not result:
@@ -339,7 +343,7 @@ def _idx_changed(self, idx):
 
         This will make the disassembler track with the PC/IP of the trace reader. 
         """
-        disassembler[self].navigate(self.reader.ip)
+        disassembler[self].navigate(self.reader.rebased_ip)
 
     def _select_trace_file(self):
         """

plugins/tenet/integration/api/ida_api.py

@@ -21,6 +21,7 @@
 import ida_idaapi
 import ida_diskio
 import ida_kernwin
+import ida_segment
 
 from .api import DisassemblerCoreAPI, DisassemblerContextAPI
 from ...util.qt import *
@@ -204,6 +205,33 @@ def is_call_insn(self, address):
             return True
         return False
 
+    def get_instruction_addresses(self):
+        """
+        Return all instruction addresses from the executable.
+        """
+        instruction_addresses = []
+
+        for seg_address in idautils.Segments():
+
+            # fetch code segments
+            seg = ida_segment.getseg(seg_address)
+            if seg.sclass != ida_segment.SEG_CODE:
+                continue
+
+            current_address = seg_address
+            end_address = seg.end_ea
+
+            # save the address of each instruction in the segment
+            while current_address < end_address:
+                current_address = ida_bytes.next_head(current_address, end_address)
+                if ida_bytes.is_code(ida_bytes.get_flags(current_address)):
+                    instruction_addresses.append(current_address)
+
+        #    print(f"Seg {seg.start_ea:08X} --> {seg.end_ea:08X} CODE")
+        #print(f" -- {len(instruction_addresses):,} instructions found")
+
+        return instruction_addresses
+
     def is_mapped(self, address):
         return ida_bytes.is_mapped(address)
 

plugins/tenet/integration/ida_integration.py

@@ -440,18 +440,19 @@ def _highlight_disassesmbly(self, lines_out, widget, lines_in):
                 ida_color |= (0xFF - int(0xFF * percent)) << 24
 
                 # save the trail color
-                trail[address] = ida_color
+                rebased_address = ctx.reader.analysis.rebase_pointer(address)
+                trail[rebased_address] = ida_color
 
         for section in lines_in.sections_lines:
             for line in section:
                 address = line.at.toea()
 
-                if address == ctx.reader.ip:
-                    color = current_color
-                elif address in backward_trail:
+                if address in backward_trail:
                     color = backward_trail[address]
                 elif address in forward_trail:
                     color = forward_trail[address]
+                elif address == ctx.reader.rebased_ip:
+                    color = current_color
                 else:
                     continue
 

plugins/tenet/trace/reader.py

@@ -35,6 +35,7 @@
 from tenet.util.misc import register_callback, notify_callback
 from tenet.trace.file import TraceFile
 from tenet.trace.types import TraceMemory
+from tenet.trace.analysis import TraceAnalysis
 
 logger = logging.getLogger("Tenet.Trace.Reader")
 
@@ -60,6 +61,7 @@ def __init__(self, filepath, architecture, dctx=None):
 
         # load the given trace file from disk
         self.trace = TraceFile(filepath, architecture)
+        self.analysis = TraceAnalysis(self.trace, dctx)
 
         self._idx_cached_registers = -1
         self._cached_registers = {}
@@ -81,6 +83,13 @@ def ip(self):
         """
         return self.get_register(self.arch.IP)
 
+    @property
+    def rebased_ip(self):
+        """
+        Return a rebased version of the current instruction pointer (if available).
+        """
+        return self.analysis.rebase_pointer(self.ip)
+
     @property
     def sp(self):
         """
@@ -279,23 +288,26 @@ def _step_over_forward(self, n):
         Step the trace forward over n instructions / calls.
         """
         address = self.get_ip(self.idx)
+        bin_address = self.analysis.rebase_pointer(address)
 
         #
         # get the address for the linear instruction address after the
         # current instruction
         #
 
-        next_address = self.dctx.get_next_insn(address)
-        if next_address == -1:
+        bin_next_address = self.dctx.get_next_insn(bin_address)
+        if bin_next_address == -1:
             self.seek(self.idx + 1)
             return
 
+        trace_next_address = self.analysis.rebase_pointer(bin_next_address)
+
         #
         # find the next time the instruction after this instruction is
         # executed in the trace
         #
 
-        next_idx = self.find_next_execution(next_address, self.idx)
+        next_idx = self.find_next_execution(trace_next_address, self.idx)
 
         #
         # the instruction after the call does not appear in the trace,
@@ -313,8 +325,9 @@ def _step_over_backward(self, n):
         Step the trace backward over n instructions / calls.
         """
         address = self.get_ip(self.idx)
+        bin_address = self.analysis.rebase_pointer(address)
 
-        prev_address = self.dctx.get_prev_insn(address)
+        bin_prev_address = self.dctx.get_prev_insn(bin_address)
 
         #
         # could not get the address of the instruction prior to the current
@@ -325,7 +338,7 @@ def _step_over_backward(self, n):
         # performant backend than the python prototype that powers this
         #
 
-        if prev_address == -1:
+        if bin_prev_address == -1:
             self.seek(self.idx - 1)
             return
 
@@ -335,7 +348,7 @@ def _step_over_backward(self, n):
         # and also pretty tricky to handle...
         #
 
-        if self.dctx.is_call_insn(prev_address):
+        if self.dctx.is_call_insn(bin_prev_address):
 
             # get the previous stack pointer address
             sp = self.get_register(self.arch.SP, self.idx - 1)
@@ -364,7 +377,9 @@ def _step_over_backward(self, n):
                 self.seek(self.idx - 1)
                 return
 
-        prev_idx = self.find_prev_execution(prev_address, self.idx)
+        trace_prev_address = self.analysis.rebase_pointer(bin_prev_address)
+
+        prev_idx = self.find_prev_execution(trace_prev_address, self.idx)
         if prev_idx == -1:
             self.seek(self.idx - 1)
             return
@@ -747,16 +762,17 @@ def get_prev_ips(self, n, step_over=False):
 
         output = []
         dctx, idx = self.dctx, self.idx
-        address = self.get_ip(idx)
+        trace_address = self.get_ip(idx)
+        bin_address = self.analysis.rebase_pointer(trace_address)
 
         # (reverse) step over any call instructions
         while len(output) < n and idx > 0:
 
-            prev_address = dctx.get_prev_insn(address)
+            bin_prev_address = dctx.get_prev_insn(bin_address)
             did_step_over = False
 
             # call instruction
-            if prev_address != -1 and dctx.is_call_insn(prev_address):
+            if bin_prev_address != -1 and dctx.is_call_insn(bin_prev_address):
 
                 # get the previous stack pointer address
                 sp = self.get_register(self.arch.SP, idx - 1)
@@ -773,16 +789,17 @@ def get_prev_ips(self, n, step_over=False):
                 # we can assume that we just returned from somewhere.
                 #
                 # 99% of the time, this will have been from the call insn at
-                # prev_address, so let's just assume that is the case and
+                # bin_prev_address, so let's just assume that is the case and
                 # 'reverse step over' onto that.
                 #
                 # NOTE: technically, we can put in more checks and stuff to
                 # try and ensure this is 'correct' but, step over and reverse
                 # step over are kind of an imperfect science as is...
                 #
 
-                if maybe_ret_address == address:
-                    prev_idx = self.find_prev_execution(prev_address, idx)
+                if maybe_ret_address == trace_address:
+                    trace_prev_address = self.analysis.rebase_pointer(bin_prev_address)
+                    prev_idx = self.find_prev_execution(trace_prev_address, idx)
                     did_step_over = bool(prev_idx != -1)
 
             #
@@ -795,7 +812,8 @@ def get_prev_ips(self, n, step_over=False):
             #
 
             if not did_step_over:
-                prev_idx = self.find_prev_execution(prev_address, idx)
+                trace_prev_address = self.analysis.rebase_pointer(bin_prev_address)
+                prev_idx = self.find_prev_execution(trace_prev_address, idx)
 
             #
             # uh, wow okay we're pretty lost and have no idea if there is
@@ -806,15 +824,16 @@ def get_prev_ips(self, n, step_over=False):
             if prev_idx == -1:
                 prev_idx = idx - 1
 
-            prev_address = self.get_ip(prev_idx)
+            trace_prev_address = self.get_ip(prev_idx)
 
             # no address was returned, so the end of trace was reached
-            if prev_address == -1:
+            if trace_prev_address == -1:
                 break
 
             # save the results and continue looping
-            output.append(prev_address)
-            address = prev_address
+            output.append(trace_prev_address)
+            trace_address = trace_prev_address
+            bin_address = self.analysis.rebase_pointer(trace_address)
             idx = prev_idx
 
         # return the list of addresses to be 'executed' next
@@ -836,7 +855,8 @@ def get_next_ips(self, n, step_over=False):
 
         output = []
         dctx, idx = self.dctx, self.idx
-        address = self.get_ip(idx)
+        trace_address = self.get_ip(idx)
+        bin_address = self.analysis.rebase_pointer(trace_address)
 
         # step over any call instructions
         while len(output) < n and idx < (self.trace.length - 1):
@@ -846,15 +866,16 @@ def get_next_ips(self, n, step_over=False):
             # current (call) instruction
             #
 
-            next_address = dctx.get_next_insn(address)
+            bin_next_address = dctx.get_next_insn(bin_address)
 
             #
             # find the next time the instruction after this instruction is
             # executed in the trace
             #
 
-            if next_address != -1:
-                next_idx = self.find_next_execution(next_address, idx)
+            if bin_next_address != -1:
+                trace_next_address = self.analysis.rebase_pointer(bin_next_address)
+                next_idx = self.find_next_execution(trace_next_address, idx)
             else:
                 next_idx = -1
 
@@ -871,15 +892,15 @@ def get_next_ips(self, n, step_over=False):
             # our stepping behavior
             #
 
-            next_address = self.get_ip(next_idx)
+            trace_next_address = self.get_ip(next_idx)
 
             # no address was returned, so the end of trace was reached
-            if next_address == -1:
+            if trace_next_address == -1:
                 break
 
             # save the results and continue looping
-            output.append(next_address)
-            address = next_address
+            output.append(trace_next_address)
+            bin_address = self.analysis.rebase_pointer(trace_next_address)
             idx = next_idx
 
         # return the list of addresses to be 'executed' next

plugins/tenet/ui/trace_view.py

@@ -866,9 +866,10 @@ def _draw_code_trace(self, painter):
 
             # get the executed/code address for the current idx that will represent this line
             address = self.reader.get_ip(idx)
+            rebased_address = self.reader.analysis.rebase_pointer(address)
 
             # select the color for instructions that can be viewed with Tenet
-            if dctx.is_mapped(address):
+            if dctx.is_mapped(rebased_address):
                 painter.setPen(self.pctx.palette.trace_instruction)
 
             # unexplorable parts of the trace are 'greyed' out (eg, not in IDB)
@@ -924,9 +925,10 @@ def _draw_code_cells(self, painter):
 
             # get the executed/code address for the current idx that will represent this cell
             address = self.reader.get_ip(idx)
+            rebased_address = self.reader.analysis.rebase_pointer(address)
 
             # select the color for instructions that can be viewed with Tenet
-            if dctx.is_mapped(address):
+            if dctx.is_mapped(rebased_address):
                 painter.setBrush(cell_color)
 
             # unexplorable parts of the trace are 'greyed' out (eg, not in IDB)