From 46f263be5f58e605e5afa813ebff659fb350b097 Mon Sep 17 00:00:00 2001 From: VakarisZ Date: Mon, 27 Sep 2021 16:51:59 +0300 Subject: [PATCH] Separate the telemetry document from telemetry_dal, also extracted external interface into __init__.py files --- monkey/monkey_island/cc/models/report.py | 9 ++-- .../cc/models/telemetries/__init__.py | 2 +- .../cc/models/telemetries/telemetry.py | 45 ------------------ .../cc/models/telemetries/telemetry_dal.py | 46 +++++++++++++++++++ .../monkey_island/cc/models/utils/__init__.py | 0 .../models/utils/field_encryptors/__init__.py | 0 .../mimikatz_results_encryptor.py | 21 --------- .../blackbox/telemetry_blackbox_endpoint.py | 4 +- .../monkey_island/cc/resources/telemetry.py | 6 +-- .../cc/services/reporting/report.py | 4 +- monkey/monkey_island/cc/utils/__init__.py | 1 + .../dict_encryptor.py} | 2 +- .../cc/utils/field_encryptors/__init__.py | 3 ++ .../field_encryptors/i_field_encryptor.py | 0 .../mimikatz_results_encryptor.py | 34 ++++++++++++++ .../field_encryptors/string_list_encryptor.py | 2 +- ...lemetry_model.py => test_telemetry_dal.py} | 25 +++++----- .../cc/models/test_report_model.py | 4 +- .../test_string_list_encryptor.py | 2 +- .../cc/services/reporting/test_report.py | 17 ++++--- 20 files changed, 123 insertions(+), 104 deletions(-) create mode 100644 monkey/monkey_island/cc/models/telemetries/telemetry_dal.py delete mode 100644 monkey/monkey_island/cc/models/utils/__init__.py delete mode 100644 monkey/monkey_island/cc/models/utils/field_encryptors/__init__.py delete mode 100644 monkey/monkey_island/cc/models/utils/field_encryptors/mimikatz_results_encryptor.py create mode 100644 monkey/monkey_island/cc/utils/__init__.py rename monkey/monkey_island/cc/{models/utils/document_encryptor.py => utils/dict_encryptor.py} (93%) create mode 100644 monkey/monkey_island/cc/utils/field_encryptors/__init__.py rename monkey/monkey_island/cc/{models => }/utils/field_encryptors/i_field_encryptor.py (100%) create mode 100644 monkey/monkey_island/cc/utils/field_encryptors/mimikatz_results_encryptor.py rename monkey/monkey_island/cc/{models => }/utils/field_encryptors/string_list_encryptor.py (81%) rename monkey/tests/unit_tests/monkey_island/cc/models/telemetries/{test_telemetry_model.py => test_telemetry_dal.py} (74%) diff --git a/monkey/monkey_island/cc/models/report.py b/monkey/monkey_island/cc/models/report.py index fb1ab3cea..cbdde0dcb 100644 --- a/monkey/monkey_island/cc/models/report.py +++ b/monkey/monkey_island/cc/models/report.py @@ -3,9 +3,8 @@ from __future__ import annotations from bson import json_util from mongoengine import DictField, Document -from monkey_island.cc.models.utils import document_encryptor -from monkey_island.cc.models.utils.document_encryptor import SensitiveField -from monkey_island.cc.models.utils.field_encryptors.string_list_encryptor import StringListEncryptor +from monkey_island.cc.utils import SensitiveField, dict_encryptor +from monkey_island.cc.utils.field_encryptors import StringListEncryptor sensitive_fields = [ SensitiveField(path="overview.config_passwords", field_encryptor=StringListEncryptor) @@ -24,7 +23,7 @@ class Report(Document): @staticmethod def save_report(report_dict: dict): report_dict = _encode_dot_char_before_mongo_insert(report_dict) - report_dict = document_encryptor.encrypt(sensitive_fields, report_dict) + report_dict = dict_encryptor.encrypt(sensitive_fields, report_dict) Report.objects.delete() Report( overview=report_dict["overview"], @@ -37,7 +36,7 @@ class Report(Document): def get_report() -> dict: report_dict = Report.objects.first().to_mongo() return _decode_dot_char_before_mongo_insert( - document_encryptor.decrypt(sensitive_fields, report_dict) + dict_encryptor.decrypt(sensitive_fields, report_dict) ) diff --git a/monkey/monkey_island/cc/models/telemetries/__init__.py b/monkey/monkey_island/cc/models/telemetries/__init__.py index b3421bbd4..6c00785f7 100644 --- a/monkey/monkey_island/cc/models/telemetries/__init__.py +++ b/monkey/monkey_island/cc/models/telemetries/__init__.py @@ -1 +1 @@ -from .telemetry import Telemetry # noqa: F401 +from .telemetry_dal import save_telemetry, get_telemetry_by_query diff --git a/monkey/monkey_island/cc/models/telemetries/telemetry.py b/monkey/monkey_island/cc/models/telemetries/telemetry.py index 371484b85..3da3d2b10 100644 --- a/monkey/monkey_island/cc/models/telemetries/telemetry.py +++ b/monkey/monkey_island/cc/models/telemetries/telemetry.py @@ -1,21 +1,6 @@ -from __future__ import annotations - -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 -from monkey_island.cc.models.utils.field_encryptors.mimikatz_results_encryptor import ( - MimikatzResultsEncryptor, -) - -sensitive_fields = [ - SensitiveField("data.credentials", MimikatzResultsEncryptor), - SensitiveField("data.mimikatz", MimikatzResultsEncryptor), -] class Telemetry(Document): @@ -27,33 +12,3 @@ class Telemetry(Document): command_control_channel = EmbeddedDocumentField(CommandControlChannel) meta = {"strict": False} - - @staticmethod - def save_telemetry(telemetry_dict: dict): - try: - telemetry_dict = document_encryptor.encrypt(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() - - @staticmethod - 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/monkey_island/cc/models/telemetries/telemetry_dal.py b/monkey/monkey_island/cc/models/telemetries/telemetry_dal.py new file mode 100644 index 000000000..88e617725 --- /dev/null +++ b/monkey/monkey_island/cc/models/telemetries/telemetry_dal.py @@ -0,0 +1,46 @@ +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.utils import FieldNotFoundError, SensitiveField, dict_encryptor +from monkey_island.cc.utils.field_encryptors import MimikatzResultsEncryptor + +sensitive_fields = [ + SensitiveField("data.credentials", MimikatzResultsEncryptor), + SensitiveField("data.mimikatz", MimikatzResultsEncryptor), +] + + +def save_telemetry(telemetry_dict: dict): + try: + telemetry_dict = dict_encryptor.encrypt(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(dict_encryptor.decrypt(sensitive_fields, telemetry)) + except FieldNotFoundError: + decrypted_list.append(telemetry) + return decrypted_list diff --git a/monkey/monkey_island/cc/models/utils/__init__.py b/monkey/monkey_island/cc/models/utils/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/monkey/monkey_island/cc/models/utils/field_encryptors/__init__.py b/monkey/monkey_island/cc/models/utils/field_encryptors/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/monkey/monkey_island/cc/models/utils/field_encryptors/mimikatz_results_encryptor.py b/monkey/monkey_island/cc/models/utils/field_encryptors/mimikatz_results_encryptor.py deleted file mode 100644 index c924a64e9..000000000 --- a/monkey/monkey_island/cc/models/utils/field_encryptors/mimikatz_results_encryptor.py +++ /dev/null @@ -1,21 +0,0 @@ -from monkey_island.cc.models.utils.field_encryptors.i_field_encryptor import IFieldEncryptor -from monkey_island.cc.server_utils.encryptor import get_encryptor - - -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: - credentials[secret_type] = get_encryptor().enc(credentials[secret_type]) - return results - - @staticmethod - def decrypt(results: dict) -> dict: - for _, credentials in results.items(): - for secret_type in MimikatzResultsEncryptor.secret_types: - credentials[secret_type] = get_encryptor().dec(credentials[secret_type]) - return results diff --git a/monkey/monkey_island/cc/resources/blackbox/telemetry_blackbox_endpoint.py b/monkey/monkey_island/cc/resources/blackbox/telemetry_blackbox_endpoint.py index f784bc323..f1e958e3e 100644 --- a/monkey/monkey_island/cc/resources/blackbox/telemetry_blackbox_endpoint.py +++ b/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.models.telemetries import Telemetry +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(Telemetry.get_telemetry_by_query(find_query))} + return {"results": list(get_telemetry_by_query(find_query))} diff --git a/monkey/monkey_island/cc/resources/telemetry.py b/monkey/monkey_island/cc/resources/telemetry.py index 3ab2c1242..1158e82f0 100644 --- a/monkey/monkey_island/cc/resources/telemetry.py +++ b/monkey/monkey_island/cc/resources/telemetry.py @@ -9,7 +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.telemetry import Telemetry as TelemetryModel +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 @@ -38,7 +38,7 @@ class Telemetry(flask_restful.Resource): find_filter["timestamp"] = {"$gt": dateutil.parser.parse(timestamp)} result["objects"] = self.telemetry_to_displayed_telemetry( - TelemetryModel.get_telemetry_by_query(query=find_filter) + get_telemetry_by_query(query=find_filter) ) return result @@ -61,7 +61,7 @@ class Telemetry(flask_restful.Resource): process_telemetry(telemetry_json) - TelemetryModel.save_telemetry(telemetry_json) + save_telemetry(telemetry_json) return {}, 201 diff --git a/monkey/monkey_island/cc/services/reporting/report.py b/monkey/monkey_island/cc/services/reporting/report.py index 32675f1d5..17e6ce037 100644 --- a/monkey/monkey_island/cc/services/reporting/report.py +++ b/monkey/monkey_island/cc/services/reporting/report.py @@ -15,7 +15,7 @@ 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.telemetries import Telemetry +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, @@ -166,7 +166,7 @@ class ReportService: @staticmethod def _get_credentials_from_system_info_telems(): formatted_creds = [] - for telem in Telemetry.get_telemetry_by_query( + for telem in get_telemetry_by_query( {"telem_category": "system_info", "data.credentials": {"$exists": True}}, {"data.credentials": 1, "monkey_guid": 1}, ): diff --git a/monkey/monkey_island/cc/utils/__init__.py b/monkey/monkey_island/cc/utils/__init__.py new file mode 100644 index 000000000..abe040645 --- /dev/null +++ b/monkey/monkey_island/cc/utils/__init__.py @@ -0,0 +1 @@ +from .dict_encryptor import FieldNotFoundError, SensitiveField diff --git a/monkey/monkey_island/cc/models/utils/document_encryptor.py b/monkey/monkey_island/cc/utils/dict_encryptor.py similarity index 93% rename from monkey/monkey_island/cc/models/utils/document_encryptor.py rename to monkey/monkey_island/cc/utils/dict_encryptor.py index afec95afa..9a6d1d3d0 100644 --- a/monkey/monkey_island/cc/models/utils/document_encryptor.py +++ b/monkey/monkey_island/cc/utils/dict_encryptor.py @@ -3,7 +3,7 @@ from typing import Callable, List, Type import dpath.util -from monkey_island.cc.models.utils.field_encryptors.i_field_encryptor import IFieldEncryptor +from monkey_island.cc.utils.field_encryptors import IFieldEncryptor class FieldNotFoundError(Exception): diff --git a/monkey/monkey_island/cc/utils/field_encryptors/__init__.py b/monkey/monkey_island/cc/utils/field_encryptors/__init__.py new file mode 100644 index 000000000..7c938d25b --- /dev/null +++ b/monkey/monkey_island/cc/utils/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 diff --git a/monkey/monkey_island/cc/models/utils/field_encryptors/i_field_encryptor.py b/monkey/monkey_island/cc/utils/field_encryptors/i_field_encryptor.py similarity index 100% rename from monkey/monkey_island/cc/models/utils/field_encryptors/i_field_encryptor.py rename to monkey/monkey_island/cc/utils/field_encryptors/i_field_encryptor.py diff --git a/monkey/monkey_island/cc/utils/field_encryptors/mimikatz_results_encryptor.py b/monkey/monkey_island/cc/utils/field_encryptors/mimikatz_results_encryptor.py new file mode 100644 index 000000000..6708ec40c --- /dev/null +++ b/monkey/monkey_island/cc/utils/field_encryptors/mimikatz_results_encryptor.py @@ -0,0 +1,34 @@ +import logging + +from monkey_island.cc.server_utils.encryption import get_datastore_encryptor +from monkey_island.cc.utils.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 diff --git a/monkey/monkey_island/cc/models/utils/field_encryptors/string_list_encryptor.py b/monkey/monkey_island/cc/utils/field_encryptors/string_list_encryptor.py similarity index 81% rename from monkey/monkey_island/cc/models/utils/field_encryptors/string_list_encryptor.py rename to monkey/monkey_island/cc/utils/field_encryptors/string_list_encryptor.py index 089155289..f939c0e22 100644 --- a/monkey/monkey_island/cc/models/utils/field_encryptors/string_list_encryptor.py +++ b/monkey/monkey_island/cc/utils/field_encryptors/string_list_encryptor.py @@ -1,7 +1,7 @@ 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.utils.field_encryptors import IFieldEncryptor class StringListEncryptor(IFieldEncryptor): 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_dal.py similarity index 74% rename from monkey/tests/unit_tests/monkey_island/cc/models/telemetries/test_telemetry_model.py rename to monkey/tests/unit_tests/monkey_island/cc/models/telemetries/test_telemetry_dal.py index 860ae05fb..f25e8ffb3 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_dal.py @@ -4,11 +4,10 @@ from datetime import datetime import mongoengine import pytest -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, -) +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.utils import SensitiveField +from monkey_island.cc.utils.field_encryptors import MimikatzResultsEncryptor MOCK_CREDENTIALS = { "Vakaris": { @@ -57,7 +56,7 @@ MOCK_SENSITIVE_FIELDS = [ @pytest.fixture(autouse=True) def patch_sensitive_fields(monkeypatch): monkeypatch.setattr( - "monkey_island.cc.models.telemetries.telemetry.sensitive_fields", + "monkey_island.cc.models.telemetries.telemetry_dal.sensitive_fields", MOCK_SENSITIVE_FIELDS, ) @@ -65,13 +64,13 @@ 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) + monkeypatch.setattr("monkey_island.cc.models.telemetries.telemetry_dal.mongo", mongo) @pytest.mark.usefixtures("uses_database", "uses_encryptor") -def test_telemetry_encryption(monkeypatch): +def test_telemetry_encryption(): - Telemetry.save_telemetry(MOCK_TELEMETRY) + save_telemetry(MOCK_TELEMETRY) assert ( not Telemetry.objects.first()["data"]["credentials"]["user"]["password"] == MOCK_CREDENTIALS["user"]["password"] @@ -81,16 +80,16 @@ def test_telemetry_encryption(monkeypatch): == MOCK_CREDENTIALS["Vakaris"]["ntlm_hash"] ) assert ( - Telemetry.get_telemetry_by_query({})[0]["data"]["credentials"]["user"]["password"] + get_telemetry_by_query({})[0]["data"]["credentials"]["user"]["password"] == MOCK_CREDENTIALS["user"]["password"] ) assert ( - Telemetry.get_telemetry_by_query({})[0]["data"]["mimikatz"]["Vakaris"]["ntlm_hash"] + 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(monkeypatch, data_for_tests_dir): +def test_no_encryption_needed(): # Make sure telemetry save doesn't break when telemetry doesn't need encryption - Telemetry.save_telemetry(MOCK_NO_ENCRYPTION_NEEDED_TELEMETRY) + save_telemetry(MOCK_NO_ENCRYPTION_NEEDED_TELEMETRY) 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 index b88b7d0b0..0c8fd90de 100644 --- 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 @@ -4,8 +4,8 @@ from typing import List import pytest from monkey_island.cc.models import Report -from monkey_island.cc.models.utils.document_encryptor import SensitiveField -from monkey_island.cc.models.utils.field_encryptors.i_field_encryptor import IFieldEncryptor +from monkey_island.cc.utils import SensitiveField +from monkey_island.cc.utils.field_encryptors import IFieldEncryptor MOCK_SENSITIVE_FIELD_CONTENTS = ["the_string", "the_string2"] MOCK_REPORT_DICT = { diff --git a/monkey/tests/unit_tests/monkey_island/cc/models/utils/field_encryptors/test_string_list_encryptor.py b/monkey/tests/unit_tests/monkey_island/cc/models/utils/field_encryptors/test_string_list_encryptor.py index a93397392..ac46898c0 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/models/utils/field_encryptors/test_string_list_encryptor.py +++ b/monkey/tests/unit_tests/monkey_island/cc/models/utils/field_encryptors/test_string_list_encryptor.py @@ -1,7 +1,7 @@ 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.utils.field_encryptors import StringListEncryptor MOCK_STRING_LIST = ["test_1", "test_2"] EMPTY_LIST = [] diff --git a/monkey/tests/unit_tests/monkey_island/cc/services/reporting/test_report.py b/monkey/tests/unit_tests/monkey_island/cc/services/reporting/test_report.py index 9f845b263..6829965f2 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/services/reporting/test_report.py +++ b/monkey/tests/unit_tests/monkey_island/cc/services/reporting/test_report.py @@ -5,7 +5,7 @@ import mongoengine import pytest from bson import ObjectId -from monkey_island.cc.models.telemetries import Telemetry +from monkey_island.cc.models.telemetries import save_telemetry from monkey_island.cc.services.reporting.report import ReportService TELEM_ID = { @@ -138,12 +138,13 @@ NODE_DICT_FAILED_EXPLOITS["exploits"][1]["result"] = False def fake_mongo(monkeypatch): mongo = mongoengine.connection.get_connection() monkeypatch.setattr("monkey_island.cc.services.reporting.report.mongo", mongo) - monkeypatch.setattr("monkey_island.cc.models.telemetries.telemetry.mongo", mongo) + monkeypatch.setattr("monkey_island.cc.models.telemetries.telemetry_dal.mongo", mongo) monkeypatch.setattr("monkey_island.cc.services.node.mongo", mongo) return mongo -def test_get_stolen_creds_exploit(fake_mongo, uses_database): +@pytest.mark.usefixtures("uses_database") +def test_get_stolen_creds_exploit(fake_mongo): fake_mongo.db.telemetry.insert_one(EXPLOIT_TELEMETRY_TELEM) stolen_creds_exploit = ReportService.get_stolen_creds() @@ -155,9 +156,10 @@ def test_get_stolen_creds_exploit(fake_mongo, uses_database): assert expected_stolen_creds_exploit == stolen_creds_exploit -def test_get_stolen_creds_system_info(fake_mongo, uses_database): +@pytest.mark.usefixtures("uses_database") +def test_get_stolen_creds_system_info(fake_mongo): fake_mongo.db.monkey.insert_one(MONKEY_TELEM) - Telemetry.save_telemetry(SYSTEM_INFO_TELEMETRY_TELEM) + save_telemetry(SYSTEM_INFO_TELEMETRY_TELEM) stolen_creds_system_info = ReportService.get_stolen_creds() expected_stolen_creds_system_info = [ @@ -169,9 +171,10 @@ def test_get_stolen_creds_system_info(fake_mongo, uses_database): assert expected_stolen_creds_system_info == stolen_creds_system_info -def test_get_stolen_creds_no_creds(fake_mongo, uses_database): +@pytest.mark.usefixtures("uses_database") +def test_get_stolen_creds_no_creds(fake_mongo): fake_mongo.db.monkey.insert_one(MONKEY_TELEM) - Telemetry.save_telemetry(NO_CREDS_TELEMETRY_TELEM) + save_telemetry(NO_CREDS_TELEMETRY_TELEM) stolen_creds_no_creds = ReportService.get_stolen_creds() expected_stolen_creds_no_creds = []