Separate the telemetry document from telemetry_dal, also extracted external interface into __init__.py files

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.utils import SensitiveField, dict_encryptor
-from monkey_island.cc.models.utils.document_encryptor import SensitiveField
+from monkey_island.cc.utils.field_encryptors import StringListEncryptor
-from monkey_island.cc.models.utils.field_encryptors.string_list_encryptor 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)
         )
 
 

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

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

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

monkey/monkey_island/cc/models/utils/field_encryptors/mimikatz_results_encryptor.py

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

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

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
 

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},
         ):

monkey/monkey_island/cc/utils/__init__.py

@@ -0,0 +1 @@
+from .dict_encryptor import FieldNotFoundError, SensitiveField

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

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

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

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

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.telemetries import get_telemetry_by_query, save_telemetry
-from monkey_island.cc.models.utils.document_encryptor import SensitiveField
+from monkey_island.cc.models.telemetries.telemetry import Telemetry
-from monkey_island.cc.models.utils.field_encryptors.mimikatz_results_encryptor import (
+from monkey_island.cc.utils import SensitiveField
-    MimikatzResultsEncryptor,
+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)

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.utils import SensitiveField
-from monkey_island.cc.models.utils.field_encryptors.i_field_encryptor import IFieldEncryptor
+from monkey_island.cc.utils.field_encryptors import IFieldEncryptor
 
 MOCK_SENSITIVE_FIELD_CONTENTS = ["the_string", "the_string2"]
 MOCK_REPORT_DICT = {

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 = []

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 = []