distinguish between analyze_gadget and analyze_addr

angrop/gadget_finder/__init__.py

@@ -1,8 +1,8 @@
 import re
 import logging
 import itertools
+from functools import partial
 from multiprocessing import Pool
-from collections import defaultdict
 
 import tqdm
 
@@ -36,8 +36,11 @@ def _set_global_gadget_analyzer(rop_gadget_analyzer):
     _global_gadget_analyzer = rop_gadget_analyzer
     _disable_loggers()
 
-def run_worker(addr):
-    res = _global_gadget_analyzer.analyze_gadget(addr)
+def run_worker(addr, allow_cond_branch=None):
+    if allow_cond_branch is None:
+        res = _global_gadget_analyzer.analyze_gadget(addr)
+    else:
+        res = _global_gadget_analyzer.analyze_gadget(addr, allow_conditional_branches=allow_cond_branch)
     if res is None:
         return []
     if isinstance(res, list):
@@ -123,8 +126,8 @@ def _initialize_gadget_analyzer(self):
         self._gadget_analyzer = gadget_analyzer.GadgetAnalyzer(self.project, self.fast_mode, arch=self.arch,
                                                                kernel_mode=self.kernel_mode, stack_gsize=self.stack_gsize)
 
-    def analyze_gadget(self, addr):
-        g = self.gadget_analyzer.analyze_gadget(addr)
+    def analyze_gadget(self, addr, allow_conditional_branches=None):
+        g = self.gadget_analyzer.analyze_gadget(addr, allow_conditional_branches=allow_conditional_branches)
         if isinstance(g, list):
             for x in g:
                 x.project = self.project
@@ -141,8 +144,9 @@ def analyze_gadget_list(self, addr_list, processes=4, show_progress=True):
             iterable = tqdm.tqdm(iterable=iterable, smoothing=0, total=len(addr_list),
                                  desc="ROP", maxinterval=0.5, dynamic_ncols=True)
 
+        func = partial(run_worker, allow_cond_branch=False)
         with Pool(processes=processes, initializer=_set_global_gadget_analyzer, initargs=initargs) as pool:
-            it = pool.imap_unordered(run_worker, iterable, chunksize=1)
+            it = pool.imap_unordered(func, iterable, chunksize=1)
             for gs in it:
                 if gs:
                     gadgets += gs

angrop/rop.py

@@ -94,16 +94,28 @@ def _screen_gadgets(self):
         self.chain_builder.syscall_gadgets = self.syscall_gadgets
         self.chain_builder.update()
 
+    def analyze_addr(self, addr):
+        """
+        return a list of gadgets that starts from addr
+        this is possible because of conditional branches
+        """
+        gs = self.gadget_finder.analyze_gadget(addr, allow_conditional_branches=True)
+        if not gs:
+            return gs
+        self._all_gadgets += gs
+        self._screen_gadgets()
+        return gs
+
     def analyze_gadget(self, addr):
-        g = self.gadget_finder.analyze_gadget(addr)
+        """
+        return a gadget or None, it filters out gadgets containing conditional_branches
+        if you'd like those, use analyze_addr
+        """
+        g = self.gadget_finder.analyze_gadget(addr, allow_conditional_branches=False)
         if g is None:
             return g
 
-        if isinstance(g, list):
-            self._all_gadgets += g
-        else:
-            self._all_gadgets.append(g)
-
+        self._all_gadgets.append(g)
         self._screen_gadgets()
         return g