Lots of enhancements, and adds the "start" command

pwndbg/auxv.py

@@ -87,7 +87,7 @@ class AUXV(dict):
     def __str__(self):
         return str({k:v for k,v in self.items() if v is not None})
 
-@pwndbg.memoize.reset_on_objfile
+@pwndbg.memoize.reset_on_start
 def get():
     return use_info_auxv() or walk_stack() or AUXV()
 

pwndbg/commands/__init__.py

@@ -34,7 +34,7 @@ debug = True
 
 class Command(gdb.Command):
     def __init__(self, function):
-        function.__doc__
+        self.__doc__ = function.__doc__
         super(Command, self).__init__(function.__name__, gdb.COMMAND_USER, gdb.COMPLETE_EXPRESSION)
         self.function = function
 
@@ -77,13 +77,4 @@ def fix(arg, sloppy=False):
 
     return None
 
-def OnlyWhenRunning(func):
-    def wrapper(*a, **kw):
-        func.__doc__
-        if not pwndbg.proc.alive:
-            pass
-        else:
-            return func(*a, **kw)
-    wrapper.__name__ = func.__name__
-    wrapper.__module__ = func.__module__
-    return wrapper
+OnlyWhenRunning = pwndbg.proc.OnlyWhenRunning

pwndbg/commands/start.py

@@ -7,10 +7,25 @@ entry point.
 import gdb
 import pwndbg.commands
 import pwndbg.symbol
+import pwndbg.events
+import pwndbg.elf
 
+break_on_first_instruction = False
 
-@pwndbg.commands.ParsedCommand
+@pwndbg.events.start
+def on_start():
+    global break_on_first_instruction
+    if break_on_first_instruction:
+        spec = "*%#x" % (int(pwndbg.elf.entry()))
+        gdb.Breakpoint(spec, temporary=True)
+        break_on_first_instruction = False
+
+@pwndbg.commands.Command
 def start():
