agent: Add type hinting to commands

This commit is contained in:
Ilija Lazoroski 2021-06-21 20:31:50 +02:00 committed by Ilija Lazoroski
parent e93df01e69
commit 680b1f54d0
2 changed files with 14 additions and 6 deletions

View File

@ -1,7 +1,10 @@
from infection_monkey.model import CMD_CARRY_OUT, CMD_EXE, MONKEY_ARG from infection_monkey.model import CMD_CARRY_OUT, CMD_EXE, MONKEY_ARG
from infection_monkey.model.host import VictimHost
def build_monkey_commandline(target_host, depth, vulnerable_port, location=None): def build_monkey_commandline(
target_host: VictimHost, depth: int, vulnerable_port: str, location: str = None
) -> str:
from infection_monkey.config import GUID from infection_monkey.config import GUID
return "".join( return "".join(
@ -17,8 +20,13 @@ def build_monkey_commandline(target_host, depth, vulnerable_port, location=None)
def build_monkey_commandline_explicitly( def build_monkey_commandline_explicitly(
parent=None, tunnel=None, server=None, depth=None, location=None, vulnerable_port=None parent: str = None,
): tunnel: str = None,
server: str = None,
depth: int = None,
location: str = None,
vulnerable_port: str = None,
) -> list:
cmdline = [] cmdline = []
if parent is not None: if parent is not None:
@ -45,13 +53,13 @@ def build_monkey_commandline_explicitly(
return cmdline return cmdline
def get_monkey_commandline_windows(destination_path, monkey_options): def get_monkey_commandline_windows(destination_path: str, monkey_options: list) -> list:
monkey_cmdline = [CMD_EXE, CMD_CARRY_OUT, destination_path, MONKEY_ARG] monkey_cmdline = [CMD_EXE, CMD_CARRY_OUT, destination_path, MONKEY_ARG]
return monkey_cmdline + monkey_options return monkey_cmdline + monkey_options
def get_monkey_commandline_linux(destination_path, monkey_options): def get_monkey_commandline_linux(destination_path: str, monkey_options: list) -> list:
monkey_cmdline = [destination_path, MONKEY_ARG] monkey_cmdline = [destination_path, MONKEY_ARG]
return monkey_cmdline + monkey_options return monkey_cmdline + monkey_options

View File

@ -21,7 +21,7 @@ def test_build_monkey_commandline_explicitly_arguments():
"80", "80",
] ]
actual = build_monkey_commandline_explicitly( actual = build_monkey_commandline_explicitly(
101010, "10.10.101.10", "127.127.127.127:5000", 0, "C:\\windows\\abc", 80 "101010", "10.10.101.10", "127.127.127.127:5000", 0, "C:\\windows\\abc", "80"
) )
assert expected == actual assert expected == actual