From 006c177abdb7ffbe006ceb0ad1aa7d2346406318 Mon Sep 17 00:00:00 2001 From: Shay Nehmad Date: Wed, 2 Oct 2019 13:04:58 +0300 Subject: [PATCH] Added lock on report generation and improved the get_completed_steps method --- monkey/monkey_island/cc/resources/root.py | 47 ++++++++++++++++------- 1 file changed, 33 insertions(+), 14 deletions(-) diff --git a/monkey/monkey_island/cc/resources/root.py b/monkey/monkey_island/cc/resources/root.py index e3b3e9854..f1914b3a2 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 @@ -18,13 +19,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 +37,10 @@ class Root(flask_restful.Resource): else: return make_response(400, {'error': 'unknown action'}) - @staticmethod @jwt_required() - def get_server_info(): + def get_server_info(self): return jsonify(ip_addresses=local_ip_addresses(), mongo=str(mongo.db), - completed_steps=Root.get_completed_steps()) + completed_steps=self.get_completed_steps()) @staticmethod @jwt_required() @@ -49,17 +51,34 @@ 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: + if self.should_generate_report(): + self.generate_report() 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) + + def generate_report(self): + # Set the event - enter the critical section + self.report_generating_lock.set() + # Not using the return value, as the get_report function also saves the report in the DB for later. + _ = ReportService.get_report() + _ = AttackReportService.get_latest_report() + # Clear the event - exit the critical section + self.report_generating_lock.clear() + + def should_generate_report(self): + # If the lock is not set, that means no one is generating a report right now. + is_any_thread_generating_a_report_right_now = not self.report_generating_lock.is_set() + return is_any_thread_generating_a_report_right_now and not ReportService.is_latest_report_exists()