From 27b4a9a3140a0b6054bef6b25cec878b82feaf33 Mon Sep 17 00:00:00 2001
From: Zach Riggle <riggle@google.com>
Date: Fri, 3 Jun 2016 14:37:33 -0700
Subject: [PATCH] Add support for argparse-style commands.

pwndbg> help hexdump
usage: hexdump [-h] [address] [count]

pwndbg> hexdump -h
usage: hexdump [-h] [address] [count]

Hexdumps data at the specified address (or at $sp)

positional arguments:
  address     Address to dump
  count       Number of bytes to dump

optional arguments:
  -h, --help  show this help message and exit

pwndbg> hexdump sp+3 5
+0000 0x7fffffffd54b  f7 ff 7f 00  00                                     |....|.   |    |    |
+0005 0x7fffffffd550

pwndbg> hexdump sp-3 7
+0000 0x7fffffffd545  00 00 00 40  2a 60 f7                               |...@|*`. |    |    |
+0007 0x7fffffffd54c
---
 pwndbg/commands/__init__.py | 32 +++++++++++++++++++++++++++-
 pwndbg/commands/hexdump.py  | 42 ++++++++++++++++---------------------
 2 files changed, 49 insertions(+), 25 deletions(-)

diff --git a/pwndbg/commands/__init__.py b/pwndbg/commands/__init__.py
index 6c954e118..79466cc1b 100644
--- a/pwndbg/commands/__init__.py
+++ b/pwndbg/commands/__init__.py
@@ -1,6 +1,7 @@
 #!/usr/bin/env python
 # -*- coding: utf-8 -*-
 from __future__ import print_function
+import argparse
 import functools
 import traceback
 import gdb
@@ -56,7 +57,6 @@ class _Command(gdb.Command):
             print('%r: %s' % (self.function.__name__.strip(),
                               self.function.__doc__.strip()))
 
-
 class _ParsedCommand(_Command):
     #: Whether to return the string 'arg' if parsing fails.
     sloppy = False
@@ -100,6 +100,9 @@ def fix(arg, sloppy=False, quiet=False):
 
     return None
 
+def fix_int(*a, **kw):
+    return int(fix(*a,**kw))
+
 def OnlyWhenRunning(function):
     @functools.wraps(function)
     def _OnlyWhenRunning(*a, **kw):
@@ -126,3 +129,30 @@ def QuietSloppyParsedCommand(func):
     c.quiet = True
     c.sloppy = True
     return c
+
+class ArgparsedCommand(object):
+    """Adds documentation and offloads parsing for a Command via argparse"""
+    def __init__(self, parser):
+        self.parser = parser
+
+        # We want to run all integer and otherwise-unspecified arguments
+        # through fix() so that GDB parses it.
+        for action in parser._actions:
+            if action.dest == 'help':
+                continue
+            if action.type not in (int, None):
+                continue
+            action.type = fix_int
+
+    def __call__(self, function):
+        self.parser.prog = function.__name__
+        @functools.wraps(function)
+        def _ArgparsedCommand(*args):
+            try:
+                args = self.parser.parse_args(args)
+            except SystemExit:
+                # If passing '-h' or '--help', argparse attempts to kill the process.
+                return
+            return function(**vars(args))
+        _ArgparsedCommand.__doc__ = self.parser.format_usage().strip()
+        return Command(_ArgparsedCommand)
diff --git a/pwndbg/commands/hexdump.py b/pwndbg/commands/hexdump.py
index f5dbda4ca..cf2294f45 100644
--- a/pwndbg/commands/hexdump.py
+++ b/pwndbg/commands/hexdump.py
@@ -1,6 +1,7 @@
 #!/usr/bin/env python
 # -*- coding: utf-8 -*-
 from __future__ import print_function
+import argparse
 import pwndbg.arch
 import pwndbg.commands
 import pwndbg.config
@@ -8,38 +9,31 @@ import pwndbg.hexdump
 import pwndbg.memory
 import pwndbg.regs
 
-default_width = pwndbg.config.Parameter('hexdump-width',
-                                         16,
-                                         'line width of hexdump command')
-default_bytes = pwndbg.config.Parameter('hexdump-bytes',
-                                         16,
-                                         'number of bytes printed by hexdump command')
+pwndbg.config.Parameter('hexdump-width',
+                         16,
+                         'line width of hexdump command')
+pwndbg.config.Parameter('hexdump-bytes',
+                         64,
+                         'number of bytes printed by hexdump command')
 
-@pwndbg.commands.ParsedCommand
-@pwndbg.commands.OnlyWhenRunning
-def hexdump(address=None, count=default_bytes):
-    """
-    Hexdumps data at the specified address.
-    Optionally provide the number of bytes to dump (default is controlled
-    by the 'hexdump-bytes' config.)
+parser = argparse.ArgumentParser(description='Hexdumps data at the specified address (or at $sp)')
+parser.add_argument('address', nargs='?', default='$sp', 
+                    help='Address to dump')
+parser.add_argument('count', nargs='?', default=pwndbg.config.hexdump_bytes, 
+                    help='Number of bytes to dump')
 
-    Note that repeating rows are collapsed.
-    """
-    address = int(address if address is not None else pwndbg.regs.sp)
+@pwndbg.commands.ArgparsedCommand(parser)
+@pwndbg.commands.OnlyWhenRunning
+def hexdump(address=None, count=pwndbg.config.hexdump_bytes):
+    address = int(address)
     address &= pwndbg.arch.ptrmask
     count   = int(count)
-
-    # if None not in (address, count):
-    #     address = int(address)
-    #     count   = int(count):
+    width   = int(pwndbg.config.hexdump_width)
 
     if count > address > 0x10000:
         count -= address
 
-    # if address is None:
-    # 	address =
-
     data = pwndbg.memory.read(address, count, partial=True)
 
-    for line in pwndbg.hexdump.hexdump(data, address=address, width=int(default_width)):
+    for line in pwndbg.hexdump.hexdump(data, address=address, width=width):
         print(line)