From f46bb60da590269131db9f8f9a3937511fe91971 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Wed, 15 Dec 2021 09:00:41 -0500 Subject: [PATCH] Agent: Add block parameter to IMaster.terminate() This allows the caller to decide whether or not they're willing to wait for the master to finish shutting down. --- monkey/infection_monkey/i_master.py | 3 ++- monkey/infection_monkey/master/automated_master.py | 7 +++++-- monkey/infection_monkey/master/mock_master.py | 2 +- 3 files changed, 8 insertions(+), 4 deletions(-) diff --git a/monkey/infection_monkey/i_master.py b/monkey/infection_monkey/i_master.py index 9caa71a4d..5269cafee 100644 --- a/monkey/infection_monkey/i_master.py +++ b/monkey/infection_monkey/i_master.py @@ -10,9 +10,10 @@ class IMaster(metaclass=abc.ABCMeta): """ @abc.abstractmethod - def terminate(self) -> None: + def terminate(self, block: bool = False) -> None: """ Stop the master and interrupt any actions that are currently being executed. + :param bool block: Whether or not to block and wait for the master to terminate. """ @abc.abstractmethod diff --git a/monkey/infection_monkey/master/automated_master.py b/monkey/infection_monkey/master/automated_master.py index a7f67d87c..75a5a5dbf 100644 --- a/monkey/infection_monkey/master/automated_master.py +++ b/monkey/infection_monkey/master/automated_master.py @@ -52,12 +52,15 @@ class AutomatedMaster(IMaster): self._master_thread.join() logger.info("The simulation has been shutdown.") - def terminate(self): + def terminate(self, block: bool = False): logger.info("Stopping automated breach and attack simulation") self._stop.set() - if self._master_thread.is_alive(): + if self._master_thread.is_alive() and block: self._master_thread.join() + # We can only have confidence that the master terminated successfully if block is set + # and join() has returned. + logger.info("AutomatedMaster successfully terminated.") def _run_master_thread(self): self._simulation_thread.start() diff --git a/monkey/infection_monkey/master/mock_master.py b/monkey/infection_monkey/master/mock_master.py index 1d8791c4c..0b4f9a3f6 100644 --- a/monkey/infection_monkey/master/mock_master.py +++ b/monkey/infection_monkey/master/mock_master.py @@ -124,7 +124,7 @@ class MockMaster(IMaster): self._telemetry_messenger.send_telemetry(FileEncryptionTelem(path, success, error)) logger.info("Finished running payloads") - def terminate(self) -> None: + def terminate(self, block: bool = False) -> None: logger.info("Terminating MockMaster") def cleanup(self) -> None: