diff --git a/monkey/infection_monkey/utils/commands.py b/monkey/infection_monkey/utils/commands.py index c2c7e188c..c290b3893 100644 --- a/monkey/infection_monkey/utils/commands.py +++ b/monkey/infection_monkey/utils/commands.py @@ -1,19 +1,20 @@ +from typing import List, Optional + from infection_monkey.config import GUID from infection_monkey.exploit.tools.helpers import AGENT_BINARY_PATH_LINUX, AGENT_BINARY_PATH_WIN64 from infection_monkey.model import CMD_CARRY_OUT, CMD_EXE, MONKEY_ARG -from infection_monkey.model.host import VictimHost # Dropper target paths DROPPER_TARGET_PATH_LINUX = AGENT_BINARY_PATH_LINUX DROPPER_TARGET_PATH_WIN64 = AGENT_BINARY_PATH_WIN64 -def build_monkey_commandline(target_host: VictimHost, depth: int, location: str = None) -> str: +def build_monkey_commandline(servers: List[str], depth: int, location: Optional[str] = None) -> str: return " " + " ".join( build_monkey_commandline_explicitly( GUID, - target_host.default_server, + servers, depth, location, ) @@ -21,19 +22,19 @@ def build_monkey_commandline(target_host: VictimHost, depth: int, location: str def build_monkey_commandline_explicitly( - parent: str = None, - server: str = None, - depth: int = None, - location: str = None, -) -> list: + parent: Optional[str] = None, + servers: Optional[List[str]] = None, + depth: Optional[int] = None, + location: Optional[str] = None, +) -> List[str]: cmdline = [] if parent is not None: cmdline.append("-p") cmdline.append(str(parent)) - if server is not None: + if servers: cmdline.append("-s") - cmdline.append(str(server)) + cmdline.append(",".join(servers)) if depth is not None: cmdline.append("-d") cmdline.append(str(depth)) @@ -44,13 +45,13 @@ def build_monkey_commandline_explicitly( return cmdline -def get_monkey_commandline_windows(destination_path: str, monkey_cmd_args: list) -> list: +def get_monkey_commandline_windows(destination_path: str, monkey_cmd_args: List[str]) -> List[str]: monkey_cmdline = [CMD_EXE, CMD_CARRY_OUT, destination_path, MONKEY_ARG] return monkey_cmdline + monkey_cmd_args -def get_monkey_commandline_linux(destination_path: str, monkey_cmd_args: list) -> list: +def get_monkey_commandline_linux(destination_path: str, monkey_cmd_args: List[str]) -> List[str]: monkey_cmdline = [destination_path, MONKEY_ARG] return monkey_cmdline + monkey_cmd_args diff --git a/monkey/tests/unit_tests/infection_monkey/utils/test_commands.py b/monkey/tests/unit_tests/infection_monkey/utils/test_commands.py index f6bbcb691..a20fd9e03 100644 --- a/monkey/tests/unit_tests/infection_monkey/utils/test_commands.py +++ b/monkey/tests/unit_tests/infection_monkey/utils/test_commands.py @@ -1,5 +1,6 @@ +import pytest + from infection_monkey.config import GUID -from infection_monkey.model.host import VictimHost from infection_monkey.utils.commands import ( build_monkey_commandline, build_monkey_commandline_explicitly, @@ -13,14 +14,14 @@ def test_build_monkey_commandline_explicitly_arguments(): "-p", "101010", "-s", - "127.127.127.127:5000", + "127.127.127.127:5000,138.138.138.138:5007", "-d", "0", "-l", "C:\\windows\\abc", ] actual = build_monkey_commandline_explicitly( - "101010", "127.127.127.127:5000", 0, "C:\\windows\\abc" + "101010", ["127.127.127.127:5000", "138.138.138.138:5007"], 0, "C:\\windows\\abc" ) assert expected == actual @@ -44,13 +45,12 @@ def test_get_monkey_commandline_windows(): "m0nk3y", "-p", "101010", + "-s", + "127.127.127.127:5000,138.138.138.138:5007", ] actual = get_monkey_commandline_windows( "C:\\windows\\abc", - [ - "-p", - "101010", - ], + ["-p", "101010", "-s", "127.127.127.127:5000,138.138.138.138:5007"], ) assert expected == actual @@ -62,23 +62,30 @@ def test_get_monkey_commandline_linux(): "m0nk3y", "-p", "101010", + "-s", + "127.127.127.127:5000,138.138.138.138:5007", ] actual = get_monkey_commandline_linux( "/home/user/monkey-linux-64", - [ - "-p", - "101010", - ], + ["-p", "101010", "-s", "127.127.127.127:5000,138.138.138.138:5007"], ) assert expected == actual def test_build_monkey_commandline(): - example_host = VictimHost(ip_addr="bla") - example_host.set_island_address("101010", "5000") + servers = ["10.10.10.10:5000", "11.11.11.11:5007"] - expected = f" -p {GUID} -s 101010:5000 -d 0 -l /home/bla" - actual = build_monkey_commandline(target_host=example_host, depth=0, location="/home/bla") + expected = f" -p {GUID} -s 10.10.10.10:5000,11.11.11.11:5007 -d 0 -l /home/bla" + actual = build_monkey_commandline(servers=servers, depth=0, location="/home/bla") + + assert expected == actual + + +@pytest.mark.parametrize("servers", [None, []]) +def test_build_monkey_commandline_empty_servers(servers): + + expected = f" -p {GUID} -d 0 -l /home/bla" + actual = build_monkey_commandline(servers, depth=0, location="/home/bla") assert expected == actual