Agent: Modify agent build command line to accept list of servers

This commit is contained in:
Ilija Lazoroski 2022-09-02 14:35:13 +02:00 committed by Mike Salvatore
parent ddc4f4d836
commit 067d50f3c4
2 changed files with 35 additions and 27 deletions

View File

@ -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

View File

@ -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