From 64826905f8cbcc67cc7cca4be42347900f4fc7c4 Mon Sep 17 00:00:00 2001
From: hugsy <hugsy@blah.cat>
Date: Wed, 9 Jun 2021 00:02:08 -0700
Subject: [PATCH] * improved command argparsing * ported a few functions to use
 it * updated unit tests to reflect the syntax change

---
 gef.py            | 259 ++++++++++++++++++++--------------------------
 tests/runtests.py |   8 +-
 2 files changed, 114 insertions(+), 153 deletions(-)

diff --git a/gef.py b/gef.py
index b543303b5..5d4ce60d3 100644
--- a/gef.py
+++ b/gef.py
@@ -2594,46 +2594,72 @@ def inner_f(*args, **kwargs):
     return wrapper
 
 
+def FakeExit(*args, **kwargs):
+    raise RuntimeWarning
+
+sys.exit = FakeExit
+
 def parse_arguments(required_arguments, optional_arguments):
     """Argument parsing decorator."""
 
+    def int_wrapper(x): return int(x, 0)
+
     def decorator(f):
         def wrapper(*args, **kwargs):
-            parser = argparse.ArgumentParser(add_help=False, argument_default=argparse.SUPPRESS)
-
+            parser = argparse.ArgumentParser(prog="gef", add_help=False)
             for argname in required_arguments:
                 argvalue = required_arguments[argname]
                 argtype = type(argvalue)
+                if argtype == int:
+                    argtype = int_wrapper
+
                 if argname.startswith("-"):
+                    # optional args
                     if argtype == bool:
-                        action = "store_true" if argvalue else "store_false"
-                        parser.add_argument(argname, action=action, help="Help for {}".format(argname), required=True, default=argvalue)
+                        parser.add_argument(argname, action="store_true" if argvalue else "store_false")
                     else:
-                        parser.add_argument(argname, type=argtype, help="Help for {}".format(argname), required=True, default=argvalue)
+                        parser.add_argument(argname, type=argtype, required=True, default=argvalue)
                 else:
-                    parser.add_argument(argname, type=argtype, help="Help for {}".format(argname), default=argvalue)
+                    # positional args
+                    parser.add_argument(argname, type=argtype, default=argvalue, nargs='?')
 
             for argname in optional_arguments:
                 if not argname.startswith("-"):
                     continue
-
                 argvalue = optional_arguments[argname]
                 argtype = type(argvalue)
+                if argtype == int:
+                    argtype = int_wrapper
                 if argtype == bool:
-                    action = "store_true" if argvalue else "store_false"
-                    parser.add_argument(argname, action=action, help="Help for {}".format(argname), required=False, default=argvalue)
+                    parser.add_argument(argname, action="store_true" if argvalue else "store_false")
                 else:
-                    parser.add_argument(argname, type=argtype, help="Help for {}".format(argname), required=False, default=argvalue)
+                    parser.add_argument(argname, type=argtype, default=argvalue)
 
             obj, cmd_args = args[0], args[1:]
             parsed_args = parser.parse_args(*cmd_args)
             kwargs["arguments"] = parsed_args
-            info("executing {}".format(f.__name__))
             return f(*args, **kwargs)
         return wrapper
     return decorator
 
 
+def copy_to_clipboard(data):
+    """Helper function to submit data to the clipboard"""
+    if sys.platform == "linux":
+        xclip = which("xclip")
+        prog = [xclip, "-selection", "clipboard", "-i"]
+    elif sys.platform == "darwin":
+        pbcopy = which("pbcopy")
+        prog = [pbcopy]
+    else:
+        raise NotImplementedError("Unsupported OS")
+
+    p = subprocess.Popen(prog, stdin=subprocess.PIPE)
+    p.stdin.write(data)
+    p.stdin.close()
+    p.wait()
+    return
+
 
 def use_stdtype():
     if is_32bit(): return "uint32_t"
@@ -4293,111 +4319,75 @@ def do_invoke(self, argv):
 class PrintFormatCommand(GenericCommand):
     """Print bytes format in high level languages."""
 
