diff --git a/monkey/monkey_island/cc/models/telemetries/__init__.py b/monkey/monkey_island/cc/models/telemetries/__init__.py index e69de29bb..b3421bbd4 100644 --- a/monkey/monkey_island/cc/models/telemetries/__init__.py +++ b/monkey/monkey_island/cc/models/telemetries/__init__.py @@ -0,0 +1 @@ +from .telemetry import Telemetry # noqa: F401 diff --git a/monkey/monkey_island/cc/models/telemetries/telemetry.py b/monkey/monkey_island/cc/models/telemetries/telemetry.py index bbefeb92f..371484b85 100644 --- a/monkey/monkey_island/cc/models/telemetries/telemetry.py +++ b/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 diff --git a/monkey/tests/unit_tests/monkey_island/cc/models/telemetries/test_telemetry_model.py b/monkey/tests/unit_tests/monkey_island/cc/models/telemetries/test_telemetry_model.py index 578aff235..860ae05fb 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/models/telemetries/test_telemetry_model.py +++ b/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"] ) diff --git a/vulture_allowlist.py b/vulture_allowlist.py index 905cc74ad..d58d4ea9b 100644 --- a/vulture_allowlist.py +++ b/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)