Agent: Move ransomware readme constants to ransomware/consts.py

This commit is contained in:
Mike Salvatore 2021-07-16 09:42:13 -04:00
parent b1fe850624
commit 07937d7238
3 changed files with 10 additions and 10 deletions

View File

@ -0,0 +1,5 @@
from pathlib import Path
README_SRC = Path(__file__).parent / "ransomware_readme.txt"
README_FILE_NAME = "README.txt"
README_SHA256_HASH = "e3d9343cbcce6097c83044327b00ead14b6e8e6aa0d411160610033a856032fc"

View File

@ -2,15 +2,13 @@ import logging
from pathlib import Path
from typing import Callable, List
from infection_monkey.ransomware.consts import README_FILE_NAME, README_SRC
from infection_monkey.ransomware.ransomware_config import RansomwareConfig
from infection_monkey.telemetry.file_encryption_telem import FileEncryptionTelem
from infection_monkey.telemetry.messengers.i_telemetry_messenger import ITelemetryMessenger
LOG = logging.getLogger(__name__)
README_SRC = Path(__file__).parent / "ransomware_readme.txt"
README_DEST = "README.txt"
class RansomwarePayload:
def __init__(
@ -39,7 +37,7 @@ class RansomwarePayload:
self._encrypt_files(file_list)
if self._config.readme_enabled:
self._leave_readme(README_SRC, self._config.target_directory / README_DEST)
self._leave_readme(README_SRC, self._config.target_directory / README_FILE_NAME)
def _find_files(self) -> List[Path]:
LOG.info(f"Collecting files in {self._config.target_directory}")

View File

@ -7,12 +7,9 @@ from tests.unit_tests.infection_monkey.ransomware.ransomware_target_files import
TEST_KEYBOARD_TXT,
)
from infection_monkey.ransomware.consts import README_FILE_NAME, README_SRC
from infection_monkey.ransomware.ransomware_config import RansomwareConfig
from infection_monkey.ransomware.ransomware_payload import (
README_DEST,
README_SRC,
RansomwarePayload,
)
from infection_monkey.ransomware.ransomware_payload import RansomwarePayload
@pytest.fixture
@ -162,7 +159,7 @@ def test_readme_true(
ransomware_payload = build_ransomware_payload(ransomware_payload_config)
ransomware_payload.run_payload()
mock_leave_readme.assert_called_with(README_SRC, ransomware_test_data / README_DEST)
mock_leave_readme.assert_called_with(README_SRC, ransomware_test_data / README_FILE_NAME)
def test_no_readme_if_no_directory(