Move vmmap to ArgparsedCommand; add sloppy_gdb_parse (#285)

* Migrate vmmap command to ArgparsedCommand

* vmmap command: better msg for no mappings

* WIP: vmmap

* Review fixes

* isort fix

pwndbg/__init__.py

@@ -163,4 +163,3 @@ signal.signal(signal.SIGWINCH, lambda signum, frame: gdb.execute("set width %i"
 # After GDB gets the fix, we should disable this only for bugged GDB versions.
 if 1:
     gdb.execute('set remote search-memory-packet off')
-

pwndbg/commands/__init__.py

@@ -5,6 +5,7 @@ from __future__ import division
 from __future__ import print_function
 from __future__ import unicode_literals
 
+import argparse
 import functools
 
 import gdb
@@ -230,3 +231,19 @@ class ArgparsedCommand(object):
 
     def __call__(self, function):
         return _ArgparsedCommand(self.parser, function)
+
+
+def sloppy_gdb_parse(s):
+    """
+    This function should be used as ``argparse.ArgumentParser`` .add_argument method's `type` helper.
+    
+    This makes the type being parsed as gdb value and if that parsing fails,
+    a string is returned.
+
+    :param s: String.
+    :return: Whatever gdb.parse_and_eval returns or string.
+    """
+    try:
+        return gdb.parse_and_eval(s)
+    except (TypeError, gdb.error):
+        return s

pwndbg/commands/vmmap.py

@@ -1,13 +1,15 @@
 #!/usr/bin/env python
 # -*- coding: utf-8 -*-
 """
-Command to print the vitual memory map a la /proc/self/maps.
+Command to print the virtual memory map a la /proc/self/maps.
 """
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 from __future__ import unicode_literals
 
+import argparse
+
 import gdb
 import six
 
@@ -17,27 +19,38 @@ import pwndbg.compat
 import pwndbg.vmmap
 
 
-@pwndbg.commands.QuietSloppyParsedCommand
-@pwndbg.commands.OnlyWhenRunning
-def vmmap(map=None):
-    """
-    Print the virtal memory map, or the specific mapping for the
-    provided address / module name.
-    """
-    int_map = None
-    str_map = None
-
-    if isinstance(map, six.string_types):
-        str_map = map
-    elif isinstance(map, six.integer_types + (gdb.Value,)):
-        int_map = int(map)
+def pages_filter(s):
+    gdbval_or_str = pwndbg.commands.sloppy_gdb_parse(s)
+
+    # returns a module filter
+    if isinstance(gdbval_or_str, six.string_types):
+        module_name = gdbval_or_str
+        return lambda page: module_name in page.objfile
+
+    # returns an address filter
+    elif isinstance(gdbval_or_str, six.integer_types + (gdb.Value,)):
+        addr = gdbval_or_str
+        return lambda page: addr in page
+
+    else:
+        raise argparse.ArgumentTypeError('Unknown vmmap argument type.')
 
-    print(M.legend())
 
-    for page in pwndbg.vmmap.get():
-        if str_map and str_map not in page.objfile:
-            continue
-        if int_map and int_map not in page:
-            continue
+parser = argparse.ArgumentParser()
+parser.description = 'Print virtual memory map pages. Results can be filtered by providing address/module name.'
+parser.add_argument('pages_filter', type=pages_filter, nargs='?', default=None,
+                    help='Address or module name.')
 
+
+@pwndbg.commands.ArgparsedCommand(parser)
+@pwndbg.commands.OnlyWhenRunning
+def vmmap(pages_filter=None):
+    pages = list(filter(pages_filter, pwndbg.vmmap.get()))
+
+    if not pages:
+        print('There are no mappings for specified address or module.')
+        return
+
+    print(M.legend())
+    for page in pages:
         print(M.get(page.vaddr, text=str(page)))