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

Agent: Reset signal handlers after the Master is cleaned up 

After the Master terminates, this resets the signal handlers to the
default handlers provided by Python.
This commit is contained in:
Mike Salvatore 2021-12-16 15:59:55 -05:00
parent 019f2c1403
commit da58392050
2 changed files with 40 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
monkey/infection_monkey/monkey.py
View File

@ -32,7 +32,7 @@ from infection_monkey.telemetry.tunnel_telem import TunnelTelem
from infection_monkey.utils.environment import is_windows_os from infection_monkey.utils.environment import is_windows_os
from infection_monkey.utils.monkey_dir import get_monkey_dir_path, remove_monkey_dir from infection_monkey.utils.monkey_dir import get_monkey_dir_path, remove_monkey_dir
from infection_monkey.utils.monkey_log_path import get_monkey_log_path 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.utils.signal_handler import register_signal_handlers, reset_signal_handlers
from infection_monkey.windows_upgrader import WindowsUpgrader from infection_monkey.windows_upgrader import WindowsUpgrader
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@ -215,6 +215,8 @@ class InfectionMonkey:
if self._master: if self._master:
self._master.cleanup() self._master.cleanup()
reset_signal_handlers()
if self._monkey_inbound_tunnel: if self._monkey_inbound_tunnel:
self._monkey_inbound_tunnel.stop() self._monkey_inbound_tunnel.stop()
self._monkey_inbound_tunnel.join() self._monkey_inbound_tunnel.join()

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

@ -1,5 +1,6 @@
import logging import logging
import signal import signal
from typing import Optional
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
@ -7,18 +8,25 @@ from infection_monkey.utils.environment import is_windows_os
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
_signal_handler = None
class StopSignalHandler: class StopSignalHandler:
def __init__(self, master: IMaster): def __init__(self, master: IMaster):
self._master = master self._master = master
def handle_posix_signals(self, signum: int, _): # Windows won't let us correctly deregister a method, but Callables and closures work.
self._handle_signal(signum, False) def __call__(self, signum: int, *args) -> Optional[bool]:
if is_windows_os():
return self._handle_windows_signals(signum)
else:
self._handle_posix_signals(signum, args)
def handle_windows_signals(self, signum: int): def _handle_windows_signals(self, signum: int) -> bool:
import win32con import win32con
if signum in {win32con.CTRL_C_EVENT, win32con.CTRL_BREAK_EVENT}: if signum in {win32con.CTRL_C_EVENT, win32con.CTRL_BREAK_EVENT}:
self._handle_signal(signum, False) self._terminate_master(signum, False)
return True return True
if signum == win32con.CTRL_CLOSE_EVENT: if signum == win32con.CTRL_CLOSE_EVENT:
@ -26,25 +34,44 @@ class StopSignalHandler:
# Calling self._handle_signal() with block=True to give the master a chance to # 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 # gracefully shut down. Note that the OS has a timeout that will forcefully kill the
# process if this handler hasn't returned in time. # process if this handler hasn't returned in time.
self._handle_signal(signum, True) self._terminate_master(signum, True)
return True return True
return False return False
def _handle_signal(self, signum: int, block: bool): def _handle_posix_signals(self, signum: int, *_):
self._terminate_master(signum, False)
def _terminate_master(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(block) self._master.terminate(block)
def register_signal_handlers(master: IMaster): def register_signal_handlers(master: IMaster):
stop_signal_handler = StopSignalHandler(master) global _signal_handler
_signal_handler = StopSignalHandler(master)
if is_windows_os(): if is_windows_os():
import win32api import win32api
# 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.handle_windows_signals, True) win32api.SetConsoleCtrlHandler(_signal_handler, True)
else: else:
signal.signal(signal.SIGINT, stop_signal_handler.handle_posix_signals) signal.signal(signal.SIGINT, _signal_handler)
signal.signal(signal.SIGTERM, stop_signal_handler.handle_posix_signals) signal.signal(signal.SIGTERM, _signal_handler)
def reset_signal_handlers():
"""
Resets the signal handlers back to the default handlers provided by Python
"""
global _signal_handler
if is_windows_os():
import win32api
win32api.SetConsoleCtrlHandler(_signal_handler, False)
else:
signal.signal(signal.SIGINT, signal.SIG_DFL)
signal.signal(signal.SIGTERM, signal.SIG_DFL)