Merge pull request #1485 from guardicore/telemetry_encryption

Telemetry encryption in database

CHANGELOG.md

@@ -41,10 +41,13 @@ Changelog](https://keepachangelog.com/en/1.0.0/).
 - Overlapping Guardicore logo in the landing page. #1441
 - PBA table collapse in security report on data change. #1423
 - Unsigned Windows agent binaries in Linux packages are now signed. #1444
+- Some of the gathered credentials no longer appear in database plaintext. #1454
 
 ### Security
 - Generate a random password when creating a new user for CommunicateAsNewUser
   PBA. #1434
+- Credentials gathered from victim machines are no longer stored plaintext in the database. #1454
+
 
 ## [1.11.0] - 2021-08-13
 ### Added

monkey/monkey_island/cc/models/__init__.py

@@ -7,4 +7,4 @@ from .creds import Creds
 from .monkey import Monkey
 from .monkey_ttl import MonkeyTtl
 from .pba_results import PbaResults
-from .report import Report
+from monkey_island.cc.models.report.report import Report

monkey/monkey_island/cc/models/report.py

@@ -1,53 +0,0 @@
-from __future__ import annotations
-
-from bson import json_util
-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 = _encode_dot_char_before_mongo_insert(report_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 _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/models/report/__init__.py

@@ -0,0 +1 @@
+from .report_dal import save_report, get_report

monkey/monkey_island/cc/models/report/report.py

@@ -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}

monkey/monkey_island/cc/models/report/report_dal.py

@@ -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)

monkey/monkey_island/cc/models/telemetries/__init__.py

@@ -0,0 +1 @@
+from .telemetry_dal import save_telemetry, get_telemetry_by_query

monkey/monkey_island/cc/models/telemetries/telemetry.py

@@ -0,0 +1,14 @@
+from mongoengine import DateTimeField, Document, DynamicField, EmbeddedDocumentField, StringField
+
+from monkey_island.cc.models import CommandControlChannel
+
+
+class Telemetry(Document):
+
+    data = DynamicField(required=True)
+    timestamp = DateTimeField(required=True)
+    monkey_guid = StringField(required=True)
+    telem_category = StringField(required=True)
+    command_control_channel = EmbeddedDocumentField(CommandControlChannel)
+
+    meta = {"strict": False}

monkey/monkey_island/cc/models/telemetries/telemetry_dal.py

@@ -0,0 +1,51 @@
+from __future__ import annotations
+
+from typing import List
+
+from monkey_island.cc.database import mongo
+from monkey_island.cc.models import CommandControlChannel
+from monkey_island.cc.models.telemetries.telemetry import Telemetry
+from monkey_island.cc.server_utils.encryption import (
+    FieldNotFoundError,
+    MimikatzResultsEncryptor,
+    SensitiveField,
+    decrypt_dict,
+    encrypt_dict,
+)
+
+sensitive_fields = [
+    SensitiveField("data.credentials", MimikatzResultsEncryptor),
+    SensitiveField("data.mimikatz", MimikatzResultsEncryptor),
+]
+
+
+def save_telemetry(telemetry_dict: dict):
+    try:
+        telemetry_dict = encrypt_dict(sensitive_fields, telemetry_dict)
+    except FieldNotFoundError:
+        pass  # Not all telemetries require encryption
+
+    cc_channel = CommandControlChannel(
+        src=telemetry_dict["command_control_channel"]["src"],
+        dst=telemetry_dict["command_control_channel"]["dst"],
+    )
+    Telemetry(
+        data=telemetry_dict["data"],
+        timestamp=telemetry_dict["timestamp"],
+        monkey_guid=telemetry_dict["monkey_guid"],
+        telem_category=telemetry_dict["telem_category"],
+        command_control_channel=cc_channel,
+    ).save()
+
+
+# A lot of codebase is using queries for telemetry collection and document field encryption is
+# not yet implemented in mongoengine. To avoid big time investment, queries are used for now.
+def get_telemetry_by_query(query: dict, output_fields=None) -> List[dict]:
+    telemetries = mongo.db.telemetry.find(query, output_fields)
+    decrypted_list = []
+    for telemetry in telemetries:
+        try:
+            decrypted_list.append(decrypt_dict(sensitive_fields, telemetry))
+        except FieldNotFoundError:
+            decrypted_list.append(telemetry)
+    return decrypted_list

