Convert aslr command to argparse style

pwndbg/commands/aslr.py

@@ -1,30 +1,30 @@
 from __future__ import print_function
+import argparse
 import gdb
 import pwndbg.color
 import pwndbg.commands
 import pwndbg.proc
 import pwndbg.vmmap
 
+options = {'on':'off', 'off':'on'}
 
-@pwndbg.commands.Command
-@pwndbg.commands.OnlyWhenRunning
-def aslr(on_or_off=None):
+parser = argparse.ArgumentParser(description='Inspect or modify ASLR status')
+parser.add_argument('state', nargs='?', type=str, choices=options,
+                    help="Turn ASLR on or off (takes effect when target is started)")
+
+@pwndbg.commands.ArgparsedCommand(parser)
+def aslr(state=None):
     """
     Check the current ASLR status, or turn it on/off.
 
     Does not take effect until the program is restarted.
     """
-    options = {'on':'off', 'off':'on'}
-
-    if on_or_off is not None:
-        on_or_off = on_or_off.lower()
-        if on_or_off not in options:
-            print('Valid options are %s' % ', '.join(map(repr, options.keys())))
-        else:
-            gdb.execute('set disable-randomization %s' % options[on_or_off], from_tty=False, to_string=True)
+    if state:
+        gdb.execute('set disable-randomization %s' % options[state], 
+                    from_tty=False, to_string=True)
 
-            if pwndbg.proc.alive:
-                print("Change will take effect when the process restarts")
+        if pwndbg.proc.alive:
+            print("Change will take effect when the process restarts")
 
     aslr = pwndbg.vmmap.check_aslr()
     status = pwndbg.color.red('OFF')

pwndbg/vmmap.py

@@ -359,15 +359,16 @@ def check_aslr():
         pass
 
     # Check the personality of the process
-    try:
-        data = pwndbg.file.get('/proc/%i/personality' % pwndbg.proc.pid)
-        personality = int(data, 16)
-        if personality & 0x40000 == 0:
-            vmmap.aslr = True
-        return vmmap.aslr
-    except:
-        print("Could not check ASLR: Couldn't get personality")
-        pass
+    if pwndbg.proc.alive:
+        try:
+            data = pwndbg.file.get('/proc/%i/personality' % pwndbg.proc.pid)
+            personality = int(data, 16)
+            if personality & 0x40000 == 0:
+                vmmap.aslr = True
+            return vmmap.aslr
+        except:
+            print("Could not check ASLR: Couldn't get personality")
+            pass
 
     # Just go with whatever GDB says it did.
     #