diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index 4160a36e0..09eef703d 100644 --- a/monkey/infection_monkey/monkey.py +++ b/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() diff --git a/monkey/infection_monkey/utils/signal_handler.py b/monkey/infection_monkey/utils/signal_handler.py new file mode 100644 index 000000000..d75b08f10 --- /dev/null +++ b/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)