p41275903
/
monkey
forked from p15670423/monkey
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Agent: Move _create_daemon_thread to threading_utils.py

This commit is contained in:
Mike Salvatore 2021-12-09 15:26:35 -05:00
parent 86203c8138
commit 3f7dbbccc2
2 changed files with 17 additions and 13 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

24
monkey/infection_monkey/master/automated_master.py
View File

@ -16,6 +16,8 @@ from infection_monkey.telemetry.scan_telem import ScanTelem
from infection_monkey.telemetry.system_info_telem import SystemInfoTelem from infection_monkey.telemetry.system_info_telem import SystemInfoTelem
from infection_monkey.utils.timer import Timer from infection_monkey.utils.timer import Timer
from .threading_utils import create_daemon_thread
CHECK_ISLAND_FOR_STOP_COMMAND_INTERVAL_SEC = 5 CHECK_ISLAND_FOR_STOP_COMMAND_INTERVAL_SEC = 5
CHECK_FOR_TERMINATE_INTERVAL_SEC = CHECK_ISLAND_FOR_STOP_COMMAND_INTERVAL_SEC / 5 CHECK_FOR_TERMINATE_INTERVAL_SEC = CHECK_ISLAND_FOR_STOP_COMMAND_INTERVAL_SEC / 5
SHUTDOWN_TIMEOUT = 5 SHUTDOWN_TIMEOUT = 5
@ -36,8 +38,8 @@ class AutomatedMaster(IMaster):
self._control_channel = control_channel self._control_channel = control_channel
self._stop = threading.Event() self._stop = threading.Event()
self._master_thread = _create_daemon_thread(target=self._run_master_thread) self._master_thread = create_daemon_thread(target=self._run_master_thread)
self._simulation_thread = _create_daemon_thread(target=self._run_simulation) self._simulation_thread = create_daemon_thread(target=self._run_simulation)
def start(self): def start(self):
logger.info("Starting automated breach and attack simulation") logger.info("Starting automated breach and attack simulation")
@ -93,7 +95,7 @@ class AutomatedMaster(IMaster):
def _run_simulation(self): def _run_simulation(self):
config = self._control_channel.get_config() config = self._control_channel.get_config()
system_info_collector_thread = _create_daemon_thread( system_info_collector_thread = create_daemon_thread(
target=self._run_plugins, target=self._run_plugins,
args=( args=(
config["system_info_collector_classes"], config["system_info_collector_classes"],
@ -101,7 +103,7 @@ class AutomatedMaster(IMaster):
self._collect_system_info, self._collect_system_info,
), ),
) )
pba_thread = _create_daemon_thread( pba_thread = create_daemon_thread(
target=self._run_plugins, target=self._run_plugins,
args=(config["post_breach_actions"].items(), "post-breach action", self._run_pba), args=(config["post_breach_actions"].items(), "post-breach action", self._run_pba),
) )
@ -116,11 +118,11 @@ class AutomatedMaster(IMaster):
# system_info_collector_thread.join() # system_info_collector_thread.join()
if self._can_propagate(): if self._can_propagate():
propagation_thread = _create_daemon_thread(target=self._propagate, args=(config,)) propagation_thread = create_daemon_thread(target=self._propagate, args=(config,))
propagation_thread.start() propagation_thread.start()
propagation_thread.join() propagation_thread.join()
payload_thread = _create_daemon_thread( payload_thread = create_daemon_thread(
target=self._run_plugins, target=self._run_plugins,
args=(config["payloads"].items(), "payload", self._run_payload), args=(config["payloads"].items(), "payload", self._run_payload),
) )
@ -159,10 +161,10 @@ class AutomatedMaster(IMaster):
hosts_to_exploit = Queue() hosts_to_exploit = Queue()
scan_thread = _create_daemon_thread( scan_thread = create_daemon_thread(
target=self._scan_network, args=(config["network_scan"], hosts_to_exploit) target=self._scan_network, args=(config["network_scan"], hosts_to_exploit)
) )
exploit_thread = _create_daemon_thread( exploit_thread = create_daemon_thread(
target=self._exploit_targets, args=(hosts_to_exploit, scan_thread) target=self._exploit_targets, args=(hosts_to_exploit, scan_thread)
) )
@ -188,7 +190,7 @@ class AutomatedMaster(IMaster):
scan_threads = [] scan_threads = []
for i in range(0, NUM_SCAN_THREADS): for i in range(0, NUM_SCAN_THREADS):
t = _create_daemon_thread( t = create_daemon_thread(
target=self._scan_ips, args=(ips_to_scan, scan_config, hosts_to_exploit) target=self._scan_ips, args=(ips_to_scan, scan_config, hosts_to_exploit)
) )
t.start() t.start()
@ -263,7 +265,3 @@ class AutomatedMaster(IMaster):
def cleanup(self): def cleanup(self):
pass pass
def _create_daemon_thread(target: Callable[[Any], None], args: Tuple[Any] = ()):
return Thread(target=target, args=args, daemon=True)

6
monkey/infection_monkey/master/threading_utils.py Normal file
View File

@ -0,0 +1,6 @@
from threading import Thread
from typing import Any, Callable, Tuple
def create_daemon_thread(target: Callable[[Any], None], args: Tuple[Any] = ()):
return Thread(target=target, args=args, daemon=True)