Merge pull request #449 from guardicore/440/bugfix/api-status-queries-trigger-multiple-report-generations

442/bugfix/api status queries trigger multiple report generations

monkey/monkey_island/cc/resources/root.py

@@ -1,5 +1,6 @@
 from datetime import datetime
 import logging
+import threading
 
 import flask_restful
 from flask import request, make_response, jsonify
@@ -9,6 +10,7 @@ from monkey_island.cc.database import mongo
 from monkey_island.cc.services.node import NodeService
 from monkey_island.cc.services.reporting.report import ReportService
 from monkey_island.cc.services.attack.attack_report import AttackReportService
+from monkey_island.cc.services.reporting.report_generation_synchronisation import is_report_being_generated, safe_generate_reports
 from monkey_island.cc.utils import local_ip_addresses
 from monkey_island.cc.services.database import Database
 
@@ -18,13 +20,15 @@ logger = logging.getLogger(__name__)
 
 
 class Root(flask_restful.Resource):
+    def __init__(self):
+        self.report_generating_lock = threading.Event()
 
     def get(self, action=None):
         if not action:
             action = request.args.get('action')
 
         if not action:
-            return Root.get_server_info()
+            return self.get_server_info()
         elif action == "reset":
             return jwt_required()(Database.reset_db)()
         elif action == "killall":
@@ -34,11 +38,12 @@ class Root(flask_restful.Resource):
         else:
             return make_response(400, {'error': 'unknown action'})
 
-    @staticmethod
     @jwt_required()
-    def get_server_info():
-        return jsonify(ip_addresses=local_ip_addresses(), mongo=str(mongo.db),
-                       completed_steps=Root.get_completed_steps())
+    def get_server_info(self):
+        return jsonify(
+            ip_addresses=local_ip_addresses(),
+            mongo=str(mongo.db),
+            completed_steps=self.get_completed_steps())
 
     @staticmethod
     @jwt_required()
@@ -49,17 +54,22 @@ class Root(flask_restful.Resource):
         logger.info('Kill all monkeys was called')
         return jsonify(status='OK')
 
-    @staticmethod
     @jwt_required()
-    def get_completed_steps():
+    def get_completed_steps(self):
         is_any_exists = NodeService.is_any_monkey_exists()
         infection_done = NodeService.is_monkey_finished_running()
-        if not infection_done:
-            report_done = False
-        else:
-            if is_any_exists:
-                ReportService.get_report()
-                AttackReportService.get_latest_report()
+
+        if infection_done:
+            # Checking is_report_being_generated here, because we don't want to wait to generate a report; rather,
+            # we want to skip and reply.
+            if not is_report_being_generated() and not ReportService.is_latest_report_exists():
+                safe_generate_reports()
             report_done = ReportService.is_report_generated()
-        return dict(run_server=True, run_monkey=is_any_exists, infection_done=infection_done,
-                    report_done=report_done)
+        else:  # Infection is not done
+            report_done = False
+
+        return dict(
+            run_server=True,
+            run_monkey=is_any_exists,
+            infection_done=infection_done,
+            report_done=report_done)

monkey/monkey_island/cc/services/attack/attack_report.py

@@ -6,6 +6,7 @@ from monkey_island.cc.services.attack.technique_reports import T1145, T1105, T10
 from monkey_island.cc.services.attack.technique_reports import T1090, T1041, T1222, T1005, T1018, T1016, T1021, T1064
 from monkey_island.cc.services.attack.attack_config import AttackConfig
 from monkey_island.cc.database import mongo
+from monkey_island.cc.services.reporting.report_generation_synchronisation import safe_generate_attack_report
 
 __author__ = "VakarisZ"
 
@@ -88,7 +89,8 @@ class AttackReportService:
             report_modifytime = latest_report['meta']['latest_monkey_modifytime']
             if monkey_modifytime and report_modifytime and monkey_modifytime == report_modifytime:
                 return latest_report
