From c7e91c5784c8a54eccc4827d8fa6e0d364febf14 Mon Sep 17 00:00:00 2001
From: VakarisZ <vakarisz@yahoo.com>
Date: Tue, 21 Sep 2021 10:37:57 +0300
Subject: [PATCH] Add report model and a unit test for it's encryption

---
 monkey/monkey_island/cc/models/__init__.py    |  1 +
 monkey/monkey_island/cc/models/report.py      | 31 ++++++++++++++++++
 .../cc/models/test_report_model.py            | 32 +++++++++++++++++++
 vulture_allowlist.py                          |  8 ++++-
 4 files changed, 71 insertions(+), 1 deletion(-)
 create mode 100644 monkey/monkey_island/cc/models/report.py
 create mode 100644 monkey/tests/unit_tests/monkey_island/cc/models/test_report_model.py

diff --git a/monkey/monkey_island/cc/models/__init__.py b/monkey/monkey_island/cc/models/__init__.py
index 50dec3f95..f99444210 100644
--- a/monkey/monkey_island/cc/models/__init__.py
+++ b/monkey/monkey_island/cc/models/__init__.py
@@ -7,3 +7,4 @@ from .creds import Creds  # noqa: F401, E402
 from .monkey import Monkey  # noqa: F401, E402
 from .monkey_ttl import MonkeyTtl  # noqa: F401, E402
 from .pba_results import PbaResults  # noqa: F401, E402
+from .report import Report  # noqa: F401, E402
diff --git a/monkey/monkey_island/cc/models/report.py b/monkey/monkey_island/cc/models/report.py
new file mode 100644
index 000000000..74b68e0be
--- /dev/null
+++ b/monkey/monkey_island/cc/models/report.py
@@ -0,0 +1,31 @@
+from __future__ import annotations
+
+from mongoengine import DictField, Document
+
+from monkey_island.cc.models.utils import report_encryptor
+
+
+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 = report_encryptor.encrypt(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 report_encryptor.decrypt(report_dict)
diff --git a/monkey/tests/unit_tests/monkey_island/cc/models/test_report_model.py b/monkey/tests/unit_tests/monkey_island/cc/models/test_report_model.py
new file mode 100644
index 000000000..b719f95a5
--- /dev/null
+++ b/monkey/tests/unit_tests/monkey_island/cc/models/test_report_model.py
@@ -0,0 +1,32 @@
+import pytest
+
+from monkey_island.cc.models import Report
+from monkey_island.cc.models.utils.field_types.string_list import StringList
+from monkey_island.cc.models.utils.report_encryptor import SensitiveField
+from monkey_island.cc.server_utils.encryptor import initialize_encryptor
+
+MOCK_SENSITIVE_FIELD_CONTENTS = ["the_string", "the_string2"]
+MOCK_REPORT_DICT = {
+    "overview": {"foo": {"the_key": MOCK_SENSITIVE_FIELD_CONTENTS, "other_key": "other_value"}},
+    "glance": {"foo": "bar"},
+    "recommendations": {"foo": "bar"},
+    "meta_info": {"foo": "bar"},
+}
+
+MOCK_SENSITIVE_FIELDS = [SensitiveField("overview.foo.the_key", StringList)]
+
+
+@pytest.mark.usefixtures("uses_database")
+def test_report_encryption(monkeypatch, data_for_tests_dir):
+    initialize_encryptor(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 not Report.objects.first()["overview"]["foo"]["the_key"] == MOCK_SENSITIVE_FIELD_CONTENTS
+    assert (
+        not Report.objects.first()["overview"]["foo"]["the_key"][1]
+        == MOCK_SENSITIVE_FIELD_CONTENTS[1]
+    )
+    assert Report.get_report()["overview"]["foo"]["the_key"] == MOCK_SENSITIVE_FIELD_CONTENTS
diff --git a/vulture_allowlist.py b/vulture_allowlist.py
index bd6d53c36..905cc74ad 100644
--- a/vulture_allowlist.py
+++ b/vulture_allowlist.py
@@ -3,7 +3,7 @@ Everything in this file is what Vulture found as dead code but either isn't real
 dead or is kept deliberately. Referencing these in a file like this makes sure that
 Vulture doesn't mark these as dead again.
 """
-
+from monkey_island.cc.models import Report
 
 fake_monkey_dir_path  # unused variable (monkey/tests/infection_monkey/post_breach/actions/test_users_custom_pba.py:37)
 set_os_linux  # unused variable (monkey/tests/infection_monkey/post_breach/actions/test_users_custom_pba.py:37)
@@ -176,6 +176,12 @@ import_status  # monkey_island\cc\resources\configuration_import.py:19
 config_schema  # monkey_island\cc\resources\configuration_import.py:25
 exception_stream  # unused attribute (monkey_island/cc/server_setup.py:104)
 ADVANCED  # unused attribute (monkey/monkey_island/cc/services/mode/mode_enum.py:6:)
+Report.overview
+Report.recommendations
+Report.glance
+Report.meta_info
+Report.meta
+Report.save_report
 
 # these are not needed for it to work, but may be useful extra information to understand what's going on
 WINDOWS_PBA_TYPE  # unused variable (monkey/monkey_island/cc/resources/pba_file_upload.py:23)