Agent: Fix batching telemetry thread shutdown

monkey/infection_monkey/monkey.py

@@ -3,6 +3,7 @@ import logging
 import os
 import subprocess
 import sys
+import threading
 from pathlib import Path, WindowsPath
 from typing import List
 
@@ -66,9 +67,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 +357,7 @@ class InfectionMonkey:
 
         puppet.load_plugin(
             "ransomware",
-            RansomwarePayload(BatchingTelemetryMessenger(self._telemetry_messenger)),
+            RansomwarePayload(self._telemetry_messenger),
             PluginType.PAYLOAD,
         )
 
@@ -415,6 +413,12 @@ class InfectionMonkey:
                 InfectionMonkey._self_delete()
 
         logger.info("Monkey is shutting down")
+        for t in threading.enumerate():
+            logger.info(t.name)
+
+        import faulthandler
+
+        faulthandler.dump_traceback()
 
     def _close_tunnel(self):
         tunnel_address = (

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),
     )