+    format_matrix = {
+        8:  (endian_str() + "B", "char", "db"),
+        16: (endian_str() + "H", "short", "dw"),
+        32: (endian_str() + "I", "int", "dd"),
+        64: (endian_str() + "Q", "long long", "dq"),
+    }
+
+    valid_formats = ["py", "c", "js", "asm"]
+
     _cmdline_ = "print-format"
-    _syntax_  = "{:s} [-f FORMAT] [-b BITSIZE] [-l LENGTH] [-c] [-h] LOCATION".format(_cmdline_)
     _aliases_ = ["pf",]
-    _example_ = "{0:s} -f py -b 8 -l 256 $rsp".format(_cmdline_)
+    _syntax_  = """{} [--lang LANG] [--bitlen SIZE] [--length LENGTH] [--clip] LOCATION
+\t--lang LANG specifies the output format for programming language (available: {}, default 'py').
+\t--bitlen SIZE specifies size of bit (possible values: {}, default is 8).
+\t--length LENGTH specifies length of array (default is 256).
+\t--clip The output data will be copied to clipboard
+\tLOCATION specifies where the address of bytes is stored.""".format(_cmdline_, str(valid_formats), str(format_matrix.keys()))
+    _example_ = "{} --lang py -l 16 $rsp".format(_cmdline_)
 
-    bitformat = {8: "<B", 16: "<H", 32: "<I", 64: "<Q"}
-    c_type = {8: "char", 16: "short", 32: "int", 64: "long long"}
-    asm_type = {8: "db", 16: "dw", 32: "dd", 64: "dq"}
 
     def __init__(self):
         super().__init__(complete=gdb.COMPLETE_LOCATION)
         return
 
-    def usage(self):
-        h = self._syntax_
-        h += "\n\t-f FORMAT specifies the output format for programming language, avaliable value is py, c, js, asm (default py).\n"
-        h += "\t-b BITSIZE sepecifies size of bit, avaliable values is 8, 16, 32, 64 (default is 8).\n"
-        h += "\t-l LENGTH specifies length of array (default is 256).\n"
-        h += "\t-c The result of data will copied to clipboard\n"
-        h += "\tLOCATION specifies where the address of bytes is stored."
-        info(h)
-        return
-
-    def clip(self, data):
-        if sys.platform == "linux":
-            xclip = which("xclip")
-            prog = [xclip, "-selection", "clipboard", "-i"]  # For linux
-        elif sys.platform == "darwin":
-            pbcopy = which("pbcopy")
-            prog = [pbcopy]  # For OSX
-        else:
-            warn("Can't copy to clipboard, platform not supported")
-            return False
-
-        try:
-            p = subprocess.Popen(prog, stdin=subprocess.PIPE)
-        except Exception:
-            warn("Can't copy to clipboard, Something went wrong while copying")
-            return False
-
-        p.stdin.write(data)
-        p.stdin.close()
-        p.wait()
-        return True
-
     @only_if_gdb_running
-    def do_invoke(self, argv):
+    @parse_arguments({"location": "$pc", }, {"--length": 256, "--bitlen": 0, "--lang": "py", "--clip": True,})
+    def do_invoke(self, argv, *args, **kwargs):
         """Default value for print-format command."""
-        lang = "py"
-        length = 256
-        bitlen = 8
-        copy_to_clipboard = False
-        supported_formats = ["py", "c", "js", "asm"]
-
-        opts, args = getopt.getopt(argv, "f:l:b:ch")
-        for o, a in opts:
-            if o == "-f": lang = a
-            elif o == "-l": length = int(gdb.parse_and_eval(a))
-            elif o == "-b": bitlen = int(a)
-            elif o == "-c": copy_to_clipboard = True
-            elif o == "-h":
-                self.usage()
-                return
-
-        if not args:
-            err("No address specified")
-            return
-
-        start_addr = int(gdb.parse_and_eval(args[0]))
+        args = kwargs["arguments"]
+        args.bitlen = args.bitlen or current_arch.ptrsize
 
