From e6ad125be97d93aa8c09a0cd84596240c55ffb6b Mon Sep 17 00:00:00 2001 From: VakarisZ Date: Fri, 24 Sep 2021 12:42:31 +0300 Subject: [PATCH] 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 --- .../cc/models/telemetries/__init__.py | 1 + .../cc/models/telemetries/telemetry.py | 19 ++++++++++++++----- .../telemetries/test_telemetry_model.py | 13 ++++++++++--- vulture_allowlist.py | 3 +++ 4 files changed, 28 insertions(+), 8 deletions(-) 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)