From 0d959e891a35ed4ea1f234a7ddc12246d38edc13 Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Fri, 16 Sep 2022 13:24:33 +0000 Subject: [PATCH 01/22] Island: Add functions to encrypt agent events --- .../cc/repository/agent_event_encryption.py | 44 ++++++++++++++ .../repository/test_agent_event_encryption.py | 57 +++++++++++++++++++ 2 files changed, 101 insertions(+) create mode 100644 monkey/monkey_island/cc/repository/agent_event_encryption.py create mode 100644 monkey/tests/unit_tests/monkey_island/cc/repository/test_agent_event_encryption.py diff --git a/monkey/monkey_island/cc/repository/agent_event_encryption.py b/monkey/monkey_island/cc/repository/agent_event_encryption.py new file mode 100644 index 000000000..83469613a --- /dev/null +++ b/monkey/monkey_island/cc/repository/agent_event_encryption.py @@ -0,0 +1,44 @@ +import json +from typing import Callable, Iterable + +from common.agent_event_serializers import JSONSerializable +from common.agent_events import AbstractAgentEvent + +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 diff --git a/monkey/tests/unit_tests/monkey_island/cc/repository/test_agent_event_encryption.py b/monkey/tests/unit_tests/monkey_island/cc/repository/test_agent_event_encryption.py new file mode 100644 index 000000000..0ebd0a5c6 --- /dev/null +++ b/monkey/tests/unit_tests/monkey_island/cc/repository/test_agent_event_encryption.py @@ -0,0 +1,57 @@ +import uuid + +import pytest + +from common.agent_event_serializers import PydanticAgentEventSerializer +from common.agent_events import AbstractAgentEvent +from monkey_island.cc.repository.agent_event_encryption import ( + decrypt_event, + encrypt_event, + get_fields_to_encrypt, +) +from monkey_island.cc.server_utils.encryption import RepositoryEncryptor + + +class FakeAgentEvent(AbstractAgentEvent): + data: str + + +@pytest.fixture +def key_file(tmp_path): + return tmp_path / "test_key.bin" + + +@pytest.fixture +def encryptor(key_file): + encryptor = RepositoryEncryptor(key_file) + encryptor.unlock(b"password") + return encryptor + + +@pytest.fixture +def serializer(): + return PydanticAgentEventSerializer() + + +def test_agent_event_encryption__encrypts(encryptor, serializer): + event = FakeAgentEvent(source=uuid.uuid4(), data="foo") + data = serializer.serialize(event) + fields = get_fields_to_encrypt(event) + encrypted_data = encrypt_event(encryptor.encrypt, data, fields) + + # Encrypted fields have the "encrypted_" prefix + assert "encrypted_data" in encrypted_data + assert encrypted_data["encrypted_data"] is not event.data + + +def test_agent_event_encryption__decrypts(encryptor, serializer): + event = FakeAgentEvent(source=uuid.uuid4(), data="foo") + + data = serializer.serialize(event) + fields = get_fields_to_encrypt(event) + encrypted_data = encrypt_event(encryptor.encrypt, data, fields) + + decrypted_data = decrypt_event(encryptor.decrypt, encrypted_data) + deserialized_event = serializer.deserialize(decrypted_data) + + assert deserialized_event == event From 54fe2a6dcaacd4ddb6ad2794b52e67efb998f465 Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Fri, 16 Sep 2022 13:33:44 +0000 Subject: [PATCH 02/22] 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}) From c19e50b7f120f41845b4841892a4701e31d7472b Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Fri, 16 Sep 2022 13:43:10 +0000 Subject: [PATCH 03/22] UT: Fix tests for event encryption --- .../repository/test_agent_event_encryption.py | 2 +- .../repository/test_mongo_event_repository.py | 21 +++++++++++++++---- 2 files changed, 18 insertions(+), 5 deletions(-) diff --git a/monkey/tests/unit_tests/monkey_island/cc/repository/test_agent_event_encryption.py b/monkey/tests/unit_tests/monkey_island/cc/repository/test_agent_event_encryption.py index 0ebd0a5c6..ff3b78b89 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/repository/test_agent_event_encryption.py +++ b/monkey/tests/unit_tests/monkey_island/cc/repository/test_agent_event_encryption.py @@ -30,7 +30,7 @@ def encryptor(key_file): @pytest.fixture def serializer(): - return PydanticAgentEventSerializer() + return PydanticAgentEventSerializer(FakeAgentEvent) def test_agent_event_encryption__encrypts(encryptor, serializer): diff --git a/monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_event_repository.py b/monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_event_repository.py index 99a780b3d..576f4a3de 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_event_repository.py +++ b/monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_event_repository.py @@ -18,6 +18,7 @@ from monkey_island.cc.repository import ( RetrievalError, StorageError, ) +from monkey_island.cc.server_utils.encryption import RepositoryEncryptor class FakeAgentEvent(AbstractAgentEvent): @@ -54,8 +55,20 @@ def mongo_client(event_serializer_registry): @pytest.fixture -def mongo_repository(mongo_client, event_serializer_registry) -> IAgentEventRepository: - return MongoEventRepository(mongo_client, event_serializer_registry) +def key_file(tmp_path): + return tmp_path / "test_key.bin" + + +@pytest.fixture +def encryptor(key_file): + encryptor = RepositoryEncryptor(key_file) + encryptor.unlock(b"password") + return encryptor + + +@pytest.fixture +def mongo_repository(mongo_client, event_serializer_registry, encryptor) -> IAgentEventRepository: + return MongoEventRepository(mongo_client, event_serializer_registry, encryptor) @pytest.fixture @@ -75,9 +88,9 @@ def error_raising_mongo_client(mongo_client) -> mongomock.MongoClient: @pytest.fixture def error_raising_mongo_repository( - error_raising_mongo_client, event_serializer_registry + error_raising_mongo_client, event_serializer_registry, encryptor ) -> IAgentEventRepository: - return MongoEventRepository(error_raising_mongo_client, event_serializer_registry) + return MongoEventRepository(error_raising_mongo_client, event_serializer_registry, encryptor) def assert_same_contents(a, b): From 71110c61a78f7ed012285a9e960a7fb2837d6adb Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Fri, 16 Sep 2022 14:00:14 +0000 Subject: [PATCH 04/22] Island: Add docstrings to agent event encryption --- .../cc/repository/agent_event_encryption.py | 23 +++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/monkey/monkey_island/cc/repository/agent_event_encryption.py b/monkey/monkey_island/cc/repository/agent_event_encryption.py index 83469613a..dd0d813cf 100644 --- a/monkey/monkey_island/cc/repository/agent_event_encryption.py +++ b/monkey/monkey_island/cc/repository/agent_event_encryption.py @@ -8,6 +8,9 @@ ENCRYPTED_PREFIX = "encrypted_" def get_fields_to_encrypt(event: AbstractAgentEvent): + """ + Get the fields of the event that are not part of the base AbstractAgentEvent. + """ return set(vars(AbstractAgentEvent)["__fields__"].keys()) ^ set(event.dict().keys()) @@ -16,6 +19,18 @@ def encrypt_event( event_data: JSONSerializable, fields: Iterable[str] = [], ) -> JSONSerializable: + """ + Encrypt a serialized AbstractAgentEvent + + The data is expected to be a dict. The encrypted fields will be given the + prefix "encrypted_". + + :param encrypt: Callable used to encrypt data + :param event_data: Serialized event to encrypt + :param fields: Fields to encrypt + :return: Serialized event with the fields encrypted + :raises TypeError: If the serialized data is not a dict + """ if not isinstance(event_data, dict): raise TypeError("Event encryption only supported for dict") @@ -31,6 +46,14 @@ def encrypt_event( def decrypt_event( decrypt: Callable[[bytes], bytes], event_data: JSONSerializable ) -> JSONSerializable: + """ + Decrypt a serialized AbstractEventData + + :param encrypt: Callable used to decrypt data + :param event_data: Serialized event to decrypt + :return: Serialized event with the fields decrypted + :raises TypeError: If the serialized data is not a dict + """ if not isinstance(event_data, dict): raise TypeError("Event decryption only supported for dict") From dc1eeefbc19c2fdec08b3cb5075833fb20bc156d Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Fri, 16 Sep 2022 14:36:40 +0000 Subject: [PATCH 05/22] Island: Add MongoEventRepository to DI container --- monkey/monkey_island/cc/server_setup.py | 14 +------------ .../monkey_island/cc/services/initialize.py | 20 +++++++++++++++---- vulture_allowlist.py | 1 - 3 files changed, 17 insertions(+), 18 deletions(-) diff --git a/monkey/monkey_island/cc/server_setup.py b/monkey/monkey_island/cc/server_setup.py index 710d45465..aca537411 100644 --- a/monkey/monkey_island/cc/server_setup.py +++ b/monkey/monkey_island/cc/server_setup.py @@ -23,11 +23,6 @@ if str(MONKEY_ISLAND_DIR_BASE_PATH) not in sys.path: sys.path.insert(0, MONKEY_ISLAND_DIR_BASE_PATH) from common import DIContainer # noqa: E402 -from common.agent_event_serializers import ( # noqa: E402 - AgentEventSerializerRegistry, - register_common_agent_event_serializers, -) -from common.network.network_utils import get_my_ip_addresses # noqa: E402 from common.version import get_version # noqa: E402 from monkey_island.cc.app import init_app # noqa: E402 from monkey_island.cc.arg_parser import IslandCmdArgs # noqa: E402 @@ -45,6 +40,7 @@ from monkey_island.cc.setup import ( # noqa: E402 setup_agent_event_handlers, setup_island_event_handlers, ) +from common.network.network_utils import get_my_ip_addresses # noqa: E402 from monkey_island.cc.setup.data_dir import IncompatibleDataDirectory, setup_data_dir # noqa: E402 from monkey_island.cc.setup.gevent_hub_error_handler import GeventHubErrorHandler # noqa: E402 from monkey_island.cc.setup.island_config_options import IslandConfigOptions # noqa: E402 @@ -71,7 +67,6 @@ def run_monkey_island(): container = _initialize_di_container(ip_addresses, version, config_options.data_dir) setup_island_event_handlers(container) setup_agent_event_handlers(container) - _setup_agent_event_serializers(container) _start_island_server(ip_addresses, island_args.setup_only, config_options, container) @@ -141,13 +136,6 @@ def _initialize_di_container( return container -def _setup_agent_event_serializers(container: DIContainer): - agent_event_serializer_registry = AgentEventSerializerRegistry() - register_common_agent_event_serializers(agent_event_serializer_registry) - - container.register_instance(AgentEventSerializerRegistry, agent_event_serializer_registry) - - def _initialize_mongodb_connection(start_mongodb: bool, data_dir: Path): mongo_db_process = None if start_mongodb: diff --git a/monkey/monkey_island/cc/services/initialize.py b/monkey/monkey_island/cc/services/initialize.py index 695db13c7..70556d811 100644 --- a/monkey/monkey_island/cc/services/initialize.py +++ b/monkey/monkey_island/cc/services/initialize.py @@ -12,6 +12,10 @@ from common.agent_configuration import ( ) from common.aws import AWSInstance from common.event_queue import IAgentEventQueue, PyPubSubAgentEventQueue +from common.agent_event_serializers import ( + AgentEventSerializerRegistry, + register_common_agent_event_serializers, +) from common.utils.file_utils import get_binary_io_sha256_hash from monkey_island.cc.event_queue import IIslandEventQueue, PyPubSubIslandEventQueue from monkey_island.cc.repository import ( @@ -35,10 +39,10 @@ from monkey_island.cc.repository import ( LocalStorageFileRepository, MongoAgentRepository, MongoCredentialsRepository, + MongoAgentEventRepository, MongoMachineRepository, MongoNodeRepository, RetrievalError, - StubbedEventRepository, ) from monkey_island.cc.server_utils.consts import MONKEY_ISLAND_ABS_PATH from monkey_island.cc.server_utils.encryption import ILockableEncryptor, RepositoryEncryptor @@ -68,6 +72,7 @@ def initialize_services(container: DIContainer, data_dir: Path): container.register_instance(IAgentEventQueue, container.resolve(PyPubSubAgentEventQueue)) container.register_instance(IIslandEventQueue, container.resolve(PyPubSubIslandEventQueue)) + _setup_agent_event_serializers(container) _register_repositories(container, data_dir) _register_services(container) @@ -107,8 +112,8 @@ def _register_repositories(container: DIContainer, data_dir: Path): ) container.register_instance(IUserRepository, container.resolve(JSONFileUserRepository)) - # TODO: Replace with MongoEventRepository - container.register_instance(IAgentEventRepository, StubbedEventRepository()) + # TODO: Figure out how to manage encryptor locking for MongoEventRepository + container.register_instance(IAgentEventRepository, container.resolve(MongoAgentEventRepository)) container.register_instance(INodeRepository, container.resolve(MongoNodeRepository)) container.register_instance(IMachineRepository, container.resolve(MongoMachineRepository)) @@ -130,9 +135,16 @@ def _build_agent_binary_repository(): return agent_binary_repository +def _setup_agent_event_serializers(container: DIContainer): + agent_event_serializer_registry = AgentEventSerializerRegistry() + register_common_agent_event_serializers(agent_event_serializer_registry) + + container.register_instance(AgentEventSerializerRegistry, agent_event_serializer_registry) + + def _log_agent_binary_hashes(agent_binary_repository: IAgentBinaryRepository): """ - Logs all the hashes of the agent executables for debbuging ease + Logs all the hashes of the agent executables for debugging ease :param agent_binary_repository: Used to retrieve the agent binaries """ diff --git a/vulture_allowlist.py b/vulture_allowlist.py index 53001be59..2843b4efb 100644 --- a/vulture_allowlist.py +++ b/vulture_allowlist.py @@ -304,7 +304,6 @@ IAgentEventRepository.save_event IAgentEventRepository.get_events_by_type IAgentEventRepository.get_events_by_tag IAgentEventRepository.get_events_by_source -MongoEventRepository # pydantic base models From c704d4a37ba0e89f20a3b89c455facf41607061d Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Fri, 16 Sep 2022 14:39:36 +0000 Subject: [PATCH 06/22] Island: Remove StubbedEventRepository --- .../monkey_island/cc/repository/__init__.py | 1 - .../cc/repository/stubbed_event_repository.py | 29 ------------------- vulture_allowlist.py | 7 +---- 3 files changed, 1 insertion(+), 36 deletions(-) delete mode 100644 monkey/monkey_island/cc/repository/stubbed_event_repository.py diff --git a/monkey/monkey_island/cc/repository/__init__.py b/monkey/monkey_island/cc/repository/__init__.py index f299ec503..2d46d352a 100644 --- a/monkey/monkey_island/cc/repository/__init__.py +++ b/monkey/monkey_island/cc/repository/__init__.py @@ -26,5 +26,4 @@ from .mongo_credentials_repository import MongoCredentialsRepository from .mongo_machine_repository import MongoMachineRepository from .mongo_agent_repository import MongoAgentRepository from .mongo_node_repository import MongoNodeRepository -from .stubbed_event_repository import StubbedEventRepository from .mongo_event_repository import MongoEventRepository diff --git a/monkey/monkey_island/cc/repository/stubbed_event_repository.py b/monkey/monkey_island/cc/repository/stubbed_event_repository.py deleted file mode 100644 index c4be4d058..000000000 --- a/monkey/monkey_island/cc/repository/stubbed_event_repository.py +++ /dev/null @@ -1,29 +0,0 @@ -from typing import Sequence, Type, TypeVar - -from common.agent_events import AbstractAgentEvent -from common.types import AgentID - -from . import IAgentEventRepository - -T = TypeVar("T", bound=AbstractAgentEvent) - - -# TODO: Remove this class after #2180 is complete -class StubbedEventRepository(IAgentEventRepository): - def save_event(self, event: AbstractAgentEvent): - return - - def get_events(self) -> Sequence[AbstractAgentEvent]: - return [] - - def get_events_by_type(self, event_type: Type[T]) -> Sequence[T]: - return [] - - def get_events_by_tag(self, tag: str) -> Sequence[AbstractAgentEvent]: - return [] - - def get_events_by_source(self, source: AgentID) -> Sequence[AbstractAgentEvent]: - return [] - - def reset(self): - return diff --git a/vulture_allowlist.py b/vulture_allowlist.py index 2843b4efb..40a2f5cdb 100644 --- a/vulture_allowlist.py +++ b/vulture_allowlist.py @@ -12,11 +12,7 @@ from infection_monkey.exploit.log4shell_utils.ldap_server import LDAPServerFacto from monkey_island.cc.event_queue import IslandEventTopic, PyPubSubIslandEventQueue from monkey_island.cc.models import Report from monkey_island.cc.models.networkmap import Arc, NetworkMap -from monkey_island.cc.repository import ( - MongoAgentRepository, - MongoMachineRepository, - StubbedEventRepository, -) +from monkey_island.cc.repository import MongoAgentRepository, MongoMachineRepository from monkey_island.cc.repository.attack.IMitigationsRepository import IMitigationsRepository from monkey_island.cc.repository.i_agent_event_repository import IAgentEventRepository from monkey_island.cc.repository.i_agent_repository import IAgentRepository @@ -283,7 +279,6 @@ IEventRepository.get_events IFindingRepository.get_findings MongoAgentRepository MongoMachineRepository -StubbedEventRepository key_list simulation netmap From 3f8c40a92c0cad11b9cd96b16821296b214e222e Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Fri, 16 Sep 2022 10:48:34 -0400 Subject: [PATCH 07/22] Island: Rename MongoEventRepository -> MongoAgentEventRepository --- .../monkey_island/cc/repository/__init__.py | 2 +- ...ory.py => mongo_agent_event_repository.py} | 2 +- .../monkey_island/cc/services/initialize.py | 6 ++-- ...y => test_mongo_agent_event_repository.py} | 34 ++++++++++--------- vulture_allowlist.py | 7 +++- 5 files changed, 29 insertions(+), 22 deletions(-) rename monkey/monkey_island/cc/repository/{mongo_event_repository.py => mongo_agent_event_repository.py} (98%) rename monkey/tests/unit_tests/monkey_island/cc/repository/{test_mongo_event_repository.py => test_mongo_agent_event_repository.py} (79%) diff --git a/monkey/monkey_island/cc/repository/__init__.py b/monkey/monkey_island/cc/repository/__init__.py index 2d46d352a..a172017f7 100644 --- a/monkey/monkey_island/cc/repository/__init__.py +++ b/monkey/monkey_island/cc/repository/__init__.py @@ -26,4 +26,4 @@ from .mongo_credentials_repository import MongoCredentialsRepository from .mongo_machine_repository import MongoMachineRepository from .mongo_agent_repository import MongoAgentRepository from .mongo_node_repository import MongoNodeRepository -from .mongo_event_repository import MongoEventRepository +from .mongo_agent_event_repository import MongoAgentEventRepository diff --git a/monkey/monkey_island/cc/repository/mongo_event_repository.py b/monkey/monkey_island/cc/repository/mongo_agent_event_repository.py similarity index 98% rename from monkey/monkey_island/cc/repository/mongo_event_repository.py rename to monkey/monkey_island/cc/repository/mongo_agent_event_repository.py index 8053d5634..2d807d896 100644 --- a/monkey/monkey_island/cc/repository/mongo_event_repository.py +++ b/monkey/monkey_island/cc/repository/mongo_agent_event_repository.py @@ -56,7 +56,7 @@ def decrypt_event( return event_data -class MongoEventRepository(IAgentEventRepository): +class MongoAgentEventRepository(IAgentEventRepository): """A repository for storing and retrieving events in MongoDB""" def __init__( diff --git a/monkey/monkey_island/cc/services/initialize.py b/monkey/monkey_island/cc/services/initialize.py index 70556d811..e240a0594 100644 --- a/monkey/monkey_island/cc/services/initialize.py +++ b/monkey/monkey_island/cc/services/initialize.py @@ -10,12 +10,12 @@ from common.agent_configuration import ( DEFAULT_RANSOMWARE_AGENT_CONFIGURATION, AgentConfiguration, ) -from common.aws import AWSInstance -from common.event_queue import IAgentEventQueue, PyPubSubAgentEventQueue from common.agent_event_serializers import ( AgentEventSerializerRegistry, register_common_agent_event_serializers, ) +from common.aws import AWSInstance +from common.event_queue import IAgentEventQueue, PyPubSubAgentEventQueue from common.utils.file_utils import get_binary_io_sha256_hash from monkey_island.cc.event_queue import IIslandEventQueue, PyPubSubIslandEventQueue from monkey_island.cc.repository import ( @@ -37,9 +37,9 @@ from monkey_island.cc.repository import ( IUserRepository, JSONFileUserRepository, LocalStorageFileRepository, + MongoAgentEventRepository, MongoAgentRepository, MongoCredentialsRepository, - MongoAgentEventRepository, MongoMachineRepository, MongoNodeRepository, RetrievalError, diff --git a/monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_event_repository.py b/monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_agent_event_repository.py similarity index 79% rename from monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_event_repository.py rename to monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_agent_event_repository.py index 576f4a3de..cebccb8f4 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_event_repository.py +++ b/monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_agent_event_repository.py @@ -13,7 +13,7 @@ from common.agent_event_serializers import ( from common.agent_events import AbstractAgentEvent from monkey_island.cc.repository import ( IAgentEventRepository, - MongoEventRepository, + MongoAgentEventRepository, RemovalError, RetrievalError, StorageError, @@ -68,7 +68,7 @@ def encryptor(key_file): @pytest.fixture def mongo_repository(mongo_client, event_serializer_registry, encryptor) -> IAgentEventRepository: - return MongoEventRepository(mongo_client, event_serializer_registry, encryptor) + return MongoAgentEventRepository(mongo_client, event_serializer_registry, encryptor) @pytest.fixture @@ -88,9 +88,9 @@ def error_raising_mongo_client(mongo_client) -> mongomock.MongoClient: @pytest.fixture def error_raising_mongo_repository( - error_raising_mongo_client, event_serializer_registry, encryptor + error_raising_mongo_client, event_serializer_registry ) -> IAgentEventRepository: - return MongoEventRepository(error_raising_mongo_client, event_serializer_registry, encryptor) + return MongoAgentEventRepository(error_raising_mongo_client, event_serializer_registry, encryptor) def assert_same_contents(a, b): @@ -99,7 +99,7 @@ def assert_same_contents(a, b): assert item in b -def test_mongo_event_repository__save_event(mongo_repository: IAgentEventRepository): +def test_mongo_agent_event_repository__save_event(mongo_repository: IAgentEventRepository): event = FakeAgentEvent(source=uuid.uuid4()) mongo_repository.save_event(event) events = mongo_repository.get_events() @@ -107,7 +107,7 @@ def test_mongo_event_repository__save_event(mongo_repository: IAgentEventReposit assert event in events -def test_mongo_event_repository__save_event_raises( +def test_mongo_agent_event_repository__save_event_raises( error_raising_mongo_repository: IAgentEventRepository, ): event = FakeAgentEvent(source=uuid.uuid4()) @@ -116,48 +116,50 @@ def test_mongo_event_repository__save_event_raises( error_raising_mongo_repository.save_event(event) -def test_mongo_event_repository__get_events(mongo_repository: IAgentEventRepository): +def test_mongo_agent_event_repository__get_events(mongo_repository: IAgentEventRepository): events = mongo_repository.get_events() assert_same_contents(events, EVENTS) -def test_mongo_event_repository__get_events_raises( +def test_mongo_agent_event_repository__get_events_raises( error_raising_mongo_repository: IAgentEventRepository, ): with pytest.raises(RetrievalError): error_raising_mongo_repository.get_events() -def test_mongo_event_repository__get_events_by_type(mongo_repository: IAgentEventRepository): +def test_mongo_agent_event_repository__get_events_by_type(mongo_repository: IAgentEventRepository): events = mongo_repository.get_events_by_type(FakeAgentItemEvent) expected_events = [EVENTS[3]] assert_same_contents(events, expected_events) -def test_mongo_event_repository__get_events_by_type_raises( +def test_mongo_agent_event_repository__get_events_by_type_raises( error_raising_mongo_repository: IAgentEventRepository, ): with pytest.raises(RetrievalError): error_raising_mongo_repository.get_events_by_type(FakeAgentItemEvent) -def test_mongo_event_repository__get_events_by_tag(mongo_repository: IAgentEventRepository): +def test_mongo_agent_event_repository__get_events_by_tag(mongo_repository: IAgentEventRepository): events = mongo_repository.get_events_by_tag("bar") expected_events = [EVENTS[1], EVENTS[2]] assert_same_contents(events, expected_events) -def test_mongo_event_repository__get_events_by_tag_raises( +def test_mongo_agent_event_repository__get_events_by_tag_raises( error_raising_mongo_repository: IAgentEventRepository, ): with pytest.raises(RetrievalError): error_raising_mongo_repository.get_events_by_tag("bar") -def test_mongo_event_repository__get_events_by_source(mongo_repository: IAgentEventRepository): +def test_mongo_agent_event_repository__get_events_by_source( + mongo_repository: IAgentEventRepository, +): source_event = EVENTS[2] events = mongo_repository.get_events_by_source(source_event.source) @@ -165,7 +167,7 @@ def test_mongo_event_repository__get_events_by_source(mongo_repository: IAgentEv assert_same_contents(events, expected_events) -def test_mongo_event_repository__get_events_by_source_raises( +def test_mongo_agent_event_repository__get_events_by_source_raises( error_raising_mongo_repository: IAgentEventRepository, ): with pytest.raises(RetrievalError): @@ -173,7 +175,7 @@ def test_mongo_event_repository__get_events_by_source_raises( error_raising_mongo_repository.get_events_by_source(source_event.source) -def test_mongo_event_repository__reset(mongo_repository: IAgentEventRepository): +def test_mongo_agent_event_repository__reset(mongo_repository: IAgentEventRepository): initial_events = mongo_repository.get_events() assert initial_events @@ -183,7 +185,7 @@ def test_mongo_event_repository__reset(mongo_repository: IAgentEventRepository): assert not events -def test_mongo_event_repository__reset_raises( +def test_mongo_agent_event_repository__reset_raises( error_raising_mongo_repository: IAgentEventRepository, ): with pytest.raises(RemovalError): diff --git a/vulture_allowlist.py b/vulture_allowlist.py index 40a2f5cdb..e2acb8949 100644 --- a/vulture_allowlist.py +++ b/vulture_allowlist.py @@ -12,7 +12,11 @@ from infection_monkey.exploit.log4shell_utils.ldap_server import LDAPServerFacto from monkey_island.cc.event_queue import IslandEventTopic, PyPubSubIslandEventQueue from monkey_island.cc.models import Report from monkey_island.cc.models.networkmap import Arc, NetworkMap -from monkey_island.cc.repository import MongoAgentRepository, MongoMachineRepository +from monkey_island.cc.repository import ( + MongoAgentEventRepository, + MongoAgentRepository, + MongoMachineRepository, +) from monkey_island.cc.repository.attack.IMitigationsRepository import IMitigationsRepository from monkey_island.cc.repository.i_agent_event_repository import IAgentEventRepository from monkey_island.cc.repository.i_agent_repository import IAgentRepository @@ -299,6 +303,7 @@ IAgentEventRepository.save_event IAgentEventRepository.get_events_by_type IAgentEventRepository.get_events_by_tag IAgentEventRepository.get_events_by_source +MongoAgentEventRepository # pydantic base models From 7634e00737b685f16a7c49d3932196776522352b Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Mon, 19 Sep 2022 06:56:41 -0400 Subject: [PATCH 08/22] Project: Remove MongoAgentEventRepository from vulture_allowlist.py --- vulture_allowlist.py | 7 +------ 1 file changed, 1 insertion(+), 6 deletions(-) diff --git a/vulture_allowlist.py b/vulture_allowlist.py index e2acb8949..40a2f5cdb 100644 --- a/vulture_allowlist.py +++ b/vulture_allowlist.py @@ -12,11 +12,7 @@ from infection_monkey.exploit.log4shell_utils.ldap_server import LDAPServerFacto from monkey_island.cc.event_queue import IslandEventTopic, PyPubSubIslandEventQueue from monkey_island.cc.models import Report from monkey_island.cc.models.networkmap import Arc, NetworkMap -from monkey_island.cc.repository import ( - MongoAgentEventRepository, - MongoAgentRepository, - MongoMachineRepository, -) +from monkey_island.cc.repository import MongoAgentRepository, MongoMachineRepository from monkey_island.cc.repository.attack.IMitigationsRepository import IMitigationsRepository from monkey_island.cc.repository.i_agent_event_repository import IAgentEventRepository from monkey_island.cc.repository.i_agent_repository import IAgentRepository @@ -303,7 +299,6 @@ IAgentEventRepository.save_event IAgentEventRepository.get_events_by_type IAgentEventRepository.get_events_by_tag IAgentEventRepository.get_events_by_source -MongoAgentEventRepository # pydantic base models From 2e63f476063f7177ec351bc6c4696ecdd6e98568 Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Mon, 19 Sep 2022 13:19:48 +0000 Subject: [PATCH 09/22] Island: Copy dict when encrypting event data --- .../cc/repository/mongo_agent_event_repository.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/monkey/monkey_island/cc/repository/mongo_agent_event_repository.py b/monkey/monkey_island/cc/repository/mongo_agent_event_repository.py index 2d807d896..336039319 100644 --- a/monkey/monkey_island/cc/repository/mongo_agent_event_repository.py +++ b/monkey/monkey_island/cc/repository/mongo_agent_event_repository.py @@ -31,13 +31,14 @@ def encrypt_event( if not isinstance(event_data, dict): raise TypeError("Event encryption only supported for dict") + data = event_data.copy() for field in fields: - event_data[ENCRYPTED_PREFIX + field] = str( + data[ENCRYPTED_PREFIX + field] = str( encrypt(json.dumps(event_data[field]).encode()), "utf-8" ) - del event_data[field] + del data[field] - return event_data + return data def decrypt_event( @@ -46,14 +47,15 @@ def decrypt_event( if not isinstance(event_data, dict): raise TypeError("Event decryption only supported for dict") + data = event_data.copy() for field in event_data.keys(): if field.startswith("encrypted_"): - event_data[field[len(ENCRYPTED_PREFIX) :]] = json.loads( + data[field[len(ENCRYPTED_PREFIX) :]] = json.loads( str(decrypt(event_data[field].encode()), "utf-8") ) - del event_data[field] + del data[field] - return event_data + return data class MongoAgentEventRepository(IAgentEventRepository): From a9f7262582f4f8f3fefdbe765feb2682ec84e0a2 Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Mon, 19 Sep 2022 13:27:58 +0000 Subject: [PATCH 10/22] UT: Fix formatting --- .../cc/repository/test_mongo_agent_event_repository.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_agent_event_repository.py b/monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_agent_event_repository.py index cebccb8f4..19a724971 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_agent_event_repository.py +++ b/monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_agent_event_repository.py @@ -90,7 +90,9 @@ def error_raising_mongo_client(mongo_client) -> mongomock.MongoClient: def error_raising_mongo_repository( error_raising_mongo_client, event_serializer_registry ) -> IAgentEventRepository: - return MongoAgentEventRepository(error_raising_mongo_client, event_serializer_registry, encryptor) + return MongoAgentEventRepository( + error_raising_mongo_client, event_serializer_registry, encryptor + ) def assert_same_contents(a, b): From aa0d5f2e197e589a94608fd30ced0e674cb91028 Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Mon, 19 Sep 2022 13:28:53 +0000 Subject: [PATCH 11/22] Island: Fix import ordering --- monkey/monkey_island/cc/server_setup.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/monkey/monkey_island/cc/server_setup.py b/monkey/monkey_island/cc/server_setup.py index aca537411..58e5adb40 100644 --- a/monkey/monkey_island/cc/server_setup.py +++ b/monkey/monkey_island/cc/server_setup.py @@ -23,6 +23,7 @@ if str(MONKEY_ISLAND_DIR_BASE_PATH) not in sys.path: sys.path.insert(0, MONKEY_ISLAND_DIR_BASE_PATH) from common import DIContainer # noqa: E402 +from common.network.network_utils import get_my_ip_addresses # noqa: E402 from common.version import get_version # noqa: E402 from monkey_island.cc.app import init_app # noqa: E402 from monkey_island.cc.arg_parser import IslandCmdArgs # noqa: E402 @@ -34,13 +35,12 @@ from monkey_island.cc.server_utils.consts import ( # noqa: E402 ) from monkey_island.cc.server_utils.island_logger import reset_logger, setup_logging # noqa: E402 from monkey_island.cc.services.initialize import initialize_services # noqa: E402 -from monkey_island.cc.setup import island_config_options_validator # noqa: E402 from monkey_island.cc.setup import ( # noqa: E402 PyWSGILoggingFilter, + island_config_options_validator, setup_agent_event_handlers, setup_island_event_handlers, ) -from common.network.network_utils import get_my_ip_addresses # noqa: E402 from monkey_island.cc.setup.data_dir import IncompatibleDataDirectory, setup_data_dir # noqa: E402 from monkey_island.cc.setup.gevent_hub_error_handler import GeventHubErrorHandler # noqa: E402 from monkey_island.cc.setup.island_config_options import IslandConfigOptions # noqa: E402 From f635c2cd5f606b821ebcc291d8139292dc2eb655 Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Mon, 19 Sep 2022 14:25:17 +0000 Subject: [PATCH 12/22] BB: Removed unneccessary logs from ZerologonAnalyzer --- envs/monkey_zoo/blackbox/analyzers/zerologon_analyzer.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/envs/monkey_zoo/blackbox/analyzers/zerologon_analyzer.py b/envs/monkey_zoo/blackbox/analyzers/zerologon_analyzer.py index b19be0273..3feb8e2ee 100644 --- a/envs/monkey_zoo/blackbox/analyzers/zerologon_analyzer.py +++ b/envs/monkey_zoo/blackbox/analyzers/zerologon_analyzer.py @@ -28,9 +28,7 @@ class ZerologonAnalyzer(Analyzer): def _analyze_credential_gathering(self) -> bool: propagation_credentials = self.island_client.get_propagation_credentials() - self.log.add_entry(f"Credentials from endpoint: {propagation_credentials}") credentials_on_island = ZerologonAnalyzer._get_relevant_credentials(propagation_credentials) - self.log.add_entry(f"Relevant credentials: {credentials_on_island}") return self._is_all_credentials_in_list(credentials_on_island) @staticmethod From ccadfccf5e20ca4248fdeba66b3a16a610ef6136 Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Mon, 19 Sep 2022 14:45:49 +0000 Subject: [PATCH 13/22] Island: Remove duplicate encryption code --- .../cc/repository/agent_event_encryption.py | 14 ++--- .../mongo_agent_event_repository.py | 51 ++----------------- 2 files changed, 11 insertions(+), 54 deletions(-) diff --git a/monkey/monkey_island/cc/repository/agent_event_encryption.py b/monkey/monkey_island/cc/repository/agent_event_encryption.py index dd0d813cf..8904773ed 100644 --- a/monkey/monkey_island/cc/repository/agent_event_encryption.py +++ b/monkey/monkey_island/cc/repository/agent_event_encryption.py @@ -34,13 +34,14 @@ def encrypt_event( if not isinstance(event_data, dict): raise TypeError("Event encryption only supported for dict") + data = event_data.copy() for field in fields: - event_data[ENCRYPTED_PREFIX + field] = str( + data[ENCRYPTED_PREFIX + field] = str( encrypt(json.dumps(event_data[field]).encode()), "utf-8" ) - del event_data[field] + del data[field] - return event_data + return data def decrypt_event( @@ -57,11 +58,12 @@ def decrypt_event( if not isinstance(event_data, dict): raise TypeError("Event decryption only supported for dict") + data = event_data.copy() for field in event_data.keys(): if field.startswith("encrypted_"): - event_data[field[len(ENCRYPTED_PREFIX) :]] = json.loads( + data[field[len(ENCRYPTED_PREFIX) :]] = json.loads( str(decrypt(event_data[field].encode()), "utf-8") ) - del event_data[field] + del data[field] - return event_data + return data diff --git a/monkey/monkey_island/cc/repository/mongo_agent_event_repository.py b/monkey/monkey_island/cc/repository/mongo_agent_event_repository.py index 336039319..0194f3d3d 100644 --- a/monkey/monkey_island/cc/repository/mongo_agent_event_repository.py +++ b/monkey/monkey_island/cc/repository/mongo_agent_event_repository.py @@ -1,62 +1,17 @@ -import json -from typing import Any, Callable, Dict, Iterable, MutableMapping, Sequence, Type +from typing import Any, Dict, MutableMapping, Sequence, Type from pymongo import MongoClient -from common.agent_event_serializers import ( - EVENT_TYPE_FIELD, - AgentEventSerializerRegistry, - JSONSerializable, -) +from common.agent_event_serializers import EVENT_TYPE_FIELD, AgentEventSerializerRegistry 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 .agent_event_encryption import decrypt_event, encrypt_event, get_fields_to_encrypt 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") - - data = event_data.copy() - for field in fields: - data[ENCRYPTED_PREFIX + field] = str( - encrypt(json.dumps(event_data[field]).encode()), "utf-8" - ) - del data[field] - - return 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") - - data = event_data.copy() - for field in event_data.keys(): - if field.startswith("encrypted_"): - data[field[len(ENCRYPTED_PREFIX) :]] = json.loads( - str(decrypt(event_data[field].encode()), "utf-8") - ) - del data[field] - - return data - class MongoAgentEventRepository(IAgentEventRepository): """A repository for storing and retrieving events in MongoDB""" From 7038179711cd81f268826625b82e21e72bba0861 Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Mon, 19 Sep 2022 14:47:10 +0000 Subject: [PATCH 14/22] UT: Add tests for event encryption exceptions --- .../cc/repository/test_agent_event_encryption.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/monkey/tests/unit_tests/monkey_island/cc/repository/test_agent_event_encryption.py b/monkey/tests/unit_tests/monkey_island/cc/repository/test_agent_event_encryption.py index ff3b78b89..07f0300ce 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/repository/test_agent_event_encryption.py +++ b/monkey/tests/unit_tests/monkey_island/cc/repository/test_agent_event_encryption.py @@ -55,3 +55,17 @@ def test_agent_event_encryption__decrypts(encryptor, serializer): deserialized_event = serializer.deserialize(decrypted_data) assert deserialized_event == event + + +def test_agent_event_encryption__encryption_throws(encryptor): + data = "Not a dict." + + with pytest.raises(TypeError): + encrypt_event(encryptor.encrypt, data, fields=[]) + + +def test_agent_event_encryption__decryption_throws(encryptor): + data = "Not a dict." + + with pytest.raises(TypeError): + decrypt_event(encryptor.decrypt, data) From 875a54aa8fab13435a3963e3e6318a1300437e3c Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Tue, 20 Sep 2022 11:35:10 +0000 Subject: [PATCH 15/22] Island: Remove stale comment --- monkey/monkey_island/cc/services/initialize.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/monkey/monkey_island/cc/services/initialize.py b/monkey/monkey_island/cc/services/initialize.py index e240a0594..d8caa96d5 100644 --- a/monkey/monkey_island/cc/services/initialize.py +++ b/monkey/monkey_island/cc/services/initialize.py @@ -111,8 +111,6 @@ def _register_repositories(container: DIContainer, data_dir: Path): ICredentialsRepository, container.resolve(MongoCredentialsRepository) ) container.register_instance(IUserRepository, container.resolve(JSONFileUserRepository)) - - # TODO: Figure out how to manage encryptor locking for MongoEventRepository container.register_instance(IAgentEventRepository, container.resolve(MongoAgentEventRepository)) container.register_instance(INodeRepository, container.resolve(MongoNodeRepository)) From e0cebd144a261aededad17fc34291a52404c51a9 Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Tue, 20 Sep 2022 11:37:11 +0000 Subject: [PATCH 16/22] UT: Use mock encryptor --- .../unit_tests/monkey_island/cc/conftest.py | 18 +++++++++++++++++- .../test_mongo_agent_event_repository.py | 18 ++++++------------ .../test_mongo_credentials_repository.py | 14 -------------- 3 files changed, 23 insertions(+), 27 deletions(-) diff --git a/monkey/tests/unit_tests/monkey_island/cc/conftest.py b/monkey/tests/unit_tests/monkey_island/cc/conftest.py index 4eef545d5..35d6fa5d2 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/conftest.py +++ b/monkey/tests/unit_tests/monkey_island/cc/conftest.py @@ -1,10 +1,26 @@ +from unittest.mock import MagicMock + import pytest from tests.unit_tests.monkey_island.cc.mongomock_fixtures import * # noqa: F401,F403,E402 -from monkey_island.cc.server_utils.encryption import unlock_datastore_encryptor +from monkey_island.cc.server_utils.encryption import ILockableEncryptor, unlock_datastore_encryptor @pytest.fixture def uses_encryptor(data_for_tests_dir): secret = "m0nk3y_u53r:3cr3t_p455w0rd" unlock_datastore_encryptor(data_for_tests_dir, secret) + + +def reverse(data: bytes) -> bytes: + return bytes(reversed(data)) + + +@pytest.fixture +def repository_encryptor(): + # NOTE: Tests will fail if any inputs to this mock encryptor are palindromes. + repository_encryptor = MagicMock(spec=ILockableEncryptor) + repository_encryptor.encrypt = MagicMock(side_effect=reverse) + repository_encryptor.decrypt = MagicMock(side_effect=reverse) + + return repository_encryptor diff --git a/monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_agent_event_repository.py b/monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_agent_event_repository.py index 19a724971..baaa37d30 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_agent_event_repository.py +++ b/monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_agent_event_repository.py @@ -18,7 +18,6 @@ from monkey_island.cc.repository import ( RetrievalError, StorageError, ) -from monkey_island.cc.server_utils.encryption import RepositoryEncryptor class FakeAgentEvent(AbstractAgentEvent): @@ -60,15 +59,10 @@ def key_file(tmp_path): @pytest.fixture -def encryptor(key_file): - encryptor = RepositoryEncryptor(key_file) - encryptor.unlock(b"password") - return encryptor - - -@pytest.fixture -def mongo_repository(mongo_client, event_serializer_registry, encryptor) -> IAgentEventRepository: - return MongoAgentEventRepository(mongo_client, event_serializer_registry, encryptor) +def mongo_repository( + mongo_client, event_serializer_registry, repository_encryptor +) -> IAgentEventRepository: + return MongoAgentEventRepository(mongo_client, event_serializer_registry, repository_encryptor) @pytest.fixture @@ -88,10 +82,10 @@ def error_raising_mongo_client(mongo_client) -> mongomock.MongoClient: @pytest.fixture def error_raising_mongo_repository( - error_raising_mongo_client, event_serializer_registry + error_raising_mongo_client, event_serializer_registry, repository_encryptor ) -> IAgentEventRepository: return MongoAgentEventRepository( - error_raising_mongo_client, event_serializer_registry, encryptor + error_raising_mongo_client, event_serializer_registry, repository_encryptor ) diff --git a/monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_credentials_repository.py b/monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_credentials_repository.py index f921c9df5..a88cba05c 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_credentials_repository.py +++ b/monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_credentials_repository.py @@ -22,20 +22,6 @@ CONFIGURED_CREDENTIALS = CREDENTIALS[0:3] STOLEN_CREDENTIALS = CREDENTIALS[3:] -def reverse(data: bytes) -> bytes: - return bytes(reversed(data)) - - -@pytest.fixture -def repository_encryptor(): - # NOTE: Tests will fail if any inputs to this mock encryptor are palindromes. - repository_encryptor = MagicMock(spec=ILockableEncryptor) - repository_encryptor.encrypt = MagicMock(side_effect=reverse) - repository_encryptor.decrypt = MagicMock(side_effect=reverse) - - return repository_encryptor - - @pytest.fixture def mongo_client(): return mongomock.MongoClient() From 2d03e497e9149f7b2ab8436ecea56ec2ef6b2516 Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Tue, 20 Sep 2022 11:38:17 +0000 Subject: [PATCH 17/22] UT: Show seq and dict are encrypted in events --- .../repository/test_agent_event_encryption.py | 25 ++++++++++++------- 1 file changed, 16 insertions(+), 9 deletions(-) diff --git a/monkey/tests/unit_tests/monkey_island/cc/repository/test_agent_event_encryption.py b/monkey/tests/unit_tests/monkey_island/cc/repository/test_agent_event_encryption.py index 07f0300ce..947c50f80 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/repository/test_agent_event_encryption.py +++ b/monkey/tests/unit_tests/monkey_island/cc/repository/test_agent_event_encryption.py @@ -1,4 +1,5 @@ import uuid +from typing import Dict, Sequence import pytest @@ -14,6 +15,13 @@ from monkey_island.cc.server_utils.encryption import RepositoryEncryptor class FakeAgentEvent(AbstractAgentEvent): data: str + list_data: Sequence[str] + dict_data: Dict[str, str] + + +EVENT = FakeAgentEvent( + source=uuid.uuid4(), data="foo", list_data=["abc", "def"], dict_data={"abc": "def"} +) @pytest.fixture @@ -34,27 +42,26 @@ def serializer(): def test_agent_event_encryption__encrypts(encryptor, serializer): - event = FakeAgentEvent(source=uuid.uuid4(), data="foo") - data = serializer.serialize(event) - fields = get_fields_to_encrypt(event) + data = serializer.serialize(EVENT) + fields = get_fields_to_encrypt(EVENT) encrypted_data = encrypt_event(encryptor.encrypt, data, fields) # Encrypted fields have the "encrypted_" prefix assert "encrypted_data" in encrypted_data - assert encrypted_data["encrypted_data"] is not event.data + assert encrypted_data["encrypted_data"] is not EVENT.data + assert encrypted_data["encrypted_list_data"] is not EVENT.list_data + assert encrypted_data["encrypted_dict_data"] is not EVENT.dict_data def test_agent_event_encryption__decrypts(encryptor, serializer): - event = FakeAgentEvent(source=uuid.uuid4(), data="foo") - - data = serializer.serialize(event) - fields = get_fields_to_encrypt(event) + data = serializer.serialize(EVENT) + fields = get_fields_to_encrypt(EVENT) encrypted_data = encrypt_event(encryptor.encrypt, data, fields) decrypted_data = decrypt_event(encryptor.decrypt, encrypted_data) deserialized_event = serializer.deserialize(decrypted_data) - assert deserialized_event == event + assert deserialized_event == EVENT def test_agent_event_encryption__encryption_throws(encryptor): From 644f3628a5061d8809ca9bf3a9df08c1202adf03 Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Tue, 20 Sep 2022 13:47:12 +0000 Subject: [PATCH 18/22] UT: Move mongo functions into a module --- .../monkey_island/cc/repository/mongo.py | 24 +++++++++++++++++ .../test_mongo_credentials_repository.py | 26 ++----------------- 2 files changed, 26 insertions(+), 24 deletions(-) create mode 100644 monkey/tests/unit_tests/monkey_island/cc/repository/mongo.py diff --git a/monkey/tests/unit_tests/monkey_island/cc/repository/mongo.py b/monkey/tests/unit_tests/monkey_island/cc/repository/mongo.py new file mode 100644 index 000000000..26fa6340b --- /dev/null +++ b/monkey/tests/unit_tests/monkey_island/cc/repository/mongo.py @@ -0,0 +1,24 @@ +from typing import Iterable + +from pymongo import MongoClient +from pymongo.collection import Collection +from pymongo.database import Database + + +def get_all_collections_in_mongo(mongo_client: MongoClient) -> Iterable[Collection]: + collections = [ + collection + for db in get_all_databases_in_mongo(mongo_client) + for collection in get_all_collections_in_database(db) + ] + + assert len(collections) > 0 + return collections + + +def get_all_databases_in_mongo(mongo_client) -> Iterable[Database]: + return (mongo_client[db_name] for db_name in mongo_client.list_database_names()) + + +def get_all_collections_in_database(db: Database) -> Iterable[Collection]: + return (db[collection_name] for collection_name in db.list_collection_names()) diff --git a/monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_credentials_repository.py b/monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_credentials_repository.py index a88cba05c..bf745bead 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_credentials_repository.py +++ b/monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_credentials_repository.py @@ -4,9 +4,8 @@ from unittest.mock import MagicMock import mongomock import pytest from pymongo import MongoClient -from pymongo.collection import Collection -from pymongo.database import Database from tests.data_for_tests.propagation_credentials import CREDENTIALS +from tests.unit_tests.monkey_island.cc.repository.mongo import get_all_collections_in_mongo from common.credentials import Credentials from monkey_island.cc.repository import ( @@ -166,9 +165,7 @@ def check_if_stored_credentials_encrypted(mongo_client: MongoClient, original_cr assert "***" not in value.decode() -def get_all_credentials_in_mongo( - mongo_client: MongoClient, -) -> Iterable[Mapping[str, Mapping[str, Any]]]: +def get_all_credentials_in_mongo(mongo_client: MongoClient) -> Iterable[Mapping[str, Any]]: encrypted_credentials = [] # Loop through all databases and collections and search for credentials. We don't want the tests @@ -180,22 +177,3 @@ def get_all_credentials_in_mongo( encrypted_credentials.append(mc) return encrypted_credentials - - -def get_all_collections_in_mongo(mongo_client: MongoClient) -> Iterable[Collection]: - collections = [ - collection - for db in get_all_databases_in_mongo(mongo_client) - for collection in get_all_collections_in_database(db) - ] - - assert len(collections) > 0 - return collections - - -def get_all_databases_in_mongo(mongo_client) -> Iterable[Database]: - return (mongo_client[db_name] for db_name in mongo_client.list_database_names()) - - -def get_all_collections_in_database(db: Database) -> Iterable[Collection]: - return (db[collection_name] for collection_name in db.list_collection_names()) From c0869aebbace26861ade09d78e7b0da11bf8f231 Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Tue, 20 Sep 2022 13:59:54 +0000 Subject: [PATCH 19/22] Island: Simplify agent encryption calls --- .../cc/repository/agent_event_encryption.py | 16 +++++----------- .../repository/mongo_agent_event_repository.py | 5 ++--- .../cc/repository/test_agent_event_encryption.py | 12 +++--------- 3 files changed, 10 insertions(+), 23 deletions(-) diff --git a/monkey/monkey_island/cc/repository/agent_event_encryption.py b/monkey/monkey_island/cc/repository/agent_event_encryption.py index 8904773ed..3e16855fe 100644 --- a/monkey/monkey_island/cc/repository/agent_event_encryption.py +++ b/monkey/monkey_island/cc/repository/agent_event_encryption.py @@ -1,23 +1,17 @@ import json -from typing import Callable, Iterable +from typing import Callable from common.agent_event_serializers import JSONSerializable from common.agent_events import AbstractAgentEvent ENCRYPTED_PREFIX = "encrypted_" - - -def get_fields_to_encrypt(event: AbstractAgentEvent): - """ - Get the fields of the event that are not part of the base AbstractAgentEvent. - """ - return set(vars(AbstractAgentEvent)["__fields__"].keys()) ^ set(event.dict().keys()) +ABSTRACT_AGENT_EVENT_FIELDS = vars(AbstractAgentEvent)["__fields__"].keys() +SERIALIZED_EVENT_FIELDS = set(ABSTRACT_AGENT_EVENT_FIELDS) | set(["type"]) def encrypt_event( encrypt: Callable[[bytes], bytes], event_data: JSONSerializable, - fields: Iterable[str] = [], ) -> JSONSerializable: """ Encrypt a serialized AbstractAgentEvent @@ -27,7 +21,6 @@ def encrypt_event( :param encrypt: Callable used to encrypt data :param event_data: Serialized event to encrypt - :param fields: Fields to encrypt :return: Serialized event with the fields encrypted :raises TypeError: If the serialized data is not a dict """ @@ -35,7 +28,8 @@ def encrypt_event( raise TypeError("Event encryption only supported for dict") data = event_data.copy() - for field in fields: + fields_to_encrypt = SERIALIZED_EVENT_FIELDS ^ set(event_data.keys()) + for field in fields_to_encrypt: data[ENCRYPTED_PREFIX + field] = str( encrypt(json.dumps(event_data[field]).encode()), "utf-8" ) diff --git a/monkey/monkey_island/cc/repository/mongo_agent_event_repository.py b/monkey/monkey_island/cc/repository/mongo_agent_event_repository.py index 0194f3d3d..f483ac3d2 100644 --- a/monkey/monkey_island/cc/repository/mongo_agent_event_repository.py +++ b/monkey/monkey_island/cc/repository/mongo_agent_event_repository.py @@ -9,7 +9,7 @@ from monkey_island.cc.repository import IAgentEventRepository from monkey_island.cc.server_utils.encryption import ILockableEncryptor from . import RemovalError, RetrievalError, StorageError -from .agent_event_encryption import decrypt_event, encrypt_event, get_fields_to_encrypt +from .agent_event_encryption import decrypt_event, encrypt_event from .consts import MONGO_OBJECT_ID_KEY @@ -30,8 +30,7 @@ class MongoAgentEventRepository(IAgentEventRepository): try: serializer = self._serializers[type(event)] serialized_event = serializer.serialize(event) - fields = get_fields_to_encrypt(event) - encrypted_event = encrypt_event(self._encryptor.encrypt, serialized_event, fields) + encrypted_event = encrypt_event(self._encryptor.encrypt, serialized_event) self._events_collection.insert_one(encrypted_event) except Exception as err: raise StorageError(f"Error saving event: {err}") diff --git a/monkey/tests/unit_tests/monkey_island/cc/repository/test_agent_event_encryption.py b/monkey/tests/unit_tests/monkey_island/cc/repository/test_agent_event_encryption.py index 947c50f80..35873a6cb 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/repository/test_agent_event_encryption.py +++ b/monkey/tests/unit_tests/monkey_island/cc/repository/test_agent_event_encryption.py @@ -5,11 +5,7 @@ import pytest from common.agent_event_serializers import PydanticAgentEventSerializer from common.agent_events import AbstractAgentEvent -from monkey_island.cc.repository.agent_event_encryption import ( - decrypt_event, - encrypt_event, - get_fields_to_encrypt, -) +from monkey_island.cc.repository.agent_event_encryption import decrypt_event, encrypt_event from monkey_island.cc.server_utils.encryption import RepositoryEncryptor @@ -43,8 +39,7 @@ def serializer(): def test_agent_event_encryption__encrypts(encryptor, serializer): data = serializer.serialize(EVENT) - fields = get_fields_to_encrypt(EVENT) - encrypted_data = encrypt_event(encryptor.encrypt, data, fields) + encrypted_data = encrypt_event(encryptor.encrypt, data) # Encrypted fields have the "encrypted_" prefix assert "encrypted_data" in encrypted_data @@ -55,8 +50,7 @@ def test_agent_event_encryption__encrypts(encryptor, serializer): def test_agent_event_encryption__decrypts(encryptor, serializer): data = serializer.serialize(EVENT) - fields = get_fields_to_encrypt(EVENT) - encrypted_data = encrypt_event(encryptor.encrypt, data, fields) + encrypted_data = encrypt_event(encryptor.encrypt, data) decrypted_data = decrypt_event(encryptor.decrypt, encrypted_data) deserialized_event = serializer.deserialize(decrypted_data) From dea7b4f74aeb148827a46a1fe15b39bd96aed50d Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Tue, 20 Sep 2022 14:03:10 +0000 Subject: [PATCH 20/22] UT: Use encryptor fixture in agent encryption tests --- .../repository/test_agent_event_encryption.py | 26 +++++++------------ 1 file changed, 9 insertions(+), 17 deletions(-) diff --git a/monkey/tests/unit_tests/monkey_island/cc/repository/test_agent_event_encryption.py b/monkey/tests/unit_tests/monkey_island/cc/repository/test_agent_event_encryption.py index 35873a6cb..156fe9571 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/repository/test_agent_event_encryption.py +++ b/monkey/tests/unit_tests/monkey_island/cc/repository/test_agent_event_encryption.py @@ -6,7 +6,6 @@ import pytest from common.agent_event_serializers import PydanticAgentEventSerializer from common.agent_events import AbstractAgentEvent from monkey_island.cc.repository.agent_event_encryption import decrypt_event, encrypt_event -from monkey_island.cc.server_utils.encryption import RepositoryEncryptor class FakeAgentEvent(AbstractAgentEvent): @@ -25,21 +24,14 @@ def key_file(tmp_path): return tmp_path / "test_key.bin" -@pytest.fixture -def encryptor(key_file): - encryptor = RepositoryEncryptor(key_file) - encryptor.unlock(b"password") - return encryptor - - @pytest.fixture def serializer(): return PydanticAgentEventSerializer(FakeAgentEvent) -def test_agent_event_encryption__encrypts(encryptor, serializer): +def test_agent_event_encryption__encrypts(repository_encryptor, serializer): data = serializer.serialize(EVENT) - encrypted_data = encrypt_event(encryptor.encrypt, data) + encrypted_data = encrypt_event(repository_encryptor.encrypt, data) # Encrypted fields have the "encrypted_" prefix assert "encrypted_data" in encrypted_data @@ -48,25 +40,25 @@ def test_agent_event_encryption__encrypts(encryptor, serializer): assert encrypted_data["encrypted_dict_data"] is not EVENT.dict_data -def test_agent_event_encryption__decrypts(encryptor, serializer): +def test_agent_event_encryption__decrypts(repository_encryptor, serializer): data = serializer.serialize(EVENT) - encrypted_data = encrypt_event(encryptor.encrypt, data) + encrypted_data = encrypt_event(repository_encryptor.encrypt, data) - decrypted_data = decrypt_event(encryptor.decrypt, encrypted_data) + decrypted_data = decrypt_event(repository_encryptor.decrypt, encrypted_data) deserialized_event = serializer.deserialize(decrypted_data) assert deserialized_event == EVENT -def test_agent_event_encryption__encryption_throws(encryptor): +def test_agent_event_encryption__encryption_throws(repository_encryptor): data = "Not a dict." with pytest.raises(TypeError): - encrypt_event(encryptor.encrypt, data, fields=[]) + encrypt_event(repository_encryptor.encrypt, data, fields=[]) -def test_agent_event_encryption__decryption_throws(encryptor): +def test_agent_event_encryption__decryption_throws(repository_encryptor): data = "Not a dict." with pytest.raises(TypeError): - decrypt_event(encryptor.decrypt, data) + decrypt_event(repository_encryptor.decrypt, data) From a83186f532bcaab8ef0744a8baa2d1629e902f8b Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Tue, 20 Sep 2022 14:05:51 +0000 Subject: [PATCH 21/22] UT: Test that stored events are encrypted --- .../test_mongo_agent_event_repository.py | 73 ++++++++++++++++++- 1 file changed, 72 insertions(+), 1 deletion(-) diff --git a/monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_agent_event_repository.py b/monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_agent_event_repository.py index baaa37d30..a97ca8209 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_agent_event_repository.py +++ b/monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_agent_event_repository.py @@ -1,10 +1,12 @@ import uuid -from typing import List +from typing import Any, Iterable, List, Mapping from unittest.mock import MagicMock import mongomock import pytest from pydantic import Field +from pymongo import MongoClient +from tests.unit_tests.monkey_island.cc.repository.mongo import get_all_collections_in_mongo from common.agent_event_serializers import ( AgentEventSerializerRegistry, @@ -18,6 +20,11 @@ from monkey_island.cc.repository import ( RetrievalError, StorageError, ) +from monkey_island.cc.repository.agent_event_encryption import ( + ENCRYPTED_PREFIX, + SERIALIZED_EVENT_FIELDS, +) +from monkey_island.cc.repository.consts import MONGO_OBJECT_ID_KEY class FakeAgentEvent(AbstractAgentEvent): @@ -103,6 +110,15 @@ def test_mongo_agent_event_repository__save_event(mongo_repository: IAgentEventR assert event in events +def test_mongo_agent_event_repository__saved_events_are_encrypted( + mongo_repository: IAgentEventRepository, mongo_client +): + event = FakeAgentEvent(source=uuid.uuid4()) + mongo_repository.save_event(event) + + assert_events_are_encrypted(mongo_client, [event]) + + def test_mongo_agent_event_repository__save_event_raises( error_raising_mongo_repository: IAgentEventRepository, ): @@ -186,3 +202,58 @@ def test_mongo_agent_event_repository__reset_raises( ): with pytest.raises(RemovalError): error_raising_mongo_repository.reset() + + +def get_all_events_in_mongo( + mongo_client: MongoClient, +) -> Iterable[Mapping[str, Mapping[str, Any]]]: + events = [] + + for collection in get_all_collections_in_mongo(mongo_client): + mongo_events = collection.find({}, {MONGO_OBJECT_ID_KEY: False}) + for mongo_event in mongo_events: + events.append(mongo_event) + + return events + + +def is_encrypted_event(original_event: AbstractAgentEvent, other_event) -> bool: + """ + Checks if an event is an encrypted version of the original + + - The number of fields match + - The AbstractAgentEvent fields match + - The remaining fields have a matching encrypted_ prefix + - The remaining fields are the encrypted version of the original fields + """ + + event = original_event.dict(simplify=True) + + # Note: The serializer adds a "type" field + event["type"] = type(original_event).__name__ + + if len(event.keys()) != len(other_event.keys()): + return False + + for field in event.keys(): + if field in SERIALIZED_EVENT_FIELDS: + if event[field] != other_event[field]: + return False + else: + encrypted_field = ENCRYPTED_PREFIX + field + if ( + encrypted_field not in other_event.keys() + or event[field] == other_event[encrypted_field] + ): + return False + + return True + + +def assert_events_are_encrypted( + mongo_client: MongoClient, original_events: Iterable[AbstractAgentEvent] +): + stored_events = get_all_events_in_mongo(mongo_client) + + for event in original_events: + assert any([is_encrypted_event(event, se) for se in stored_events]) From b27dea54e7005ef1a70886fb43c151997172927a Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Tue, 20 Sep 2022 14:08:37 +0000 Subject: [PATCH 22/22] Island: Fix typo in docstring --- monkey/monkey_island/cc/repository/agent_event_encryption.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/monkey/monkey_island/cc/repository/agent_event_encryption.py b/monkey/monkey_island/cc/repository/agent_event_encryption.py index 3e16855fe..fc347b46f 100644 --- a/monkey/monkey_island/cc/repository/agent_event_encryption.py +++ b/monkey/monkey_island/cc/repository/agent_event_encryption.py @@ -44,7 +44,7 @@ def decrypt_event( """ Decrypt a serialized AbstractEventData - :param encrypt: Callable used to decrypt data + :param decrypt: Callable used to decrypt data :param event_data: Serialized event to decrypt :return: Serialized event with the fields decrypted :raises TypeError: If the serialized data is not a dict