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.
This commit is contained in:
Mike Salvatore 2021-12-17 09:08:30 -05:00
parent 0a4ff25843
commit 2299c029d7
4 changed files with 49 additions and 61 deletions

View File

@ -11,7 +11,7 @@ from .ransomware_config import RansomwareConfig
logger = logging.getLogger(__name__)
class RansomwarePayload:
class Ransomware:
def __init__(
self,
config: RansomwareConfig,

View File

@ -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,

View File

@ -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"]

View File

@ -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()