From a0b43a17a2d37d4b5e72d9e86d738632c3329667 Mon Sep 17 00:00:00 2001
From: Mike Salvatore <mike.s.salvatore@gmail.com>
Date: Thu, 24 Jun 2021 12:19:43 -0400
Subject: [PATCH] agent: Implement IBatchableTelem in RansomwareTelem

This allows encryption attempt telmetries to be batched into one
telemetry object so they can be sent to the island in batches.
---
 .../telemetry/ransomware_telem.py             | 13 ++++++++-----
 .../ransomware/test_ransomware_payload.py     | 12 ++++++------
 .../telemetry/test_ransomware_telem.py        | 19 +++++++++----------
 3 files changed, 23 insertions(+), 21 deletions(-)

diff --git a/monkey/infection_monkey/telemetry/ransomware_telem.py b/monkey/infection_monkey/telemetry/ransomware_telem.py
index c56e8337c..64cce13c2 100644
--- a/monkey/infection_monkey/telemetry/ransomware_telem.py
+++ b/monkey/infection_monkey/telemetry/ransomware_telem.py
@@ -1,11 +1,13 @@
-from typing import List, Tuple
+from typing import Tuple
 
 from common.common_consts.telem_categories import TelemCategoryEnum
 from infection_monkey.telemetry.base_telem import BaseTelem
+from infection_monkey.telemetry.batchable_telem_mixin import BatchableTelemMixin
+from infection_monkey.telemetry.i_batchable_telem import IBatchableTelem
 
 
-class RansomwareTelem(BaseTelem):
-    def __init__(self, attempts: List[Tuple[str, str]]):
+class RansomwareTelem(BatchableTelemMixin, IBatchableTelem, BaseTelem):
+    def __init__(self, entry: Tuple[str, str]):
         """
         Ransomware telemetry constructor
         :param attempts: List of tuples with each tuple containing the path
@@ -14,9 +16,10 @@ class RansomwareTelem(BaseTelem):
                          containing the directory path and error string.
         """
         super().__init__()
-        self.attempts = attempts
+
+        self._telemetry_entries.append(entry)
 
     telem_category = TelemCategoryEnum.RANSOMWARE
 
     def get_data(self):
-        return {"ransomware_attempts": self.attempts}
+        return {"ransomware_attempts": self._telemetry_entries}
diff --git a/monkey/tests/unit_tests/infection_monkey/ransomware/test_ransomware_payload.py b/monkey/tests/unit_tests/infection_monkey/ransomware/test_ransomware_payload.py
index 35aef048c..86fb5c336 100644
--- a/monkey/tests/unit_tests/infection_monkey/ransomware/test_ransomware_payload.py
+++ b/monkey/tests/unit_tests/infection_monkey/ransomware/test_ransomware_payload.py
@@ -148,10 +148,10 @@ def test_telemetry_success(ransomware_payload, telemetry_messenger_spy):
     telem_1 = telemetry_messenger_spy.telemetries[0]
     telem_2 = telemetry_messenger_spy.telemetries[1]
 
-    assert ALL_ZEROS_PDF in telem_1.get_data()["ransomware_attempts"][0]
-    assert telem_1.get_data()["ransomware_attempts"][1] == ""
-    assert TEST_KEYBOARD_TXT in telem_2.get_data()["ransomware_attempts"][0]
-    assert telem_2.get_data()["ransomware_attempts"][1] == ""
+    assert ALL_ZEROS_PDF in telem_1.get_data()["ransomware_attempts"][0][0]
+    assert telem_1.get_data()["ransomware_attempts"][0][1] == ""
+    assert TEST_KEYBOARD_TXT in telem_2.get_data()["ransomware_attempts"][0][0]
+    assert telem_2.get_data()["ransomware_attempts"][0][1] == ""
 
 
 def test_telemetry_failure(monkeypatch, ransomware_payload, telemetry_messenger_spy):
@@ -164,5 +164,5 @@ def test_telemetry_failure(monkeypatch, ransomware_payload, telemetry_messenger_
     ransomware_payload.run_payload()
     telem_1 = telemetry_messenger_spy.telemetries[0]
 
-    assert "/file/not/exist" in telem_1.get_data()["ransomware_attempts"][0]
-    assert "No such file or directory" in telem_1.get_data()["ransomware_attempts"][1]
+    assert "/file/not/exist" in telem_1.get_data()["ransomware_attempts"][0][0]
+    assert "No such file or directory" in telem_1.get_data()["ransomware_attempts"][0][1]
diff --git a/monkey/tests/unit_tests/infection_monkey/telemetry/test_ransomware_telem.py b/monkey/tests/unit_tests/infection_monkey/telemetry/test_ransomware_telem.py
index 4994c9287..e2e674ecd 100644
--- a/monkey/tests/unit_tests/infection_monkey/telemetry/test_ransomware_telem.py
+++ b/monkey/tests/unit_tests/infection_monkey/telemetry/test_ransomware_telem.py
@@ -1,20 +1,19 @@
 import json
 
-import pytest
-
 from infection_monkey.telemetry.ransomware_telem import RansomwareTelem
 
-ATTEMPTS = [("<file>", "<encryption attempt result>")]
+ENCRYPTION_ATTEMPTS = [("<file1>", "<encryption attempt result>"), ("<file2>", "")]
 
 
-@pytest.fixture
-def ransomware_telem_test_instance():
-    return RansomwareTelem(ATTEMPTS)
+def test_ransomware_telem_send(spy_send_telemetry):
+    ransomware_telem_1 = RansomwareTelem(ENCRYPTION_ATTEMPTS[0])
+    ransomware_telem_2 = RansomwareTelem(ENCRYPTION_ATTEMPTS[1])
 
+    ransomware_telem_1.add_telemetry_to_batch(ransomware_telem_2)
+
+    ransomware_telem_1.send()
+    expected_data = {"ransomware_attempts": ENCRYPTION_ATTEMPTS}
+    expected_data = json.dumps(expected_data, cls=ransomware_telem_1.json_encoder)
 
-def test_ransomware_telem_send(ransomware_telem_test_instance, spy_send_telemetry):
-    ransomware_telem_test_instance.send()
-    expected_data = {"ransomware_attempts": ATTEMPTS}
-    expected_data = json.dumps(expected_data, cls=ransomware_telem_test_instance.json_encoder)
     assert spy_send_telemetry.data == expected_data
     assert spy_send_telemetry.telem_category == "ransomware"