Agent: Call IMaster.terminate() from signal handler

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
@@ -108,7 +113,8 @@ class InfectionMonkey(object):
             logger.info("Monkey is starting...")
 
             logger.debug("Starting the setup phase.")
-            register_signal_handlers()
+            mock_master = MockMaster(MockPuppet(), LegacyTelemetryMessengerAdapter())
+            register_signal_handlers(mock_master)
             input()
             # Sets island's IP and port for monkey to communicate to
             self.set_default_server()

monkey/infection_monkey/utils/signal_handler.py

@@ -1,19 +1,25 @@
 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__)
 
 
-def stop_signal_handler(_, __):
-    # IMaster.cleanup()
-    logger.debug("Some kind of interrupt signal was sent to the Monkey Agent")
-    raise PlannedShutdownException("Monkey Agent got an interrupt signal")
+class StopSignalHandler:
+    def __init__(self, master: IMaster):
+        self._master = master
+
+    def __call__(self, _, __):
+        self._master.terminate()
+        logger.debug("Some kind of interrupt signal was sent to the Monkey Agent")
+        raise PlannedShutdownException("Monkey Agent got an interrupt signal")
 
 
-def register_signal_handlers():
+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)