diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index 444bde452..02463e988 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -4,6 +4,7 @@ import os import subprocess import sys import time +from threading import Thread import infection_monkey.tunnel as tunnel from common.network.network_utils import get_host_from_network_location @@ -138,9 +139,9 @@ class InfectionMonkey(object): StateTelem(is_done=False, version=get_version()).send() TunnelTelem().send() - LOG.debug("Starting the post-breach phase.") - self.collect_system_info_if_configured() - PostBreach().execute_all_configured() + LOG.debug("Starting the post-breach phase asynchronously.") + post_breach_phase = Thread(target=self.start_post_breach_phase) + post_breach_phase.start() LOG.debug("Starting the propagation phase.") self.shutdown_by_max_depth_reached() @@ -230,10 +231,17 @@ class InfectionMonkey(object): if monkey_tunnel: monkey_tunnel.stop() monkey_tunnel.join() + + post_breach_phase.join() + except PlannedShutdownException: LOG.info("A planned shutdown of the Monkey occurred. Logging the reason and finishing execution.") LOG.exception("Planned shutdown, reason:") + def start_post_breach_phase(self): + self.collect_system_info_if_configured() + PostBreach().execute_all_configured() + def shutdown_by_max_depth_reached(self): if 0 == WormConfiguration.depth: TraceTelem(MAX_DEPTH_REACHED_MESSAGE).send() diff --git a/monkey/infection_monkey/post_breach/post_breach_handler.py b/monkey/infection_monkey/post_breach/post_breach_handler.py index 33eb55fc6..e20727ccf 100644 --- a/monkey/infection_monkey/post_breach/post_breach_handler.py +++ b/monkey/infection_monkey/post_breach/post_breach_handler.py @@ -1,4 +1,5 @@ import logging +from multiprocessing.dummy import Pool from typing import Sequence from infection_monkey.post_breach.pba import PBA @@ -24,12 +25,8 @@ class PostBreach(object): """ Executes all post breach actions. """ - for pba in self.pba_list: - try: - LOG.debug("Executing PBA: '{}'".format(pba.name)) - pba.run() - except Exception as e: - LOG.error("PBA {} failed. Error info: {}".format(pba.name, e)) + pool = Pool(5) + pool.map(self.run_pba, self.pba_list) LOG.info("All PBAs executed. Total {} executed.".format(len(self.pba_list))) @staticmethod @@ -38,3 +35,10 @@ class PostBreach(object): :return: A list of PBA objects. """ return PBA.get_instances() + + def run_pba(self, pba): + try: + LOG.debug("Executing PBA: '{}'".format(pba.name)) + pba.run() + except Exception as e: + LOG.error("PBA {} failed. Error info: {}".format(pba.name, e))