From 7310cc0b15cdd8cc061492b6fef3f057e698229f Mon Sep 17 00:00:00 2001
From: vakarisz <vakarisz@yahoo.com>
Date: Mon, 30 May 2022 15:58:35 +0300
Subject: [PATCH] Agent: Improve the commandline argument parsing

---
 monkey/infection_monkey/dropper.py |  3 +--
 monkey/infection_monkey/main.py    | 22 ++++++++++++----------
 monkey/infection_monkey/monkey.py  |  2 +-
 3 files changed, 14 insertions(+), 13 deletions(-)

diff --git a/monkey/infection_monkey/dropper.py b/monkey/infection_monkey/dropper.py
index 4bbc272ea..bf5ace48a 100644
--- a/monkey/infection_monkey/dropper.py
+++ b/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]),
diff --git a/monkey/infection_monkey/main.py b/monkey/infection_monkey/main.py
index 1ee8023a3..3c71c9721 100644
--- a/monkey/infection_monkey/main.py
+++ b/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()
diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py
index 7f09adde4..4d2807194 100644
--- a/monkey/infection_monkey/monkey.py
+++ b/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