From 224861d0848cfd91f5c0fbd34b9a3ffe6c86fd02 Mon Sep 17 00:00:00 2001 From: Shay Nehmad Date: Tue, 31 Mar 2020 20:15:34 +0300 Subject: [PATCH] Added "-s" flag for setup-only, useful for building Islands --- monkey/monkey_island.py | 14 +++++++++++++- monkey/monkey_island/cc/main.py | 6 +++++- monkey/monkey_island/cc/setup.py | 10 ++++++++-- 3 files changed, 26 insertions(+), 4 deletions(-) diff --git a/monkey/monkey_island.py b/monkey/monkey_island.py index 8f31c83f8..b3b74e6b3 100644 --- a/monkey/monkey_island.py +++ b/monkey/monkey_island.py @@ -1,4 +1,16 @@ from monkey_island.cc.main import main + +def parse_cli_args(): + import argparse + parser = argparse.ArgumentParser(description="Infection Monkey Island CnC Server. See https://infectionmonkey.com") + 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.") + args = parser.parse_args() + return args.setup_only + + if "__main__" == __name__: - main() + is_setup_only = parse_cli_args() + main(is_setup_only) diff --git a/monkey/monkey_island/cc/main.py b/monkey/monkey_island/cc/main.py index 045f5452e..7a49b447e 100644 --- a/monkey/monkey_island/cc/main.py +++ b/monkey/monkey_island/cc/main.py @@ -29,7 +29,7 @@ from common.version import get_version from monkey_island.cc.setup import setup -def main(): +def main(should_setup_only): from tornado.wsgi import WSGIContainer from tornado.httpserver import HTTPServer from tornado.ioloop import IOLoop @@ -46,6 +46,10 @@ def main(): setup() + if should_setup_only: + logger.warning("Setup only flag passed. Exiting.") + return + if env.is_debug(): app.run(host='0.0.0.0', debug=True, ssl_context=(crt_path, key_path)) else: diff --git a/monkey/monkey_island/cc/setup.py b/monkey/monkey_island/cc/setup.py index e1b53affa..4953f950c 100644 --- a/monkey/monkey_island/cc/setup.py +++ b/monkey/monkey_island/cc/setup.py @@ -1,10 +1,16 @@ +import logging + from monkey_island.cc.services.attack.mitre_api_interface import MitreApiInterface -from cc.models.attack.attack_mitigations import AttackMitigations +from monkey_island.cc.models.attack.attack_mitigations import AttackMitigations from monkey_island.cc.database import mongo from pymongo import errors +logger = logging.getLogger(__name__) + + def setup(): + logger.info("Setting up the Monkey Island, this might take a while...") try_store_mitigations_on_mongo() @@ -13,7 +19,7 @@ def try_store_mitigations_on_mongo(): try: mongo.db.validate_collection(mitigation_collection_name) if mongo.db.attack_mitigations.count() == 0: - raise errors.OperationFailure("Mitigation collection empty") + raise errors.OperationFailure("Mitigation collection empty. Try dropping the collection and running again") except errors.OperationFailure: try: mongo.db.create_collection(mitigation_collection_name)