Agent, Common: Add positive_int argument type

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

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)

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)