From 0b953c8cff6455a4d137795f59d7068409be311f Mon Sep 17 00:00:00 2001
From: Mike Salvatore <mike.s.salvatore@gmail.com>
Date: Tue, 22 Jun 2021 13:23:14 -0400
Subject: [PATCH] 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().
---
 .../ransomware/ransomware_payload.py          | 12 +++++-
 monkey/infection_monkey/ransomware/utils.py   | 14 ------
 .../infection_monkey/ransomware/test_utils.py | 43 -------------------
 3 files changed, 10 insertions(+), 59 deletions(-)
 delete mode 100644 monkey/infection_monkey/ransomware/utils.py
 delete mode 100644 monkey/tests/unit_tests/infection_monkey/ransomware/test_utils.py

diff --git a/monkey/infection_monkey/ransomware/ransomware_payload.py b/monkey/infection_monkey/ransomware/ransomware_payload.py
index 779181992..a47e3ba18 100644
--- a/monkey/infection_monkey/ransomware/ransomware_payload.py
+++ b/monkey/infection_monkey/ransomware/ransomware_payload.py
@@ -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:
diff --git a/monkey/infection_monkey/ransomware/utils.py b/monkey/infection_monkey/ransomware/utils.py
deleted file mode 100644
index dde0ff973..000000000
--- a/monkey/infection_monkey/ransomware/utils.py
+++ /dev/null
@@ -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))
diff --git a/monkey/tests/unit_tests/infection_monkey/ransomware/test_utils.py b/monkey/tests/unit_tests/infection_monkey/ransomware/test_utils.py
deleted file mode 100644
index 37a7f059d..000000000
--- a/monkey/tests/unit_tests/infection_monkey/ransomware/test_utils.py
+++ /dev/null
@@ -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