From 2299c029d7c36706990bd2716db6bbc7866fe6ec Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Fri, 17 Dec 2021 09:08:30 -0500 Subject: [PATCH] Agent: Rename RansomwarePayload to Ransomware A payload adheres to a specific IPayload interface. The class that is now called RansomwarePayload is just a concrete ransomware. A new RansomwarePayload will be introduced to wrap the build and execute of the Ransomware. --- .../{ransomware_payload.py => ransomware.py} | 2 +- ...yload_builder.py => ransomware_builder.py} | 8 +- .../payload/ransomware/test_file_selectors.py | 2 +- ...nsomware_payload.py => test_ransomware.py} | 98 ++++++++----------- 4 files changed, 49 insertions(+), 61 deletions(-) rename monkey/infection_monkey/payload/ransomware/{ransomware_payload.py => ransomware.py} (99%) rename monkey/infection_monkey/payload/ransomware/{ransomware_payload_builder.py => ransomware_builder.py} (89%) rename monkey/tests/unit_tests/infection_monkey/payload/ransomware/{test_ransomware_payload.py => test_ransomware.py} (55%) diff --git a/monkey/infection_monkey/payload/ransomware/ransomware_payload.py b/monkey/infection_monkey/payload/ransomware/ransomware.py similarity index 99% rename from monkey/infection_monkey/payload/ransomware/ransomware_payload.py rename to monkey/infection_monkey/payload/ransomware/ransomware.py index 227829d10..2f09e386f 100644 --- a/monkey/infection_monkey/payload/ransomware/ransomware_payload.py +++ b/monkey/infection_monkey/payload/ransomware/ransomware.py @@ -11,7 +11,7 @@ from .ransomware_config import RansomwareConfig logger = logging.getLogger(__name__) -class RansomwarePayload: +class Ransomware: def __init__( self, config: RansomwareConfig, diff --git a/monkey/infection_monkey/payload/ransomware/ransomware_payload_builder.py b/monkey/infection_monkey/payload/ransomware/ransomware_builder.py similarity index 89% rename from monkey/infection_monkey/payload/ransomware/ransomware_payload_builder.py rename to monkey/infection_monkey/payload/ransomware/ransomware_builder.py index c5eb034f0..671f20f1f 100644 --- a/monkey/infection_monkey/payload/ransomware/ransomware_payload_builder.py +++ b/monkey/infection_monkey/payload/ransomware/ransomware_builder.py @@ -12,8 +12,8 @@ from infection_monkey.utils.bit_manipulators import flip_bits from . import readme_dropper from .file_selectors import ProductionSafeTargetFileSelector from .in_place_file_encryptor import InPlaceFileEncryptor +from .ransomware import Ransomware from .ransomware_config import RansomwareConfig -from .ransomware_payload import RansomwarePayload from .targeted_file_extensions import TARGETED_FILE_EXTENSIONS EXTENSION = ".m0nk3y" @@ -22,8 +22,8 @@ CHUNK_SIZE = 4096 * 24 logger = logging.getLogger(__name__) -def build_ransomware_payload(config: dict): - logger.debug(f"Ransomware payload configuration:\n{pformat(config)}") +def build_ransomware(config: dict): + logger.debug(f"Ransomware configuration:\n{pformat(config)}") ransomware_config = RansomwareConfig(config) file_encryptor = _build_file_encryptor() @@ -31,7 +31,7 @@ def build_ransomware_payload(config: dict): leave_readme = _build_leave_readme() telemetry_messenger = _build_telemetry_messenger() - return RansomwarePayload( + return Ransomware( ransomware_config, file_encryptor, file_selector, diff --git a/monkey/tests/unit_tests/infection_monkey/payload/ransomware/test_file_selectors.py b/monkey/tests/unit_tests/infection_monkey/payload/ransomware/test_file_selectors.py index 635b39587..f779b733e 100644 --- a/monkey/tests/unit_tests/infection_monkey/payload/ransomware/test_file_selectors.py +++ b/monkey/tests/unit_tests/infection_monkey/payload/ransomware/test_file_selectors.py @@ -13,7 +13,7 @@ from tests.unit_tests.infection_monkey.payload.ransomware.ransomware_target_file from tests.utils import is_user_admin from infection_monkey.payload.ransomware.file_selectors import ProductionSafeTargetFileSelector -from infection_monkey.payload.ransomware.ransomware_payload import README_SRC +from infection_monkey.payload.ransomware.ransomware import README_SRC TARGETED_FILE_EXTENSIONS = [".pdf", ".txt"] diff --git a/monkey/tests/unit_tests/infection_monkey/payload/ransomware/test_ransomware_payload.py b/monkey/tests/unit_tests/infection_monkey/payload/ransomware/test_ransomware.py similarity index 55% rename from monkey/tests/unit_tests/infection_monkey/payload/ransomware/test_ransomware_payload.py rename to monkey/tests/unit_tests/infection_monkey/payload/ransomware/test_ransomware.py index 39ca057ee..a7e9f8a90 100644 --- a/monkey/tests/unit_tests/infection_monkey/payload/ransomware/test_ransomware_payload.py +++ b/monkey/tests/unit_tests/infection_monkey/payload/ransomware/test_ransomware.py @@ -8,17 +8,17 @@ from tests.unit_tests.infection_monkey.payload.ransomware.ransomware_target_file ) from infection_monkey.payload.ransomware.consts import README_FILE_NAME, README_SRC +from infection_monkey.payload.ransomware.ransomware import Ransomware from infection_monkey.payload.ransomware.ransomware_config import RansomwareConfig -from infection_monkey.payload.ransomware.ransomware_payload import RansomwarePayload @pytest.fixture -def ransomware_payload(build_ransomware_payload, ransomware_payload_config): - return build_ransomware_payload(ransomware_payload_config) +def ransomware(build_ransomware, ransomware_config): + return build_ransomware(ransomware_config) @pytest.fixture -def build_ransomware_payload( +def build_ransomware( mock_file_encryptor, mock_file_selector, mock_leave_readme, telemetry_messenger_spy ): def inner( @@ -27,7 +27,7 @@ def build_ransomware_payload( file_selector=mock_file_selector, leave_readme=mock_leave_readme, ): - return RansomwarePayload( + return Ransomware( config, file_encryptor, file_selector, @@ -39,7 +39,7 @@ def build_ransomware_payload( @pytest.fixture -def ransomware_payload_config(ransomware_test_data): +def ransomware_config(ransomware_test_data): class RansomwareConfigStub(RansomwareConfig): def __init__(self, encryption_enabled, readme_enabled, target_directory): self.encryption_enabled = encryption_enabled @@ -69,18 +69,16 @@ def mock_leave_readme(): def test_files_selected_from_target_dir( - ransomware_payload, - ransomware_payload_config, + ransomware, + ransomware_config, mock_file_selector, ): - ransomware_payload.run_payload() - mock_file_selector.assert_called_with(ransomware_payload_config.target_directory) + ransomware.run_payload() + mock_file_selector.assert_called_with(ransomware_config.target_directory) -def test_all_selected_files_encrypted( - ransomware_test_data, ransomware_payload, mock_file_encryptor -): - ransomware_payload.run_payload() +def test_all_selected_files_encrypted(ransomware_test_data, ransomware, mock_file_encryptor): + ransomware.run_payload() assert mock_file_encryptor.call_count == 2 mock_file_encryptor.assert_any_call(ransomware_test_data / ALL_ZEROS_PDF) @@ -88,30 +86,30 @@ def test_all_selected_files_encrypted( def test_encryption_skipped_if_configured_false( - build_ransomware_payload, ransomware_payload_config, mock_file_encryptor + build_ransomware, ransomware_config, mock_file_encryptor ): - ransomware_payload_config.encryption_enabled = False + ransomware_config.encryption_enabled = False - ransomware_payload = build_ransomware_payload(ransomware_payload_config) - ransomware_payload.run_payload() + ransomware = build_ransomware(ransomware_config) + ransomware.run_payload() assert mock_file_encryptor.call_count == 0 def test_encryption_skipped_if_no_directory( - build_ransomware_payload, ransomware_payload_config, mock_file_encryptor + build_ransomware, ransomware_config, mock_file_encryptor ): - ransomware_payload_config.encryption_enabled = True - ransomware_payload_config.target_directory = None + ransomware_config.encryption_enabled = True + ransomware_config.target_directory = None - ransomware_payload = build_ransomware_payload(ransomware_payload_config) - ransomware_payload.run_payload() + ransomware = build_ransomware(ransomware_config) + ransomware.run_payload() assert mock_file_encryptor.call_count == 0 -def test_telemetry_success(ransomware_payload, telemetry_messenger_spy): - ransomware_payload.run_payload() +def test_telemetry_success(ransomware, telemetry_messenger_spy): + ransomware.run_payload() assert len(telemetry_messenger_spy.telemetries) == 2 telem_1 = telemetry_messenger_spy.telemetries[0] @@ -125,19 +123,15 @@ def test_telemetry_success(ransomware_payload, telemetry_messenger_spy): assert telem_2.get_data()["files"][0]["error"] == "" -def test_telemetry_failure( - build_ransomware_payload, ransomware_payload_config, telemetry_messenger_spy -): +def test_telemetry_failure(build_ransomware, ransomware_config, telemetry_messenger_spy): file_not_exists = "/file/not/exist" mfe = MagicMock( side_effect=FileNotFoundError(f"[Errno 2] No such file or directory: '{file_not_exists}'") ) mfs = MagicMock(return_value=[PurePosixPath(file_not_exists)]) - ransomware_payload = build_ransomware_payload( - config=ransomware_payload_config, file_encryptor=mfe, file_selector=mfs - ) + ransomware = build_ransomware(config=ransomware_config, file_encryptor=mfe, file_selector=mfs) - ransomware_payload.run_payload() + ransomware.run_payload() telem = telemetry_messenger_spy.telemetries[0] assert file_not_exists in telem.get_data()["files"][0]["path"] @@ -145,42 +139,36 @@ def test_telemetry_failure( assert "No such file or directory" in telem.get_data()["files"][0]["error"] -def test_readme_false(build_ransomware_payload, ransomware_payload_config, mock_leave_readme): - ransomware_payload_config.readme_enabled = False - ransomware_payload = build_ransomware_payload(ransomware_payload_config) +def test_readme_false(build_ransomware, ransomware_config, mock_leave_readme): + ransomware_config.readme_enabled = False + ransomware = build_ransomware(ransomware_config) - ransomware_payload.run_payload() + ransomware.run_payload() mock_leave_readme.assert_not_called() -def test_readme_true( - build_ransomware_payload, ransomware_payload_config, mock_leave_readme, ransomware_test_data -): - ransomware_payload_config.readme_enabled = True - ransomware_payload = build_ransomware_payload(ransomware_payload_config) +def test_readme_true(build_ransomware, ransomware_config, mock_leave_readme, ransomware_test_data): + ransomware_config.readme_enabled = True + ransomware = build_ransomware(ransomware_config) - ransomware_payload.run_payload() + ransomware.run_payload() mock_leave_readme.assert_called_with(README_SRC, ransomware_test_data / README_FILE_NAME) -def test_no_readme_if_no_directory( - build_ransomware_payload, ransomware_payload_config, mock_leave_readme -): - ransomware_payload_config.target_directory = None - ransomware_payload_config.readme_enabled = True +def test_no_readme_if_no_directory(build_ransomware, ransomware_config, mock_leave_readme): + ransomware_config.target_directory = None + ransomware_config.readme_enabled = True - ransomware_payload = build_ransomware_payload(ransomware_payload_config) + ransomware = build_ransomware(ransomware_config) - ransomware_payload.run_payload() + ransomware.run_payload() mock_leave_readme.assert_not_called() -def test_leave_readme_exceptions_handled(build_ransomware_payload, ransomware_payload_config): +def test_leave_readme_exceptions_handled(build_ransomware, ransomware_config): leave_readme = MagicMock(side_effect=Exception("Test exception when leaving README")) - ransomware_payload_config.readme_enabled = True - ransomware_payload = build_ransomware_payload( - config=ransomware_payload_config, leave_readme=leave_readme - ) + ransomware_config.readme_enabled = True + ransomware = build_ransomware(config=ransomware_config, leave_readme=leave_readme) # Test will fail if exception is raised and not handled - ransomware_payload.run_payload() + ransomware.run_payload()