agent: Add explicit "success" field to FileEncryptionTelem

This commit is contained in:
Mike Salvatore 2021-06-28 13:37:00 -04:00
parent c1af3f8165
commit dbd6dedb95
4 changed files with 17 additions and 10 deletions

View File

@ -56,10 +56,10 @@ class RansomewarePayload:
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
@ -67,6 +67,6 @@ class RansomewarePayload:
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)

View File

@ -7,7 +7,7 @@ from infection_monkey.telemetry.i_batchable_telem import IBatchableTelem
class FileEncryptionTelem(BatchableTelemMixin, IBatchableTelem, BaseTelem):
def __init__(self, filepath: Path, error: 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
@ -17,7 +17,7 @@ class FileEncryptionTelem(BatchableTelemMixin, IBatchableTelem, BaseTelem):
"""
super().__init__()
self._telemetry_entries.append({"path": filepath, "error": error})
self._telemetry_entries.append({"path": filepath, "success": success, "error": error})
telem_category = TelemCategoryEnum.FILE_ENCRYPTION

View File

@ -134,8 +134,10 @@ def test_telemetry_success(ransomware_payload, telemetry_messenger_spy):
telem_2 = telemetry_messenger_spy.telemetries[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"] == ""
@ -150,4 +152,5 @@ def test_telemetry_failure(monkeypatch, ransomware_payload, telemetry_messenger_
telem_1 = telemetry_messenger_spy.telemetries[0]
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"]

View File

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