From bcc0c86526733e0bd1993d5e72806c4b54a10509 Mon Sep 17 00:00:00 2001
From: Kekoa Kaaikala <kekoa.kaaikala@gmail.com>
Date: Wed, 24 Aug 2022 19:51:19 +0000
Subject: [PATCH] Agent: Fix batching telemetry thread shutdown

---
 monkey/infection_monkey/monkey.py                    | 12 ++++++++----
 .../payload/ransomware/ransomware_builder.py         |  5 ++++-
 2 files changed, 12 insertions(+), 5 deletions(-)

diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py
index 98b2cab78..e212ef2d4 100644
--- a/monkey/infection_monkey/monkey.py
+++ b/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 = (
diff --git a/monkey/infection_monkey/payload/ransomware/ransomware_builder.py b/monkey/infection_monkey/payload/ransomware/ransomware_builder.py
index 1a2649060..e52f52cab 100644
--- a/monkey/infection_monkey/payload/ransomware/ransomware_builder.py
+++ b/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),
     )