p15670423
/
monkey
1
0
Fork 5
Code Issues Pull Requests 3 Projects Releases Wiki Activity

Agent: Improve Windows signal handler

This commit is contained in:
Mike Salvatore 2021-11-24 12:51:33 -05:00
parent e04e8d3177
commit d31fd2c811
1 changed files with 24 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -12,22 +12,39 @@ class StopSignalHandler:
def __init__(self, master: IMaster): def __init__(self, master: IMaster):
self._master = master self._master = master
def __call__(self, signum, _=None): def handle_posix_signals(self, signum, _):
self._handle_signal(signum)
# Windows signal handlers must return boolean. Only raising this exception for POSIX
# signals.
raise PlannedShutdownException("Monkey Agent got an interrupt signal")
def handle_windows_signals(self, signum):
import win32con
# TODO: This signal handler gets called for a CTRL_CLOSE_EVENT, but the system immediately
# kills the process after the handler returns. After the master is implemented and the
# setup/teardown of the Agent is fully refactored, revisit this signal handler and
# 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}:
self._handle_signal(signum)
return True
return False
def _handle_signal(self, signum):
logger.info(f"The Monkey Agent received signal {signum}") logger.info(f"The Monkey Agent received signal {signum}")
self._master.terminate() self._master.terminate()
raise PlannedShutdownException("Monkey Agent got an interrupt signal")
def register_signal_handlers(master: IMaster): def register_signal_handlers(master: IMaster):
stop_signal_handler = StopSignalHandler(master) stop_signal_handler = StopSignalHandler(master)
signal.signal(signal.SIGINT, stop_signal_handler)
signal.signal(signal.SIGTERM, stop_signal_handler)
if is_windows_os(): if is_windows_os():
import win32api import win32api
signal.signal(signal.SIGBREAK, stop_signal_handler)
# CTRL_CLOSE_EVENT signal has a timeout of 5000ms, # CTRL_CLOSE_EVENT signal has a timeout of 5000ms,
# after that OS will forcefully kill the process # after that OS will forcefully kill the process
win32api.SetConsoleCtrlHandler(stop_signal_handler, True) win32api.SetConsoleCtrlHandler(stop_signal_handler.handle_windows_signals, True)
else:
signal.signal(signal.SIGINT, stop_signal_handler.handle_posix_signals)
signal.signal(signal.SIGTERM, stop_signal_handler.handle_posix_signals)