From 7c108e1f2ee5921c9ba4e376d3783156d5990bdc Mon Sep 17 00:00:00 2001 From: Shreya Date: Fri, 7 Aug 2020 13:47:05 +0530 Subject: [PATCH] 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))