import signal

SIGNALS = [signal.SIGINT, signal.SIGTERM]


class SignalProtector:
    def __init__(self):
        self._sig = None
        self._frame = None
        self._old_handlers = None

    def __enter__(self):
        self._sig = None
        self._frame = None
        self._old_handlers = [signal.signal(sig, self._handler) for sig in SIGNALS]

    def __exit__(self, exc_type, exc_val, exc_tb):
        for sig, handler in zip(SIGNALS, self._old_handlers):
            signal.signal(sig, handler)
        for sig, handler in zip(SIGNALS, self._old_handlers):
            if self._sig == sig:
                handler(self._sig, self._frame)

    def _handler(self, sig, frame):
        self._sig = sig
        self._frame = frame