On-disk gadget cache (#23, part 1)

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] [--interactive]
+               [--output {text,json,csv}] [--op <op>] [--dst <reg>] [--src <reg>] [--ropchain <file>] [--exhaustive | --no-exhaustive] [--interactive] [--cache] [--cache-dir <dir>]
 
 This tool allows you to search for gadgets, operations, and ROP chains using a backtracking algorithm in a tree-like structure
 
@@ -73,6 +73,17 @@ options:
   --exhaustive, --no-exhaustive
                         exhaustive search for ROP chains
   --interactive         scan the binary once and drop into an interactive prompt
+  --cache               cache discovered gadgets on disk and reuse them on repeated runs over the same file and options
+  --cache-dir <dir>     directory for the gadget cache (default: $XDG_CACHE_HOME/rop3)
+```
+
+### Gadget cache
+
+With `--cache`, the gadgets discovered for a binary are stored on disk and reused on later runs over the same file and options, skipping the scan. The cache key binds the file content hash and every option that affects the result, so a changed binary or option misses cleanly. This is especially handy for large binaries and for the interactive mode.
+
+```Shell
+$ python rop3.py --binary libc.so.6 --cache        # first run scans and caches
+$ python rop3.py --binary libc.so.6 --cache --op mov --dst rdi --src rax   # reuses the cache
 ```
 
 ### Interactive mode

rop3/api.py

@@ -38,7 +38,8 @@ class Rop3:
     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):
+                 badchars=None, badchar_bytes=None, arch=None, symbols=False,
+                 cache=False, cache_dir=None):
         self.binaries = [binaries] if isinstance(binaries, str) else list(binaries)
         self.base = base
         self.badchars = badchars
@@ -62,7 +63,7 @@ def __init__(self, binaries, *, depth=gadfinder.DEPTH, rop=True, jop=False,
         if avoid_canary:
             flags |= gadfinder.AVOID_CANARY
 
-        self._finder = GadFinder(depth, flags)
+        self._finder = GadFinder(depth, flags, cache=cache, cache_dir=cache_dir)
         self._gadgets = None
 
     @classmethod
@@ -76,7 +77,8 @@ def from_args(cls, args):
         self.badchar_bytes = args.badchar_bytes
         self.arch = args.arch
         self.symbols = args.symbols
-        self._finder = GadFinder(args.depth, args.flags)
+        self._finder = GadFinder(args.depth, args.flags,
+                                 cache=args.cache, cache_dir=args.cache_dir)
         self._gadgets = None
         return self
 

rop3/args.py

@@ -51,6 +51,8 @@ def __init__(self):
         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')
+        self.argparser.add_argument('--cache', action='store_true', default=False, help='cache discovered gadgets on disk and reuse them on repeated runs over the same file and options')
+        self.argparser.add_argument('--cache-dir', type=str, metavar='<dir>', default=None, help='directory for the gadget cache (default: $XDG_CACHE_HOME/rop3)')
 
     def parse_args(self, arguments):
         args = self.argparser.parse_args(arguments)

rop3/cache.py

@@ -0,0 +1,77 @@
+'''
+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 os
+import json
+import hashlib
+import tempfile
+
+import rop3.debug as debug
+
+''' Bump when the on-disk record format changes to invalidate old entries '''
+CACHE_VERSION = 1
+
+
+def default_cache_dir():
+    base = os.environ.get('XDG_CACHE_HOME') or os.path.join(
+        os.path.expanduser('~'), '.cache')
+    return os.path.join(base, 'rop3')
+
+
+class GadgetCache:
+    '''
+    Caches the raw gadget records (vaddr, hex-bytes) discovered for a binary so
+    repeated runs over the same file and parameters skip the scan. The key
+    binds the file content hash and every parameter that affects the record
+    set, so a changed file or option misses cleanly.
+    '''
+
+    def __init__(self, cache_dir=None):
+        self.cache_dir = cache_dir or default_cache_dir()
+
+    def key(self, file_hash, params: dict) -> str:
+        material = json.dumps(
+            {'version': CACHE_VERSION, 'hash': file_hash, 'params': params},
+            sort_keys=True)
+        return hashlib.sha256(material.encode()).hexdigest()
+
+    @staticmethod
+    def file_hash(data: bytes) -> str:
+        return hashlib.sha256(data).hexdigest()
+
+    def _path(self, key: str) -> str:
+        return os.path.join(self.cache_dir, key + '.json')
+
+    def load(self, key: str):
+        ''' Return the cached records (list of [vaddr, hex]) or None on miss. '''
+        try:
+            with open(self._path(key), 'r') as f:
+                return json.load(f)
+        except (FileNotFoundError, ValueError):
+            return None
+
+    def store(self, key: str, records) -> None:
+        ''' Persist records atomically; failures are non-fatal (just a warning). '''
+        try:
+            os.makedirs(self.cache_dir, exist_ok=True)
+            path = self._path(key)
+            fd, tmp = tempfile.mkstemp(dir=self.cache_dir, suffix='.tmp')
+            with os.fdopen(fd, 'w') as f:
+                json.dump(records, f)
+            os.replace(tmp, path)
+        except OSError as exc:
+            debug.warning(f'could not write gadget cache: {exc}')

rop3/gadfinder.py

@@ -20,6 +20,7 @@
 import bisect
 import capstone
 
+from rop3.cache import GadgetCache
 import rop3.utils as utils
 import rop3.debug as debug
 import rop3.binary
@@ -52,9 +53,10 @@ class GadFinder:
     '''
     Class to search gadgets in a binary
     '''
-    def __init__(self, depth=DEPTH, flags=DEFAULT):
+    def __init__(self, depth=DEPTH, flags=DEFAULT, cache=False, cache_dir=None):
         self.depth = depth
         self.flags = flags
+        self._cache = GadgetCache(cache_dir) if cache else None
 
     def find(self, filenames: list[str], base=None, badchars=None,
              badchar_bytes=None, arch=None, symbols=False) -> list[Gadget]:
@@ -155,6 +157,39 @@ def find_op_from_gadgets(self, gadgets, op, dst=None, src=None):
         return op_obj.filter_gadgets(gadgets)
 
     def _search_gadgets(self, binary, badchars, badchar_bytes=None, symbol_table=None):
+        '''
+        Yield the gadgets of a binary, building them either from the on-disk
+        cache (when enabled and warm) or from a fresh scan. The cache only
+        stores the raw (vaddr, bytes) records; everything address/disassembly
+        derived (decodes, symbol) is rebuilt here.
+        '''
+        if self._cache is not None:
+            key = self._cache.key(
+                self._cache.file_hash(binary.raw_data),
+                self._record_params(binary, badchars, badchar_bytes))
+            cached = self._cache.load(key)
+            if cached is not None:
+                debug.info(f'{os.path.basename(binary.filename)}: '
+                           f'{len(cached)} gadgets from cache')
+                yield from self._reconstruct(binary, cached, symbol_table)
+                return
+
+        records = [] if self._cache is not None else None
+        arch = arch_singleton.arch.arch
+        mode = arch_singleton.arch.mode
+        for vaddr, raw, decodes in self._scan(binary, badchars, badchar_bytes):
+            if records is not None:
+                records.append([vaddr, raw.hex()])
+            symbol = self._nearest_symbol(vaddr, symbol_table) if symbol_table else None
+            yield Gadget(filename=binary.filename, arch=arch, mode=mode,
+                         vaddr=vaddr, decodes=decodes, bytes=raw, symbol=symbol)
+
+        if records is not None:
+            self._cache.store(key, records)
+
+    def _scan(self, binary, badchars, badchar_bytes):
+        ''' Single pass over the executable sections; yields the raw
+            (vaddr, bytes, decodes) of every valid gadget (one disassembly). '''
         sections = binary.get_exec_sections()
         arch = arch_singleton.arch.arch
         mode = arch_singleton.arch.mode
@@ -181,17 +216,34 @@ def _search_gadgets(self, binary, badchars, badchar_bytes=None, symbol_table=Non
                                 continue
                             decodes = list(md.disasm(bytes_, vaddr))
                             if self._is_valid_gadget(decodes):
-                                symbol = self._nearest_symbol(vaddr, symbol_table) \
-                                    if symbol_table else None
-                                yield Gadget(
-                                    filename=binary.filename,
-                                    arch=arch,
-                                    mode=mode,
-                                    vaddr=vaddr,
-                                    decodes=decodes,
-                                    bytes=bytes_,
-                                    symbol=symbol,
-                                )
+                                yield (vaddr, bytes_, decodes)
+
+    def _reconstruct(self, binary, records, symbol_table):
+        ''' Rebuild Gadget objects from cached (vaddr, hex-bytes) records. '''
+        arch = arch_singleton.arch.arch
+        mode = arch_singleton.arch.mode
+        md = capstone.Cs(arch, mode)
+        md.detail = True
+        for vaddr, hexbytes in records:
+            raw = bytes.fromhex(hexbytes)
+            decodes = list(md.disasm(raw, vaddr))
+            symbol = self._nearest_symbol(vaddr, symbol_table) if symbol_table else None
+            yield Gadget(filename=binary.filename, arch=arch, mode=mode,
+                         vaddr=vaddr, decodes=decodes, bytes=raw, symbol=symbol)
+
+    def _record_params(self, binary, badchars, badchar_bytes) -> dict:
+        ''' Everything (besides file content) that changes the raw record set,
+            so a different option misses the cache cleanly. '''
+        arch = arch_singleton.arch
+        return {
+            'depth': self.depth,
+            'flags': int(self.flags),
+            'arch': [arch.arch, arch.mode],
+            'badchars': sorted(badchars) if badchars else None,
+            'badchar_bytes': sorted(badchar_bytes) if badchar_bytes else None,
+            'sections': [[s['vaddr'], len(s['opcodes'])]
+                         for s in binary.get_exec_sections()],
+        }
 
     def _gad_terminations(self):
         ret = []

tests/test_cache.py

@@ -0,0 +1,76 @@
+'''
+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/>.
+'''
+
+from rop3.cache import GadgetCache
+from rop3 import Rop3
+
+from conftest import build_minimal_elf, EM_X86_64, ET_DYN
+
+TEXT = b'\x58\xc3\x5b\xc3\x90\xc3'   # pop rax;ret  pop rbx;ret  nop;ret
+
+
+def test_key_changes_with_content(tmp_path):
+    c = GadgetCache(str(tmp_path))
+    h1 = c.file_hash(b'aaaa')
+    h2 = c.file_hash(b'bbbb')
+    assert c.key(h1, {'depth': 5}) != c.key(h2, {'depth': 5})
+    assert c.key(h1, {'depth': 5}) == c.key(h1, {'depth': 5})
+
+
+def test_key_changes_with_params(tmp_path):
+    c = GadgetCache(str(tmp_path))
+    h = c.file_hash(b'aaaa')
+    assert c.key(h, {'depth': 5}) != c.key(h, {'depth': 6})
+
+
+def test_store_and_load_roundtrip(tmp_path):
+    c = GadgetCache(str(tmp_path))
+    key = c.key('hash', {'depth': 5})
+    assert c.load(key) is None
+    records = [[0x1000, '58c3'], [0x1010, '5bc3']]
+    c.store(key, records)
+    assert c.load(key) == records
+
+
+def _elf(tmp_path):
+    data = build_minimal_elf(64, EM_X86_64, TEXT, 0x1000, ET_DYN)
+    path = tmp_path / 'sample.elf'
+    path.write_bytes(data)
+    return str(path)
+
+
+def test_cache_produces_same_gadgets(tmp_path):
+    ''' A warm cache yields the same gadgets as a cold scan. '''
+    elf = _elf(tmp_path)
+    cache_dir = str(tmp_path / 'cache')
+
+    cold = Rop3(elf, cache=True, cache_dir=cache_dir)
+    first = sorted((g.vaddr, g.text_repr) for g in cold.gadgets())
+
+    warm = Rop3(elf, cache=True, cache_dir=cache_dir)   # fresh instance, warm cache
+    second = sorted((g.vaddr, g.text_repr) for g in warm.gadgets())
+
+    assert first == second
+    assert first == sorted((g.vaddr, g.text_repr)
+                           for g in Rop3(elf).gadgets())   # == uncached
+
+
+def test_cache_file_written(tmp_path):
+    elf = _elf(tmp_path)
+    cache_dir = tmp_path / 'cache'
+    Rop3(elf, cache=True, cache_dir=str(cache_dir)).gadgets()
+    assert list(cache_dir.glob('*.json'))