diff --git a/monkey/common/utils/argparse_types.py b/monkey/common/utils/argparse_types.py new file mode 100644 index 000000000..a920ef1b4 --- /dev/null +++ b/monkey/common/utils/argparse_types.py @@ -0,0 +1,9 @@ +import argparse + + +def positive_int(_input: str): + int_value = int(_input) + if int_value <= 0: + raise argparse.ArgumentTypeError(f"{_input} is not a positive integer") + + return int_value diff --git a/monkey/infection_monkey/dropper.py b/monkey/infection_monkey/dropper.py index 8e5fb0f5c..48e5bb86d 100644 --- a/monkey/infection_monkey/dropper.py +++ b/monkey/infection_monkey/dropper.py @@ -9,6 +9,7 @@ import sys import time from pathlib import PosixPath, WindowsPath +from common.utils.argparse_types import positive_int from common.utils.attack_utils import UsageEnum from infection_monkey.utils.commands import ( build_monkey_commandline_explicitly, @@ -45,7 +46,7 @@ class MonkeyDrops(object): arg_parser.add_argument("-p", "--parent") arg_parser.add_argument("-t", "--tunnel") arg_parser.add_argument("-s", "--server") - arg_parser.add_argument("-d", "--depth", type=int) + arg_parser.add_argument("-d", "--depth", type=positive_int) arg_parser.add_argument("-l", "--location") arg_parser.add_argument("-vp", "--vulnerable-port") self.opts = arg_parser.parse_args(args) diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index 4eeec8b76..6e3d9b0f0 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -8,6 +8,7 @@ from typing import List import infection_monkey.tunnel as tunnel from common.network.network_utils import address_to_ip_port +from common.utils.argparse_types import positive_int from common.utils.attack_utils import ScanStatus, UsageEnum from common.version import get_version from infection_monkey.config import GUID @@ -104,7 +105,7 @@ class InfectionMonkey: arg_parser.add_argument("-p", "--parent") arg_parser.add_argument("-t", "--tunnel") arg_parser.add_argument("-s", "--server") - arg_parser.add_argument("-d", "--depth", type=int) + arg_parser.add_argument("-d", "--depth", type=positive_int) opts = arg_parser.parse_args(args) InfectionMonkey._log_arguments(opts)