monkey/monkey_island/cc/models/utils/report_encryptor.py

@@ -1,50 +0,0 @@
-from dataclasses import dataclass
-from typing import Callable, Type
-
-import dpath.util
-
-from monkey_island.cc.models.utils.field_encryptors.i_field_encryptor import IFieldEncryptor
-from monkey_island.cc.models.utils.field_encryptors.string_list_encryptor import StringListEncryptor
-
-
-@dataclass
-class SensitiveField:
-    path: str
-    path_separator = "."
-    field_type: Type[IFieldEncryptor]
-
-
-sensitive_fields = [
-    SensitiveField(path="overview.config_passwords", field_type=StringListEncryptor)
-]
-
-
-def encrypt(report: dict) -> dict:
-    for sensitive_field in sensitive_fields:
-        _apply_operation_to_report_field(
-            report, sensitive_field, sensitive_field.field_type.encrypt
-        )
-
-    return report
-
-
-def decrypt(report: dict) -> dict:
-    for sensitive_field in sensitive_fields:
-        _apply_operation_to_report_field(
-            report, sensitive_field, sensitive_field.field_type.decrypt
-        )
-    return report
-
-
-def _apply_operation_to_report_field(
-    report: dict, sensitive_field: SensitiveField, operation: Callable
-):
-    field_value = dpath.util.get(report, sensitive_field.path, sensitive_field.path_separator, None)
-    if field_value is None:
-        raise Exception(
-            f"Can't encrypt object because the path {sensitive_field.path} doesn't exist."
-        )
-
-    modified_value = operation(field_value)
-
-    dpath.util.set(report, sensitive_field.path, modified_value, sensitive_field.path_separator)

monkey/monkey_island/cc/resources/blackbox/telemetry_blackbox_endpoint.py

@@ -2,7 +2,7 @@ import flask_restful
 from bson import json_util
 from flask import request
 
-from monkey_island.cc.database import mongo
+from monkey_island.cc.models.telemetries import get_telemetry_by_query
 from monkey_island.cc.resources.auth.auth import jwt_required
 
 
@@ -10,4 +10,4 @@ class TelemetryBlackboxEndpoint(flask_restful.Resource):
     @jwt_required
     def get(self, **kw):
         find_query = json_util.loads(request.args.get("find_query"))
-        return {"results": list(mongo.db.telemetry.find(find_query))}
+        return {"results": list(get_telemetry_by_query(find_query))}

monkey/monkey_island/cc/resources/telemetry.py

@@ -9,6 +9,7 @@ from flask import request
 from common.common_consts.telem_categories import TelemCategoryEnum
 from monkey_island.cc.database import mongo
 from monkey_island.cc.models.monkey import Monkey
+from monkey_island.cc.models.telemetries import get_telemetry_by_query, save_telemetry
 from monkey_island.cc.resources.auth.auth import jwt_required
 from monkey_island.cc.resources.blackbox.utils.telem_store import TestTelemStore
 from monkey_island.cc.services.node import NodeService
@@ -37,7 +38,7 @@ class Telemetry(flask_restful.Resource):
             find_filter["timestamp"] = {"$gt": dateutil.parser.parse(timestamp)}
 
         result["objects"] = self.telemetry_to_displayed_telemetry(
-            mongo.db.telemetry.find(find_filter)
+            get_telemetry_by_query(query=find_filter)
         )
         return result
 
@@ -60,8 +61,9 @@ class Telemetry(flask_restful.Resource):
 
         process_telemetry(telemetry_json)
 
