agent: Add explicit fields to FileEncryptionTelem

This commit is contained in:
Mike Salvatore 2021-06-28 13:30:53 -04:00
parent 49eb1cd996
commit df1f3cda78
4 changed files with 20 additions and 13 deletions

View File

@ -68,5 +68,5 @@ class RansomewarePayload:
filepath.rename(new_filepath)
def _send_telemetry(self, filepath: Path, error: str):
encryption_attempt = FileEncryptionTelem((str(filepath), str(error)))
encryption_attempt = FileEncryptionTelem(str(filepath), str(error))
self._telemetry_messenger.send_telemetry(encryption_attempt)

View File

@ -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,7 +7,7 @@ from infection_monkey.telemetry.i_batchable_telem import IBatchableTelem
class FileEncryptionTelem(BatchableTelemMixin, IBatchableTelem, BaseTelem):
def __init__(self, entry: Tuple[str, str]):
def __init__(self, filepath: Path, error: str):
"""
File Encryption telemetry constructor
:param attempts: List of tuples with each tuple containing the path
@ -17,7 +17,7 @@ class FileEncryptionTelem(BatchableTelemMixin, IBatchableTelem, BaseTelem):
"""
super().__init__()
self._telemetry_entries.append(entry)
self._telemetry_entries.append({"path": filepath, "error": error})
telem_category = TelemCategoryEnum.FILE_ENCRYPTION

View File

@ -133,10 +133,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()["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]["error"] == ""
assert TEST_KEYBOARD_TXT in telem_2.get_data()["files"][0]["path"]
assert telem_2.get_data()["files"][0]["error"] == ""
def test_telemetry_failure(monkeypatch, ransomware_payload, telemetry_messenger_spy):
@ -149,5 +149,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()["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 "No such file or directory" in telem_1.get_data()["files"][0]["error"]

View File

@ -2,12 +2,19 @@ import json
from infection_monkey.telemetry.file_encryption_telem import FileEncryptionTelem
ENCRYPTION_ATTEMPTS = [("<file1>", "<encryption attempt result>"), ("<file2>", "")]
ENCRYPTION_ATTEMPTS = [
{"path": "<file1>", "error": "<encryption attempt result>"},
{"path": "<file2>", "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]["error"]
)
file_encryption_telem_2 = FileEncryptionTelem(
ENCRYPTION_ATTEMPTS[1]["path"], ENCRYPTION_ATTEMPTS[1]["error"]
)
file_encryption_telem_1.add_telemetry_to_batch(file_encryption_telem_2)