From 3450ac93a31cef8101c08bccc0eedb04327e6a3c Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Mon, 24 Jan 2022 14:01:03 -0500 Subject: [PATCH] Agent: Extract code from try/except in _leave_readme_in_target_directory --- monkey/infection_monkey/payload/ransomware/ransomware.py | 8 ++++---- .../payload/ransomware/test_ransomware.py | 7 ++++++- 2 files changed, 10 insertions(+), 5 deletions(-) diff --git a/monkey/infection_monkey/payload/ransomware/ransomware.py b/monkey/infection_monkey/payload/ransomware/ransomware.py index febdfe025..c4351acaf 100644 --- a/monkey/infection_monkey/payload/ransomware/ransomware.py +++ b/monkey/infection_monkey/payload/ransomware/ransomware.py @@ -71,11 +71,11 @@ class Ransomware: self._telemetry_messenger.send_telemetry(encryption_attempt) 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 + if interrupt.is_set(): + logger.debug("Received a stop signal, skipping leave readme") + return + try: 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 adffe6f88..365f9fecd 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 @@ -117,7 +117,12 @@ def test_interrupt_while_encrypting( mfe.assert_any_call(ransomware_test_data / HELLO_TXT) -def test_no_readme_after_interrupt(ransomware, interrupt, mock_leave_readme): +def test_no_readme_after_interrupt( + ransomware_options, build_ransomware, interrupt, mock_leave_readme +): + ransomware_options.readme_enabled = True + ransomware = build_ransomware(ransomware_options) + interrupt.set() ransomware.run(interrupt)