From afbc313a7cff0d16e9b3c24cf33d2b9deba5a672 Mon Sep 17 00:00:00 2001
From: Ilija Lazoroski <ilija.la@live.com>
Date: Fri, 17 Dec 2021 16:10:42 +0100
Subject: [PATCH] Agent: Handle interrupts in ransomware

---
 .../payload/ransomware/ransomware.py          | 12 +++---
 .../payload/ransomware/test_ransomware.py     | 38 +++++++++++++++++++
 2 files changed, 44 insertions(+), 6 deletions(-)

diff --git a/monkey/infection_monkey/payload/ransomware/ransomware.py b/monkey/infection_monkey/payload/ransomware/ransomware.py
index fcc0a53e7..003112cc3 100644
--- a/monkey/infection_monkey/payload/ransomware/ransomware.py
+++ b/monkey/infection_monkey/payload/ransomware/ransomware.py
@@ -43,12 +43,8 @@ class Ransomware:
             file_list = self._find_files()
             self._encrypt_files(file_list, interrupt)
 
-        if interrupt.is_set():
-            logger.debug("Received a stop signal, skipping remaining tasks of ransomware payload")
-            return
-
         if self._config.readme_enabled:
-            self._leave_readme_in_target_directory()
+            self._leave_readme_in_target_directory(interrupt)
 
     def _find_files(self) -> List[Path]:
         logger.info(f"Collecting files in {self._target_directory}")
@@ -76,8 +72,12 @@ class Ransomware:
         encryption_attempt = FileEncryptionTelem(str(filepath), success, error)
         self._telemetry_messenger.send_telemetry(encryption_attempt)
 
-    def _leave_readme_in_target_directory(self):
+    def _leave_readme_in_target_directory(self, interrupt: threading.Event):
         try:
+            if interrupt.is_set():
+                logger.debug("Received a stop signal, skipping leave readme")
+                return
+
             self._leave_readme(README_SRC, self._readme_file_path)
         except Exception as ex:
             logger.warning(f"An error occurred while attempting to leave a README.txt file: {ex}")
diff --git a/monkey/tests/unit_tests/infection_monkey/payload/ransomware/test_ransomware.py b/monkey/tests/unit_tests/infection_monkey/payload/ransomware/test_ransomware.py
index 94de5aabc..adffe6f88 100644
--- a/monkey/tests/unit_tests/infection_monkey/payload/ransomware/test_ransomware.py
+++ b/monkey/tests/unit_tests/infection_monkey/payload/ransomware/test_ransomware.py
@@ -5,6 +5,7 @@ from unittest.mock import MagicMock
 import pytest
 from tests.unit_tests.infection_monkey.payload.ransomware.ransomware_target_files import (
     ALL_ZEROS_PDF,
+    HELLO_TXT,
     TEST_KEYBOARD_TXT,
 )
 
@@ -69,6 +70,11 @@ def mock_leave_readme():
     return MagicMock()
 
 
+@pytest.fixture
+def interrupt():
+    return threading.Event()
+
+
 def test_files_selected_from_target_dir(
     ransomware,
     ransomware_options,
@@ -86,6 +92,38 @@ def test_all_selected_files_encrypted(ransomware_test_data, ransomware, mock_fil
     mock_file_encryptor.assert_any_call(ransomware_test_data / TEST_KEYBOARD_TXT)
 
 
+def test_interrupt_while_encrypting(
+    ransomware_test_data, interrupt, ransomware_options, build_ransomware
+):
+    selected_files = [
+        ransomware_test_data / ALL_ZEROS_PDF,
+        ransomware_test_data / HELLO_TXT,
+        ransomware_test_data / TEST_KEYBOARD_TXT,
+    ]
+    mfs = MagicMock(return_value=selected_files)
+
+    def _callback(file_path, *_):
+        # Block all threads here until 2 threads reach this barrier, then set stop
+        # and test that neither thread continues to scan.
+        if file_path.name == HELLO_TXT:
+            interrupt.set()
+
+    mfe = MagicMock(side_effect=_callback)
+
+    build_ransomware(ransomware_options, mfe, mfs).run(interrupt)
+
+    assert mfe.call_count == 2
+    mfe.assert_any_call(ransomware_test_data / ALL_ZEROS_PDF)
+    mfe.assert_any_call(ransomware_test_data / HELLO_TXT)
+
+
+def test_no_readme_after_interrupt(ransomware, interrupt, mock_leave_readme):
+    interrupt.set()
+    ransomware.run(interrupt)
+
+    mock_leave_readme.assert_not_called()
+
+
 def test_encryption_skipped_if_configured_false(
     build_ransomware, ransomware_options, mock_file_encryptor
 ):