diff --git a/monkey/infection_monkey/exploit/tools/helpers.py b/monkey/infection_monkey/exploit/tools/helpers.py index 84bd2c52b..87f5636eb 100644 --- a/monkey/infection_monkey/exploit/tools/helpers.py +++ b/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-.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) diff --git a/monkey/tests/unit_tests/infection_monkey/exploit/tools/test_helpers.py b/monkey/tests/unit_tests/infection_monkey/exploit/tools/test_helpers.py new file mode 100644 index 000000000..2c1441961 --- /dev/null +++ b/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")