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 791fe384c..09f079c19 100644 --- a/monkey/monkey_island/cc/main.py +++ b/monkey/monkey_island/cc/main.py @@ -31,17 +31,17 @@ from monkey_island.cc.bootloader_server import BootloaderHttpServer from monkey_island.cc.setup import setup -def main(): +def main(should_setup_only): logger.info("Starting bootloader server") mongo_url = os.environ.get('MONGO_URL', env.get_mongo_url()) bootloader_server_thread = Thread(target=BootloaderHttpServer(mongo_url).serve_forever, daemon=True) bootloader_server_thread.start() - start_island_server() + start_island_server(should_setup_only) bootloader_server_thread.join() -def start_island_server(): +def start_island_server(should_setup_only): from tornado.wsgi import WSGIContainer from tornado.httpserver import HTTPServer from tornado.ioloop import IOLoop @@ -58,6 +58,10 @@ def start_island_server(): 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 c37c24f3b..007fca6ef 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)