From ef1ef3475ba8ec721ec04faea4566bb2e50f6167 Mon Sep 17 00:00:00 2001 From: VakarisZ Date: Wed, 24 Feb 2021 09:52:34 +0200 Subject: [PATCH] Extracted island argument parsing into a separate file --- monkey/monkey_island.py | 44 ++++----------------------- monkey/monkey_island/cc/arg_parser.py | 42 +++++++++++++++++++++++++ 2 files changed, 48 insertions(+), 38 deletions(-) create mode 100644 monkey/monkey_island/cc/arg_parser.py diff --git a/monkey/monkey_island.py b/monkey/monkey_island.py index 837f6f714..e2323a7c3 100644 --- a/monkey/monkey_island.py +++ b/monkey/monkey_island.py @@ -1,5 +1,7 @@ from gevent import monkey as gevent_monkey +from monkey_island.cc.arg_parser import parse_cli_args + gevent_monkey.patch_all() import json # noqa: E402 @@ -11,51 +13,17 @@ from monkey_island.cc.consts import ( from monkey_island.cc.island_logger import json_setup_logging # noqa: E402 -def parse_cli_args(): - import argparse - - parser = argparse.ArgumentParser( - description="Infection Monkey Island CnC Server. See https://infectionmonkey.com", - formatter_class=argparse.ArgumentDefaultsHelpFormatter, - ) - parser.add_argument( - "-s", - "--setup-only", - action="store_true", - help="Pass this flag to cause the Island to setup and exit without actually starting. " - "This is useful for preparing Island to boot faster later-on, so for " - "compiling/packaging Islands.", - ) - parser.add_argument( - "--server-config", - action="store", - help="The path to the server configuration file.", - default=DEFAULT_SERVER_CONFIG_PATH, - ) - parser.add_argument( - "--logger-config", - action="store", - help="The path to the logging configuration file.", - default=DEFAULT_LOGGER_CONFIG_PATH, - ) - args = parser.parse_args() - - return (args.setup_only, args.server_config, args.logger_config) - - if "__main__" == __name__: - # TODO: Address https://github.com/guardicore/monkey/pull/963#discussion_r575022748 - # before merging appimage PR - (is_setup_only, server_config, logger_config) = parse_cli_args() + island_args = parse_cli_args() # This is here in order to catch EVERYTHING, some functions are being called on # imports, so the log init needs to be first. try: - json_setup_logging(logger_config) - except (json.JSONDecodeError) as ex: + json_setup_logging(island_args.logger_config) + except json.JSONDecodeError as ex: print(f"Error loading logging config: {ex}") exit(1) from monkey_island.cc.main import main # noqa: E402 - main(is_setup_only, server_config) + main(island_args.setup_only, island_args.server_config) diff --git a/monkey/monkey_island/cc/arg_parser.py b/monkey/monkey_island/cc/arg_parser.py new file mode 100644 index 000000000..8716a9b4a --- /dev/null +++ b/monkey/monkey_island/cc/arg_parser.py @@ -0,0 +1,42 @@ +from dataclasses import dataclass + +from monkey_island.cc.consts import DEFAULT_SERVER_CONFIG_PATH, DEFAULT_LOGGER_CONFIG_PATH + + +@dataclass +class IslandArgs: + setup_only: bool + server_config: str + logger_config: str + + +def parse_cli_args() -> IslandArgs: + import argparse + + parser = argparse.ArgumentParser( + description="Infection Monkey Island CnC Server. See https://infectionmonkey.com", + formatter_class=argparse.ArgumentDefaultsHelpFormatter + ) + parser.add_argument( + "-s", + "--setup-only", + action="store_true", + help="Pass this flag to cause the Island to setup and exit without actually starting. " + "This is useful for preparing Island to boot faster later-on, so for " + "compiling/packaging Islands.", + ) + parser.add_argument( + "--server-config", + action="store", + help="The path to the server configuration file.", + default=DEFAULT_SERVER_CONFIG_PATH, + ) + parser.add_argument( + "--logger-config", + action="store", + help="The path to the logging configuration file.", + default=DEFAULT_LOGGER_CONFIG_PATH, + ) + args = parser.parse_args() + + return IslandArgs(args.setup_only, args.server_config, args.logger_config)