diff --git a/monkey/infection_monkey/master/automated_master.py b/monkey/infection_monkey/master/automated_master.py index 88841f401..0ef31129c 100644 --- a/monkey/infection_monkey/master/automated_master.py +++ b/monkey/infection_monkey/master/automated_master.py @@ -9,7 +9,6 @@ from infection_monkey.i_master import IMaster from infection_monkey.i_puppet import IPuppet from infection_monkey.model import VictimHostFactory from infection_monkey.network import NetworkInterface -from infection_monkey.post_breach.custom_pba import CustomPBA from infection_monkey.telemetry.credentials_telem import CredentialsTelem from infection_monkey.telemetry.messengers.i_telemetry_messenger import ITelemetryMessenger from infection_monkey.telemetry.post_breach_telem import PostBreachTelem @@ -17,6 +16,7 @@ from infection_monkey.utils.threading import create_daemon_thread, interruptible from infection_monkey.utils.timer import Timer from . import Exploiter, IPScanner, Propagator +from .option_parsing import custom_pba_is_enabled CHECK_ISLAND_FOR_STOP_COMMAND_INTERVAL_SEC = 5 CHECK_FOR_TERMINATE_INTERVAL_SEC = CHECK_ISLAND_FOR_STOP_COMMAND_INTERVAL_SEC / 5 @@ -214,7 +214,7 @@ class AutomatedMaster(IMaster): ): self._run_plugins(plugins, "post-breach action", callback) - if CustomPBA.should_run(custom_pba_options): + if custom_pba_is_enabled(custom_pba_options): self._run_plugins([("CustomPBA", custom_pba_options)], "post-breach action", callback) def _run_plugins( diff --git a/monkey/infection_monkey/master/option_parsing.py b/monkey/infection_monkey/master/option_parsing.py new file mode 100644 index 000000000..c35bf6303 --- /dev/null +++ b/monkey/infection_monkey/master/option_parsing.py @@ -0,0 +1,13 @@ +from typing import Dict + +from infection_monkey.utils.environment import is_windows_os + + +def custom_pba_is_enabled(pba_options: Dict) -> bool: + if not is_windows_os(): + if pba_options["linux_command"]: + return True + else: + if pba_options["windows_command"]: + return True + return False diff --git a/monkey/infection_monkey/post_breach/custom_pba/custom_pba.py b/monkey/infection_monkey/post_breach/custom_pba/custom_pba.py index 51e9ef7a1..453dfb6ed 100644 --- a/monkey/infection_monkey/post_breach/custom_pba/custom_pba.py +++ b/monkey/infection_monkey/post_breach/custom_pba/custom_pba.py @@ -65,16 +65,6 @@ class CustomPBA(PBA): self.download_pba_file(get_monkey_dir_path(), self.filename) return super(CustomPBA, self)._execute_default() - @staticmethod - def should_run(options): - if not is_windows_os(): - if options["linux_filename"] or options["linux_command"]: - return True - else: - if options["windows_filename"] or options["windows_command"]: - return True - return False - def download_pba_file(self, dst_dir, filename): """ Handles post breach action file download