Repository: Gallopsled/pwntools

Issue:
atexit module can't handle multiple instances
The `atexit` module uses a dictionary to keep track of the routines to call.  This means that no routines can be registered more than once, for example with different arguments.

Additionally, this means that `pwnlib.atexit` does not guarantee order.

Per the [Python documentation](https://docs.python.org/2/library/atexit.html#module-atexit) for `atexit`:

_atexit runs these functions in the reverse order in which they were registered; if you register A, B, and C, at interpreter termination time they will be run in the order C, B, A._

However, multiple registration of routines breaks `unregister` in its current incarnation.  The suggested remedy for this is to return a unique ID for each `register` call that is passed to `unregister`, rather than passing in the function to `unregister`.



Discussion:
@br0ns has a local fix. He will do a PR when he has time.


Resolving patch:
diff --git a/pwnlib/atexit.py b/pwnlib/atexit.py
--- a/pwnlib/atexit.py
+++ b/pwnlib/atexit.py
@@ -8,58 +8,76 @@
 printed twice when the standard :mod:`atexit` is used.
 """
 
-import sys, traceback
+import sys, traceback, threading
+from .context import context
 
 __all__ = ['register', 'unregister']
 
+_lock = threading.Lock()
+_ident = 0
 _handlers = {}
 
 def register(func, *args, **kwargs):
     """register(func, *args, **kwargs)
 
-    Registers a function to be called when an unhandled exception occurs.  The
-    function will be called with positional arguments `args` and keyword
-    arguments `kwargs`, i.e. ``func(*args, **kwargs)``.
+    Registers a function to be called on program termination.  The function will
+    be called with positional arguments `args` and keyword arguments `kwargs`,
+    i.e. ``func(*args, **kwargs)``.  The current `context` is recorded and will
+    be the one used when the handler is run.
 
-    If `func` is already registered then `args` and `kwargs` will be updated.
+    E.g. to suppress logging output from an exit-handler one could write::
 
-    This function can be used as a decorator::
+      with context.local(log_level = 'error'):
+        atexit.register(handler)
 
-      def f():
-        ...
-      atexit.register(f)
+    An identifier is returned which can be used to unregister the exit-handler.
 
-    is equivalent to this::
+    This function can be used as a decorator::
 
       @atexit.register
-      def f():
+      def handler():
         ...
 
-    """
-    _handlers[func] = (args, kwargs)
-    return func
+    Notice however that this will bind ``handler`` to the identifier and not the
+    actual exit-handler.  The exit-handler can then be unregistered with::
+
+      atexit.unregister(handler)
 
-def unregister(func):
-    """unregister(func)
+    This function is thread safe.
 
-    Remove `func` from the collection of registered functions.  If `func` isn't
-    registered this is a no-op.
     """
-    if func in _handlers:
-        del _handlers[func]
+    global _ident
+    with _lock:
+        ident = _ident
+        _ident += 1
+    _handlers[ident] = (func, args, kwargs, vars(context))
+    return ident
+
+def unregister(ident):
+    """unregister(ident)
+
+    Remove the exit-handler identified by `ident` from the list of registered
+    handlers.  If `ident` isn't registered this is a no-op.
+    """
+    if ident in _handlers:
+        del _handlers[ident]
 
 def _run_handlers():
     """_run_handlers()
 
-    Run registered exit handlers.  The order is arbitrary.
+    Run registered exit-handlers.  They run in the reverse order of which they
+    were registered.
 
     If a handler raises an exception, it will be printed but nothing else
     happens, i.e. other handlers will be run and `sys.excepthook` will not be
     called for that reason.
     """
-    for func, (args, kwargs) in _handlers.items():
+    context.clear()
+    for _ident, (func, args, kwargs, ctx) in \
+        sorted(_handlers.items(), reverse = True):
         try:
-            func(*args, **kwargs)
+            with context.local(**ctx):
+                func(*args, **kwargs)
         except SystemExit:
             pass
         except:


Test patch: