From 28f2d32579924b99dec4fd54ec3aa84f0ee4d40d Mon Sep 17 00:00:00 2001
From: Shay Nehmad <shay.nehmad@guardicore.com>
Date: Wed, 5 Jun 2019 19:25:09 +0300
Subject: [PATCH] Added assertion of mongo db version.

Also refactoed a bit to make the main shorter.
---
 monkey/monkey_island/cc/database.py | 11 ++++++++++
 monkey/monkey_island/cc/main.py     | 32 ++++++++++++++++++++++++-----
 2 files changed, 38 insertions(+), 5 deletions(-)

diff --git a/monkey/monkey_island/cc/database.py b/monkey/monkey_island/cc/database.py
index 8fb3b120b..082553e5f 100644
--- a/monkey/monkey_island/cc/database.py
+++ b/monkey/monkey_island/cc/database.py
@@ -25,3 +25,14 @@ def is_db_server_up(mongo_url):
         return True
     except ServerSelectionTimeoutError:
         return False
+
+
+def get_db_version(mongo_url):
+    """
+    Return the mongo db version
+    :param mongo_url: Which mongo to check.
+    :return: version as a tuple (e.g. `(u'4', u'0', u'8')`)
+    """
+    client = MongoClient(mongo_url, serverSelectionTimeoutMS=100)
+    server_version = tuple(client.server_info()['version'].split('.'))
+    return server_version
diff --git a/monkey/monkey_island/cc/main.py b/monkey/monkey_island/cc/main.py
index 412c3c399..103280a05 100644
--- a/monkey/monkey_island/cc/main.py
+++ b/monkey/monkey_island/cc/main.py
@@ -6,6 +6,8 @@ import sys
 import time
 import logging
 
+MINIMUM_MONGO_DB_VERSION_REQUIRED = "4.0.0"
+
 BASE_PATH = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
 
 if BASE_PATH not in sys.path:
@@ -22,7 +24,7 @@ from monkey_island.cc.app import init_app
 from monkey_island.cc.exporter_init import populate_exporter_list
 from monkey_island.cc.utils import local_ip_addresses
 from monkey_island.cc.environment.environment import env
-from monkey_island.cc.database import is_db_server_up
+from monkey_island.cc.database import is_db_server_up, get_db_version
 
 
 def main():
@@ -31,10 +33,8 @@ def main():
     from tornado.ioloop import IOLoop
 
     mongo_url = os.environ.get('MONGO_URL', env.get_mongo_url())
-
-    while not is_db_server_up(mongo_url):
-        logger.info('Waiting for MongoDB server')
-        time.sleep(1)
+    wait_for_mongo_db_server(mongo_url)
+    assert_mongo_db_version(mongo_url)
 
     populate_exporter_list()
     app = init_app(mongo_url)
@@ -55,5 +55,27 @@ def main():
         IOLoop.instance().start()
 
 
+def wait_for_mongo_db_server(mongo_url):
+    while not is_db_server_up(mongo_url):
+        logger.info('Waiting for MongoDB server on {0}'.format(mongo_url))
+        time.sleep(1)
+
+
+def assert_mongo_db_version(mongo_url):
+    """
+    Checks if the mongodb version is new enough for running the app.
+    If the DB is too old, quits.
+    :param mongo_url: URL to the mongo the Island will use
+    """
+    required_version = tuple(MINIMUM_MONGO_DB_VERSION_REQUIRED.split("."))
+    server_version = get_db_version(mongo_url)
+    if server_version < required_version:
+        logger.error(
+            'Mongo DB version too old. {0} is required, but got {1}'.format(str(required_version), str(server_version)))
+        sys.exit(-1)
+    else:
+        logger.info('Mongo DB version OK. Got {0}'.format(str(server_version)))
+
+
 if __name__ == '__main__':
     main()