Run post-breach phase in separate thread

This commit is contained in:
Shreya 2020-08-03 15:56:09 +05:30
parent f0428d09d2
commit c0bff448c4
1 changed files with 10 additions and 2 deletions

View File

@ -4,6 +4,7 @@ import os
import subprocess import subprocess
import sys import sys
import time import time
from threading import Thread
import infection_monkey.tunnel as tunnel import infection_monkey.tunnel as tunnel
from common.network.network_utils import get_host_from_network_location from common.network.network_utils import get_host_from_network_location
@ -139,8 +140,8 @@ class InfectionMonkey(object):
TunnelTelem().send() TunnelTelem().send()
LOG.debug("Starting the post-breach phase.") LOG.debug("Starting the post-breach phase.")
self.collect_system_info_if_configured() post_breach_phase = Thread(target=self.start_post_breach_phase)
PostBreach().execute_all_configured() post_breach_phase.start()
LOG.debug("Starting the propagation phase.") LOG.debug("Starting the propagation phase.")
self.shutdown_by_max_depth_reached() self.shutdown_by_max_depth_reached()
@ -230,10 +231,17 @@ class InfectionMonkey(object):
if monkey_tunnel: if monkey_tunnel:
monkey_tunnel.stop() monkey_tunnel.stop()
monkey_tunnel.join() monkey_tunnel.join()
post_breach_phase.join()
except PlannedShutdownException: except PlannedShutdownException:
LOG.info("A planned shutdown of the Monkey occurred. Logging the reason and finishing execution.") LOG.info("A planned shutdown of the Monkey occurred. Logging the reason and finishing execution.")
LOG.exception("Planned shutdown, reason:") 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): def shutdown_by_max_depth_reached(self):
if 0 == WormConfiguration.depth: if 0 == WormConfiguration.depth:
TraceTelem(MAX_DEPTH_REACHED_MESSAGE).send() TraceTelem(MAX_DEPTH_REACHED_MESSAGE).send()