Agent: Add a method that appends random string to filename in path

This method will be used to avoid duplication in destination file paths and will avoid clashes of exploiters writing to same files

monkey/infection_monkey/exploit/tools/helpers.py

@@ -1,22 +1,38 @@
 import logging
 import random
 import string
+from pathlib import Path
 from typing import Any, Mapping
 
 from infection_monkey.model import VictimHost
 
 logger = logging.getLogger(__name__)
 
+RAND_SUFFIX_LEN = 8
+
 
 def get_random_file_suffix() -> str:
     character_set = list(string.ascii_letters + string.digits + "_" + "-")
     # random.SystemRandom can block indefinitely in Linux
-    random_string = "".join(random.choices(character_set, k=8))  # noqa: DUO102
+    random_string = "".join(random.choices(character_set, k=RAND_SUFFIX_LEN))  # noqa: DUO102
     return random_string
 
 
-def get_agent_dest_path(host: VictimHost, options: Mapping[str, Any]) -> str:
+def get_agent_dest_path(host: VictimHost, options: Mapping[str, Any]) -> Path:
     if host.os["type"] == "windows":
-        return options["dropper_target_path_win_64"]
+        path = Path(options["dropper_target_path_win_64"])
     else:
-        return options["dropper_target_path_linux"]
+        path = Path(options["dropper_target_path_linux"])
+
+    return _add_random_suffix(path)
+
+
+#  Turns C:\\monkey.exe into C:\\monkey-<random_string>.exe
+#  Useful to avoid duplicate file paths
+def _add_random_suffix(path: Path) -> Path:
+    stem = path.name.split(".")[0]
+    suffixes = path.suffixes
+    stem = f"{stem}-{get_random_file_suffix()}"
+    rand_filename = "".join([stem, *suffixes])
+    path = path.with_name(rand_filename)
+    return path

monkey/tests/unit_tests/infection_monkey/exploit/tools/test_helpers.py

@@ -0,0 +1,40 @@
+from unittest.mock import Mock
+
+import pytest
+
+from infection_monkey.exploit.tools.helpers import RAND_SUFFIX_LEN, get_agent_dest_path
+
+
+def _get_host_and_options(os, path):
+    host = Mock()
+    host.os = {"type": os}
+    options = {"dropper_target_path_win_64": path, "dropper_target_path_linux": path}
+    return host, options
+
+
+@pytest.mark.parametrize("os", ["windows", "linux"])
+@pytest.mark.parametrize("path", ["C:\\monkey.exe", "/tmp/monkey-linux-64", "mon.key.exe"])
+def test_get_agent_dest_path(os, path):
+    host, options = _get_host_and_options(os, path)
+    rand_path = get_agent_dest_path(host, options)
+
+    # Assert that filename got longer by RAND_SUFFIX_LEN and one dash
+    assert len(str(rand_path)) == (len(str(path)) + RAND_SUFFIX_LEN + 1)
+
+
+def test_get_agent_dest_path_randomness():
+    host, options = _get_host_and_options("windows", "monkey.exe")
+
+    path1 = get_agent_dest_path(host, options)
+    path2 = get_agent_dest_path(host, options)
+
+    assert not path1 == path2
+
+
+def test_get_agent_dest_path_str_place():
+    host, options = _get_host_and_options("windows", "C:\\abc\\monkey.exe")
+
+    rand_path = get_agent_dest_path(host, options)
+
+    assert str(rand_path).startswith("C:\\abc\\monkey-")
+    assert str(rand_path).endswith(".exe")