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
|
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
|
# SCHEMA
|
||||||
name = StringField(required=True)
|
name = StringField(required=True)
|
||||||
time = DateTimeField(required=True)
|
time = DateTimeField(required=True)
|
|
@ -6,7 +6,7 @@ from os import mkdir, path
|
||||||
|
|
||||||
from flask import request
|
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
|
from monkey_island.cc.services.config import ConfigService
|
||||||
|
|
||||||
TELEM_SAMPLE_DIR = "./telem_sample"
|
TELEM_SAMPLE_DIR = "./telem_sample"
|
||||||
|
@ -20,7 +20,7 @@ class TestTelemStore:
|
||||||
TELEMS_EXPORTED = False
|
TELEMS_EXPORTED = False
|
||||||
|
|
||||||
@staticmethod
|
@staticmethod
|
||||||
def store_test_telem(f):
|
def store_exported_telem(f):
|
||||||
@wraps(f)
|
@wraps(f)
|
||||||
def decorated_function(*args, **kwargs):
|
def decorated_function(*args, **kwargs):
|
||||||
if ConfigService.is_test_telem_export_enabled():
|
if ConfigService.is_test_telem_export_enabled():
|
||||||
|
@ -35,7 +35,7 @@ class TestTelemStore:
|
||||||
.replace(">", "_")
|
.replace(">", "_")
|
||||||
.replace(":", "_")
|
.replace(":", "_")
|
||||||
)
|
)
|
||||||
TestTelem(
|
ExportedTelem(
|
||||||
name=name, method=method, endpoint=endpoint, content=content, time=time
|
name=name, method=method, endpoint=endpoint, content=content, time=time
|
||||||
).save()
|
).save()
|
||||||
return f(*args, **kwargs)
|
return f(*args, **kwargs)
|
||||||
|
@ -43,7 +43,7 @@ class TestTelemStore:
|
||||||
return decorated_function
|
return decorated_function
|
||||||
|
|
||||||
@staticmethod
|
@staticmethod
|
||||||
def export_test_telems():
|
def export_telems():
|
||||||
logger.info(f"Exporting all telemetries to {TELEM_SAMPLE_DIR}")
|
logger.info(f"Exporting all telemetries to {TELEM_SAMPLE_DIR}")
|
||||||
try:
|
try:
|
||||||
mkdir(TELEM_SAMPLE_DIR)
|
mkdir(TELEM_SAMPLE_DIR)
|
||||||
|
@ -51,9 +51,9 @@ class TestTelemStore:
|
||||||
logger.info("Deleting all previous telemetries.")
|
logger.info("Deleting all previous telemetries.")
|
||||||
shutil.rmtree(TELEM_SAMPLE_DIR)
|
shutil.rmtree(TELEM_SAMPLE_DIR)
|
||||||
mkdir(TELEM_SAMPLE_DIR)
|
mkdir(TELEM_SAMPLE_DIR)
|
||||||
for test_telem in TestTelem.objects():
|
for test_telem in ExportedTelem.objects():
|
||||||
with open(
|
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",
|
"w",
|
||||||
) as file:
|
) as file:
|
||||||
file.write(test_telem.to_json(indent=2))
|
file.write(test_telem.to_json(indent=2))
|
||||||
|
@ -61,8 +61,8 @@ class TestTelemStore:
|
||||||
logger.info("Telemetries exported!")
|
logger.info("Telemetries exported!")
|
||||||
|
|
||||||
@staticmethod
|
@staticmethod
|
||||||
def get_unique_file_path_for_test_telem(target_dir: str, test_telem: TestTelem):
|
def get_unique_file_path_for_export_telem(target_dir: str, test_telem: ExportedTelem):
|
||||||
telem_filename = TestTelemStore._get_filename_by_test_telem(test_telem)
|
telem_filename = TestTelemStore._get_filename_by_export_telem(test_telem)
|
||||||
for i in range(MAX_SAME_CATEGORY_TELEMS):
|
for i in range(MAX_SAME_CATEGORY_TELEMS):
|
||||||
potential_filepath = path.join(target_dir, (telem_filename + str(i)))
|
potential_filepath = path.join(target_dir, (telem_filename + str(i)))
|
||||||
if path.exists(potential_filepath):
|
if path.exists(potential_filepath):
|
||||||
|
@ -73,10 +73,10 @@ class TestTelemStore:
|
||||||
)
|
)
|
||||||
|
|
||||||
@staticmethod
|
@staticmethod
|
||||||
def _get_filename_by_test_telem(test_telem: TestTelem):
|
def _get_filename_by_export_telem(test_telem: ExportedTelem):
|
||||||
endpoint_part = test_telem.name
|
endpoint_part = test_telem.name
|
||||||
return endpoint_part + "_" + test_telem.method
|
return endpoint_part + "_" + test_telem.method
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
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))
|
return LogService.log_exists(ObjectId(exists_monkey_id))
|
||||||
|
|
||||||
# Used by monkey. can't secure.
|
# Used by monkey. can't secure.
|
||||||
@TestTelemStore.store_test_telem
|
@TestTelemStore.store_exported_telem
|
||||||
def post(self):
|
def post(self):
|
||||||
telemetry_json = json.loads(request.data)
|
telemetry_json = json.loads(request.data)
|
||||||
|
|
||||||
|
|
|
@ -35,7 +35,7 @@ class Monkey(flask_restful.Resource):
|
||||||
return {}
|
return {}
|
||||||
|
|
||||||
# Used by monkey. can't secure.
|
# Used by monkey. can't secure.
|
||||||
@TestTelemStore.store_test_telem
|
@TestTelemStore.store_exported_telem
|
||||||
def patch(self, guid):
|
def patch(self, guid):
|
||||||
monkey_json = json.loads(request.data)
|
monkey_json = json.loads(request.data)
|
||||||
update = {"$set": {"modifytime": datetime.now()}}
|
update = {"$set": {"modifytime": datetime.now()}}
|
||||||
|
@ -60,7 +60,7 @@ class Monkey(flask_restful.Resource):
|
||||||
|
|
||||||
# Used by monkey. can't secure.
|
# Used by monkey. can't secure.
|
||||||
# Called on monkey wakeup to initialize local configuration
|
# Called on monkey wakeup to initialize local configuration
|
||||||
@TestTelemStore.store_test_telem
|
@TestTelemStore.store_exported_telem
|
||||||
def post(self, **kw):
|
def post(self, **kw):
|
||||||
monkey_json = json.loads(request.data)
|
monkey_json = json.loads(request.data)
|
||||||
monkey_json["creds"] = []
|
monkey_json["creds"] = []
|
||||||
|
|
|
@ -44,7 +44,7 @@ class Telemetry(flask_restful.Resource):
|
||||||
return result
|
return result
|
||||||
|
|
||||||
# Used by monkey. can't secure.
|
# Used by monkey. can't secure.
|
||||||
@TestTelemStore.store_test_telem
|
@TestTelemStore.store_exported_telem
|
||||||
def post(self):
|
def post(self):
|
||||||
telemetry_json = json.loads(request.data)
|
telemetry_json = json.loads(request.data)
|
||||||
telemetry_json["data"] = json.loads(telemetry_json["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():
|
if not is_report_being_generated() and not ReportService.is_latest_report_exists():
|
||||||
safe_generate_reports()
|
safe_generate_reports()
|
||||||
if ConfigService.is_test_telem_export_enabled() and not TestTelemStore.TELEMS_EXPORTED:
|
if ConfigService.is_test_telem_export_enabled() and not TestTelemStore.TELEMS_EXPORTED:
|
||||||
TestTelemStore.export_test_telems()
|
TestTelemStore.export_telems()
|
||||||
|
|
Loading…
Reference in New Issue