From fd8ef7f39ef1c91951220c523d2cba5d792544c4 Mon Sep 17 00:00:00 2001 From: VakarisZ Date: Thu, 22 Apr 2021 17:19:20 +0300 Subject: [PATCH] Refactored test_telems to exported_telems and moved them from the test folder, because they are generated in production --- .../cc/models/exported_telem.py} | 7 +++++-- .../resources/blackbox/utils/telem_store.py | 20 +++++++++---------- monkey/monkey_island/cc/resources/log.py | 2 +- monkey/monkey_island/cc/resources/monkey.py | 4 ++-- .../monkey_island/cc/resources/telemetry.py | 2 +- .../cc/services/infection_lifecycle.py | 2 +- 6 files changed, 20 insertions(+), 17 deletions(-) rename monkey/{tests/monkey_island/cc/models/test_telem.py => monkey_island/cc/models/exported_telem.py} (50%) diff --git a/monkey/tests/monkey_island/cc/models/test_telem.py b/monkey/monkey_island/cc/models/exported_telem.py similarity index 50% rename from monkey/tests/monkey_island/cc/models/test_telem.py rename to monkey/monkey_island/cc/models/exported_telem.py index 8dd1cb658..6df2296fb 100644 --- a/monkey/tests/monkey_island/cc/models/test_telem.py +++ b/monkey/monkey_island/cc/models/exported_telem.py @@ -1,10 +1,13 @@ """ -Define a Document Schema for the TestTelem document. +Define a Document Schema for the TelemForExport document. """ from mongoengine import DateTimeField, Document, StringField -class TestTelem(Document): +# This document describes exported telemetry. +# These telemetries are used to mock monkeys sending telemetries to the island. +# This way we can replicate island state without running monkeys. +class ExportedTelem(Document): # SCHEMA name = StringField(required=True) time = DateTimeField(required=True) diff --git a/monkey/monkey_island/cc/resources/blackbox/utils/telem_store.py b/monkey/monkey_island/cc/resources/blackbox/utils/telem_store.py index 574712cda..2130cef5a 100644 --- a/monkey/monkey_island/cc/resources/blackbox/utils/telem_store.py +++ b/monkey/monkey_island/cc/resources/blackbox/utils/telem_store.py @@ -6,7 +6,7 @@ from os import mkdir, path from flask import request -from monkey_island.cc.models.test_telem import TestTelem +from monkey_island.cc.models.exported_telem import ExportedTelem from monkey_island.cc.services.config import ConfigService TELEM_SAMPLE_DIR = "./telem_sample" @@ -20,7 +20,7 @@ class TestTelemStore: TELEMS_EXPORTED = False @staticmethod - def store_test_telem(f): + def store_exported_telem(f): @wraps(f) def decorated_function(*args, **kwargs): if ConfigService.is_test_telem_export_enabled(): @@ -35,7 +35,7 @@ class TestTelemStore: .replace(">", "_") .replace(":", "_") ) - TestTelem( + ExportedTelem( name=name, method=method, endpoint=endpoint, content=content, time=time ).save() return f(*args, **kwargs) @@ -43,7 +43,7 @@ class TestTelemStore: return decorated_function @staticmethod - def export_test_telems(): + def export_telems(): logger.info(f"Exporting all telemetries to {TELEM_SAMPLE_DIR}") try: mkdir(TELEM_SAMPLE_DIR) @@ -51,9 +51,9 @@ class TestTelemStore: logger.info("Deleting all previous telemetries.") shutil.rmtree(TELEM_SAMPLE_DIR) mkdir(TELEM_SAMPLE_DIR) - for test_telem in TestTelem.objects(): + for test_telem in ExportedTelem.objects(): with open( - TestTelemStore.get_unique_file_path_for_test_telem(TELEM_SAMPLE_DIR, test_telem), + TestTelemStore.get_unique_file_path_for_export_telem(TELEM_SAMPLE_DIR, test_telem), "w", ) as file: file.write(test_telem.to_json(indent=2)) @@ -61,8 +61,8 @@ class TestTelemStore: logger.info("Telemetries exported!") @staticmethod - def get_unique_file_path_for_test_telem(target_dir: str, test_telem: TestTelem): - telem_filename = TestTelemStore._get_filename_by_test_telem(test_telem) + def get_unique_file_path_for_export_telem(target_dir: str, test_telem: ExportedTelem): + telem_filename = TestTelemStore._get_filename_by_export_telem(test_telem) for i in range(MAX_SAME_CATEGORY_TELEMS): potential_filepath = path.join(target_dir, (telem_filename + str(i))) if path.exists(potential_filepath): @@ -73,10 +73,10 @@ class TestTelemStore: ) @staticmethod - def _get_filename_by_test_telem(test_telem: TestTelem): + def _get_filename_by_export_telem(test_telem: ExportedTelem): endpoint_part = test_telem.name return endpoint_part + "_" + test_telem.method if __name__ == "__main__": - TestTelemStore.export_test_telems() + TestTelemStore.export_telems() diff --git a/monkey/monkey_island/cc/resources/log.py b/monkey/monkey_island/cc/resources/log.py index 126e1f697..298674d3b 100644 --- a/monkey/monkey_island/cc/resources/log.py +++ b/monkey/monkey_island/cc/resources/log.py @@ -24,7 +24,7 @@ class Log(flask_restful.Resource): return LogService.log_exists(ObjectId(exists_monkey_id)) # Used by monkey. can't secure. - @TestTelemStore.store_test_telem + @TestTelemStore.store_exported_telem def post(self): telemetry_json = json.loads(request.data) diff --git a/monkey/monkey_island/cc/resources/monkey.py b/monkey/monkey_island/cc/resources/monkey.py index 7d8a18a98..2b4042622 100644 --- a/monkey/monkey_island/cc/resources/monkey.py +++ b/monkey/monkey_island/cc/resources/monkey.py @@ -35,7 +35,7 @@ class Monkey(flask_restful.Resource): return {} # Used by monkey. can't secure. - @TestTelemStore.store_test_telem + @TestTelemStore.store_exported_telem def patch(self, guid): monkey_json = json.loads(request.data) update = {"$set": {"modifytime": datetime.now()}} @@ -60,7 +60,7 @@ class Monkey(flask_restful.Resource): # Used by monkey. can't secure. # Called on monkey wakeup to initialize local configuration - @TestTelemStore.store_test_telem + @TestTelemStore.store_exported_telem def post(self, **kw): monkey_json = json.loads(request.data) monkey_json["creds"] = [] diff --git a/monkey/monkey_island/cc/resources/telemetry.py b/monkey/monkey_island/cc/resources/telemetry.py index 15e7139a0..d2f662c1d 100644 --- a/monkey/monkey_island/cc/resources/telemetry.py +++ b/monkey/monkey_island/cc/resources/telemetry.py @@ -44,7 +44,7 @@ class Telemetry(flask_restful.Resource): return result # Used by monkey. can't secure. - @TestTelemStore.store_test_telem + @TestTelemStore.store_exported_telem def post(self): telemetry_json = json.loads(request.data) telemetry_json["data"] = json.loads(telemetry_json["data"]) diff --git a/monkey/monkey_island/cc/services/infection_lifecycle.py b/monkey/monkey_island/cc/services/infection_lifecycle.py index 476bc2a0e..5529cc70d 100644 --- a/monkey/monkey_island/cc/services/infection_lifecycle.py +++ b/monkey/monkey_island/cc/services/infection_lifecycle.py @@ -54,4 +54,4 @@ class InfectionLifecycle: if not is_report_being_generated() and not ReportService.is_latest_report_exists(): safe_generate_reports() if ConfigService.is_test_telem_export_enabled() and not TestTelemStore.TELEMS_EXPORTED: - TestTelemStore.export_test_telems() + TestTelemStore.export_telems()