From c0bff448c4f884d99b78c8fcfb640d49a46fd9c8 Mon Sep 17 00:00:00 2001 From: Shreya Date: Mon, 3 Aug 2020 15:56:09 +0530 Subject: [PATCH 1/3] Run post-breach phase in separate thread --- monkey/infection_monkey/monkey.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index 444bde452..be97332a4 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 @@ -139,8 +140,8 @@ class InfectionMonkey(object): TunnelTelem().send() LOG.debug("Starting the post-breach phase.") - self.collect_system_info_if_configured() - PostBreach().execute_all_configured() + 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() From 7c108e1f2ee5921c9ba4e376d3783156d5990bdc Mon Sep 17 00:00:00 2001 From: Shreya Date: Fri, 7 Aug 2020 13:47:05 +0530 Subject: [PATCH 2/3] Make PBAs run parallely --- .../post_breach/post_breach_handler.py | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/monkey/infection_monkey/post_breach/post_breach_handler.py b/monkey/infection_monkey/post_breach/post_breach_handler.py index 33eb55fc6..86f519c9e 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(4) + 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)) From 444c2cb7ddb9c2ec1fd86689567a7f68e1140516 Mon Sep 17 00:00:00 2001 From: Shreya Date: Mon, 10 Aug 2020 11:29:37 +0530 Subject: [PATCH 3/3] Change max threads from 4 to 5 & modify log message --- monkey/infection_monkey/monkey.py | 2 +- monkey/infection_monkey/post_breach/post_breach_handler.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index be97332a4..02463e988 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -139,7 +139,7 @@ class InfectionMonkey(object): StateTelem(is_done=False, version=get_version()).send() TunnelTelem().send() - LOG.debug("Starting the post-breach phase.") + LOG.debug("Starting the post-breach phase asynchronously.") post_breach_phase = Thread(target=self.start_post_breach_phase) post_breach_phase.start() diff --git a/monkey/infection_monkey/post_breach/post_breach_handler.py b/monkey/infection_monkey/post_breach/post_breach_handler.py index 86f519c9e..e20727ccf 100644 --- a/monkey/infection_monkey/post_breach/post_breach_handler.py +++ b/monkey/infection_monkey/post_breach/post_breach_handler.py @@ -25,7 +25,7 @@ class PostBreach(object): """ Executes all post breach actions. """ - pool = Pool(4) + pool = Pool(5) pool.map(self.run_pba, self.pba_list) LOG.info("All PBAs executed. Total {} executed.".format(len(self.pba_list)))