Extract DAL interface for report model into a separate report_dal.py file

This commit is contained in:
VakarisZ 2021-09-28 10:58:23 +03:00
parent 1160ac6af0
commit a24eb841c1
4 changed files with 69 additions and 63 deletions

View File

@ -1,60 +0,0 @@
from __future__ import annotations
from bson import json_util
from mongoengine import DictField, Document
from monkey_island.cc.utils import SensitiveField, dict_encryptor
from monkey_island.cc.utils.field_encryptors import StringListEncryptor
sensitive_fields = [
SensitiveField(path="overview.config_passwords", field_encryptor=StringListEncryptor)
]
class Report(Document):
overview = DictField(required=True)
glance = DictField(required=True)
recommendations = DictField(required=True)
meta_info = DictField(required=True)
meta = {"strict": False}
@staticmethod
def save_report(report_dict: dict):
report_dict = _encode_dot_char_before_mongo_insert(report_dict)
report_dict = dict_encryptor.encrypt(sensitive_fields, report_dict)
Report.objects.delete()
Report(
overview=report_dict["overview"],
glance=report_dict["glance"],
recommendations=report_dict["recommendations"],
meta_info=report_dict["meta_info"],
).save()
@staticmethod
def get_report() -> dict:
report_dict = Report.objects.first().to_mongo()
return _decode_dot_char_before_mongo_insert(
dict_encryptor.decrypt(sensitive_fields, report_dict)
)
def _encode_dot_char_before_mongo_insert(report_dict):
"""
mongodb doesn't allow for '.' and '$' in a key's name, this function replaces the '.'
char with the unicode
,,, combo instead.
:return: dict with formatted keys with no dots.
"""
report_as_json = json_util.dumps(report_dict).replace(".", ",,,")
return json_util.loads(report_as_json)
def _decode_dot_char_before_mongo_insert(report_dict):
"""
this function replaces the ',,,' combo with the '.' char instead.
:return: report dict with formatted keys (',,,' -> '.')
"""
report_as_json = json_util.dumps(report_dict).replace(",,,", ".")
return json_util.loads(report_as_json)

View File

@ -0,0 +1,13 @@
from __future__ import annotations
from mongoengine import DictField, Document
class Report(Document):
overview = DictField(required=True)
glance = DictField(required=True)
recommendations = DictField(required=True)
meta_info = DictField(required=True)
meta = {"strict": False}

View File

@ -0,0 +1,52 @@
from __future__ import annotations
from bson import json_util
from monkey_island.cc.models.report.report import Report
from monkey_island.cc.server_utils.encryption import (
SensitiveField,
StringListEncryptor,
decrypt_dict,
encrypt_dict,
)
sensitive_fields = [
SensitiveField(path="overview.config_passwords", field_encryptor=StringListEncryptor)
]
def save_report(report_dict: dict):
report_dict = _encode_dot_char_before_mongo_insert(report_dict)
report_dict = encrypt_dict(sensitive_fields, report_dict)
Report.objects.delete()
Report(
overview=report_dict["overview"],
glance=report_dict["glance"],
recommendations=report_dict["recommendations"],
meta_info=report_dict["meta_info"],
).save()
def get_report() -> dict:
report_dict = Report.objects.first().to_mongo()
return _decode_dot_char_before_mongo_insert(decrypt_dict(sensitive_fields, report_dict))
def _encode_dot_char_before_mongo_insert(report_dict):
"""
mongodb doesn't allow for '.' and '$' in a key's name, this function replaces the '.'
char with the unicode
,,, combo instead.
:return: dict with formatted keys with no dots.
"""
report_as_json = json_util.dumps(report_dict).replace(".", ",,,")
return json_util.loads(report_as_json)
def _decode_dot_char_before_mongo_insert(report_dict):
"""
this function replaces the ',,,' combo with the '.' char instead.
:return: report dict with formatted keys (',,,' -> '.')
"""
report_as_json = json_util.dumps(report_dict).replace(",,,", ".")
return json_util.loads(report_as_json)

View File

@ -14,7 +14,8 @@ from common.config_value_paths import (
from common.network.network_range import NetworkRange
from common.network.segmentation_utils import get_ip_in_src_and_not_in_dst
from monkey_island.cc.database import mongo
from monkey_island.cc.models import Monkey, Report
from monkey_island.cc.models import Monkey
from monkey_island.cc.models.report import get_report, save_report
from monkey_island.cc.models.telemetries import get_telemetry_by_query
from monkey_island.cc.services.config import ConfigService
from monkey_island.cc.services.configuration.utils import (
@ -635,7 +636,7 @@ class ReportService:
"meta_info": {"latest_monkey_modifytime": monkey_latest_modify_time},
}
ReportExporterManager().export(report)
Report.save_report(report)
save_report(report)
return report
@staticmethod
@ -697,4 +698,4 @@ class ReportService:
if not ReportService.is_latest_report_exists():
return safe_generate_regular_report()
return Report.get_report()
return get_report()