Expose rop3 as a Python library and add an interactive mode (#24)

README.md

@@ -35,7 +35,7 @@ Now, you can install dependencies in [requirements.txt](requirements.txt):
 ```
 usage: rop3.py [-h] [-v] [--depth <bytes>] [--all] [--rop | --no-rop] [--retf | --no-retf] [--jop | --no-jop] [--allow-undeterministic-gadgets] [--allow-complex-memory-ops] [--verbose]
                [--binary <file> [<file> ...]] [--badchar <hex> [<hex> ...]] [--badchar-bytes <hex> [<hex> ...]] [--keep-canary-address] [--base <hex> [<hex> ...]] [--arch <name>] [--symbols]
-               [--output {text,json,csv}] [--op <op>] [--dst <reg>] [--src <reg>] [--ropchain <file>] [--exhaustive | --no-exhaustive]
+               [--output {text,json,csv}] [--op <op>] [--dst <reg>] [--src <reg>] [--ropchain <file>] [--exhaustive | --no-exhaustive] [--interactive]
 
 This tool allows you to search for gadgets, operations, and ROP chains using a backtracking algorithm in a tree-like structure
 
@@ -72,6 +72,41 @@ options:
   --ropchain <file>     plain text file with a ROP chain
   --exhaustive, --no-exhaustive
                         exhaustive search for ROP chains
+  --interactive         scan the binary once and drop into an interactive prompt
+```
+
+### Interactive mode
+
+With `--interactive`, rop3 scans the binary once and drops into a prompt so you can explore its gadgets without re-scanning on every query:
+
+```Shell
+$ python rop3.py --binary /bin/ls --interactive
+Loaded 71 gadgets from /bin/ls
+rop3> count
+71
+rop3> search pop rbp
+[ls @ 0x100000777]: pop rbp ; ret (x29)
+...
+rop3> op mov rdi rax
+rop3> chain chain.txt
+rop3> quit
+```
+
+Commands: `gadgets`/`search [substring]`, `count`, `op <name> [dst] [src]`, `chain <file>`, `help`, `quit`.
+
+### Use as a library
+
+rop3 can also be used programmatically through the `Rop3` class. Gadgets are scanned once and cached on the instance:
+
+```python
+from rop3 import Rop3
+
+r = Rop3("libc.so.6", base="0x7f0000000000", symbols=True)
+for gadget in r.gadgets():
+    print(gadget)
+
+r.find_op("mov", dst="rdi", src="rax")   # list of matching gadgets
+r.ropchain("chain.txt")                  # iterator over ROP chains
 ```
 
 In the work that we presented in [15th IEEE Workshop on Offensive Technologies (WOOT21)](https://www.ieee-security.org/TC/SP2021/SPW2021/WOOT21/), we used rop3 to evaluate the executional power of Return Oriented Programming in a [subset of most common Windows DLLs](https://drive.google.com/file/d/1gOxUolzrw-xlaW6K-fhzZ7Z-sqxiaZeZ/view?usp=sharing>). Check the [paper](https://drive.google.com/file/d/1Pe7s7bLhJ_20MC-duQ7YiLP-Rx5VCjFK/view?usp=sharing) for further details.

rop3/__init__.py

@@ -25,6 +25,8 @@
 import rop3.ropchain
 import rop3.gadfinder as gadfinder
 
+from rop3.api import Rop3
+
 def main():
     args = rop3.args.ArgumentParser().parse_args(sys.argv[1:])
 
@@ -35,21 +37,17 @@ def main():
         utils.show_version()
         sys.exit(0)
     elif args.binary:
-        finder = gadfinder.GadFinder(args.depth, args.flags)
+        rop = Rop3.from_args(args)
 
         try:
-            if args.ropchain:
-                ropchain = rop3.ropchain.RopChain(finder)
-                result = ropchain.search_from_files(args.binary, args.ropchain,
-                                         base=args.base, badchars=args.badchar,
-                                         badchar_bytes=args.badchar_bytes,
-                                         arch=args.arch, symbols=args.symbols)
+            if args.interactive:
+                from rop3.interactive import Rop3Shell
+                Rop3Shell(rop).cmdloop()
+            elif args.ropchain:
+                result = rop.ropchain(args.ropchain)
                 utils.output_ropchains(result, args.output, exhaustive=args.exhaustive)
             elif args.op:
-                result = finder.find_op(args.binary, args.op, args.dst, args.src,
-                                        base=args.base, badchars=args.badchar,
-                                        badchar_bytes=args.badchar_bytes,
-                                        arch=args.arch, symbols=args.symbols)
+                result = rop.find_op(args.op, args.dst, args.src)
                 if result and isinstance(result[0], list):
                     ''' Composite operation: a list of chains '''
                     if args.output == 'text':
@@ -61,10 +59,7 @@ def main():
                 else:
                     utils.output_gadgets(result, args.output)
             else:
-                gadgets = finder.find(args.binary, base=args.base, badchars=args.badchar,
-                                      badchar_bytes=args.badchar_bytes,
-                                      arch=args.arch, symbols=args.symbols)
-                utils.output_gadgets(gadgets, args.output)
+                utils.output_gadgets(rop.gadgets(), args.output)
         except parser.ParserException as exc:
             debug.error(str(exc))
         except rop3.ropchain.RopChainNotFound as exc:

rop3/api.py

@@ -0,0 +1,102 @@
+'''
+This file is part of rop3 (https://github.com/reverseame/rop3).
+
+rop3 is free software: you can redistribute it and/or modify
+it under the terms of the GNU General Public License as published by
+the Free Software Foundation, either version 3 of the License, or
+(at your option) any later version.
+
+rop3 is distributed in the hope that it will be useful,
+but WITHOUT ANY WARRANTY; without even the implied warranty of
+MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+GNU General Public License for more details.
+
+You should have received a copy of the GNU General Public License
+along with rop3. If not, see <https://www.gnu.org/licenses/>.
+'''
+
+import rop3.gadfinder as gadfinder
+from rop3.gadfinder import GadFinder
+from rop3.ropchain import RopChain
+
+
+class Rop3:
+    '''
+    High-level, reusable entry point to rop3.
+
+    Example:
+        from rop3 import Rop3
+        r = Rop3("libc.so.6", base="0x7f0000000000")
+        r.gadgets()                       # list[Gadget]
+        r.find_op("mov", dst="rdi", src="rax")
+        r.ropchain("chain.txt")
+
+    The discovered gadgets are scanned once and cached on the instance, so
+    repeated queries (and the interactive mode) do not re-scan the binary.
+    '''
+
+    def __init__(self, binaries, *, depth=gadfinder.DEPTH, rop=True, jop=False,
+                 retf=False, all=False, allow_undeterministic=False,
+                 allow_complex_mem=False, avoid_canary=True, base=None,
+                 badchars=None, badchar_bytes=None, arch=None, symbols=False):
+        self.binaries = [binaries] if isinstance(binaries, str) else list(binaries)
+        self.base = base
+        self.badchars = badchars
+        self.badchar_bytes = badchar_bytes
+        self.arch = arch
+        self.symbols = symbols
+
+        flags = 0
+        if all:
+            flags |= gadfinder.KEEP_DUPLICATES
+        if jop:
+            flags |= gadfinder.JOP
+        if rop:
+            flags |= gadfinder.ROP
+        if retf:
+            flags |= gadfinder.RETF
+        if allow_undeterministic:
+            flags |= gadfinder.ALLOW_UNDETERMINISTIC
+        if allow_complex_mem:
+            flags |= gadfinder.ALLOW_COMPLEX_MEM
+        if avoid_canary:
+            flags |= gadfinder.AVOID_CANARY
+
+        self._finder = GadFinder(depth, flags)
+        self._gadgets = None
+
+    @classmethod
+    def from_args(cls, args):
+        ''' Build an instance from a parsed argparse namespace (reuses the
+            already-computed flags), as the CLI does. '''
+        self = cls.__new__(cls)
+        self.binaries = args.binary
+        self.base = args.base
+        self.badchars = args.badchar
+        self.badchar_bytes = args.badchar_bytes
+        self.arch = args.arch
+        self.symbols = args.symbols
+        self._finder = GadFinder(args.depth, args.flags)
+        self._gadgets = None
+        return self
+
+    @property
+    def finder(self) -> GadFinder:
+        return self._finder
+
+    def gadgets(self, refresh=False):
+        ''' All gadgets in the binaries (scanned once, then cached). '''
+        if self._gadgets is None or refresh:
+            self._gadgets = self._finder.find(
+                self.binaries, base=self.base, badchars=self.badchars,
+                badchar_bytes=self.badchar_bytes, arch=self.arch,
+                symbols=self.symbols)
+        return self._gadgets
+
+    def find_op(self, op, dst=None, src=None):
+        ''' Gadgets (or ROP chains, for composite ops) implementing `op`. '''
+        return self._finder.find_op_from_gadgets(self.gadgets(), op, dst, src)
+
+    def ropchain(self, ropfile):
+        ''' Iterator over ROP chains satisfying the operations in `ropfile`. '''
+        return RopChain(self._finder).search_from_gadgets(self.gadgets(), ropfile)

rop3/args.py

@@ -50,6 +50,7 @@ def __init__(self):
         self.argparser.add_argument('--src', type=str, metavar='<reg>', help='specify a source register for the operation')
         self.argparser.add_argument('--ropchain', type=str, metavar='<file>', help='plain text file with a ROP chain')
         self.argparser.add_argument('--exhaustive', action=argparse.BooleanOptionalAction, help="exhaustive search for ROP chains", default=False)
+        self.argparser.add_argument('--interactive', action='store_true', default=False, help='scan the binary once and drop into an interactive prompt')
 
     def parse_args(self, arguments):
         args = self.argparser.parse_args(arguments)

rop3/interactive.py

@@ -0,0 +1,108 @@
+'''
+This file is part of rop3 (https://github.com/reverseame/rop3).
+
+rop3 is free software: you can redistribute it and/or modify
+it under the terms of the GNU General Public License as published by
+the Free Software Foundation, either version 3 of the License, or
+(at your option) any later version.
+
+rop3 is distributed in the hope that it will be useful,
+but WITHOUT ANY WARRANTY; without even the implied warranty of
+MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+GNU General Public License for more details.
+
+You should have received a copy of the GNU General Public License
+along with rop3. If not, see <https://www.gnu.org/licenses/>.
+'''
+
+import cmd
+import shlex
+
+import rop3.utils as utils
+import rop3.parser as parser
+import rop3.ropchain
+
+
+class Rop3Shell(cmd.Cmd):
+    '''
+    Interactive REPL over a Rop3 instance. The binary is scanned once (its
+    gadgets are cached on the Rop3 instance), so each command is cheap.
+    '''
+    intro = ('rop3 interactive mode. The binary is scanned once; '
+             'type "help" or "?" for commands.')
+    prompt = 'rop3> '
+
+    def __init__(self, rop3):
+        super().__init__()
+        self.rop3 = rop3
+
+    def preloop(self):
+        count = len(self.rop3.gadgets())
+        print(f'Loaded {count} gadgets from {", ".join(self.rop3.binaries)}')
+
+    def do_gadgets(self, arg):
+        'gadgets [substring]: list gadgets, optionally filtered by a substring'
+        needle = arg.strip()
+        shown = 0
+        for gadget in self.rop3.gadgets():
+            if not needle or needle in str(gadget):
+                utils.print_gadget(gadget)
+                shown += 1
+        print(f'({shown} gadgets)')
+
+    do_search = do_gadgets
+
+    def do_count(self, arg):
+        'count: print the number of gadgets found'
+        print(len(self.rop3.gadgets()))
+
+    def do_op(self, arg):
+        'op <name> [dst] [src]: search for an operation'
+        parts = shlex.split(arg)
+        if not parts:
+            print('usage: op <name> [dst] [src]')
+            return
+        op = parts[0]
+        dst = parts[1] if len(parts) > 1 else None
+        src = parts[2] if len(parts) > 2 else None
+        try:
+            result = self.rop3.find_op(op, dst, src)
+        except parser.ParserException as exc:
+            print(str(exc))
+            return
+        if result and isinstance(result[0], list):
+            for chain in result:
+                for gadget in chain:
+                    utils.print_gadget(gadget)
+        else:
+            for gadget in result:
+                utils.print_gadget(gadget)
+
+    def do_chain(self, arg):
+        'chain <file>: build a ROP chain from a plain-text ROP file'
+        ropfile = arg.strip()
+        if not ropfile:
+            print('usage: chain <file>')
+            return
+        try:
+            chains = self.rop3.ropchain(ropfile)
+            first = next(iter(chains), None)
+            if first is None:
+                print('No ROP chain found')
+            else:
+                utils.print_ropchain(first)
+        except rop3.ropchain.RopChainNotFound as exc:
+            print(f'No ROP chain found: {exc}')
+
+    def do_quit(self, arg):
+        'quit: exit the interactive mode'
+        return True
+
+    do_exit = do_quit
+
+    def do_EOF(self, arg):
+        print()
+        return True
+
+    def emptyline(self):
+        pass