Agent: Improve the commandline argument parsing

monkey/infection_monkey/dropper.py

@@ -46,8 +46,7 @@ class MonkeyDrops(object):
         arg_parser.add_argument("-d", "--depth", type=int)
         arg_parser.add_argument("-l", "--location")
         arg_parser.add_argument("-vp", "--vulnerable-port")
-        self.monkey_args = args[1:]
-        self.opts, _ = arg_parser.parse_known_args(args)
+        self.opts = arg_parser.parse_args(args)
 
         self._config = {
             "source_path": os.path.abspath(sys.argv[0]),

monkey/infection_monkey/main.py

@@ -44,25 +44,27 @@ LOG_CONFIG = {
 def main():
     global logger
 
-    if 2 > len(sys.argv):
-        return True
     freeze_support()  # required for multiprocessing + pyinstaller on windows
-    monkey_mode = sys.argv[1]
-
-    if not (monkey_mode in [MONKEY_ARG, DROPPER_ARG]):
-        return True
 
     arg_parser = argparse.ArgumentParser()
-    _, monkey_args = arg_parser.parse_known_args(sys.argv[2:])
+    arg_parser.add_argument(
+        "mode",
+        choices=[MONKEY_ARG, DROPPER_ARG],
+        help=f"'{MONKEY_ARG}' mode will run the agent in the current session/terminal."
+        f"'{DROPPER_ARG}' will detach the agent from the current session "
+        f"and will start it on a separate process.",
+    )
+    mode_args, mode_specific_args = arg_parser.parse_known_args()
+    mode = mode_args.mode
 
     formatted_config = pformat(WormConfiguration.hide_sensitive_info(WormConfiguration.as_dict()))
     print(f"Loaded Configuration:\n{formatted_config}")
 
     try:
-        if MONKEY_ARG == monkey_mode:
+        if MONKEY_ARG == mode:
             log_path = get_agent_log_path()
             monkey_cls = InfectionMonkey
-        elif DROPPER_ARG == monkey_mode:
+        elif DROPPER_ARG == mode:
             log_path = get_dropper_log_path()
             monkey_cls = MonkeyDrops
         else:
@@ -98,7 +100,7 @@ def main():
     logger.info(f"version: {get_version()}")
     logger.info(f"writing log file to {log_path}")
 
-    monkey = monkey_cls(monkey_args)
+    monkey = monkey_cls(mode_specific_args)
 
     try:
         monkey.start()

monkey/infection_monkey/monkey.py

@@ -102,7 +102,7 @@ class InfectionMonkey:
         arg_parser.add_argument("-t", "--tunnel")
         arg_parser.add_argument("-s", "--server")
         arg_parser.add_argument("-d", "--depth", type=int)
-        opts, _ = arg_parser.parse_known_args(args)
+        opts = arg_parser.parse_args(args)
         InfectionMonkey._log_arguments(opts)
 
         return opts