-        if bitlen not in [8, 16, 32, 64]:
-            err("Size of bit must be in 8, 16, 32, or 64")
+        valid_bitlens = self.format_matrix.keys()
+        if args.bitlen not in valid_bitlens:
+            err("Size of bit must be in: {}".format(str(valid_bitlens)))
             return
 
-        if lang not in supported_formats:
-            err("Language must be : {}".format(str(supported_formats)))
+        if args.lang not in self.valid_formats:
+            err("Language must be in: {}".format(str(self.valid_formats)))
             return
 
-        size = int(bitlen / 8)
-        end_addr = start_addr + length * size
-        bf = self.bitformat[bitlen]
+        start_addr = int(gdb.parse_and_eval(args.location))
+        size = int(args.bitlen / 8)
+        end_addr = start_addr + args.length * size
+        fmt = self.format_matrix[args.bitlen][0]
         data = []
-        out = ""
 
-        for address in range(start_addr, end_addr, size):
-            value = struct.unpack(bf, read_memory(address, size))[0]
+        for addr in range(start_addr, end_addr, size):
+            value = struct.unpack(fmt, read_memory(addr, size))[0]
             data += [value]
         sdata = ", ".join(map(hex, data))
 
-        if lang == "py":
+        if args.lang == "py":
             out = "buf = [{}]".format(sdata)
-        elif lang == "c":
-            out =  "unsigned {0} buf[{1}] = {{{2}}};".format(self.c_type[bitlen], length, sdata)
-        elif lang == "js":
+        elif args.lang == "c":
+            c_type = self.format_matrix[args.bitlen][1]
+            out = "unsigned {0} buf[{1}] = {{{2}}};".format(c_type, length, sdata)
+        elif args.lang == "js":
             out = "var buf = [{}]".format(sdata)
-        elif lang == "asm":
-            out += "buf {0} {1}".format(self.asm_type[bitlen], sdata)
+        elif args.lang == "asm":
+            asm_type = self.format_matrix[args.bitlen][2]
+            out = "buf {0} {1}".format(asm_type, sdata)
 
-        if copy_to_clipboard:
-            if self.clip(bytes(out, "utf-8")):
+        if args.clip:
+            if copy_to_clipboard(gef_pybytes(out)):
                 info("Copied to clipboard")
             else:
                 warn("There's a problem while copying")
 
-        print(out)
+        gef_print(out)
         return
 
 
