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 = [("", "")] +ENCRYPTION_ATTEMPTS = [("", ""), ("", "")] -@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"