-        return AttackReportService.generate_new_report()
+
+        return safe_generate_attack_report()
 
     @staticmethod
     def is_report_generated():

monkey/monkey_island/cc/services/reporting/report.py

@@ -1,25 +1,24 @@
-import itertools
 import functools
-
-import ipaddress
+import itertools
 import logging
 
+import ipaddress
 from bson import json_util
 from enum import Enum
-
 from six import text_type
 
+from common.network.network_range import NetworkRange
 from common.network.segmentation_utils import get_ip_in_src_and_not_in_dst
 from monkey_island.cc.database import mongo
 from monkey_island.cc.models import Monkey
-from monkey_island.cc.services.reporting.report_exporter_manager import ReportExporterManager
 from monkey_island.cc.services.config import ConfigService
 from monkey_island.cc.services.configuration.utils import get_config_network_segments_as_subnet_groups
 from monkey_island.cc.services.edge import EdgeService
 from monkey_island.cc.services.node import NodeService
-from monkey_island.cc.utils import local_ip_addresses, get_subnets
 from monkey_island.cc.services.reporting.pth_report import PTHReportService
-from common.network.network_range import NetworkRange
+from monkey_island.cc.services.reporting.report_exporter_manager import ReportExporterManager
+from monkey_island.cc.services.reporting.report_generation_synchronisation import safe_generate_regular_report
+from monkey_island.cc.utils import local_ip_addresses, get_subnets
 
 __author__ = "itay.mizeretz"
 
@@ -780,7 +779,7 @@ class ReportService:
     def get_report():
         if ReportService.is_latest_report_exists():
             return ReportService.decode_dot_char_before_mongo_insert(mongo.db.report.find_one())
-        return ReportService.generate_report()
+        return safe_generate_regular_report()
 
     @staticmethod
     def did_exploit_type_succeed(exploit_type):

monkey/monkey_island/cc/services/reporting/report_generation_synchronisation.py

@@ -0,0 +1,52 @@
+import logging
+import threading
+
+logger = logging.getLogger(__name__)
+
+# These are pseudo-singletons - global Locks. These locks will allow only one thread to generate a report at a time.
+# Report generation can be quite slow if there is a lot of data, and the UI queries the Root service often; without
+# the locks, these requests would accumulate, overload the server, eventually causing it to crash.
+logger.debug("Initializing report generation locks.")
+__report_generating_lock = threading.Semaphore()
+__attack_report_generating_lock = threading.Semaphore()
+__regular_report_generating_lock = threading.Semaphore()
+
+
+def safe_generate_reports():
+    # Entering the critical section; Wait until report generation is available.
+    __report_generating_lock.acquire()
+    report = safe_generate_regular_report()
+    attack_report = safe_generate_attack_report()
+    # Leaving the critical section.
+    __report_generating_lock.release()
+    return report, attack_report
+
+
+def safe_generate_regular_report():
+    # Local import to avoid circular imports
+    from monkey_island.cc.services.reporting.report import ReportService
+    __regular_report_generating_lock.acquire()
+    report = ReportService.generate_report()
+    __regular_report_generating_lock.release()
+    return report
+
+
+def safe_generate_attack_report():
+    # Local import to avoid circular imports
+    from monkey_island.cc.services.attack.attack_report import AttackReportService
+    __attack_report_generating_lock.acquire()
+    attack_report = AttackReportService.generate_new_report()
+    __attack_report_generating_lock.release()
+    return attack_report
+
+
+def is_report_being_generated():
+    # From https://docs.python.org/2/library/threading.html#threading.Semaphore.acquire:
+    # When invoked with blocking set to false, do not block.
+    #   If a call without an argument would block, return false immediately;
+    #   otherwise, do the same thing as when called without arguments, and return true.
+    is_report_being_generated_right_now = not __report_generating_lock.acquire(blocking=False)
+    if not is_report_being_generated_right_now:
+        # We're not using the critical resource; we just checked its state.
+        __report_generating_lock.release()
+    return is_report_being_generated_right_now