@@ -5845,9 +5835,16 @@ class UnicornEmulateCommand(GenericCommand):
     the next instruction from current PC."""
 
     _cmdline_ = "unicorn-emulate"
-    _syntax_  = "{:s} [-f LOCATION] [-t LOCATION] [-n NB_INSTRUCTION] [-s] [-o PATH] [-h]".format(_cmdline_)
+    _syntax_  = """{:s} [--start LOCATION] [--until LOCATION] [--skip-emulation] [--output-file PATH] [NB_INSTRUCTION]
+\n\t--start LOCATION specifies the start address of the emulated run (default $pc).
+\t--until LOCATION specifies the end address of the emulated run.
+\t--skip-emulation\t do not execute the script once generated.
+\t--output-file /PATH/TO/SCRIPT.py writes the persistent Unicorn script into this file.
+\tNB_INSTRUCTION indicates the number of instructions to execute
+\nAdditional options can be setup via `gef config unicorn-emulate`
+""".format(_cmdline_)
     _aliases_ = ["emulate",]
-    _example_ = "{0:s} -f $pc -n 10 -o /tmp/my-gef-emulation.py".format(_cmdline_)
+    _example_ = "{0:s} --start $pc 10 --output-file /tmp/my-gef-emulation.py".format(_cmdline_)
 
     def __init__(self):
         super().__init__(complete=gdb.COMPLETE_LOCATION)
@@ -5855,18 +5852,6 @@ def __init__(self):
         self.add_setting("show_disassembly", False, "Show every instruction executed")
         return
 
-    def help(self):
-        h = self._syntax_
-        h += "\n\t-f LOCATION specifies the start address of the emulated run (default $pc).\n"
-        h += "\t-t LOCATION specifies the end address of the emulated run.\n"
-        h += "\t-s      Script-Only: do not execute the script once generated.\n"
-        h += "\t-o /PATH/TO/SCRIPT.py writes the persistent Unicorn script into this file.\n"
-        h += "\t-n NB_INSTRUCTION indicates the number of instructions to execute (mutually exclusive with `-t` and `-g`).\n"
-        h += "\t-g NB_GADGET indicates the number of gadgets to execute (mutually exclusive with `-t` and `-n`).\n"
-        h += "\nAdditional options can be setup via `gef config unicorn-emulate`\n"
-        info(h)
-        return
-
     def pre_load(self):
         try:
             __import__("unicorn")
@@ -5882,12 +5867,12 @@ def pre_load(self):
         return
 
     @only_if_gdb_running
-    @parse_arguments({}, {"--start": 0, "--nb": 1, "--until": 0, "--dry-run": False, "--output-file": ""})
+    @parse_arguments({"nb": 1}, {"--start": 0, "--until": 0, "--skip-emulation": True, "--output-file": ""})
     def do_invoke(self, argv, *args, **kwargs):
         args = kwargs["arguments"]
         start_address = args.start or current_arch.pc
         end_address = args.until or self.get_unicorn_end_addr(start_address, args.nb)
-        self.run_unicorn(start_address, end_address, to_script_only=args.dry_run, to_file=args.output_file)
+        self.run_unicorn(start_address, end_address, skip_emulation=args.skip_emulation, to_file=args.output_file)
         return
 
     def get_unicorn_end_addr(self, start_addr, nb):
@@ -5897,7 +5882,7 @@ def get_unicorn_end_addr(self, start_addr, nb):
 
     def run_unicorn(self, start_insn_addr, end_insn_addr, *args, **kwargs):
         verbose = self.get_setting("verbose") or False
-        to_script_only = kwargs.get("to_script_only", False)
+        skip_emulation = kwargs.get("skip_emulation", False)
         arch, mode = get_unicorn_arch(to_string=True)
         unicorn_registers = get_unicorn_registers(to_string=True)
         cs_arch, cs_mode = get_capstone_arch(to_string=True)
@@ -6086,7 +6071,7 @@ def emulate(emu, start_addr, end_addr):
             info("Unicorn script generated as '{}'".format(tmp_filename))
             os.chmod(tmp_filename, 0o700)
 
-        if to_script_only:
+        if skip_emulation:
             return
 
         ok("Starting emulation: {:#x} {} {:#x}".format(start_insn_addr, RIGHT_ARROW, end_insn_addr))
@@ -6335,7 +6320,9 @@ class NopCommand(GenericCommand):
     aware."""
 
     _cmdline_ = "nop"
-    _syntax_  = "{:s} [-b NUM_BYTES] [-h] [LOCATION]".format(_cmdline_)
+    _syntax_  = """{:s} [-nb NUM_BYTES] [LOCATION]
+  LOCATION\taddress/symbol to patch
+    --nb NUM_BYTES\tInstead of writing one instruction, patch the specified number of bytes""".format(_cmdline_)
     _example_ = "{:s} $pc".format(_cmdline_)
 
     def __init__(self):
@@ -6347,30 +6334,12 @@ def get_insn_size(self, addr):
         next_insn = gef_instruction_n(addr, 2)
         return next_insn.address - cur_insn.address
 