+    """
+    Set a breakpoint at a convenient location in the binary,
+    generally 'main', 'init', or the entry point.
+    """
     symbols = ["main",
                 "_main",
                 "start",
@@ -24,5 +39,18 @@ def start():
             b = gdb.Breakpoint('*%#x' % address, temporary=True)
             gdb.execute('run', from_tty=False, to_string=True)
             break
+
     else:
-        print("Could not find a good place to start :(")
+        entry()
+
+
+@pwndbg.commands.Command
+def entry():
+    """
+    Set a breakpoint at the first instruction executed in
+    the target binary.
+    """
+    global break_on_first_instruction
+    break_on_first_instruction = True
+    print("Trying experimental breakpoint")
+    gdb.execute('run', from_tty=False, to_string=True)

pwndbg/elf.py

@@ -18,6 +18,7 @@ import gdb
 import pwndbg.auxv
 import pwndbg.events
 import pwndbg.info
+import pwndbg.proc
 import pwndbg.memoize
 import pwndbg.memory
 import pwndbg.stack
@@ -47,17 +48,17 @@ Elf64_Phdr f;
 
 subprocess.check_output('gcc -c -g %s.c -o %s.o' % (gef_elf, gef_elf), shell=True)
 
-@pwndbg.memoize.reset_on_objfile
+@pwndbg.proc.OnlyWhenRunning
+@pwndbg.memoize.reset_on_start
 def exe():
     """
     Return a loaded ELF header object pointing to the Ehdr of the
     main executable.
     """
-    elf = None
-    ptr = entry()
-    return load(ptr)
+    return load(entry())
 
-@pwndbg.memoize.reset_on_objfile
+@pwndbg.proc.OnlyWhenRunning
+@pwndbg.memoize.reset_on_start
 def entry():
     """
     Return the address of the entry point for the main executable.
@@ -127,9 +128,16 @@ def get_ehdr(pointer):
     try:
         data = pwndbg.memory.read(base, 4)
 
-        while data != b'\x7FELF':
+        # Do not search more than 4MB of memory
+        for i in range(1024):
             base -= pwndbg.memory.PAGE_SIZE
             data = pwndbg.memory.read(base, 4)
+
+            if data == b'\x7FELF':
+                break
+        else:
+            print("ERROR: Could not find ELF base!")
+            return None, None
     except gdb.MemoryError:
         return None, None
 
@@ -203,7 +211,7 @@ def map(pointer, objfile=''):
     ei_class, ehdr         = get_ehdr(pointer)
     return map_inner(ei_class, ehdr, objfile)
 
-@pwndbg.memoize.reset_on_stop
+@pwndbg.memoize.reset_on_objfile
 def map_inner(ei_class, ehdr, objfile):
     if not ehdr:
         return []

pwndbg/events.py

@@ -9,15 +9,49 @@ import traceback
 import gdb
 import sys
 
-debug = True
+debug = False
 pause = 0
 
+
+# There is no GDB way to get a notification when the binary itself
+# is loaded from disk, by the operating system, before absolutely
+# anything happens
+#
+# However, we get an Objfile event when the binary is loaded, before
+# its entry point is invoked.
+#
+# We also get an Objfile event when we load up GDB, so we need
+# to detect when the binary is running or not.
+class StartEvent(object):
+    def __init__(self):
+        self.registered = set()
+        self.running    = False
+    def connect(self, function):
+        self.registered.add(function)
+    def disconnect(self, function):
+        if function in self.registered:
+            self.registered.remove(function)
+    def on_new_objfile(self, o):
+        if self.running or not gdb.selected_thread():
+            return
+
+        self.running = True
+
+        for function in self.registered:
+            function(o)
+
+    def on_stop(self, e):
+        self.running = False
+
+gdb.events.start = StartEvent()
+
 # In order to support reloading, we must be able to re-fire
 # all 'objfile' and 'stop' events.
 registered = {gdb.events.exited: [],
               gdb.events.cont: [],
               gdb.events.new_objfile: [],
-              gdb.events.stop: []}
+              gdb.events.stop: [],
+              gdb.events.start: []}
 
 class Pause(object):
     def __enter__(self, *a, **kw):
@@ -28,6 +62,9 @@ class Pause(object):
         pause -= 1
 
 def connect(func, event_handler, name=''):
+    if debug:
+        print("Connecting", func.__name__, event_handler)
+
     def caller(*a):
         func.__doc__
         if debug: sys.stdout.write('%r %s.%s %r\n' % (name, func.__module__, func.__name__, a))
@@ -37,9 +74,10 @@ def connect(func, event_handler, name=''):
         except Exception as e:
             if debug: print(traceback.format_exc())
             raise e
-        if debug: sys.stdout.write('DONE %r %s.%s %r\n' % (name, func.__module__, func.__name__, a))
+
     registered[event_handler].append(caller)
     caller.name     = func.__name__
+    caller.__name__ = func.__name__
     event_handler.connect(caller)
     return func
 
@@ -47,6 +85,7 @@ def exit(func):        return connect(func, gdb.events.exited, 'exit')
 def cont(func):        return connect(func, gdb.events.cont, 'cont')
 def new_objfile(func): return connect(func, gdb.events.new_objfile, 'obj')
 def stop(func):        return connect(func, gdb.events.stop, 'stop')
+def start(func):       return connect(func, gdb.events.start, 'start')
 
 def after_reload():
     return
@@ -56,9 +95,16 @@ def after_reload():
     #     for f in registered[gdb.events.stop]:
     #         f()
 
-
 def on_reload():
     for event, functions in registered.items():
         for function in functions:
             event.disconnect(function)
         registered[event] = []
+
+@new_objfile
+def _start_newobjfile(o=None):
+    gdb.events.start.on_new_objfile(o)
+
+@stop
+def _start_stop(o=None):
+    gdb.events.start.on_stop(o)

pwndbg/ida.py

@@ -16,12 +16,18 @@ import pwndbg.events
 import pwndbg.memoize
 import pwndbg.memory
 import pwndbg.regs
+import pwndbg.compat
 
 try:
     import xmlrpc.client as xmlrpclib
 except:
     import xmlrpclib
 
+if pwndbg.compat.python2:
+    xmlrpclib.Marshaller.dispatch[type(0L)] = lambda _, v, w: w("<value><i8>%d</i8></value>" % v)
+    xmlrpclib.Marshaller.dispatch[type(0)] = lambda _, v, w: w("<value><i8>%d</i8></value>" % v)
+
+
 _ida = None
 
 xmlrpclib.Marshaller.dispatch[type(0)] = lambda _, v, w: w("<value><i8>%d</i8></value>" % v)

pwndbg/memoize.py

@@ -15,37 +15,51 @@ import sys
 import gdb
 import pwndbg.events
 
+debug = False
 
 class memoize(object):
     def __call__(self, *args):
+        how = None
+
         if not isinstance(args, collections.Hashable):
             print("Cannot memoize %r!", file=sys.stderr)
-            return self.func(*args)
+            how   = "Not memoizeable!"
+            value = self.func(*args)
 
         if args in self.cache:
-            return self.cache[args]
+            how   = "Cached"
+            value = self.cache[args]
 
+        else:
+            how   = "Executed"
             value = self.func(*args)
             self.cache[args] = value
 
             if isinstance(value, list):
                 print("Shouldnt cache mutable types! %r" % self.func.__name__)
 
+        if debug:
+            print("%s: %s(%r)" % (how, self, args))
+            print(".... %r" % (value,))
         return value
 
     def __repr__(self):
-        return self.func.__doc__
+        funcname = self.func.__module__ + '.' + self.func.__name__
+        return "<%s-memoized function %s>" % (self.kind, funcname)
 
     def __get__(self, obj, objtype):
         return functools.partial(self.__call__, obj)
 
     def clear(self):
+        if debug:
+            print("Clearing %s %r" % (self, self.cache))
         self.cache.clear()
 
 
 
 class reset_on_stop(memoize):
     caches = []
+    kind   = 'stop'
 
     def __init__(self, func):
         self.func  = func
@@ -60,6 +74,7 @@ class reset_on_stop(memoize):
 
 class reset_on_exit(memoize):
     caches = []
+    kind   = 'exit'
 
     def __init__(self, func):
         self.func  = func
@@ -76,6 +91,7 @@ class reset_on_exit(memoize):
 
 class reset_on_objfile(memoize):
     caches = []
+    kind   = 'objfile'
 
     def __init__(self, func):
         self.func  = func
@@ -89,3 +105,21 @@ class reset_on_objfile(memoize):
     def __reset():
         for obj in reset_on_objfile.caches:
             obj.clear()
+
+class reset_on_start(memoize):
+    caches = []
+    kind   = 'start'
+
+    def __init__(self, func):
+        self.func  = func
+        self.cache = {}
+        self.caches.append(self)
+        self.__name__ = func.__name__
+        self.__module__ = func.__module__
+
+    @staticmethod
+    @pwndbg.events.stop
+    @pwndbg.events.start
+    def __reset():
+        for obj in reset_on_start.caches:
+            obj.clear()

pwndbg/memory.py

@@ -72,25 +72,20 @@ def page_offset(address): return (address & (PAGE_SIZE-1))
 assert round_down(0xdeadbeef, 0x1000) == 0xdeadb000
 assert round_up(0xdeadbeef, 0x1000)   == 0xdeadc000
 
-def find_upper_boundary(addr):
+def find_upper_boundary(addr, max_pages=1024):
     addr = pwndbg.memory.page_align(int(addr))
     try:
-        while True:
-            import sys
-            sys.stdout.write(hex(addr) + '\n')
-            sys.stdout.flush()
+        for i in range(max_pages):
             pwndbg.memory.read(addr, 1)
             addr += pwndbg.memory.PAGE_SIZE
     except gdb.MemoryError:
         pass
     return addr
 
-def find_lower_boundary(addr):
+def find_lower_boundary(addr, max_pages=1024):
     addr = pwndbg.memory.page_align(int(addr))
     try:
-        while True:
-            sys.stdout.write(hex(addr) + '\n')
-            sys.stdout.flush()
+        for i in range(max_pages):
             pwndbg.memory.read(addr, 1)
             addr -= pwndbg.memory.PAGE_SIZE
     except gdb.MemoryError:

pwndbg/proc.py

@@ -21,7 +21,16 @@ class module(ModuleType):
 
     @property
     def alive(self):
-        return self.pid > 0
+        return gdb.selected_thread() is not None
+
+    def OnlyWhenRunning(self, func):
+        def wrapper(*a, **kw):
+            func.__doc__
+            if self.alive:
+                return func(*a, **kw)
+        wrapper.__name__ = func.__name__
+        wrapper.__module__ = func.__module__
+        return wrapper
 
 # To prevent garbage collection
 tether = sys.modules[__name__]

pwndbg/stack.py

@@ -37,9 +37,6 @@ def update():
     For each running thread, updates the known address range
     for its stack.
     """
-    # import pdb
-    # pdb.set_trace()
-
     curr_thread = gdb.selected_thread()
 
     try:
@@ -51,7 +48,7 @@ def update():
             # a new Page mapping for it.
             page = stacks.get(thread.ptid, None)
             if page is None:
-                start = sp + 0x1000 & ~(0xfff) #pwndbg.memory.find_lower_boundary(sp)
+                start = pwndbg.memory.find_lower_boundary(sp)
                 stop  = pwndbg.memory.find_upper_boundary(sp)
                 page  = pwndbg.memory.Page(start, stop-start, 6 if not is_executable() else 7, 0, '[stack]')
                 stacks[thread.ptid] = page