-        telem_id = mongo.db.telemetry.insert(telemetry_json)
+        save_telemetry(telemetry_json)
-        return mongo.db.telemetry.find_one_or_404({"_id": telem_id})
+
+        return {}, 201
 
     @staticmethod
     def telemetry_to_displayed_telemetry(telemetry):

monkey/monkey_island/cc/server_utils/encryption/__init__.py

@@ -11,3 +11,11 @@ from monkey_island.cc.server_utils.encryption.data_store_encryptor import (
     get_datastore_encryptor,
     initialize_datastore_encryptor,
 )
+from .dict_encryption.dict_encryptor import (
+    SensitiveField,
+    encrypt_dict,
+    decrypt_dict,
+    FieldNotFoundError,
+)
+from .dict_encryption.field_encryptors.mimikatz_results_encryptor import MimikatzResultsEncryptor
+from .dict_encryption.field_encryptors.string_list_encryptor import StringListEncryptor

monkey/monkey_island/cc/server_utils/encryption/dict_encryption/dict_encryptor.py

@@ -0,0 +1,50 @@
+from dataclasses import dataclass
+from typing import Callable, List, Type
+
+import dpath.util
+
+from monkey_island.cc.server_utils.encryption.dict_encryption.field_encryptors import (
+    IFieldEncryptor,
+)
+
+
+class FieldNotFoundError(Exception):
+    pass
+
+
+@dataclass
+class SensitiveField:
+    path: str
+    path_separator = "."
+    field_encryptor: Type[IFieldEncryptor]
+
+
+def encrypt_dict(sensitive_fields: List[SensitiveField], document_dict: dict) -> dict:
+    for sensitive_field in sensitive_fields:
+        _apply_operation_to_document_field(
+            document_dict, sensitive_field, sensitive_field.field_encryptor.encrypt
+        )
+
+    return document_dict
+
+
+def decrypt_dict(sensitive_fields: List[SensitiveField], document_dict: dict) -> dict:
+    for sensitive_field in sensitive_fields:
+        _apply_operation_to_document_field(
+            document_dict, sensitive_field, sensitive_field.field_encryptor.decrypt
+        )
+    return document_dict
+
+
+def _apply_operation_to_document_field(
+    report: dict, sensitive_field: SensitiveField, operation: Callable
+):
+    field_value = dpath.util.get(report, sensitive_field.path, sensitive_field.path_separator, None)
+    if field_value is None:
+        raise FieldNotFoundError(
+            f"Can't encrypt object because the path {sensitive_field.path} doesn't exist."
+        )
+
+    modified_value = operation(field_value)
+
+    dpath.util.set(report, sensitive_field.path, modified_value, sensitive_field.path_separator)

monkey/monkey_island/cc/server_utils/encryption/dict_encryption/field_encryptors/__init__.py

@@ -0,0 +1,3 @@
+from .i_field_encryptor import IFieldEncryptor
+from .mimikatz_results_encryptor import MimikatzResultsEncryptor
+from .string_list_encryptor import StringListEncryptor

monkey/monkey_island/cc/server_utils/encryption/dict_encryption/field_encryptors/mimikatz_results_encryptor.py

