Merge pull request #1802 from guardicore/1782-random-agent-filename

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

monkey/infection_monkey/exploit/tools/helpers.py

@@ -1,22 +1,36 @@
 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]
+    stem = f"{stem}-{get_random_file_suffix()}"
+    rand_filename = "".join([stem, *path.suffixes])
+    return path.with_name(rand_filename)

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 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")