agent: Remove ransomware/utils.py

The code for getting files to encrypt has become so trivial that it no
longer warrants a separate function outside of _find_files().
This commit is contained in:
Mike Salvatore 2021-06-22 13:23:14 -04:00
parent 5c1902ca73
commit 0b953c8cff
3 changed files with 10 additions and 59 deletions

View File

@ -1,6 +1,12 @@
import logging
from pathlib import Path
from infection_monkey.ransomware.utils import get_files_to_encrypt
from infection_monkey.ransomware.valid_file_extensions import VALID_FILE_EXTENSIONS_FOR_ENCRYPTION
from infection_monkey.utils.dir_utils import (
file_extension_filter,
filter_files,
get_all_files_in_directory,
)
from infection_monkey.utils.environment import is_windows_os
LOG = logging.getLogger(__name__)
@ -25,7 +31,9 @@ class RansomewarePayload:
if is_windows_os()
else self.config["linux_dir_ransom"]
)
return get_files_to_encrypt(dir_path)
all_files = get_all_files_in_directory(Path(dir_path))
return filter_files(all_files, file_extension_filter(VALID_FILE_EXTENSIONS_FOR_ENCRYPTION))
def _encrypt_files(self, file_list):
for file in file_list:

View File

@ -1,14 +0,0 @@
from pathlib import Path
from typing import List
from infection_monkey.ransomware.valid_file_extensions import VALID_FILE_EXTENSIONS_FOR_ENCRYPTION
from infection_monkey.utils.dir_utils import (
file_extension_filter,
filter_files,
get_all_files_in_directory,
)
def get_files_to_encrypt(dir_path: str) -> List[Path]:
all_files = get_all_files_in_directory(Path(dir_path))
return filter_files(all_files, file_extension_filter(VALID_FILE_EXTENSIONS_FOR_ENCRYPTION))

View File

@ -1,43 +0,0 @@
import infection_monkey.ransomware.utils
VALID_FILE_EXTENSION_1 = "file.3ds"
VALID_FILE_EXTENSION_2 = "file.jpg.zip"
INVALID_FILE_EXTENSION_1 = "file.pqr"
INVALID_FILE_EXTENSION_2 = "file.xyz"
SUBDIR_1 = "subdir1"
SUBDIR_2 = "subdir2"
def test_get_files_to_encrypt__no_files(monkeypatch):
all_files = []
monkeypatch.setattr(
"infection_monkey.ransomware.utils.get_all_files_in_directory", lambda _: all_files
)
expected_return_value = []
assert infection_monkey.ransomware.utils.get_files_to_encrypt("") == expected_return_value
def test_get_files_to_encrypt__no_valid_files(monkeypatch):
all_files = [INVALID_FILE_EXTENSION_1, INVALID_FILE_EXTENSION_2]
monkeypatch.setattr(
"infection_monkey.ransomware.utils.get_all_files_in_directory", lambda _: all_files
)
expected_return_value = []
assert infection_monkey.ransomware.utils.get_files_to_encrypt("") == expected_return_value
def test_get_files_to_encrypt__valid_files(monkeypatch):
all_files = [
VALID_FILE_EXTENSION_1,
INVALID_FILE_EXTENSION_1,
VALID_FILE_EXTENSION_2,
INVALID_FILE_EXTENSION_2,
]
monkeypatch.setattr(
"infection_monkey.ransomware.utils.get_all_files_in_directory", lambda _: all_files
)
expected_return_value = [VALID_FILE_EXTENSION_1, VALID_FILE_EXTENSION_2]
assert infection_monkey.ransomware.utils.get_files_to_encrypt("") == expected_return_value