From f5cb48e5b05987185e1872af321314daa6b3ebe7 Mon Sep 17 00:00:00 2001
From: Kyle Zeng <zengyhkyle@gmail.com>
Date: Fri, 17 Jan 2025 18:16:28 -0700
Subject: [PATCH] implement gadget.pp()

---
 angrop/gadget_finder/__init__.py | 16 ++++++++++++++--
 angrop/rop.py                    |  5 ++++-
 angrop/rop_gadget.py             |  8 +++++++-
 angrop/rop_utils.py              |  3 ---
 4 files changed, 25 insertions(+), 7 deletions(-)

diff --git a/angrop/gadget_finder/__init__.py b/angrop/gadget_finder/__init__.py
index 3b9031b..768ca90 100644
--- a/angrop/gadget_finder/__init__.py
+++ b/angrop/gadget_finder/__init__.py
@@ -119,7 +119,10 @@ def _initialize_gadget_analyzer(self):
                                                                kernel_mode=self.kernel_mode, stack_gsize=self.stack_gsize)
 
     def analyze_gadget(self, addr):
-        return self.gadget_analyzer.analyze_gadget(addr)
+        g = self.gadget_analyzer.analyze_gadget(addr)
+        if g:
+            g.project = self.project
+        return g
 
     def analyze_gadget_list(self, addr_list, processes=4, show_progress=True):
         gadgets = []
@@ -136,6 +139,9 @@ def analyze_gadget_list(self, addr_list, processes=4, show_progress=True):
                 if gs:
                     gadgets += gs
 
+        for g in gadgets:
+            g.project = self.project
+
         return sorted(gadgets, key=lambda x: x.addr)
 
     def get_duplicates(self):
@@ -145,7 +151,7 @@ def get_duplicates(self):
         cache = self._cache
         return {k:v for k,v in cache.items() if len(v) >= 2}
 
-    def find_gadgets(self, processes=4, show_progress=True):
+    def find_gadgets(self, processes=16, show_progress=True):
         self._cache = {}
 
         initargs = (self.gadget_analyzer,)
@@ -167,6 +173,9 @@ def find_gadgets(self, processes=4, show_progress=True):
                 )
             )
 
+        for g in gadgets:
+            g.project = self.project
+
         return sorted(gadgets, key=lambda x: x.addr), self.get_duplicates()
 
     def find_gadgets_single_threaded(self, show_progress=True):
@@ -178,6 +187,9 @@ def find_gadgets_single_threaded(self, show_progress=True):
         for addr in self._addresses_to_check_with_caching(show_progress):
             gadgets.extend(self.gadget_analyzer.analyze_gadget(addr, allow_conditional_branches=True))
 
+        for g in gadgets:
+            g.project = self.project
+
         return sorted(gadgets, key=lambda x: x.addr), self.get_duplicates()
 
     def _block_has_ip_relative(self, addr, bl):
diff --git a/angrop/rop.py b/angrop/rop.py
index 9798afa..285f105 100644
--- a/angrop/rop.py
+++ b/angrop/rop.py
@@ -143,11 +143,14 @@ def find_gadgets_single_threaded(self, show_progress=True):
         return self.rop_gadgets
 
     def _get_cache_tuple(self):
-        return (self._all_gadgets, self._duplicates)
+        all_gadgets = [x for x in self._all_gadgets]
+        for g in all_gadgets: g.project = None
+        return (all_gadgets, self._duplicates)
 
     def _load_cache_tuple(self, tup):
         self._all_gadgets = tup[0]
         self._duplicates = tup[1]
+        for g in self._all_gadgets: g.project = self.project
         self._screen_gadgets()
 
     def save_gadgets(self, path):
diff --git a/angrop/rop_gadget.py b/angrop/rop_gadget.py
index d01f467..f07afc9 100644
--- a/angrop/rop_gadget.py
+++ b/angrop/rop_gadget.py
@@ -1,3 +1,5 @@
+from .rop_utils import addr_to_asmstring
+
 class RopMemAccess:
     """Holds information about memory accesses
     Attributes:
@@ -93,6 +95,7 @@ class RopGadget:
     Gadget objects
     """
     def __init__(self, addr):
+        self.project = None
         self.addr = addr
         self.block_length = None
         self.stack_change = None
@@ -141,6 +144,9 @@ def has_symbolic_access(self):
         accesses = set(self.mem_reads + self.mem_writes + self.mem_changes)
         return any(x.is_symbolic_access() for x in accesses)
 
+    def pp(self):
+        print("; ".join(addr_to_asmstring(self.project, addr) for addr in self.bbl_addrs))
+
     def __str__(self):
         s = "Gadget %#x\n" % self.addr
         s += "Stack change: %#x\n" % self.stack_change
@@ -198,7 +204,7 @@ def __repr__(self):
         return "<Gadget %#x>" % self.addr
 
     def copy(self):
-        out = RopGadget(self.addr)
+        out = RopGadget(self.project, self.addr)
         out.addr = self.addr
         out.changed_regs = set(self.changed_regs)
         out.popped_regs = set(self.popped_regs)
diff --git a/angrop/rop_utils.py b/angrop/rop_utils.py
index d555ab9..d063a6f 100644
--- a/angrop/rop_utils.py
+++ b/angrop/rop_utils.py
@@ -11,9 +11,6 @@ def addr_to_asmstring(project, addr):
     block = project.factory.block(addr)
     return "; ".join(["%s %s" %(i.mnemonic, i.op_str) for i in block.capstone.insns])
 
-def gadget_to_asmstring(project, gadget):
-    return "; ".join(addr_to_asmstring(project, addr) for addr in gadget.bbl_addrs)
-
 
 def get_ast_dependency(ast):
     """