agent: Return results from RansomwareBitflipEncryptor.encrypt_files()

monkey/infection_monkey/ransomware/ransomware_bitflip_encryptor.py

@@ -1,5 +1,5 @@
 from pathlib import Path
-from typing import List
+from typing import List, Optional, Tuple
 
 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}")
         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:
-            self._encrypt_single_file_in_place(filepath)
-            self._add_extension(filepath)
+            try:
+                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):
         with open(filepath, "rb+") as f:

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

@@ -18,13 +18,42 @@ def with_extension(filename):
 
 
 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[1]) == TEST_KEYBOARD_TXT_CLEARTEXT_SHA256
 
-    encryptor = RansomwareBitflipEncryptor(".new")
+    encryptor = RansomwareBitflipEncryptor(EXTENSION)
     encryptor.encrypt_files(file_list)
 
-    assert hash_file(with_extension(file_list[0])) == ALL_ZEROS_PDF_ENCRYPTED_SHA256
-    assert hash_file(with_extension(file_list[1])) == TEST_KEYBOARD_TXT_ENCRYPTED_SHA256
+    assert hash_file(with_extension(orig_all_zeros)) == ALL_ZEROS_PDF_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])