From fae0c8ded2159e657fa97d744248ab6c4772b437 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Mon, 24 Jan 2022 08:56:04 -0500 Subject: [PATCH] Agent: Replace if checks with interruptable_iter() in for loops --- monkey/infection_monkey/master/automated_master.py | 9 +++------ monkey/infection_monkey/master/exploiter.py | 7 ++----- monkey/infection_monkey/master/ip_scanner.py | 12 +++--------- .../payload/ransomware/ransomware.py | 12 +++++------- 4 files changed, 13 insertions(+), 27 deletions(-) diff --git a/monkey/infection_monkey/master/automated_master.py b/monkey/infection_monkey/master/automated_master.py index 94a508e0c..c68e77cb7 100644 --- a/monkey/infection_monkey/master/automated_master.py +++ b/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: - if self._stop.is_set(): - logger.debug(f"Received a stop signal, skipping remaining {plugin_type}s") - return - + interrupted_message = f"Received a stop signal, skipping remaining {plugin_type}s" + for p in interruptable_iter(plugin, self._stop, interrupted_message): callback(p) logger.info(f"Finished running {plugin_type}s") diff --git a/monkey/infection_monkey/master/exploiter.py b/monkey/infection_monkey/master/exploiter.py index 8acca5ffa..09f6ebf4b 100644 --- a/monkey/infection_monkey/master/exploiter.py +++ b/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: - if stop.is_set(): - break - + for exploiter in interruptable_iter(exploiters_to_run, stop): exploiter_name = exploiter["name"] exploiter_results = self._run_exploiter(exploiter_name, victim_host, stop) results_callback(exploiter_name, victim_host, exploiter_results) diff --git a/monkey/infection_monkey/master/ip_scanner.py b/monkey/infection_monkey/master/ip_scanner.py index bae6358e2..0f7132a27 100644 --- a/monkey/infection_monkey/master/ip_scanner.py +++ b/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: - if stop.is_set(): - break - + for p in interruptable_iter(ports, stop): 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: - if stop.is_set(): - break - + for f in interruptable_iter(fingerprinters, stop): fingerprint_data[f] = self._puppet.fingerprint(f, ip, ping_scan_data, port_scan_data) return fingerprint_data diff --git a/monkey/infection_monkey/payload/ransomware/ransomware.py b/monkey/infection_monkey/payload/ransomware/ransomware.py index 003112cc3..febdfe025 100644 --- a/monkey/infection_monkey/payload/ransomware/ransomware.py +++ b/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: - if interrupt.is_set(): - logger.debug( - "Received a stop signal, skipping remaining files for encryption of " - "ransomware payload" - ) - return + interrupted_message = ( + "Received a stop signal, skipping remaining files for encryption of ransomware payload" + ) + for filepath in interruptable_iter(file_list, interrupt, interrupted_message): try: logger.debug(f"Encrypting {filepath}") self._encrypt_file(filepath)