Island: Move encode/decode dot mongo functions to Report model

monkey/monkey_island/cc/models/report.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from bson import json_util
 from mongoengine import DictField, Document
 
 from monkey_island.cc.models.utils import report_encryptor
@@ -16,6 +17,7 @@ class Report(Document):
 
     @staticmethod
     def save_report(report_dict: dict):
+        report_dict = _encode_dot_char_before_mongo_insert(report_dict)
         report_dict = report_encryptor.encrypt(report_dict)
         Report.objects.delete()
         Report(
@@ -28,4 +30,24 @@ class Report(Document):
     @staticmethod
     def get_report() -> dict:
         report_dict = Report.objects.first().to_mongo()
-        return report_encryptor.decrypt(report_dict)
+        return _decode_dot_char_before_mongo_insert(report_encryptor.decrypt(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)

monkey/monkey_island/cc/services/reporting/report.py

@@ -4,8 +4,6 @@ import itertools
 import logging
 from typing import List
 
-from bson import json_util
-
 from common.config_value_paths import (
     EXPLOITER_CLASSES_PATH,
     LOCAL_NETWORK_SCAN_PATH,
@@ -636,7 +634,6 @@ class ReportService:
             "meta_info": {"latest_monkey_modifytime": monkey_latest_modify_time},
         }
         ReportExporterManager().export(report)
-        report = ReportService.encode_dot_char_before_mongo_insert(report)
         Report.save_report(report)
         return report
 
@@ -666,17 +663,6 @@ class ReportService:
         logger.info("Issues generated for reporting")
         return issues_dict
 
-    @staticmethod
-    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)
-
     @staticmethod
     def is_latest_report_exists():
         """
@@ -705,18 +691,9 @@ class ReportService:
                 "Report cache not cleared. DeleteResult: " + delete_result.raw_result
             )
 
-    @staticmethod
-    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)
-
     @staticmethod
     def get_report():
         if not ReportService.is_latest_report_exists():
             return safe_generate_regular_report()
 
-        return ReportService.decode_dot_char_before_mongo_insert(Report.get_report())
+        return Report.get_report()

monkey/tests/unit_tests/monkey_island/cc/models/test_report_model.py

@@ -1,3 +1,4 @@
+import copy
 from typing import List
 
 import pytest
@@ -32,15 +33,30 @@ class MockFieldEncryptor(IFieldEncryptor):
         return [MockFieldEncryptor.plaintext[int(v)] for v in value]
 
 
-MOCK_SENSITIVE_FIELDS = [SensitiveField("overview.foo.the_key", MockFieldEncryptor)]
+@pytest.fixture(autouse=True)
+def patch_sensitive_fields(monkeypatch):
+    mock_sensitive_fields = [SensitiveField("overview.foo.the_key", MockFieldEncryptor)]
+    monkeypatch.setattr(
+        "monkey_island.cc.models.utils.report_encryptor.sensitive_fields", mock_sensitive_fields
+    )
 
 
 @pytest.mark.usefixtures("uses_database")
-def test_report_encryption(monkeypatch, data_for_tests_dir):
+def test_report_encryption(data_for_tests_dir):
-    monkeypatch.setattr(
-        "monkey_island.cc.models.utils.report_encryptor.sensitive_fields", MOCK_SENSITIVE_FIELDS
-    )
     Report.save_report(MOCK_REPORT_DICT)
 
     assert Report.objects.first()["overview"]["foo"]["the_key"] == ["0", "1"]
     assert Report.get_report()["overview"]["foo"]["the_key"] == MOCK_SENSITIVE_FIELD_CONTENTS
+
+
+@pytest.mark.usefixtures("uses_database")
+def test_report_dot_encoding(data_for_tests_dir):
+    mrd = copy.deepcopy(MOCK_REPORT_DICT)
+    mrd["meta_info"] = {"foo.bar": "baz"}
+    Report.save_report(mrd)
+
+    assert "foo.bar" not in Report.objects.first()["meta_info"]
+    assert "foo,,,bar" in Report.objects.first()["meta_info"]
+
+    report = Report.get_report()
+    assert "foo.bar" in report["meta_info"]