From 680b1f54d0e137d711ad7e9bd44ea4a5b8a0ef23 Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Mon, 21 Jun 2021 20:31:50 +0200 Subject: [PATCH] agent: Add type hinting to commands --- monkey/infection_monkey/utils/commands.py | 18 +++++++++++++----- .../infection_monkey/utils/test_commands.py | 2 +- 2 files changed, 14 insertions(+), 6 deletions(-) diff --git a/monkey/infection_monkey/utils/commands.py b/monkey/infection_monkey/utils/commands.py index e94f1124e..5742389f0 100644 --- a/monkey/infection_monkey/utils/commands.py +++ b/monkey/infection_monkey/utils/commands.py @@ -1,7 +1,10 @@ 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 return "".join( @@ -17,8 +20,13 @@ def build_monkey_commandline(target_host, depth, vulnerable_port, location=None) 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 = [] if parent is not None: @@ -45,13 +53,13 @@ def build_monkey_commandline_explicitly( 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] 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] return monkey_cmdline + monkey_options 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 4eebd0ee4..c43d02e41 100644 --- a/monkey/tests/unit_tests/infection_monkey/utils/test_commands.py +++ b/monkey/tests/unit_tests/infection_monkey/utils/test_commands.py @@ -21,7 +21,7 @@ def test_build_monkey_commandline_explicitly_arguments(): "80", ] 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