From 1c24411b267ff3485e9dbe58013b92fe7b71b66f Mon Sep 17 00:00:00 2001
From: Shreya Malviya <shreya.malviya@gmail.com>
Date: Tue, 29 Mar 2022 13:24:14 +0530
Subject: [PATCH] Agent: Pass telemetry messenger to PBAs for sending ATT&CK
 telem

---
 .../post_breach/actions/users_custom_pba.py     | 17 +++++++++--------
 monkey/infection_monkey/post_breach/pba.py      | 11 ++++++++---
 2 files changed, 17 insertions(+), 11 deletions(-)

diff --git a/monkey/infection_monkey/post_breach/actions/users_custom_pba.py b/monkey/infection_monkey/post_breach/actions/users_custom_pba.py
index 4c706a1c1..b1ccec85c 100644
--- a/monkey/infection_monkey/post_breach/actions/users_custom_pba.py
+++ b/monkey/infection_monkey/post_breach/actions/users_custom_pba.py
@@ -65,8 +65,7 @@ class UsersPBA(PBA):
                 return True
         return False
 
-    @staticmethod
-    def download_pba_file(dst_dir, filename):
+    def download_pba_file(self, dst_dir, filename):
         """
         Handles post breach action file download
         :param dst_dir: Destination directory
@@ -84,12 +83,14 @@ class UsersPBA(PBA):
         if not status:
             status = ScanStatus.USED
 
-        T1105Telem(
-            status,
-            WormConfiguration.current_server.split(":")[0],
-            get_interface_to_target(WormConfiguration.current_server.split(":")[0]),
-            filename,
-        ).send()
+        self._telemetry_messenger.send_telemetry(
+            T1105Telem(
+                status,
+                WormConfiguration.current_server.split(":")[0],
+                get_interface_to_target(WormConfiguration.current_server.split(":")[0]),
+                filename,
+            )
+        )
 
         if status == ScanStatus.SCANNED:
             return False
diff --git a/monkey/infection_monkey/post_breach/pba.py b/monkey/infection_monkey/post_breach/pba.py
index 8b50f08ba..769ff2de0 100644
--- a/monkey/infection_monkey/post_breach/pba.py
+++ b/monkey/infection_monkey/post_breach/pba.py
@@ -5,6 +5,8 @@ from typing import Iterable
 from common.utils.attack_utils import ScanStatus
 from infection_monkey.i_puppet.i_puppet import PostBreachData
 from infection_monkey.telemetry.attack.t1064_telem import T1064Telem
+from infection_monkey.telemetry.messengers.i_telemetry_messenger import ITelemetryMessenger
+from infection_monkey.telemetry.post_breach_telem import PostBreachTelem
 from infection_monkey.utils.environment import is_windows_os
 
 logger = logging.getLogger(__name__)
@@ -34,9 +36,12 @@ class PBA:
             exec_funct = self._execute_default
             result = exec_funct()
             if self.scripts_were_used_successfully(result):
-                T1064Telem(
-                    ScanStatus.USED, f"Scripts were used to execute {self.name} post breach action."
-                ).send()
+                self.telemetry_messenger.send_telemetry(
+                    T1064Telem(
+                        ScanStatus.USED,
+                        f"Scripts were used to execute {self.name} post breach action.",
+                    )
+                )
             self.pba_data.append(PostBreachData(self.name, self.command, result))
             return self.pba_data
         else: