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:
parent
5c1902ca73
commit
0b953c8cff
|
@ -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:
|
||||
|
|
|
@ -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))
|
|
@ -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
|
Loading…
Reference in New Issue