Merge pull request #2211 from guardicore/fix-batching-telemetry-thread-shutdown

Fix batching telemetry thread shutdown

monkey/infection_monkey/monkey.py

@@ -66,9 +66,6 @@ from infection_monkey.puppet.puppet import Puppet
 from infection_monkey.system_singleton import SystemSingleton
 from infection_monkey.telemetry.attack.t1106_telem import T1106Telem
 from infection_monkey.telemetry.attack.t1107_telem import T1107Telem
-from infection_monkey.telemetry.messengers.batching_telemetry_messenger import (
-    BatchingTelemetryMessenger,
-)
 from infection_monkey.telemetry.messengers.exploit_intercepting_telemetry_messenger import (
     ExploitInterceptingTelemetryMessenger,
 )
@@ -359,7 +356,7 @@ class InfectionMonkey:
 
         puppet.load_plugin(
             "ransomware",
-            RansomwarePayload(BatchingTelemetryMessenger(self._telemetry_messenger)),
+            RansomwarePayload(self._telemetry_messenger),
             PluginType.PAYLOAD,
         )
 

monkey/infection_monkey/payload/ransomware/ransomware_builder.py

@@ -1,6 +1,9 @@
 import logging
 from pprint import pformat
 
+from infection_monkey.telemetry.messengers.batching_telemetry_messenger import (
+    BatchingTelemetryMessenger,
+)
 from infection_monkey.telemetry.messengers.i_telemetry_messenger import ITelemetryMessenger
 from infection_monkey.utils.bit_manipulators import flip_bits
 
@@ -29,7 +32,7 @@ def build_ransomware(options: dict, telemetry_messenger: ITelemetryMessenger):
         file_encryptor,
         file_selector,
         leave_readme,
-        telemetry_messenger,
+        BatchingTelemetryMessenger(telemetry_messenger),
     )
 
 

monkey/infection_monkey/telemetry/messengers/batching_telemetry_messenger.py

@@ -47,7 +47,7 @@ class BatchingTelemetryMessenger(ITelemetryMessenger):
         def start(self):
             self._should_run_batch_thread = True
             self._manage_telemetry_batches_thread = threading.Thread(
-                target=self._manage_telemetry_batches
+                name="BatchingTelemetryMessengerThread", target=self._manage_telemetry_batches
             )
             self._manage_telemetry_batches_thread.start()