Agent: Replace if checks with interruptable_iter() in for loops

monkey/infection_monkey/master/automated_master.py

@@ -11,7 +11,7 @@ from infection_monkey.network import NetworkInterface
 from infection_monkey.telemetry.messengers.i_telemetry_messenger import ITelemetryMessenger
 from infection_monkey.telemetry.post_breach_telem import PostBreachTelem
 from infection_monkey.telemetry.system_info_telem import SystemInfoTelem
-from infection_monkey.utils.threading import create_daemon_thread
+from infection_monkey.utils.threading import create_daemon_thread, interruptable_iter
 from infection_monkey.utils.timer import Timer
 
 from . import Exploiter, IPScanner, Propagator
@@ -195,11 +195,8 @@ class AutomatedMaster(IMaster):
         logger.info(f"Running {plugin_type}s")
         logger.debug(f"Found {len(plugin)} {plugin_type}(s) to run")
 
-        for p in plugin:
+        interrupted_message = f"Received a stop signal, skipping remaining {plugin_type}s"
-            if self._stop.is_set():
+        for p in interruptable_iter(plugin, self._stop, interrupted_message):
-                logger.debug(f"Received a stop signal, skipping remaining {plugin_type}s")
-                return
-
             callback(p)
 
         logger.info(f"Finished running {plugin_type}s")

monkey/infection_monkey/master/exploiter.py

@@ -7,7 +7,7 @@ from typing import Callable, Dict, List
 
 from infection_monkey.i_puppet import ExploiterResultData, IPuppet
 from infection_monkey.model import VictimHost
-from infection_monkey.utils.threading import run_worker_threads
+from infection_monkey.utils.threading import interruptable_iter, run_worker_threads
 
 QUEUE_TIMEOUT = 2
 
@@ -74,10 +74,7 @@ class Exploiter:
         results_callback: Callback,
         stop: Event,
     ):
-        for exploiter in exploiters_to_run:
+        for exploiter in interruptable_iter(exploiters_to_run, stop):
-            if stop.is_set():
-                break
-
             exploiter_name = exploiter["name"]
             exploiter_results = self._run_exploiter(exploiter_name, victim_host, stop)
             results_callback(exploiter_name, victim_host, exploiter_results)

monkey/infection_monkey/master/ip_scanner.py

@@ -13,7 +13,7 @@ from infection_monkey.i_puppet import (
     PortStatus,
 )
 from infection_monkey.network import NetworkAddress
-from infection_monkey.utils.threading import run_worker_threads
+from infection_monkey.utils.threading import interruptable_iter, run_worker_threads
 
 from . import IPScanResults
 
@@ -85,10 +85,7 @@ class IPScanner:
     ) -> Dict[int, PortScanData]:
         port_scan_data = {}
 
-        for p in ports:
+        for p in interruptable_iter(ports, stop):
-            if stop.is_set():
-                break
-
             port_scan_data[p] = self._puppet.scan_tcp_port(ip, p, timeout)
 
         return port_scan_data
@@ -107,10 +104,7 @@ class IPScanner:
     ) -> Dict[str, FingerprintData]:
         fingerprint_data = {}
 
-        for f in fingerprinters:
+        for f in interruptable_iter(fingerprinters, stop):
-            if stop.is_set():
-                break
-
             fingerprint_data[f] = self._puppet.fingerprint(f, ip, ping_scan_data, port_scan_data)
 
         return fingerprint_data

monkey/infection_monkey/payload/ransomware/ransomware.py

@@ -5,6 +5,7 @@ from typing import Callable, List
 
 from infection_monkey.telemetry.file_encryption_telem import FileEncryptionTelem
 from infection_monkey.telemetry.messengers.i_telemetry_messenger import ITelemetryMessenger
+from infection_monkey.utils.threading import interruptable_iter
 
 from .consts import README_FILE_NAME, README_SRC
 from .ransomware_options import RansomwareOptions
@@ -53,13 +54,10 @@ class Ransomware:
     def _encrypt_files(self, file_list: List[Path], interrupt: threading.Event):
         logger.info(f"Encrypting files in {self._target_directory}")
 
-        for filepath in file_list:
+        interrupted_message = (
-            if interrupt.is_set():
+            "Received a stop signal, skipping remaining files for encryption of ransomware payload"
-                logger.debug(
+        )
-                    "Received a stop signal, skipping remaining files for encryption of "
+        for filepath in interruptable_iter(file_list, interrupt, interrupted_message):
-                    "ransomware payload"
-                )
-                return
             try:
                 logger.debug(f"Encrypting {filepath}")
                 self._encrypt_file(filepath)