forked from p15670423/monkey
Make PBAs run parallely
This commit is contained in:
parent
c0bff448c4
commit
7c108e1f2e
|
@ -1,4 +1,5 @@
|
||||||
import logging
|
import logging
|
||||||
|
from multiprocessing.dummy import Pool
|
||||||
from typing import Sequence
|
from typing import Sequence
|
||||||
|
|
||||||
from infection_monkey.post_breach.pba import PBA
|
from infection_monkey.post_breach.pba import PBA
|
||||||
|
@ -24,12 +25,8 @@ class PostBreach(object):
|
||||||
"""
|
"""
|
||||||
Executes all post breach actions.
|
Executes all post breach actions.
|
||||||
"""
|
"""
|
||||||
for pba in self.pba_list:
|
pool = Pool(4)
|
||||||
try:
|
pool.map(self.run_pba, self.pba_list)
|
||||||
LOG.debug("Executing PBA: '{}'".format(pba.name))
|
|
||||||
pba.run()
|
|
||||||
except Exception as e:
|
|
||||||
LOG.error("PBA {} failed. Error info: {}".format(pba.name, e))
|
|
||||||
LOG.info("All PBAs executed. Total {} executed.".format(len(self.pba_list)))
|
LOG.info("All PBAs executed. Total {} executed.".format(len(self.pba_list)))
|
||||||
|
|
||||||
@staticmethod
|
@staticmethod
|
||||||
|
@ -38,3 +35,10 @@ class PostBreach(object):
|
||||||
:return: A list of PBA objects.
|
:return: A list of PBA objects.
|
||||||
"""
|
"""
|
||||||
return PBA.get_instances()
|
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))
|
||||||
|
|
Loading…
Reference in New Issue