mirror of
https://github.com/Gallopsled/pwntools.git
synced 2026-10-12 02:02:27 +00:00
peace-maker
658479605b
* Fix atexception handlers in term mode The colored-traceback module registers its own `sys.excepthook` but doesn't call the original handler. Reorder the includes such that the colored-traceback is setup before and our own hooks afterwards. We always call the old excepthook and forward the exception to the colored-traceback after firing our own handlers. It worked before in non-term mode when stderr wasn't a tty and colors were disabled. * Update CHANGELOG * Import atexception in doctests
122 lines
3.7 KiB
Python
122 lines
3.7 KiB
Python
"""
|
|
Analogous to atexit, this module allows the programmer to register functions to
|
|
be run if an unhandled exception occurs.
|
|
"""
|
|
import sys
|
|
import threading
|
|
import traceback
|
|
from types import TracebackType
|
|
from typing import Any, Callable, ParamSpec, TypeVar
|
|
|
|
from pwnlib.context import context
|
|
|
|
__all__ = ['register', 'unregister']
|
|
|
|
_lock = threading.Lock()
|
|
_ident = 0
|
|
_handlers: dict[int, tuple[Callable[..., Any], Any, Any, dict[str, Any]]] = {}
|
|
|
|
_P = ParamSpec('_P')
|
|
_R = TypeVar('_R')
|
|
|
|
def register(func: Callable[_P, _R], *args: _P.args, **kwargs: _P.kwargs) -> int:
|
|
"""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)``. The current `context`
|
|
is recorded and will be the one used when the handler is run.
|
|
|
|
E.g. to suppress logging output from an exception-handler one could write::
|
|
|
|
with context.local(log_level = 'error'):
|
|
atexception.register(handler)
|
|
|
|
An identifier is returned which can be used to unregister the
|
|
exception-handler.
|
|
|
|
This function can be used as a decorator::
|
|
|
|
@atexception.register
|
|
def handler():
|
|
...
|
|
|
|
Notice however that this will bind ``handler`` to the identifier and not the
|
|
actual exception-handler. The exception-handler can then be unregistered
|
|
with::
|
|
|
|
ident = atexception.register(handler)
|
|
atexception.unregister(ident)
|
|
|
|
This function is thread safe.
|
|
|
|
>>> def handler(a, b):
|
|
... print("Exception handler", a, b)
|
|
>>> ident = atexception.register(handler, 13, 37)
|
|
>>> raise Exception("Unexpected error") # doctest: +ELLIPSIS
|
|
Traceback (most recent call last):
|
|
...
|
|
Exception: Unexpected error
|
|
Exception handler 13 37
|
|
>>> atexception.unregister(ident)
|
|
>>> raise Exception("Unexpected error") # doctest: +ELLIPSIS
|
|
Traceback (most recent call last):
|
|
...
|
|
Exception: Unexpected error
|
|
"""
|
|
global _ident
|
|
with _lock:
|
|
ident = _ident
|
|
_ident += 1
|
|
_handlers[ident] = (func, args, kwargs, vars(context))
|
|
return ident
|
|
|
|
def unregister(ident: int) -> None:
|
|
"""unregister(ident)
|
|
|
|
Remove the exception-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() -> None:
|
|
"""_run_handlers()
|
|
|
|
Run registered 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.
|
|
"""
|
|
for _ident, (func, args, kwargs, ctx) in \
|
|
sorted(_handlers.items(), reverse = True):
|
|
try:
|
|
with context.local():
|
|
context.clear()
|
|
context.update(**ctx)
|
|
func(*args, **kwargs)
|
|
except SystemExit:
|
|
pass
|
|
except Exception:
|
|
# extract the current exception and rewind the traceback to where it
|
|
# originated
|
|
typ, val, tb = sys.exc_info()
|
|
traceback.print_exception(typ, value=val, tb=tb.tb_next if tb else None)
|
|
|
|
# we rely on the existing excepthook to print exceptions
|
|
_oldhook = getattr(sys, 'excepthook', None)
|
|
|
|
def _newhook(typ: type[BaseException], val: BaseException, tb: TracebackType | None) -> None:
|
|
"""_newhook(typ, val, tb)
|
|
|
|
Our excepthook replacement. First the original hook is called to print the
|
|
exception, then each handler is called.
|
|
"""
|
|
if _oldhook:
|
|
_oldhook(typ, val, tb)
|
|
if _handlers:
|
|
_run_handlers()
|
|
|
|
sys.excepthook = _newhook
|