From 9d36f20b4238db1afe9c1a89f3dbf95e66d51142 Mon Sep 17 00:00:00 2001 From: VakarisZ Date: Tue, 23 Nov 2021 15:27:09 +0200 Subject: [PATCH 1/6] Agent: register signal handlers Agent will now handle interrupt and break signals on linux and windows --- monkey/infection_monkey/monkey.py | 3 +++ .../infection_monkey/utils/signal_handler.py | 21 +++++++++++++++++++ 2 files changed, 24 insertions(+) create mode 100644 monkey/infection_monkey/utils/signal_handler.py diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index 4160a36e0..4c5584558 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -38,6 +38,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 +108,8 @@ class InfectionMonkey(object): logger.info("Monkey is starting...") logger.debug("Starting the setup phase.") + register_signal_handlers() + input() # 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..32efa078b --- /dev/null +++ b/monkey/infection_monkey/utils/signal_handler.py @@ -0,0 +1,21 @@ +import logging +import signal + +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") + + +def register_signal_handlers(): + signal.signal(signal.SIGINT, stop_signal_handler) + signal.signal(signal.SIGTERM, stop_signal_handler) + + if is_windows_os(): + signal.signal(signal.SIGBREAK, stop_signal_handler) From 27ef06c546234e5e4f8526a7c9afb23cd0231e29 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 23 Nov 2021 12:17:17 -0500 Subject: [PATCH 2/6] Agent: Call IMaster.terminate() from signal handler --- monkey/infection_monkey/monkey.py | 8 +++++++- monkey/infection_monkey/utils/signal_handler.py | 16 +++++++++++----- 2 files changed, 18 insertions(+), 6 deletions(-) diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index 4c5584558..2fde8cf43 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 @@ -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() diff --git a/monkey/infection_monkey/utils/signal_handler.py b/monkey/infection_monkey/utils/signal_handler.py index 32efa078b..d125ad98d 100644 --- a/monkey/infection_monkey/utils/signal_handler.py +++ b/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) From 068307f0ebdd51d143e4fa16aa809fd8ef9a3ce1 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 23 Nov 2021 13:09:17 -0500 Subject: [PATCH 3/6] Agent: Handle window close event on Windows --- monkey/infection_monkey/utils/signal_handler.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/monkey/infection_monkey/utils/signal_handler.py b/monkey/infection_monkey/utils/signal_handler.py index d125ad98d..f15fded3a 100644 --- a/monkey/infection_monkey/utils/signal_handler.py +++ b/monkey/infection_monkey/utils/signal_handler.py @@ -12,7 +12,7 @@ class StopSignalHandler: def __init__(self, master: IMaster): self._master = master - def __call__(self, _, __): + def __call__(self, _, __=None): self._master.terminate() logger.debug("Some kind of interrupt signal was sent to the Monkey Agent") raise PlannedShutdownException("Monkey Agent got an interrupt signal") @@ -24,4 +24,7 @@ def register_signal_handlers(master: IMaster): signal.signal(signal.SIGTERM, stop_signal_handler) if is_windows_os(): + import win32api + signal.signal(signal.SIGBREAK, stop_signal_handler) + win32api.SetConsoleCtrlHandler(stop_signal_handler, True) From 6149ef630bc63c221fea730eff5dc5b00042297c Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 23 Nov 2021 13:18:11 -0500 Subject: [PATCH 4/6] Agent: Improve signal handler log message --- monkey/infection_monkey/utils/signal_handler.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/monkey/infection_monkey/utils/signal_handler.py b/monkey/infection_monkey/utils/signal_handler.py index f15fded3a..a2f865ce2 100644 --- a/monkey/infection_monkey/utils/signal_handler.py +++ b/monkey/infection_monkey/utils/signal_handler.py @@ -12,9 +12,9 @@ class StopSignalHandler: def __init__(self, master: IMaster): self._master = master - def __call__(self, _, __=None): + def __call__(self, signum, __=None): + logger.info(f"The Monkey Agent received signal {signum}") self._master.terminate() - logger.debug("Some kind of interrupt signal was sent to the Monkey Agent") raise PlannedShutdownException("Monkey Agent got an interrupt signal") From 73329e9729ddc309c1c9a954f3546652e97b2fde Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 23 Nov 2021 13:28:01 -0500 Subject: [PATCH 5/6] Agent: Remove input() call in monkey.py The call to input() was used to pause the execution of the agent while testing the new signal handlers. It is no longer needed. --- monkey/infection_monkey/monkey.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index 2fde8cf43..09eef703d 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -115,7 +115,7 @@ class InfectionMonkey(object): logger.debug("Starting the setup phase.") 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() self.set_default_port() From 3f7c4a8859b1d71cab4982f49489c8b06d309d50 Mon Sep 17 00:00:00 2001 From: VakarisZ Date: Wed, 24 Nov 2021 10:40:05 +0200 Subject: [PATCH 6/6] Agent: add a comment warning that windows will terminate the process 5s after CTRL_CLOSE_EVENT signal The comment will warn us that in case that particular signal is raised, the cleanup shouldn't take longer than 5s --- monkey/infection_monkey/utils/signal_handler.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/monkey/infection_monkey/utils/signal_handler.py b/monkey/infection_monkey/utils/signal_handler.py index a2f865ce2..d75b08f10 100644 --- a/monkey/infection_monkey/utils/signal_handler.py +++ b/monkey/infection_monkey/utils/signal_handler.py @@ -12,7 +12,7 @@ class StopSignalHandler: def __init__(self, master: IMaster): self._master = master - def __call__(self, signum, __=None): + 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") @@ -27,4 +27,7 @@ def register_signal_handlers(master: IMaster): 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)