From 5b64ea5151e715012fbd406fed41296e37fc002e Mon Sep 17 00:00:00 2001
From: Shreya <shreya.malviya@gmail.com>
Date: Tue, 22 Jun 2021 11:41:49 +0530
Subject: [PATCH] agent: ransomware: Iterate through files in directory and get
 list of files to encrypt

---
 monkey/infection_monkey/ransomware/utils.py | 94 +++++++++++++++++++++
 vulture_allowlist.py                        |  1 +
 2 files changed, 95 insertions(+)
 create mode 100644 monkey/infection_monkey/ransomware/utils.py

diff --git a/monkey/infection_monkey/ransomware/utils.py b/monkey/infection_monkey/ransomware/utils.py
new file mode 100644
index 000000000..d9cc7d5ce
--- /dev/null
+++ b/monkey/infection_monkey/ransomware/utils.py
@@ -0,0 +1,94 @@
+import os
+from typing import Iterator, List
+
+VALID_FILE_EXTENSIONS_FOR_ENCRYPTION = {
+    ".3ds",
+    ".7z",
+    ".accdb",
+    ".ai",
+    ".asp",
+    ".aspx",
+    ".avhd",
+    ".avi",
+    ".back",
+    ".bak",
+    ".c",
+    ".cfg",
+    ".conf",
+    ".cpp",
+    ".cs",
+    ".ctl",
+    ".dbf",
+    ".disk",
+    ".djvu",
+    ".doc",
+    ".docx",
+    ".dwg",
+    ".eml",
+    ".fdb",
+    ".giff",
+    ".gz",
+    ".h",
+    ".hdd",
+    ".jpg",
+    ".jpeg",
+    ".kdbx",
+    ".mail",
+    ".mdb",
+    ".mpg",
+    ".mpeg",
+    ".msg",
+    ".nrg",
+    ".ora",
+    ".ost",
+    ".ova",
+    ".ovf",
+    ".pdf",
+    ".php",
+    ".pmf",
+    ".png",
+    ".ppt",
+    ".pptx",
+    ".pst",
+    ".pvi",
+    ".py",
+    ".pyc",
+    ".rar",
+    ".rtf",
+    ".sln",
+    ".sql",
+    ".tar",
+    ".tiff",
+    ".txt",
+    ".vbox",
+    ".vbs",
+    ".vcb",
+    ".vdi",
+    ".vfd",
+    ".vmc",
+    ".vmdk",
+    ".vmsd",
+    ".vmx",
+    ".vsdx",
+    ".vsv",
+    ".work",
+    ".xls",
+    ".xlsx",
+    ".xvd",
+    ".zip",
+}
+
+
+def get_files_to_encrypt(dir_path: str) -> List[str]:
+    all_files = get_all_files_in_directory(dir_path)
+
+    files_to_encrypt = []
+    for file in all_files:
+        if os.path.splitext(file)[1] in VALID_FILE_EXTENSIONS_FOR_ENCRYPTION:
+            files_to_encrypt.append(file)
+
+    return files_to_encrypt
+
+
+def get_all_files_in_directory(dir_path: str) -> Iterator:
+    return filter(os.path.isfile, [os.path.join(dir_path, item) for item in os.listdir(dir_path)])
diff --git a/vulture_allowlist.py b/vulture_allowlist.py
index 2c937ee4f..304ff6f12 100644
--- a/vulture_allowlist.py
+++ b/vulture_allowlist.py
@@ -171,6 +171,7 @@ ISLAND  # unused variable (monkey/monkey_island/cc/services/utils/node_states.py
 MONKEY_LINUX_RUNNING  # unused variable (monkey/monkey_island/cc/services/utils/node_states.py:26)
 import_status  # monkey_island\cc\resources\configuration_import.py:19
 config_schema  # monkey_island\cc\resources\configuration_import.py:25
+get_files_to_encrypt  # monkey/infection_monkey/ransomware/utils.py:82
 
 # these are not needed for it to work, but may be useful extra information to understand what's going on
 WINDOWS_PBA_TYPE  # unused variable (monkey/monkey_island/cc/resources/pba_file_upload.py:23)