Agent: Move the decision if custom pba should run to master

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(

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

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