From 18aa627d316e4a52df9e9b665d0e957794b7fc31 Mon Sep 17 00:00:00 2001
From: Zach Riggle <riggle@google.com>
Date: Sat, 9 May 2015 14:53:35 -0700
Subject: [PATCH] Enable cache-while-running; add ASLR command; work around GDB
 internal aborts

---
 pwndbg/commands/aslr.py | 25 +++++++++++++++++++++++--
 pwndbg/memoize.py       | 22 +++++++++++++++++++++-
 pwndbg/remote.py        | 10 ++++++++--
 pwndbg/vmmap.py         | 20 +++++++++++++++-----
 4 files changed, 67 insertions(+), 10 deletions(-)

diff --git a/pwndbg/commands/aslr.py b/pwndbg/commands/aslr.py
index fa46b48b1..aa797199d 100644
--- a/pwndbg/commands/aslr.py
+++ b/pwndbg/commands/aslr.py
@@ -1,13 +1,34 @@
+import gdb
 import pwndbg.color
 import pwndbg.commands
+import pwndbg.proc
 import pwndbg.vmmap
 
 
 @pwndbg.commands.OnlyWhenRunning
 @pwndbg.commands.Command
-def aslr():
+def aslr(on_or_off=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 pwndbg.proc.alive:
+                print("Change will take effect when the process restarts")
+
+    aslr = pwndbg.vmmap.check_aslr()
     status = pwndbg.color.red('OFF')
-    if pwndbg.vmmap.aslr:
+
+    if aslr:
         status = pwndbg.color.green('ON')
 
     print("ASLR is %s" % status)
diff --git a/pwndbg/memoize.py b/pwndbg/memoize.py
index 896eca61c..11ec0c249 100644
--- a/pwndbg/memoize.py
+++ b/pwndbg/memoize.py
@@ -18,6 +18,8 @@ import pwndbg.events
 debug = False
 
 class memoize(object):
+    caching = True
+
     def __init__(self, func):
         self.func  = func
         self.cache = {}
@@ -32,7 +34,7 @@ class memoize(object):
             how   = "Not memoizeable!"
             value = self.func(*args)
 
-        if args in self.cache:
+        if self.caching and args in self.cache:
             how   = "Cached"
             value = self.cache[args]
 
@@ -103,3 +105,21 @@ class reset_on_start(memoize):
     def __reset():
         for obj in reset_on_start.caches:
             obj.clear()
+
+class while_running(memoize):
+    caches = []
+    kind   = 'running'
+    caching = False
+
+    @staticmethod
+    @pwndbg.events.start
+    def __start_caching():
+        while_running.caching = True
+
+    @staticmethod
+    @pwndbg.events.exit
+    def __reset():
+        for obj in while_running.caches:
+            obj.clear()
+        while_running.caching = False
+
diff --git a/pwndbg/remote.py b/pwndbg/remote.py
index 6e09b8c81..02d5e3a7a 100644
--- a/pwndbg/remote.py
+++ b/pwndbg/remote.py
@@ -5,7 +5,13 @@ Information about whether the debuggee is local (under GDB) or remote
 (under GDBSERVER or QEMU stub).
 """
 import gdb
+import pwndbg.proc
 
-
+@pwndbg.proc.OnlyWhenRunning
 def is_remote():
-    return 'serial line' in gdb.execute('info program',to_string=True)
+    # N.B.: We cannot use "info program" because of:
+    # https://sourceware.org/bugzilla/show_bug.cgi?id=18335
+    #
+    # return 'serial line' in gdb.execute('info program',to_string=True,)
+
+    return 'Remote' in gdb.execute('info file',to_string=True,from_tty=False)
diff --git a/pwndbg/vmmap.py b/pwndbg/vmmap.py
index 657b242b7..77285499b 100644
--- a/pwndbg/vmmap.py
+++ b/pwndbg/vmmap.py
@@ -317,16 +317,26 @@ def find_boundaries(addr, name=''):
 
 aslr = False
 
-@pwndbg.events.stop
-@pwndbg.memoize.reset_on_exit
+@pwndbg.events.new_objfile
+@pwndbg.memoize.while_running
 def check_aslr():
     vmmap = sys.modules[__name__]
     vmmap.aslr = False
 
+    # Check to see if ASLR is disabled on the system.
+    # if not pwndbg.remote.is_remote():
     system_aslr = True
-    data        = ''
-    try: data = pwndbg.file.get('/proc/sys/kernel/randomize_va_space')
-    except OSError: pass
+    data        = b''
+
+    try:
+        data = pwndbg.file.get('/proc/sys/kernel/randomize_va_space')
+    except Exception as e:
+        print(e)
+        pass
+
+    # Systemwide ASLR is disabled
+    if b'0' in data:
+        return
 
     output = gdb.execute('show disable-randomization', to_string=True)
     if "is off." in output: