Merge pull request #1661 from guardicore/1597-signal-handling

1597 signal handling

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

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()

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:

monkey/infection_monkey/utils/exceptions/planned_shutdown_error.py

@@ -1,2 +0,0 @@
-class PlannedShutdownError(Exception):
-    pass

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):