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

This commit is contained in:
VakarisZ 2021-09-27 16:51:59 +03:00
parent 51f6fbe356
commit 46f263be5f
20 changed files with 123 additions and 104 deletions

View File

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

View File

@ -1 +1 @@
from .telemetry import Telemetry # noqa: F401
from .telemetry_dal import save_telemetry, get_telemetry_by_query

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

@ -0,0 +1,3 @@
from .i_field_encryptor import IFieldEncryptor
from .mimikatz_results_encryptor import MimikatzResultsEncryptor
from .string_list_encryptor import StringListEncryptor

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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