diff --git a/monkey/infection_monkey/i_master.py b/monkey/infection_monkey/i_master.py index 9caa71a4d..5269cafee 100644 --- a/monkey/infection_monkey/i_master.py +++ b/monkey/infection_monkey/i_master.py @@ -10,9 +10,10 @@ class IMaster(metaclass=abc.ABCMeta): """ @abc.abstractmethod - def terminate(self) -> None: + def terminate(self, block: bool = False) -> None: """ Stop the master and interrupt any actions that are currently being executed. + :param bool block: Whether or not to block and wait for the master to terminate. """ @abc.abstractmethod diff --git a/monkey/infection_monkey/master/automated_master.py b/monkey/infection_monkey/master/automated_master.py index a7f67d87c..75a5a5dbf 100644 --- a/monkey/infection_monkey/master/automated_master.py +++ b/monkey/infection_monkey/master/automated_master.py @@ -52,12 +52,15 @@ class AutomatedMaster(IMaster): self._master_thread.join() logger.info("The simulation has been shutdown.") - def terminate(self): + def terminate(self, block: bool = False): logger.info("Stopping automated breach and attack simulation") self._stop.set() - if self._master_thread.is_alive(): + if self._master_thread.is_alive() and block: self._master_thread.join() + # We can only have confidence that the master terminated successfully if block is set + # and join() has returned. + logger.info("AutomatedMaster successfully terminated.") def _run_master_thread(self): self._simulation_thread.start() diff --git a/monkey/infection_monkey/master/mock_master.py b/monkey/infection_monkey/master/mock_master.py index 1d8791c4c..0b4f9a3f6 100644 --- a/monkey/infection_monkey/master/mock_master.py +++ b/monkey/infection_monkey/master/mock_master.py @@ -124,7 +124,7 @@ class MockMaster(IMaster): self._telemetry_messenger.send_telemetry(FileEncryptionTelem(path, success, error)) logger.info("Finished running payloads") - def terminate(self) -> None: + def terminate(self, block: bool = False) -> None: logger.info("Terminating MockMaster") def cleanup(self) -> None: diff --git a/monkey/infection_monkey/utils/exceptions/planned_shutdown_error.py b/monkey/infection_monkey/utils/exceptions/planned_shutdown_error.py deleted file mode 100644 index 885340c23..000000000 --- a/monkey/infection_monkey/utils/exceptions/planned_shutdown_error.py +++ /dev/null @@ -1,2 +0,0 @@ -class PlannedShutdownError(Exception): - pass 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):