From beea47111337249e5b914b494761459939556e58 Mon Sep 17 00:00:00 2001
From: chinggg <24590067+chinggg@users.noreply.github.com>
Date: Thu, 18 Aug 2022 12:36:25 +0800
Subject: [PATCH] feat(r2): CallStack and fuzzy backtrace hook

---
 examples/extensions/r2/hello_r2.py |  1 +
 qiling/extensions/r2/callstack.py  | 73 ++++++++++++++++++++++++++++++
 qiling/extensions/r2/r2.py         | 37 +++++++++++++++
 3 files changed, 111 insertions(+)
 create mode 100644 qiling/extensions/r2/callstack.py

diff --git a/examples/extensions/r2/hello_r2.py b/examples/extensions/r2/hello_r2.py
index ebd54c452..3b02293ea 100644
--- a/examples/extensions/r2/hello_r2.py
+++ b/examples/extensions/r2/hello_r2.py
@@ -35,6 +35,7 @@ def my_sandbox(path, rootfs):
     ql.hook_address(func, r2.functions['main'].offset)
     # enable trace powered by r2 symsmap
     # r2.enable_trace()
+    r2.bt(0x401906)
     ql.run()
 
 if __name__ == "__main__":
diff --git a/qiling/extensions/r2/callstack.py b/qiling/extensions/r2/callstack.py
new file mode 100644
index 000000000..e527cf03f
--- /dev/null
+++ b/qiling/extensions/r2/callstack.py
@@ -0,0 +1,73 @@
+from dataclasses import dataclass
+from typing import Iterator, Optional
+
+
+@dataclass
+class CallStack:
+    """See https://github.com/angr/angr/blob/master/angr/state_plugins/callstack.py"""
+    addr: int
+    sp: int
+    bp: int
+    name: str = None  # 'name + offset'
+    next: Optional['CallStack'] = None
+
+    def __iter__(self) -> Iterator['CallStack']:
+        """
+        Iterate through the callstack, from top to bottom
+        (most recent first).
+        """
+        i = self
+        while i is not None:
+            yield i
+            i = i.next
+
+    def __getitem__(self, k):
+        """
+        Returns the CallStack at index k, indexing from the top of the stack.
+        """
+        orig_k = k
+        for i in self:
+            if k == 0:
+                return i
+            k -= 1
+        raise IndexError(orig_k)
+
+    def __len__(self):
+        """
+        Get how many frames there are in the current call stack.
+
+        :return: Number of frames
+        :rtype: int
+        """
+
+        o = 0
+        for _ in self:
+            o += 1
+        return o
+
+    def __repr__(self):
+        """
+        Get a string representation.
+
+        :return: A printable representation of the CallStack object
+        :rtype: str
+        """
+        return "<CallStack (depth %d)>" % len(self)
+
+    def __str__(self):
+        return "Backtrace:\n" + "\n".join(f"Frame {i}: [{f.name}] {f.addr:#x} sp={f.sp:#x}, bp={f.bp:#x}" for i, f in enumerate(self))
+
+    def __eq__(self, other):
+        if not isinstance(other, CallStack):
+            return False
+
+        if self.addr != other.addr or self.sp != other.sp or self.bp != other.bp:
+            return False
+
+        return self.next == other.next
+
+    def __ne__(self, other):
+        return not (self == other)
+
+    def __hash__(self):
+        return hash(tuple((c.addr, c.sp, c.bp) for c in self))
diff --git a/qiling/extensions/r2/r2.py b/qiling/extensions/r2/r2.py
index 13a655b2f..e86128d74 100644
--- a/qiling/extensions/r2/r2.py
+++ b/qiling/extensions/r2/r2.py
@@ -13,6 +13,7 @@
 from qiling.const import QL_ARCH
 from qiling.extensions import trace
 from unicorn import UC_PROT_NONE, UC_PROT_READ, UC_PROT_WRITE, UC_PROT_EXEC, UC_PROT_ALL
+from .callstack import CallStack
 
 if TYPE_CHECKING:
     from qiling.core import Qiling
@@ -268,6 +269,42 @@ def dis_nbytes(self, addr: int, size: int) -> List[Instruction]:
         insts = [Instruction(**dic) for dic in self._cmdj(f"pDj {size} @ {addr}")]
         return insts
 
+    def dis_ninsts(self, addr: int, n: int=1) -> List[Instruction]:
+        insts = [Instruction(**dic) for dic in self._cmdj(f"pdj {n} @ {addr}")]
+        return insts
+
+    def _backtrace_fuzzy(self, at: int = None, depth: int = 128) -> Optional[CallStack]:
+        '''Fuzzy backtrace, see https://github.com/radareorg/radare2/blob/master/libr/debug/p/native/bt/fuzzy_all.c#L38
+        Args:
+            at: address to start walking stack, default to current SP
+            depth: limit of stack walking
+        Returns:
+            List of Frame
+        '''
+        sp = at or self.ql.arch.regs.arch_sp
+        wordsize = self.ql.arch.bits // 8
+        oldframe = None
+        cursp = oldsp = sp
+        for i in range(depth):
+            addr = self.ql.stack_read(i * wordsize)
+            inst = self.dis_ninsts(addr)[0]
+            if inst.type.lower() == 'call':
+                frame = CallStack(addr=addr, sp=cursp, bp=oldsp, name=self.at(addr), next=oldframe)
+                oldframe = frame
+                oldsp = cursp
+            cursp += wordsize
+        return oldframe
+
+    def bt(self, target: Union[int, str]):
+        '''Backtrace when reaching target'''
+        def bt_hook(ql: 'Qiling', addr: int, size: int, target):
+            if isinstance(target, str):
+                target = self.where(target)
+            if addr <= target and target <= addr + size:
+                callstack = self._backtrace_fuzzy()
+                print(callstack)
+        self.ql.hook_code(bt_hook, target)
+
     def disassembler(self, ql: 'Qiling', addr: int, size: int, filt: Pattern[str]=None) -> int:
         '''A human-friendly monkey patch of QlArchUtils.disassembler powered by r2, can be used for hook_code
             :param ql: Qiling instance