-    def do_invoke(self, argv):
-        opts, args = getopt.getopt(argv, "b:h")
-        num_bytes = 0
-        for o, a in opts:
-            if o == "-b":
-                num_bytes = int(a, 0)
-            elif o == "-h":
-                self.help()
-                return
-
-        if args:
-            loc = parse_address(args[0])
-        else:
-            loc = current_arch.pc
-
-        self.nop_bytes(loc, num_bytes)
-        return
-
-    def help(self):
-        m = self._syntax_
-        m += "\n  LOCATION\taddress/symbol to patch\n"
-        m += "  -b NUM_BYTES\tInstead of writing one instruction, patch the specified number of bytes\n"
-        m += "  -h \t\tprint this help\n"
-        info(m)
+    @parse_arguments({"address": 0}, {"--nb": 0, })
+    def do_invoke(self, argv, *args, **kwargs):
+        args = kwargs["arguments"]
+        address = parse_address(args.address) if args.address else current_arch.pc
+        number_of_bytes = args.nb or 1
+        self.nop_bytes(address, number_of_bytes)
         return
 
     @only_if_gdb_running
@@ -6406,29 +6375,21 @@ class StubCommand(GenericCommand):
     function to be called and disrupt your runtime flow (ex. fork)."""
 
     _cmdline_ = "stub"
-    _syntax_  = """{:s} [-r RETVAL] [-h] [LOCATION]
+    _syntax_  = """{:s} [--retval RETVAL] [LOCATION]
 \tLOCATION\taddress/symbol to stub out
-\t-r RETVAL\tSet the return value""".format(_cmdline_)
-    _example_ = "{:s} -r 0 fork".format(_cmdline_)
+\t--retval RETVAL\tSet the return value""".format(_cmdline_)
+    _example_ = "{:s} --retval 0 fork".format(_cmdline_)
 
     def __init__(self):
         super().__init__(complete=gdb.COMPLETE_LOCATION)
         return
 
     @only_if_gdb_running
-    def do_invoke(self, argv):
-        try:
-            opts, args = getopt.getopt(argv, "r:")
-            retval = 0
-            for o, a in opts:
-                if o == "-r":
-                    retval = int(a, 0)
-        except getopt.GetoptError:
-            self.usage()
-            return
-
-        loc = args[0] if args else "*{:#x}".format(current_arch.pc)
-        StubBreakpoint(loc, retval)
+    @parse_arguments({"address": ""}, {"--retval": 0})
+    def do_invoke(self, argv, *args, **kwargs):
+        args = kwargs["arguments"]
+        loc = args.address if args.address else "*{:#x}".format(current_arch.pc)
+        StubBreakpoint(loc, args.retval)
         return
 
 
@@ -8255,7 +8216,7 @@ def context_trace(self):
             else:
                 try:
                     insn = next(gef_disassemble(pc, 1))
-                except gdb.MemoryError:
+                except gdb.MemoryError as e:
                     break
                 items.append(Color.redify("{} {}".format(insn.mnemonic, ", ".join(insn.operands))))
 
diff --git a/tests/runtests.py b/tests/runtests.py
index 15bd1dd76..f89593d25 100755
--- a/tests/runtests.py
+++ b/tests/runtests.py
@@ -324,12 +324,12 @@ def test_cmd_print_format(self):
         res = gdb_start_silent_cmd("print-format $rsp")
         self.assertNoException(res)
         self.assertTrue("buf = [" in res)
-        res = gdb_start_silent_cmd("print-format -f js $rsp")
+        res = gdb_start_silent_cmd("print-format --lang js $rsp")
         self.assertNoException(res)
         self.assertTrue("var buf = [" in res)
-        res = gdb_start_silent_cmd("print-format -f iDontExist $rsp")
+        res = gdb_start_silent_cmd("print-format --lang iDontExist $rsp")
         self.assertNoException(res)
-        self.assertTrue("Language must be :" in res)
+        self.assertTrue("Language must be in:" in res)
         return
 
     def test_cmd_process_status(self):
@@ -465,7 +465,7 @@ def test_cmd_trace_run(self):
         return
 
     def test_cmd_unicorn_emulate(self):
-        cmd = "emu -n 1"
+        cmd = "emu 10"
         res = gdb_run_cmd(cmd)
         self.assertFailIfInactiveSession(res)