From 2e6b361a9d35ba06815522635430e639601c25c0 Mon Sep 17 00:00:00 2001 From: vakaris_zilius Date: Wed, 23 Mar 2022 10:49:25 +0000 Subject: [PATCH 1/3] 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 --- .../infection_monkey/exploit/tools/helpers.py | 24 +++++++++-- .../exploit/tools/test_helpers.py | 40 +++++++++++++++++++ 2 files changed, 60 insertions(+), 4 deletions(-) create mode 100644 monkey/tests/unit_tests/infection_monkey/exploit/tools/test_helpers.py diff --git a/monkey/infection_monkey/exploit/tools/helpers.py b/monkey/infection_monkey/exploit/tools/helpers.py index 84bd2c52b..5e9406160 100644 --- a/monkey/infection_monkey/exploit/tools/helpers.py +++ b/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-.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 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..00e44d2e4 --- /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 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") From efb0039e3405917523ddef3082d08706dfb6e944 Mon Sep 17 00:00:00 2001 From: vakaris_zilius Date: Wed, 23 Mar 2022 13:33:26 +0000 Subject: [PATCH 2/3] Agent: Make _add_random_suffix method code more concise --- monkey/infection_monkey/exploit/tools/helpers.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/monkey/infection_monkey/exploit/tools/helpers.py b/monkey/infection_monkey/exploit/tools/helpers.py index 5e9406160..87f5636eb 100644 --- a/monkey/infection_monkey/exploit/tools/helpers.py +++ b/monkey/infection_monkey/exploit/tools/helpers.py @@ -31,8 +31,6 @@ def get_agent_dest_path(host: VictimHost, options: Mapping[str, Any]) -> Path: # 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 + rand_filename = "".join([stem, *path.suffixes]) + return path.with_name(rand_filename) From c2b06f22f045ceac86e9c9f19e4d0615e44fb98c Mon Sep 17 00:00:00 2001 From: vakaris_zilius Date: Wed, 23 Mar 2022 13:37:33 +0000 Subject: [PATCH 3/3] Agent: Improve path comparison style in test_helpers.py --- .../unit_tests/infection_monkey/exploit/tools/test_helpers.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 index 00e44d2e4..2c1441961 100644 --- a/monkey/tests/unit_tests/infection_monkey/exploit/tools/test_helpers.py +++ b/monkey/tests/unit_tests/infection_monkey/exploit/tools/test_helpers.py @@ -28,7 +28,7 @@ def test_get_agent_dest_path_randomness(): path1 = get_agent_dest_path(host, options) path2 = get_agent_dest_path(host, options) - assert not path1 == path2 + assert path1 != path2 def test_get_agent_dest_path_str_place():