forked from p15670423/monkey
Agent: Replace if checks with interruptable_iter() in for loops
This commit is contained in:
parent
0c877833c5
commit
fae0c8ded2
|
@ -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")
|
||||
|
|
|
@ -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)
|
||||
|
|
|
@ -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
|
||||
|
|
|
@ -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)
|
||||
|
|
Loading…
Reference in New Issue