@@ -0,0 +1,36 @@
+import logging
+
+from monkey_island.cc.server_utils.encryption import get_datastore_encryptor
+from monkey_island.cc.server_utils.encryption.dict_encryption.field_encryptors import (
+    IFieldEncryptor,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class MimikatzResultsEncryptor(IFieldEncryptor):
+
+    secret_types = ["password", "ntlm_hash", "lm_hash"]
+
+    @staticmethod
+    def encrypt(results: dict) -> dict:
+        for _, credentials in results.items():
+            for secret_type in MimikatzResultsEncryptor.secret_types:
+                try:
+                    credentials[secret_type] = get_datastore_encryptor().enc(
+                        credentials[secret_type]
+                    )
+                except ValueError as e:
+                    logger.error(
+                        f"Failed encrypting sensitive field for "
+                        f"user {credentials['username']}! Error: {e}"
+                    )
+                    credentials[secret_type] = get_datastore_encryptor().enc("")
+        return results
+
+    @staticmethod
+    def decrypt(results: dict) -> dict:
+        for _, credentials in results.items():
+            for secret_type in MimikatzResultsEncryptor.secret_types:
+                credentials[secret_type] = get_datastore_encryptor().dec(credentials[secret_type])
+        return results

monkey/monkey_island/cc/server_utils/encryption/dict_encryption/field_encryptors/string_list_encryptor.py

@@ -1,7 +1,9 @@
 from typing import List
 
-from monkey_island.cc.models.utils.field_encryptors.i_field_encryptor import IFieldEncryptor
 from monkey_island.cc.server_utils.encryption import get_datastore_encryptor
+from monkey_island.cc.server_utils.encryption.dict_encryption.field_encryptors import (
+    IFieldEncryptor,
+)
 
 
 class StringListEncryptor(IFieldEncryptor):

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

@@ -14,7 +14,9 @@ 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 (
     get_config_network_segments_as_subnet_groups,
@@ -165,7 +167,7 @@ class ReportService:
     @staticmethod
     def _get_credentials_from_system_info_telems():
         formatted_creds = []
-        for telem in mongo.db.telemetry.find(
+        for telem in get_telemetry_by_query(
             {"telem_category": "system_info", "data.credentials": {"$exists": True}},
             {"data.credentials": 1, "monkey_guid": 1},
         ):
@@ -634,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
@@ -696,4 +698,4 @@ class ReportService:
         if not ReportService.is_latest_report_exists():
             return safe_generate_regular_report()
 
-        return Report.get_report()
+        return get_report()

monkey/tests/unit_tests/monkey_island/cc/conftest.py

@@ -10,6 +10,8 @@ from tests.unit_tests.monkey_island.cc.server_utils.encryption.test_password_bas
     STANDARD_PLAINTEXT_MONKEY_CONFIG_FILENAME,
 )
 
+from monkey_island.cc.server_utils.encryption import initialize_datastore_encryptor
+
 
 @pytest.fixture
 def monkey_config(data_for_tests_dir):
@@ -23,3 +25,8 @@ def monkey_config(data_for_tests_dir):
 @pytest.fixture
 def monkey_config_json(monkey_config):
     return json.dumps(monkey_config)
+
+
+@pytest.fixture
+def uses_encryptor(data_for_tests_dir):
+    initialize_datastore_encryptor(data_for_tests_dir)

monkey/tests/unit_tests/monkey_island/cc/models/telemetries/test_telemetry_dal.py

@@ -0,0 +1,97 @@
+from copy import deepcopy
+from datetime import datetime
+
+import mongoengine
+import pytest
+
+from monkey_island.cc.models.telemetries import get_telemetry_by_query, save_telemetry
+from monkey_island.cc.models.telemetries.telemetry import Telemetry
+from monkey_island.cc.server_utils.encryption import SensitiveField
+from monkey_island.cc.server_utils.encryption.dict_encryption.field_encryptors import (
+    MimikatzResultsEncryptor,
+)
+
+MOCK_CREDENTIALS = {
+    "Vakaris": {
+        "username": "M0nk3y",
+        "password": "",
+        "ntlm_hash": "e87f2f73e353f1d95e42ce618601b61f",
+        "lm_hash": "",
+    },
+    "user": {"username": "user", "password": "test", "ntlm_hash": "", "lm_hash": ""},
+}
+
+MOCK_DATA_DICT = {
+    "network_info": {},
+    "credentials": deepcopy(MOCK_CREDENTIALS),
+    "mimikatz": deepcopy(MOCK_CREDENTIALS),
+}
+
+MOCK_TELEMETRY = {
+    "timestamp": datetime.now(),
+    "command_control_channel": {
+        "src": "192.168.56.1",
+        "dst": "192.168.56.2",
+    },
+    "monkey_guid": "211375648895908",
+    "telem_category": "system_info",
+    "data": MOCK_DATA_DICT,
+}
+
+MOCK_NO_ENCRYPTION_NEEDED_TELEMETRY = {
+    "timestamp": datetime.now(),
+    "command_control_channel": {
+        "src": "192.168.56.1",
+        "dst": "192.168.56.2",
+    },
+    "monkey_guid": "211375648895908",
+    "telem_category": "state",
+    "data": {"done": False},
+}
+
+MOCK_SENSITIVE_FIELDS = [
+    SensitiveField("data.credentials", MimikatzResultsEncryptor),
+    SensitiveField("data.mimikatz", MimikatzResultsEncryptor),
+]
+
+
+@pytest.fixture(autouse=True)
+def patch_sensitive_fields(monkeypatch):
+    monkeypatch.setattr(
+        "monkey_island.cc.models.telemetries.telemetry_dal.sensitive_fields",
+        MOCK_SENSITIVE_FIELDS,
+    )
+
+
+@pytest.fixture(autouse=True)
+def fake_mongo(monkeypatch):
+    mongo = mongoengine.connection.get_connection()
+    monkeypatch.setattr("monkey_island.cc.models.telemetries.telemetry_dal.mongo", mongo)
+
+
+@pytest.mark.usefixtures("uses_database", "uses_encryptor")
+def test_telemetry_encryption():
+
+    save_telemetry(MOCK_TELEMETRY)
+    assert (
+        not Telemetry.objects.first()["data"]["credentials"]["user"]["password"]
+        == MOCK_CREDENTIALS["user"]["password"]
+    )
+    assert (
+        not Telemetry.objects.first()["data"]["mimikatz"]["Vakaris"]["ntlm_hash"]
+        == MOCK_CREDENTIALS["Vakaris"]["ntlm_hash"]
+    )
+    assert (
+        get_telemetry_by_query({})[0]["data"]["credentials"]["user"]["password"]
+        == MOCK_CREDENTIALS["user"]["password"]
+    )
+    assert (
+        get_telemetry_by_query({})[0]["data"]["mimikatz"]["Vakaris"]["ntlm_hash"]
+        == MOCK_CREDENTIALS["Vakaris"]["ntlm_hash"]
+    )
+
+
+@pytest.mark.usefixtures("uses_database", "uses_encryptor")
+def test_no_encryption_needed():
+    # Make sure telemetry save doesn't break when telemetry doesn't need encryption
+    save_telemetry(MOCK_NO_ENCRYPTION_NEEDED_TELEMETRY)

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

@@ -4,8 +4,11 @@ from typing import List
 import pytest
 
 from monkey_island.cc.models import Report
-from monkey_island.cc.models.utils.field_encryptors.i_field_encryptor import IFieldEncryptor
+from monkey_island.cc.models.report import get_report, save_report
-from monkey_island.cc.models.utils.report_encryptor import SensitiveField
+from monkey_island.cc.server_utils.encryption import SensitiveField
+from monkey_island.cc.server_utils.encryption.dict_encryption.field_encryptors import (
+    IFieldEncryptor,
+)
 
 MOCK_SENSITIVE_FIELD_CONTENTS = ["the_string", "the_string2"]
 MOCK_REPORT_DICT = {
@@ -19,51 +22,51 @@ MOCK_REPORT_DICT = {
 }
 
 
-class MockFieldEncryptor(IFieldEncryptor):
+class MockStringListEncryptor(IFieldEncryptor):
     plaintext = []
 
     @staticmethod
     def encrypt(value: List[str]) -> List[str]:
-        return [MockFieldEncryptor._encrypt(v) for v in value]
+        return [MockStringListEncryptor._encrypt(v) for v in value]
 
     @staticmethod
     def _encrypt(value: str) -> str:
-        MockFieldEncryptor.plaintext.append(value)
+        MockStringListEncryptor.plaintext.append(value)
-        return f"ENCRYPTED_{str(len(MockFieldEncryptor.plaintext) - 1)}"
+        return f"ENCRYPTED_{str(len(MockStringListEncryptor.plaintext) - 1)}"
 
     @staticmethod
     def decrypt(value: List[str]) -> List[str]:
-        return MockFieldEncryptor.plaintext
+        return MockStringListEncryptor.plaintext
 
 
 @pytest.fixture(autouse=True)
 def patch_sensitive_fields(monkeypatch):
     mock_sensitive_fields = [
-        SensitiveField("overview.foo.the_key", MockFieldEncryptor),
+        SensitiveField("overview.foo.the_key", MockStringListEncryptor),
-        SensitiveField("overview.bar.the_key", MockFieldEncryptor),
+        SensitiveField("overview.bar.the_key", MockStringListEncryptor),
     ]
     monkeypatch.setattr(
-        "monkey_island.cc.models.utils.report_encryptor.sensitive_fields", mock_sensitive_fields
+        "monkey_island.cc.models.report.report_dal.sensitive_fields", mock_sensitive_fields
     )
 
 
 @pytest.mark.usefixtures("uses_database")
 def test_report_encryption():
-    Report.save_report(MOCK_REPORT_DICT)
+    save_report(MOCK_REPORT_DICT)
 
     assert Report.objects.first()["overview"]["foo"]["the_key"] == ["ENCRYPTED_0", "ENCRYPTED_1"]
     assert Report.objects.first()["overview"]["bar"]["the_key"] == []
-    assert Report.get_report()["overview"]["foo"]["the_key"] == MOCK_SENSITIVE_FIELD_CONTENTS
+    assert get_report()["overview"]["foo"]["the_key"] == MOCK_SENSITIVE_FIELD_CONTENTS
 
 
 @pytest.mark.usefixtures("uses_database")
 def test_report_dot_encoding():
     mrd = copy.deepcopy(MOCK_REPORT_DICT)
     mrd["meta_info"] = {"foo.bar": "baz"}
-    Report.save_report(mrd)
+    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()
+    report = get_report()
     assert "foo.bar" in report["meta_info"]

monkey/tests/unit_tests/monkey_island/cc/models/utils/field_encryptors/test_string_list_encryptor.py

@@ -1,7 +1,9 @@
 import pytest
 
-from monkey_island.cc.models.utils.field_encryptors.string_list_encryptor import StringListEncryptor
 from monkey_island.cc.server_utils.encryption import initialize_datastore_encryptor
+from monkey_island.cc.server_utils.encryption.dict_encryption.field_encryptors import (
+    StringListEncryptor,
+)
 
 MOCK_STRING_LIST = ["test_1", "test_2"]
 EMPTY_LIST = []

monkey/tests/unit_tests/monkey_island/cc/mongomock_fixtures.py

@@ -1,32 +1,23 @@
 import mongoengine
 import pytest
 
-from monkey_island.cc.models import Monkey
+# Database name has to match the db used in the codebase,
-from monkey_island.cc.models.edge import Edge
+# else the name needs to be mocked during tests.
-from monkey_island.cc.models.zero_trust.finding import Finding
+# Currently its used like so: "mongo.db.telemetry.find()".
+MOCK_DB_NAME = "db"
 
 
 @pytest.fixture(scope="module", autouse=True)
 def change_to_mongo_mock():
     # Make sure tests are working with mongomock
     mongoengine.disconnect()
-    mongoengine.connect("mongoenginetest", host="mongomock://localhost")
+    mongoengine.connect(MOCK_DB_NAME, host="mongomock://localhost")
 
 
 @pytest.fixture(scope="function")
 def uses_database():
-    _clean_edge_db()
+    _drop_database()
-    _clean_monkey_db()
-    _clean_finding_db()
 
 
-def _clean_monkey_db():
+def _drop_database():
-    Monkey.objects().delete()
+    mongoengine.connection.get_connection().drop_database(MOCK_DB_NAME)
-
-
-def _clean_edge_db():
-    Edge.objects().delete()
-
-
-def _clean_finding_db():
-    Finding.objects().delete()

monkey/tests/unit_tests/monkey_island/cc/services/reporting/test_report.py

@@ -1,10 +1,11 @@
 import datetime
 from copy import deepcopy
 
-import mongomock
+import mongoengine
 import pytest
 from bson import ObjectId
 
+from monkey_island.cc.models.telemetries import save_telemetry
 from monkey_island.cc.services.reporting.report import ReportService
 
 TELEM_ID = {
@@ -49,6 +50,11 @@ SYSTEM_INFO_TELEMETRY_TELEM = {
     "_id": TELEM_ID["system_info_creds"],
     "monkey_guid": MONKEY_GUID,
     "telem_category": "system_info",
+    "timestamp": datetime.datetime(2021, 2, 19, 9, 0, 14, 984000),
+    "command_control_channel": {
+        "src": "192.168.56.1",
+        "dst": "192.168.56.2",
+    },
     "data": {
         "credentials": {
             USER: {
@@ -64,6 +70,11 @@ NO_CREDS_TELEMETRY_TELEM = {
     "_id": TELEM_ID["no_creds"],
     "monkey_guid": MONKEY_GUID,
     "telem_category": "exploit",
+    "timestamp": datetime.datetime(2021, 2, 19, 9, 0, 14, 984000),
+    "command_control_channel": {
+        "src": "192.168.56.1",
+        "dst": "192.168.56.2",
+    },
     "data": {
         "machine": {
             "ip_addr": VICTIM_IP,
@@ -125,12 +136,14 @@ NODE_DICT_FAILED_EXPLOITS["exploits"][1]["result"] = False
 
 @pytest.fixture
 def fake_mongo(monkeypatch):
-    mongo = mongomock.MongoClient()
+    mongo = mongoengine.connection.get_connection()
     monkeypatch.setattr("monkey_island.cc.services.reporting.report.mongo", mongo)
+    monkeypatch.setattr("monkey_island.cc.models.telemetries.telemetry_dal.mongo", mongo)
     monkeypatch.setattr("monkey_island.cc.services.node.mongo", mongo)
     return mongo
 
 
+@pytest.mark.usefixtures("uses_database")
 def test_get_stolen_creds_exploit(fake_mongo):
     fake_mongo.db.telemetry.insert_one(EXPLOIT_TELEMETRY_TELEM)
 
@@ -143,9 +156,10 @@ def test_get_stolen_creds_exploit(fake_mongo):
     assert expected_stolen_creds_exploit == stolen_creds_exploit
 
 
+@pytest.mark.usefixtures("uses_database")
 def test_get_stolen_creds_system_info(fake_mongo):
     fake_mongo.db.monkey.insert_one(MONKEY_TELEM)
-    fake_mongo.db.telemetry.insert_one(SYSTEM_INFO_TELEMETRY_TELEM)
+    save_telemetry(SYSTEM_INFO_TELEMETRY_TELEM)
 
     stolen_creds_system_info = ReportService.get_stolen_creds()
     expected_stolen_creds_system_info = [
@@ -157,8 +171,10 @@ def test_get_stolen_creds_system_info(fake_mongo):
     assert expected_stolen_creds_system_info == stolen_creds_system_info
 
 
+@pytest.mark.usefixtures("uses_database")
 def test_get_stolen_creds_no_creds(fake_mongo):
-    fake_mongo.db.telemetry.insert_one(NO_CREDS_TELEMETRY_TELEM)
+    fake_mongo.db.monkey.insert_one(MONKEY_TELEM)
+    save_telemetry(NO_CREDS_TELEMETRY_TELEM)
 
     stolen_creds_no_creds = ReportService.get_stolen_creds()
     expected_stolen_creds_no_creds = []

vulture_allowlist.py

@@ -181,7 +181,6 @@ 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)