From 16a470ee2279eb4d7093add3acd2d9d9baa4410e Mon Sep 17 00:00:00 2001
From: Itay Mizeretz <itay.mizeretz@guardicore.com>
Date: Tue, 31 Oct 2017 13:55:29 +0200
Subject: [PATCH] Add wait for mongodb

---
 monkey_island/cc/database.py | 11 +++++++++++
 monkey_island/cc/main.py     | 12 +++++++++++-
 2 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/monkey_island/cc/database.py b/monkey_island/cc/database.py
index 31828705a..fadff853d 100644
--- a/monkey_island/cc/database.py
+++ b/monkey_island/cc/database.py
@@ -1,5 +1,16 @@
 from flask_pymongo import PyMongo
+from flask_pymongo import MongoClient
+from pymongo.errors import ServerSelectionTimeoutError
 
 __author__ = 'Barak'
 
 mongo = PyMongo()
+
+
+def is_db_server_up(mongo_url):
+    client = MongoClient(mongo_url, serverSelectionTimeoutMS=100)
+    try:
+        client.server_info()
+        return True
+    except ServerSelectionTimeoutError:
+        return False
diff --git a/monkey_island/cc/main.py b/monkey_island/cc/main.py
index dd133cfd1..20885f47b 100644
--- a/monkey_island/cc/main.py
+++ b/monkey_island/cc/main.py
@@ -3,6 +3,8 @@ from __future__ import print_function  # In python 2.7
 import os
 import sys
 
+import time
+
 BASE_PATH = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
 if BASE_PATH not in sys.path:
     sys.path.insert(0, BASE_PATH)
@@ -10,13 +12,21 @@ if BASE_PATH not in sys.path:
 from cc.app import init_app
 from cc.utils import local_ip_addresses
 from cc.island_config import DEFAULT_MONGO_URL, ISLAND_PORT
+from cc.database import is_db_server_up
 
 if __name__ == '__main__':
     from tornado.wsgi import WSGIContainer
     from tornado.httpserver import HTTPServer
     from tornado.ioloop import IOLoop
 
-    app = init_app(os.environ.get('MONGO_URL', DEFAULT_MONGO_URL))
+    mongo_url = os.environ.get('MONGO_URL', DEFAULT_MONGO_URL)
+
+    while not is_db_server_up(mongo_url):
+        print('Waiting for MongoDB server')
+        time.sleep(1)
+
+    app = init_app(mongo_url)
+
     http_server = HTTPServer(WSGIContainer(app),
                              ssl_options={'certfile': os.environ.get('SERVER_CRT', 'server.crt'),
                                           'keyfile': os.environ.get('SERVER_KEY', 'server.key')})