From 8351b3b39d930b8c05a1be6ba512eaca47cc225b Mon Sep 17 00:00:00 2001
From: Zach Riggle <riggle@google.com>
Date: Wed, 1 Feb 2017 23:52:44 -0500
Subject: [PATCH] wip

---
 pwndbg/elf.py      | 26 +++++++++++++++++++++++---
 pwndbg/elftypes.py | 20 --------------------
 pwndbg/events.py   |  7 +++++++
 3 files changed, 30 insertions(+), 23 deletions(-)

diff --git a/pwndbg/elf.py b/pwndbg/elf.py
index 51d019635..c57453510 100644
--- a/pwndbg/elf.py
+++ b/pwndbg/elf.py
@@ -12,15 +12,17 @@ from __future__ import division
 from __future__ import print_function
 from __future__ import unicode_literals
 
+import ctypes
 import os
 import re
 import subprocess
+import sys
 import tempfile
 
 import gdb
 
 import pwndbg.auxv
-import pwndbg.elftypes as E
+import pwndbg.elftypes
 import pwndbg.events
 import pwndbg.info
 import pwndbg.memoize
@@ -32,6 +34,24 @@ import pwndbg.stack
 PF_X, PF_W, PF_R = 1,2,4
 ET_EXEC, ET_DYN  = 2,3
 
+
+module = sys.modules[__name__]
+
+@pwndbg.events.start
+@pwndbg.events.new_objfile
+def update():
+    if pwndbg.arch.ptrsize == 4:
+        Ehdr = pwndbg.elftypes.Elf32_Ehdr
+        Phdr = pwndbg.elftypes.Elf32_Phdr
+    else:
+        Ehdr = pwndbg.elftypes.Elf64_Ehdr
+        Phdr = pwndbg.elftypes.Elf64_Phdr
+
+    module.__dict__.update(locals())
+
+update()
+
+
 def read(typ, address, blob=None):
     size = ctypes.sizeof(typ)
 
@@ -128,7 +148,7 @@ def get_ehdr(pointer):
     ei_class = pwndbg.memory.byte(base+4)
 
     # Find out where the section headers start
-    Elfhdr   = read(E.Ehdr, base)
+    Elfhdr   = read(Ehdr, base)
     return ei_class, Elfhdr
 
 def get_phdrs(pointer):
@@ -146,7 +166,7 @@ def get_phdrs(pointer):
     phoff     = Elfhdr.e_phoff
     phentsize = Elfhdr.e_phentsize
 
-    x = (phnum, phentsize, read(E.Phdr, Elfhdr.address + phoff))
+    x = (phnum, phentsize, read(Phdr, Elfhdr.address + phoff))
     return x
 
 def iter_phdrs(ehdr):
diff --git a/pwndbg/elftypes.py b/pwndbg/elftypes.py
index 1512cdda8..7ce5e5e08 100644
--- a/pwndbg/elftypes.py
+++ b/pwndbg/elftypes.py
@@ -313,23 +313,3 @@ class Elf64_Phdr(pwndbg.ctypes.Structure):
                 ("p_filesz", Elf64_Xword),
                 ("p_memsz", Elf64_Xword),
                 ("p_align", Elf64_Xword),]
-
-module = sys.modules[__name__]
-
-@pwndbg.events.start
-@pwndbg.events.new_objfile
-def reload():
-    six.moves.reload_module(module)
-    module.update()
-
-def update():
-    if pwndbg.arch.ptrsize == 4:
-        Ehdr = Elf32_Ehdr
-        Phdr = Elf32_Phdr
-    else:
-        Ehdr = Elf64_Ehdr
-        Phdr = Elf64_Phdr
-
-    module.__dict__.update(locals())
-
-update()
diff --git a/pwndbg/events.py b/pwndbg/events.py
index 14878e6ef..b86b00103 100644
--- a/pwndbg/events.py
+++ b/pwndbg/events.py
@@ -122,6 +122,13 @@ def connect(func, event_handler, name=''):
                 traceback.print_exc()
                 raise e
 
+    registry = registered[event_handler]
+
+    for event in registry.copy():
+        if event.__name__ == func.__name__ \
+        and event.__module__ == func.__module__:
+            print("Re-registered %s" % func)
+
     registered[event_handler].append(caller)
     event_handler.connect(caller)
     return func