From d8946feb69b5d26f1298df0d6b5bd6f89fe00703 Mon Sep 17 00:00:00 2001
From: Itay Mizeretz <itay.mizeretz@guardicore.com>
Date: Tue, 27 Feb 2018 19:13:28 +0200
Subject: [PATCH] Fix CR

---
 monkey_island/cc/app.py            |  3 ++-
 monkey_island/cc/database.py       | 15 +++++++++++++--
 monkey_island/cc/resources/log.py  |  2 +-
 monkey_island/cc/resources/root.py |  3 +--
 monkey_island/cc/services/log.py   | 18 +++++++-----------
 5 files changed, 24 insertions(+), 17 deletions(-)

diff --git a/monkey_island/cc/app.py b/monkey_island/cc/app.py
index 2d8041eb0..1d632f6fa 100644
--- a/monkey_island/cc/app.py
+++ b/monkey_island/cc/app.py
@@ -6,7 +6,7 @@ from bson.json_util import dumps
 from flask import Flask, send_from_directory, make_response
 from werkzeug.exceptions import NotFound
 
-from cc.database import mongo
+from cc.database import mongo, database
 from cc.resources.client_run import ClientRun
 from cc.resources.edge import Edge
 from cc.resources.local_run import LocalRun
@@ -75,6 +75,7 @@ def init_app(mongo_url):
     mongo.init_app(app)
 
     with app.app_context():
+        database.init()
         ConfigService.init_config()
 
     app.add_url_rule('/', 'serve_home', serve_home)
diff --git a/monkey_island/cc/database.py b/monkey_island/cc/database.py
index fadff853d..8fb3b120b 100644
--- a/monkey_island/cc/database.py
+++ b/monkey_island/cc/database.py
@@ -1,5 +1,5 @@
-from flask_pymongo import PyMongo
-from flask_pymongo import MongoClient
+import gridfs
+from flask_pymongo import MongoClient, PyMongo
 from pymongo.errors import ServerSelectionTimeoutError
 
 __author__ = 'Barak'
@@ -7,6 +7,17 @@ __author__ = 'Barak'
 mongo = PyMongo()
 
 
+class Database:
+    def __init__(self):
+        self.gridfs = None
+
+    def init(self):
+        self.gridfs = gridfs.GridFS(mongo.db)
+
+
+database = Database()
+
+
 def is_db_server_up(mongo_url):
     client = MongoClient(mongo_url, serverSelectionTimeoutMS=100)
     try:
diff --git a/monkey_island/cc/resources/log.py b/monkey_island/cc/resources/log.py
index e49c7024d..76e802f0a 100644
--- a/monkey_island/cc/resources/log.py
+++ b/monkey_island/cc/resources/log.py
@@ -24,7 +24,7 @@ class Log(flask_restful.Resource):
         telemetry_json = json.loads(request.data)
 
         monkey_id = NodeService.get_monkey_by_guid(telemetry_json['monkey_guid'])['_id']
-        # This is base64 so no data will be lost. this'll take 2 time less space.
+        # This shouldn't contain any unicode characters. this'll take 2 time less space.
         log_data = str(telemetry_json['log'])
         log_id = LogService.add_log(monkey_id, log_data)
 
diff --git a/monkey_island/cc/resources/root.py b/monkey_island/cc/resources/root.py
index f725163aa..3acb870cf 100644
--- a/monkey_island/cc/resources/root.py
+++ b/monkey_island/cc/resources/root.py
@@ -33,8 +33,7 @@ class Root(flask_restful.Resource):
 
     @staticmethod
     def reset_db():
-        [mongo.db[x].drop() for x in
-         ['config', 'monkey', 'telemetry', 'node', 'edge', 'report', 'log', 'fs.chunks', 'fs.files']]
+        [mongo.db[x].drop() for x in mongo.db.collection_names()]
         ConfigService.init_config()
         return jsonify(status='OK')
 
diff --git a/monkey_island/cc/services/log.py b/monkey_island/cc/services/log.py
index 2783a1dfa..81603e62e 100644
--- a/monkey_island/cc/services/log.py
+++ b/monkey_island/cc/services/log.py
@@ -1,9 +1,7 @@
 from datetime import datetime
 
-import gridfs
-
 import cc.services.node
-from cc.database import mongo
+from cc.database import mongo, database
 
 __author__ = "itay.mizeretz"
 
@@ -16,8 +14,7 @@ class LogService:
     def get_log_by_monkey_id(monkey_id):
         log = mongo.db.log.find_one({'monkey_id': monkey_id})
         if log:
-            fs = gridfs.GridFS(mongo.db)
-            log_file = fs.get(log['file_id'])
+            log_file = database.gridfs.get(log['file_id'])
             monkey_label = cc.services.node.NodeService.get_monkey_label(
                 cc.services.node.NodeService.get_monkey_by_id(log['monkey_id']))
             return \
@@ -29,16 +26,15 @@ class LogService:
 
     @staticmethod
     def remove_logs_by_monkey_id(monkey_id):
-        fs = gridfs.GridFS(mongo.db)
-        for log in mongo.db.log.find({'monkey_id': monkey_id}):
-            fs.delete(log['file_id'])
-        mongo.db.log.delete_many({'monkey_id': monkey_id})
+        log = mongo.db.log.find_one({'monkey_id': monkey_id})
+        if log is not None:
+            database.gridfs.delete(log['file_id'])
+            mongo.db.log.delete_one({'monkey_id': monkey_id})
 
     @staticmethod
     def add_log(monkey_id, log_data, timestamp=datetime.now()):
         LogService.remove_logs_by_monkey_id(monkey_id)
-        fs = gridfs.GridFS(mongo.db)
-        file_id = fs.put(log_data)
+        file_id = database.gridfs.put(log_data)
         return mongo.db.log.insert(
             {
                 'monkey_id': monkey_id,