diff --git a/monkey/monkey_island/cc/resources/root.py b/monkey/monkey_island/cc/resources/root.py index e3b3e9854..d7cae8bd7 100644 --- a/monkey/monkey_island/cc/resources/root.py +++ b/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) diff --git a/monkey/monkey_island/cc/services/attack/attack_report.py b/monkey/monkey_island/cc/services/attack/attack_report.py index c04e6870f..c7457c2f6 100644 --- a/monkey/monkey_island/cc/services/attack/attack_report.py +++ b/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(): diff --git a/monkey/monkey_island/cc/services/reporting/report.py b/monkey/monkey_island/cc/services/reporting/report.py index f00fbc22c..03d0f5011 100644 --- a/monkey/monkey_island/cc/services/reporting/report.py +++ b/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): diff --git a/monkey/monkey_island/cc/services/reporting/report_generation_synchronisation.py b/monkey/monkey_island/cc/services/reporting/report_generation_synchronisation.py new file mode 100644 index 000000000..9025ff68f --- /dev/null +++ b/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