p20319658
/
monkey
forked from p34709852/monkey
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Agent: Wait for master to terminate on windows CTRL_CLOSE_EVENT

This commit is contained in:
Mike Salvatore 2021-12-15 09:02:44 -05:00
parent f46bb60da5
commit 3f9bd24228
1 changed files with 15 additions and 15 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

30
monkey/infection_monkey/utils/signal_handler.py
View File

@ -3,7 +3,6 @@ import signal
from infection_monkey.i_master import IMaster from infection_monkey.i_master import IMaster
from infection_monkey.utils.environment import is_windows_os from infection_monkey.utils.environment import is_windows_os
from infection_monkey.utils.exceptions.planned_shutdown_error import PlannedShutdownError
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@ -12,28 +11,29 @@ class StopSignalHandler:
def __init__(self, master: IMaster): def __init__(self, master: IMaster):
self._master = master self._master = master
def handle_posix_signals(self, signum, _): def handle_posix_signals(self, signum: int, _):
self._handle_signal(signum) self._handle_signal(signum, False)
# Windows signal handlers must return boolean. Only raising this exception for POSIX
# signals.
raise PlannedShutdownError("Monkey Agent got an interrupt signal")
def handle_windows_signals(self, signum): def handle_windows_signals(self, signum: int):
import win32con import win32con
# TODO: This signal handler gets called for a CTRL_CLOSE_EVENT, but the system immediately if signum in {win32con.CTRL_C_EVENT, win32con.CTRL_BREAK_EVENT}:
# kills the process after the handler returns. After the master is implemented and the self._handle_signal(signum, False)
# setup/teardown of the Agent is fully refactored, revisit this signal handler and return True
# 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}: if signum == win32con.CTRL_CLOSE_EVENT:
self._handle_signal(signum) # 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 True
return False return False
def _handle_signal(self, signum): def _handle_signal(self, signum: int, block: bool):
logger.info(f"The Monkey Agent received signal {signum}") logger.info(f"The Monkey Agent received signal {signum}")
self._master.terminate() self._master.terminate(block)
def register_signal_handlers(master: IMaster): def register_signal_handlers(master: IMaster):