forked from p15670423/monkey
Refactored test_telems to exported_telems and moved them from the test folder, because they are generated in production
This commit is contained in:
parent
c3f31c0c78
commit
8121f08aa9
|
@ -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)
|
|
@ -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()
|
||||
|
|
|
@ -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)
|
||||
|
||||
|
|
|
@ -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"] = []
|
||||
|
|
|
@ -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"])
|
||||
|
|
|
@ -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()
|
||||
|
|
Loading…
Reference in New Issue