From 91d68113f9bf13fc2e3f31cab4235c6e1c321bcb Mon Sep 17 00:00:00 2001
From: VakarisZ <vakarisz@yahoo.com>
Date: Thu, 23 Apr 2020 19:46:05 +0300
Subject: [PATCH] Refactored root endpoint and added service hook for things
 that should happen after exploitation is done

---
 monkey/monkey_island/cc/resources/root.py     | 37 ++------------
 .../cc/services/infection_lifecycle.py        | 51 +++++++++++++++++++
 2 files changed, 54 insertions(+), 34 deletions(-)
 create mode 100644 monkey/monkey_island/cc/services/infection_lifecycle.py

diff --git a/monkey/monkey_island/cc/resources/root.py b/monkey/monkey_island/cc/resources/root.py
index 59a8fbe7c..216329905 100644
--- a/monkey/monkey_island/cc/resources/root.py
+++ b/monkey/monkey_island/cc/resources/root.py
@@ -1,6 +1,5 @@
 import logging
 import threading
-from datetime import datetime
 
 import flask_restful
 from flask import request, make_response, jsonify
@@ -8,10 +7,7 @@ from flask import request, make_response, jsonify
 from monkey_island.cc.auth import jwt_required
 from monkey_island.cc.database import mongo
 from monkey_island.cc.services.database import Database
-from monkey_island.cc.services.node import NodeService
-from monkey_island.cc.services.reporting.report import ReportService
-from monkey_island.cc.services.reporting.report_generation_synchronisation import is_report_being_generated, \
-    safe_generate_reports
+from monkey_island.cc.services.infection_lifecycle import InfectionLifecycle
 from monkey_island.cc.utils import local_ip_addresses
 
 __author__ = 'Barak'
@@ -32,7 +28,7 @@ class Root(flask_restful.Resource):
         elif action == "reset":
             return jwt_required()(Database.reset_db)()
         elif action == "killall":
-            return Root.kill_all()
+            return jwt_required()(InfectionLifecycle.kill_all)()
         elif action == "is-up":
             return {'is-up': True}
         else:
@@ -43,33 +39,6 @@ class Root(flask_restful.Resource):
         return jsonify(
             ip_addresses=local_ip_addresses(),
             mongo=str(mongo.db),
-            completed_steps=self.get_completed_steps())
+            completed_steps=InfectionLifecycle.get_completed_steps())
 
-    @staticmethod
-    @jwt_required()
-    def kill_all():
-        mongo.db.monkey.update({'dead': False}, {'$set': {'config.alive': False, 'modifytime': datetime.now()}},
-                               upsert=False,
-                               multi=True)
-        logger.info('Kill all monkeys was called')
-        return jsonify(status='OK')
 
-    @jwt_required()
-    def get_completed_steps(self):
-        is_any_exists = NodeService.is_any_monkey_exists()
-        infection_done = NodeService.is_monkey_finished_running()
-
-        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()
-        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/infection_lifecycle.py b/monkey/monkey_island/cc/services/infection_lifecycle.py
new file mode 100644
index 000000000..e79cfe947
--- /dev/null
+++ b/monkey/monkey_island/cc/services/infection_lifecycle.py
@@ -0,0 +1,51 @@
+import logging
+from datetime import datetime
+
+from flask import jsonify
+
+from monkey_island.cc.database import mongo
+from monkey_island.cc.resources.test.utils.telem_store import TestTelemStore
+from monkey_island.cc.services.config import ConfigService
+from monkey_island.cc.services.node import NodeService
+from monkey_island.cc.services.reporting.report import ReportService
+from monkey_island.cc.services.reporting.report_generation_synchronisation import is_report_being_generated, \
+    safe_generate_reports
+
+logger = logging.getLogger(__name__)
+
+
+class InfectionLifecycle:
+
+    @staticmethod
+    def kill_all():
+        mongo.db.monkey.update({'dead': False}, {'$set': {'config.alive': False, 'modifytime': datetime.now()}},
+                               upsert=False,
+                               multi=True)
+        logger.info('Kill all monkeys was called')
+        return jsonify(status='OK')
+
+    @staticmethod
+    def get_completed_steps():
+        is_any_exists = NodeService.is_any_monkey_exists()
+        infection_done = NodeService.is_monkey_finished_running()
+
+        if infection_done:
+            InfectionLifecycle._on_finished_infection()
+            report_done = ReportService.is_report_generated()
+        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)
+
+    @staticmethod
+    def _on_finished_infection():
+        # 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()
+        if ConfigService.is_test_telem_export_enabled():
+            TestTelemStore.export_test_telems()