Merge branch 'add-fields-to-file-encryption-telemetry' into develop

Author: mssalvatore

monkey/infection_monkey/ransomware/ransomware_payload.py

@@ -64,19 +64,19 @@ def _encrypt_files(self, file_list: List[Path]) -> List[Tuple[Path, Optional[Exc
                 LOG.debug(f"Encrypting {filepath}")
                 self._encryptor.encrypt_file_in_place(filepath)
                 self._add_extension(filepath)
-                self._send_telemetry(filepath, "")
+                self._send_telemetry(filepath, True, "")
             except Exception as ex:
                 LOG.warning(f"Error encrypting {filepath}: {ex}")
-                self._send_telemetry(filepath, str(ex))
+                self._send_telemetry(filepath, False, str(ex))
 
         return results
 
     def _add_extension(self, filepath: Path):
         new_filepath = filepath.with_suffix(f"{filepath.suffix}{self._new_file_extension}")
         filepath.rename(new_filepath)
 
-    def _send_telemetry(self, filepath: Path, error: str):
-        encryption_attempt = FileEncryptionTelem((str(filepath), str(error)))
+    def _send_telemetry(self, filepath: Path, success: bool, error: str):
+        encryption_attempt = FileEncryptionTelem(str(filepath), success, error)
         self._telemetry_messenger.send_telemetry(encryption_attempt)
 
     def _leave_readme(self):

monkey/infection_monkey/telemetry/file_encryption_telem.py

@@ -1,4 +1,4 @@
-from typing import Tuple
+from pathlib import Path
 
 from common.common_consts.telem_categories import TelemCategoryEnum
 from infection_monkey.telemetry.base_telem import BaseTelem
@@ -7,17 +7,16 @@
 
 
 class FileEncryptionTelem(BatchableTelemMixin, IBatchableTelem, BaseTelem):
-    def __init__(self, entry: Tuple[str, str]):
+    def __init__(self, filepath: Path, success: bool, error: str):
         """
         File Encryption telemetry constructor
-        :param attempts: List of tuples with each tuple containing the path
-                         of a file it tried encrypting and its result.
-                         If ransomware fails completely - list of one tuple
-                         containing the directory path and error string.
+        :param filepath: The path to the file that monkey attempted to encrypt
+        :param success: True if encryption was successful, false otherwise
+        :param error: An error message describing the failure. Empty unless success == False
         """
         super().__init__()
 
-        self._telemetry_entries.append(entry)
+        self._telemetry_entries.append({"path": filepath, "success": success, "error": error})
 
     telem_category = TelemCategoryEnum.FILE_ENCRYPTION
 

monkey/tests/unit_tests/infection_monkey/ransomware/test_ransomware_payload.py

@@ -1,5 +1,5 @@
 import os
-from pathlib import Path, PurePath
+from pathlib import Path, PurePosixPath
 
 import pytest
 from tests.unit_tests.infection_monkey.ransomware.ransomware_target_files import (
@@ -138,24 +138,27 @@ 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()["files"][0][0]
-    assert telem_1.get_data()["files"][0][1] == ""
-    assert TEST_KEYBOARD_TXT in telem_2.get_data()["files"][0][0]
-    assert telem_2.get_data()["files"][0][1] == ""
+    assert ALL_ZEROS_PDF in telem_1.get_data()["files"][0]["path"]
+    assert telem_1.get_data()["files"][0]["success"]
+    assert telem_1.get_data()["files"][0]["error"] == ""
+    assert TEST_KEYBOARD_TXT in telem_2.get_data()["files"][0]["path"]
+    assert telem_2.get_data()["files"][0]["success"]
+    assert telem_2.get_data()["files"][0]["error"] == ""
 
 
 def test_telemetry_failure(monkeypatch, ransomware_payload, telemetry_messenger_spy):
     monkeypatch.setattr(
         ransomware_payload_module,
         "select_production_safe_target_files",
-        lambda a, b: [PurePath("/file/not/exist")],
+        lambda a, b: [PurePosixPath("/file/not/exist")],
     ),
 
     ransomware_payload.run_payload()
     telem_1 = telemetry_messenger_spy.telemetries[0]
 
-    assert "/file/not/exist" in telem_1.get_data()["files"][0][0]
-    assert "No such file or directory" in telem_1.get_data()["files"][0][1]
+    assert "/file/not/exist" in telem_1.get_data()["files"][0]["path"]
+    assert not telem_1.get_data()["files"][0]["success"]
+    assert "No such file or directory" in telem_1.get_data()["files"][0]["error"]
 
 
 def test_readme_false(ransomware_payload_config, ransomware_target, telemetry_messenger_spy):

monkey/tests/unit_tests/infection_monkey/telemetry/test_file_encryption_telem.py

@@ -2,12 +2,23 @@
 
 from infection_monkey.telemetry.file_encryption_telem import FileEncryptionTelem
 
-ENCRYPTION_ATTEMPTS = [("<file1>", "<encryption attempt result>"), ("<file2>", "")]
+ENCRYPTION_ATTEMPTS = [
+    {"path": "<file1>", "success": False, "error": "<encryption attempt result>"},
+    {"path": "<file2>", "success": True, "error": ""},
+]
 
 
 def test_file_encryption_telem_send(spy_send_telemetry):
-    file_encryption_telem_1 = FileEncryptionTelem(ENCRYPTION_ATTEMPTS[0])
-    file_encryption_telem_2 = FileEncryptionTelem(ENCRYPTION_ATTEMPTS[1])
+    file_encryption_telem_1 = FileEncryptionTelem(
+        ENCRYPTION_ATTEMPTS[0]["path"],
+        ENCRYPTION_ATTEMPTS[0]["success"],
+        ENCRYPTION_ATTEMPTS[0]["error"],
+    )
+    file_encryption_telem_2 = FileEncryptionTelem(
+        ENCRYPTION_ATTEMPTS[1]["path"],
+        ENCRYPTION_ATTEMPTS[1]["success"],
+        ENCRYPTION_ATTEMPTS[1]["error"],
+    )
 
     file_encryption_telem_1.add_telemetry_to_batch(file_encryption_telem_2)