Change the telemetry model to have a method for fetching the telemetries based on queries.

Telemetry code mainly uses queries and mongoengine has no good way of field encryption, that's why this method prefers to handle queries rather than Telemetry models

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

@@ -0,0 +1 @@
+from .telemetry import Telemetry  # noqa: F401

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

@@ -1,7 +1,10 @@
 from __future__ import annotations
 
-from mongoengine import DateTimeField, DictField, Document, EmbeddedDocumentField, StringField
+from typing import List
 
+from mongoengine import DateTimeField, Document, DynamicField, EmbeddedDocumentField, StringField
+
+from monkey_island.cc.database import mongo
 from monkey_island.cc.models import CommandControlChannel
 from monkey_island.cc.models.utils import document_encryptor
 from monkey_island.cc.models.utils.document_encryptor import FieldNotFoundError, SensitiveField
@@ -17,7 +20,7 @@ sensitive_fields = [
 
 class Telemetry(Document):
 
-    data = DictField(required=True)
+    data = DynamicField(required=True)
     timestamp = DateTimeField(required=True)
     monkey_guid = StringField(required=True)
     telem_category = StringField(required=True)
@@ -45,6 +48,12 @@ class Telemetry(Document):
         ).save()
 
     @staticmethod
-    def get_telemetry() -> dict:
-        telemetry_dict = Telemetry.objects.first().to_mongo()
-        return document_encryptor.decrypt(sensitive_fields, telemetry_dict)
+    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(document_encryptor.decrypt(sensitive_fields, telemetry))
+            except FieldNotFoundError:
+                decrypted_list.append(telemetry)
+        return decrypted_list

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

@@ -1,9 +1,10 @@
 from copy import deepcopy
 from datetime import datetime
 
+import mongoengine
 import pytest
 
-from monkey_island.cc.models.telemetries.telemetry import Telemetry
+from monkey_island.cc.models.telemetries import Telemetry
 from monkey_island.cc.models.utils.document_encryptor import SensitiveField
 from monkey_island.cc.models.utils.field_encryptors.mimikatz_results_encryptor import (
     MimikatzResultsEncryptor,
@@ -61,6 +62,12 @@ def patch_sensitive_fields(monkeypatch):
     )
 
 
+@pytest.fixture(autouse=True)
+def fake_mongo(monkeypatch):
+    mongo = mongoengine.connection.get_connection()
+    monkeypatch.setattr("monkey_island.cc.models.telemetries.telemetry.mongo", mongo)
+
+
 @pytest.mark.usefixtures("uses_database", "uses_encryptor")
 def test_telemetry_encryption(monkeypatch):
 
@@ -74,11 +81,11 @@ def test_telemetry_encryption(monkeypatch):
         == MOCK_CREDENTIALS["Vakaris"]["ntlm_hash"]
     )
     assert (
-        Telemetry.get_telemetry()["data"]["credentials"]["user"]["password"]
+        Telemetry.get_telemetry_by_query({})[0]["data"]["credentials"]["user"]["password"]
         == MOCK_CREDENTIALS["user"]["password"]
     )
     assert (
-        Telemetry.get_telemetry()["data"]["mimikatz"]["Vakaris"]["ntlm_hash"]
+        Telemetry.get_telemetry_by_query({})[0]["data"]["mimikatz"]["Vakaris"]["ntlm_hash"]
         == MOCK_CREDENTIALS["Vakaris"]["ntlm_hash"]
     )
 

vulture_allowlist.py

@@ -4,6 +4,7 @@ dead or is kept deliberately. Referencing these in a file like this makes sure t
 Vulture doesn't mark these as dead again.
 """
 from monkey_island.cc.models import Report
+from monkey_island.cc.models.telemetries import Telemetry
 
 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)
@@ -182,6 +183,8 @@ Report.glance
 Report.meta_info
 Report.meta
 Report.save_report
+Telemetry.save_telemetry
+Telemetry.get_telemetry_by_query
 
 # 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)