diff --git a/monkey/infection_monkey/utils/signal_handler.py b/monkey/infection_monkey/utils/signal_handler.py index 6fda3bc12..831b31441 100644 --- a/monkey/infection_monkey/utils/signal_handler.py +++ b/monkey/infection_monkey/utils/signal_handler.py @@ -3,7 +3,6 @@ import signal from infection_monkey.i_master import IMaster from infection_monkey.utils.environment import is_windows_os -from infection_monkey.utils.exceptions.planned_shutdown_error import PlannedShutdownError logger = logging.getLogger(__name__) @@ -12,28 +11,29 @@ class StopSignalHandler: def __init__(self, master: IMaster): self._master = master - def handle_posix_signals(self, signum, _): - self._handle_signal(signum) - # Windows signal handlers must return boolean. Only raising this exception for POSIX - # signals. - raise PlannedShutdownError("Monkey Agent got an interrupt signal") + def handle_posix_signals(self, signum: int, _): + self._handle_signal(signum, False) - def handle_windows_signals(self, signum): + def handle_windows_signals(self, signum: int): import win32con - # TODO: This signal handler gets called for a CTRL_CLOSE_EVENT, but the system immediately - # kills the process after the handler returns. After the master is implemented and the - # setup/teardown of the Agent is fully refactored, revisit this signal handler and - # modify as necessary to more gracefully handle CTRL_CLOSE_EVENT signals. - if signum in {win32con.CTRL_C_EVENT, win32con.CTRL_BREAK_EVENT, win32con.CTRL_CLOSE_EVENT}: - self._handle_signal(signum) + if signum in {win32con.CTRL_C_EVENT, win32con.CTRL_BREAK_EVENT}: + self._handle_signal(signum, False) + return True + + if signum == win32con.CTRL_CLOSE_EVENT: + # After the signal handler returns True, the OS will forcefully kill the process. + # Calling self._handle_signal() with block=True to give the master a chance to + # gracefully shut down. Note that the OS has a timeout that will forcefully kill the + # process if this handler hasn't returned in time. + self._handle_signal(signum, True) return True return False - def _handle_signal(self, signum): + def _handle_signal(self, signum: int, block: bool): logger.info(f"The Monkey Agent received signal {signum}") - self._master.terminate() + self._master.terminate(block) def register_signal_handlers(master: IMaster):