From 20890e51ec1509e64c3c416b5900b81e47574f21 Mon Sep 17 00:00:00 2001
From: Ilija Lazoroski <ilija.la@live.com>
Date: Wed, 15 Dec 2021 12:44:48 +0100
Subject: [PATCH] Agent: Remove ransomware cleanup function

---
 .../ransomware/ransomware_payload.py          | 21 ---------------
 .../ransomware/test_ransomware_payload.py     | 27 +------------------
 2 files changed, 1 insertion(+), 47 deletions(-)

diff --git a/monkey/infection_monkey/ransomware/ransomware_payload.py b/monkey/infection_monkey/ransomware/ransomware_payload.py
index ff2a89d64..a1e052970 100644
--- a/monkey/infection_monkey/ransomware/ransomware_payload.py
+++ b/monkey/infection_monkey/ransomware/ransomware_payload.py
@@ -30,7 +30,6 @@ class RansomwarePayload:
         self._readme_file_path = (
             self._target_directory / README_FILE_NAME if self._target_directory else None
         )
-        self._readme_incomplete = False
 
     def run_payload(self):
         if not self._target_directory:
@@ -67,26 +66,6 @@ class RansomwarePayload:
 
     def _leave_readme_in_target_directory(self):
         try:
-            self._readme_incomplete = True
             self._leave_readme(README_SRC, self._readme_file_path)
-            self._readme_incomplete = False
         except Exception as ex:
             logger.warning(f"An error occurred while attempting to leave a README.txt file: {ex}")
-
-    def cleanup(self):
-        # This cleanup function is only concerned with cleaning up and replacing *incomplete*
-        # README.txt files; its goal is not to ensure the existence of a README file. Therefore,
-        # only retry if a README.txt file actually exists.
-        if self._readme_incomplete and self._readme_file_path.exists():
-            logger.info(
-                "The process of leaving a README.txt was interrupted. Removing the corrupt file "
-                "and trying again."
-            )
-            try:
-                self._readme_file_path.unlink()
-                self._leave_readme_in_target_directory()
-            except Exception as ex:
-                logger.error(
-                    "An error occurred while trying to remove the corrupt or incomplete README.txt "
-                    f"file: {ex}"
-                )
diff --git a/monkey/tests/unit_tests/infection_monkey/ransomware/test_ransomware_payload.py b/monkey/tests/unit_tests/infection_monkey/ransomware/test_ransomware_payload.py
index 09a330553..24eb8443d 100644
--- a/monkey/tests/unit_tests/infection_monkey/ransomware/test_ransomware_payload.py
+++ b/monkey/tests/unit_tests/infection_monkey/ransomware/test_ransomware_payload.py
@@ -1,4 +1,4 @@
-from pathlib import Path, PurePosixPath
+from pathlib import PurePosixPath
 from unittest.mock import MagicMock
 
 import pytest
@@ -184,28 +184,3 @@ def test_leave_readme_exceptions_handled(build_ransomware_payload, ransomware_pa
 
     # Test will fail if exception is raised and not handled
     ransomware_payload.run_payload()
-    ransomware_payload.cleanup()
-
-
-def test_cleanup_incomplete_readme(build_ransomware_payload, ransomware_payload_config):
-    def leave_readme(_: Path, dest: Path):
-        if leave_readme.i == 0:
-            dest.touch()
-
-        leave_readme.i += 1
-
-        raise Exception("Test exception when leaving README")
-
-    leave_readme.i = 0
-
-    ransomware_payload_config.readme_enabled = True
-    ransomware_payload = build_ransomware_payload(
-        config=ransomware_payload_config, leave_readme=leave_readme
-    )
-
-    ransomware_payload.run_payload()
-    assert (ransomware_payload_config.target_directory / README_FILE_NAME).exists()
-
-    ransomware_payload.cleanup()
-    assert not (ransomware_payload_config.target_directory / README_FILE_NAME).exists()
-    assert leave_readme.i == 2