forked from p15670423/monkey
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:
parent
0a4ff25843
commit
2299c029d7
|
@ -11,7 +11,7 @@ from .ransomware_config import RansomwareConfig
|
|||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class RansomwarePayload:
|
||||
class Ransomware:
|
||||
def __init__(
|
||||
self,
|
||||
config: RansomwareConfig,
|
|
@ -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,
|
|
@ -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"]
|
||||
|
||||
|
|
|
@ -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()
|
Loading…
Reference in New Issue