Refactored test_telems to exported_telems and moved them from the test folder, because they are generated in production

This commit is contained in:
VakarisZ 2021-04-22 17:19:20 +03:00 committed by VakarisZ
parent c3f31c0c78
commit 8121f08aa9
6 changed files with 20 additions and 17 deletions

View File

@ -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)

View File

@ -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()

View File

@ -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)

View File

@ -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"] = []

View File

@ -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"])

View File

@ -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()