Agent: Add _create_daemon_thread() utility function to AutomatedMaster

This commit is contained in:
Mike Salvatore 2021-12-06 19:13:53 -05:00
parent 261826fc78
commit e8de38881c
1 changed files with 10 additions and 11 deletions

View File

@ -30,8 +30,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 = threading.Thread(target=self._run_master_thread, daemon=True) self._master_thread = _create_daemon_thread(target=self._run_master_thread)
self._simulation_thread = threading.Thread(target=self._run_simulation, daemon=True) 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")
@ -87,19 +87,17 @@ 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 = threading.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"],
"system info collector", "system info collector",
self._collect_system_info, self._collect_system_info,
), ),
daemon=True,
) )
pba_thread = threading.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),
daemon=True,
) )
system_info_collector_thread.start() system_info_collector_thread.start()
@ -112,16 +110,13 @@ class AutomatedMaster(IMaster):
system_info_collector_thread.join() system_info_collector_thread.join()
if self._can_propagate(): if self._can_propagate():
propagation_thread = threading.Thread( propagation_thread = _create_daemon_thread(target=self._propagate, args=(config,))
target=self._propagate, args=(config,), daemon=True
)
propagation_thread.start() propagation_thread.start()
propagation_thread.join() propagation_thread.join()
payload_thread = threading.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),
daemon=True,
) )
payload_thread.start() payload_thread.start()
payload_thread.join() payload_thread.join()
@ -177,3 +172,7 @@ class AutomatedMaster(IMaster):
def cleanup(self): def cleanup(self):
pass pass
def _create_daemon_thread(target: Callable[[Any], None], args: Tuple[Any] = ()):
return threading.Thread(target=target, args=args, daemon=True)