From 54fe2a6dcaacd4ddb6ad2794b52e67efb998f465 Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Fri, 16 Sep 2022 13:33:44 +0000 Subject: [PATCH] Island: Add encryption to MongoEventRepository --- .../cc/repository/mongo_event_repository.py | 62 +++++++++++++++++-- 1 file changed, 57 insertions(+), 5 deletions(-) diff --git a/monkey/monkey_island/cc/repository/mongo_event_repository.py b/monkey/monkey_island/cc/repository/mongo_event_repository.py index 8a0538357..8053d5634 100644 --- a/monkey/monkey_island/cc/repository/mongo_event_repository.py +++ b/monkey/monkey_island/cc/repository/mongo_event_repository.py @@ -1,30 +1,81 @@ -from typing import Any, Dict, MutableMapping, Sequence, Type +import json +from typing import Any, Callable, Dict, Iterable, MutableMapping, Sequence, Type from pymongo import MongoClient -from common.agent_event_serializers import EVENT_TYPE_FIELD, AgentEventSerializerRegistry +from common.agent_event_serializers import ( + EVENT_TYPE_FIELD, + AgentEventSerializerRegistry, + JSONSerializable, +) from common.agent_events import AbstractAgentEvent from common.types import AgentID from monkey_island.cc.repository import IAgentEventRepository +from monkey_island.cc.server_utils.encryption import ILockableEncryptor from . import RemovalError, RetrievalError, StorageError from .consts import MONGO_OBJECT_ID_KEY +ENCRYPTED_PREFIX = "encrypted_" + + +def get_fields_to_encrypt(event: AbstractAgentEvent): + return set(vars(AbstractAgentEvent)["__fields__"].keys()) ^ set(event.dict().keys()) + + +def encrypt_event( + encrypt: Callable[[bytes], bytes], + event_data: JSONSerializable, + fields: Iterable[str] = [], +) -> JSONSerializable: + if not isinstance(event_data, dict): + raise TypeError("Event encryption only supported for dict") + + for field in fields: + event_data[ENCRYPTED_PREFIX + field] = str( + encrypt(json.dumps(event_data[field]).encode()), "utf-8" + ) + del event_data[field] + + return event_data + + +def decrypt_event( + decrypt: Callable[[bytes], bytes], event_data: JSONSerializable +) -> JSONSerializable: + if not isinstance(event_data, dict): + raise TypeError("Event decryption only supported for dict") + + for field in event_data.keys(): + if field.startswith("encrypted_"): + event_data[field[len(ENCRYPTED_PREFIX) :]] = json.loads( + str(decrypt(event_data[field].encode()), "utf-8") + ) + del event_data[field] + + return event_data + class MongoEventRepository(IAgentEventRepository): """A repository for storing and retrieving events in MongoDB""" def __init__( - self, mongo_client: MongoClient, serializer_registry: AgentEventSerializerRegistry + self, + mongo_client: MongoClient, + serializer_registry: AgentEventSerializerRegistry, + encryptor: ILockableEncryptor, ): self._events_collection = mongo_client.monkey_island.events self._serializers = serializer_registry + self._encryptor = encryptor def save_event(self, event: AbstractAgentEvent): try: serializer = self._serializers[type(event)] serialized_event = serializer.serialize(event) - self._events_collection.insert_one(serialized_event) + fields = get_fields_to_encrypt(event) + encrypted_event = encrypt_event(self._encryptor.encrypt, serialized_event, fields) + self._events_collection.insert_one(encrypted_event) except Exception as err: raise StorageError(f"Error saving event: {err}") @@ -61,9 +112,10 @@ class MongoEventRepository(IAgentEventRepository): raise RemovalError(f"Error resetting the repository: {err}") def _deserialize(self, mongo_record: MutableMapping[str, Any]) -> AbstractAgentEvent: + decrypted_event = decrypt_event(self._encryptor.decrypt, mongo_record) event_type = mongo_record[EVENT_TYPE_FIELD] serializer = self._serializers[event_type] - return serializer.deserialize(mongo_record) + return serializer.deserialize(decrypted_event) def _query_events(self, query: Dict[Any, Any]) -> Sequence[AbstractAgentEvent]: serialized_events = self._events_collection.find(query, {MONGO_OBJECT_ID_KEY: False})