Added "-s" flag for setup-only, useful for building Islands

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)

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:

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)