Merge pull request #1627 from guardicore/1594-signal-handlers

Agent: register signal handlers

monkey/infection_monkey/monkey.py

@@ -13,18 +13,23 @@ from common.version import get_version
 from infection_monkey.config import WormConfiguration
 from infection_monkey.control import ControlClient
 from infection_monkey.exploit.HostExploiter import HostExploiter
+from infection_monkey.master.mock_master import MockMaster
 from infection_monkey.model import DELAY_DELETE_CMD
 from infection_monkey.network.firewall import app as firewall
 from infection_monkey.network.HostFinger import HostFinger
 from infection_monkey.network.network_scanner import NetworkScanner
 from infection_monkey.network.tools import get_interface_to_target, is_running_on_island
 from infection_monkey.post_breach.post_breach_handler import PostBreach
+from infection_monkey.puppet.mock_puppet import MockPuppet
 from infection_monkey.ransomware.ransomware_payload_builder import build_ransomware_payload
 from infection_monkey.system_info import SystemInfoCollector
 from infection_monkey.system_singleton import SystemSingleton
 from infection_monkey.telemetry.attack.t1106_telem import T1106Telem
 from infection_monkey.telemetry.attack.t1107_telem import T1107Telem
 from infection_monkey.telemetry.attack.victim_host_telem import VictimHostTelem
+from infection_monkey.telemetry.messengers.legacy_telemetry_messenger_adapter import (
+    LegacyTelemetryMessengerAdapter,
+)
 from infection_monkey.telemetry.scan_telem import ScanTelem
 from infection_monkey.telemetry.state_telem import StateTelem
 from infection_monkey.telemetry.system_info_telem import SystemInfoTelem
@@ -38,6 +43,7 @@ from infection_monkey.utils.monkey_dir import (
     remove_monkey_dir,
 )
 from infection_monkey.utils.monkey_log_path import get_monkey_log_path
+from infection_monkey.utils.signal_handler import register_signal_handlers
 from infection_monkey.windows_upgrader import WindowsUpgrader
 
 MAX_DEPTH_REACHED_MESSAGE = "Reached max depth, skipping propagation phase."
@@ -107,6 +113,9 @@ class InfectionMonkey(object):
             logger.info("Monkey is starting...")
 
             logger.debug("Starting the setup phase.")
+            mock_master = MockMaster(MockPuppet(), LegacyTelemetryMessengerAdapter())
+            register_signal_handlers(mock_master)
+
             # Sets island's IP and port for monkey to communicate to
             self.set_default_server()
             self.set_default_port()

monkey/infection_monkey/utils/signal_handler.py

@@ -0,0 +1,33 @@
+import logging
+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_exception import PlannedShutdownException
+
+logger = logging.getLogger(__name__)
+
+
+class StopSignalHandler:
+    def __init__(self, master: IMaster):
+        self._master = master
+
+    def __call__(self, signum, _=None):
+        logger.info(f"The Monkey Agent received signal {signum}")
+        self._master.terminate()
+        raise PlannedShutdownException("Monkey Agent got an interrupt signal")
+
+
+def register_signal_handlers(master: IMaster):
+    stop_signal_handler = StopSignalHandler(master)
+    signal.signal(signal.SIGINT, stop_signal_handler)
+    signal.signal(signal.SIGTERM, stop_signal_handler)
+
+    if is_windows_os():
+        import win32api
+
+        signal.signal(signal.SIGBREAK, stop_signal_handler)
+
+        # CTRL_CLOSE_EVENT signal has a timeout of 5000ms,
+        # after that OS will forcefully kill the process
+        win32api.SetConsoleCtrlHandler(stop_signal_handler, True)