diff --git a/CHANGELOG.md b/CHANGELOG.md index 27bf5e27f..e1d37d08b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -73,6 +73,7 @@ Changelog](https://keepachangelog.com/en/1.0.0/). - Island log download button from "Telemetries"(previously called "Logs") page. #1640 - "/api/client-monkey" endpoint. #1889 - "+dev" from version numbers. #1553 +- agent's "--config" argument. #906 ### Fixed - A bug in network map page that caused delay of telemetry log loading. #1545 diff --git a/monkey/infection_monkey/config.py b/monkey/infection_monkey/config.py index 1f8c46311..ddf224b10 100644 --- a/monkey/infection_monkey/config.py +++ b/monkey/infection_monkey/config.py @@ -1,12 +1,8 @@ -import os -import sys import uuid from abc import ABCMeta GUID = str(uuid.getnode()) -EXTERNAL_CONFIG_FILE = os.path.join(os.path.abspath(os.path.dirname(sys.argv[0])), "monkey.bin") - SENSITIVE_FIELDS = [ "exploit_password_list", "exploit_user_list", 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 74961e0ad..3c71c9721 100644 --- a/monkey/infection_monkey/main.py +++ b/monkey/infection_monkey/main.py @@ -1,5 +1,4 @@ import argparse -import json import logging import logging.config import os @@ -12,7 +11,7 @@ from pprint import pformat # noinspection PyUnresolvedReferences import infection_monkey.post_breach # noqa: F401 from common.version import get_version -from infection_monkey.config import EXTERNAL_CONFIG_FILE, WormConfiguration +from infection_monkey.config import WormConfiguration from infection_monkey.dropper import MonkeyDrops from infection_monkey.model import DROPPER_ARG, MONKEY_ARG from infection_monkey.monkey import InfectionMonkey @@ -45,44 +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 - - config_file = EXTERNAL_CONFIG_FILE arg_parser = argparse.ArgumentParser() - arg_parser.add_argument("-c", "--config") - opts, monkey_args = arg_parser.parse_known_args(sys.argv[2:]) - if opts.config: - config_file = opts.config - if os.path.isfile(config_file): - # using print because config can also change log locations - print("Loading config from %s." % config_file) - try: - with open(config_file) as config_fo: - json_dict = json.load(config_fo) - WormConfiguration.from_kv(json_dict) - except ValueError as e: - print("Error loading config: %s, using default" % (e,)) - else: - print( - "Config file wasn't supplied and default path: %s wasn't found, using internal " - "default" % (config_file,) - ) + 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: @@ -118,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