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 diff --git a/monkey/monkey_island/cc/repository/__init__.py b/monkey/monkey_island/cc/repository/__init__.py index f299ec503..a172017f7 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 +from .mongo_agent_event_repository import MongoAgentEventRepository 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..fc347b46f --- /dev/null +++ b/monkey/monkey_island/cc/repository/agent_event_encryption.py @@ -0,0 +1,63 @@ +import json +from typing import Callable + +from common.agent_event_serializers import JSONSerializable +from common.agent_events import AbstractAgentEvent + +ENCRYPTED_PREFIX = "encrypted_" +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, +) -> 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 + :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") + + data = event_data.copy() + 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" + ) + del data[field] + + return data + + +def decrypt_event( + decrypt: Callable[[bytes], bytes], event_data: JSONSerializable +) -> JSONSerializable: + """ + Decrypt a serialized AbstractEventData + + :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 + """ + 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 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 79% rename from monkey/monkey_island/cc/repository/mongo_event_repository.py rename to monkey/monkey_island/cc/repository/mongo_agent_event_repository.py index 8a0538357..f483ac3d2 100644 --- a/monkey/monkey_island/cc/repository/mongo_event_repository.py +++ b/monkey/monkey_island/cc/repository/mongo_agent_event_repository.py @@ -6,25 +6,32 @@ from common.agent_event_serializers import EVENT_TYPE_FIELD, AgentEventSerialize 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 from .consts import MONGO_OBJECT_ID_KEY -class MongoEventRepository(IAgentEventRepository): +class MongoAgentEventRepository(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) + 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}") @@ -61,9 +68,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}) 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/monkey/monkey_island/cc/server_setup.py b/monkey/monkey_island/cc/server_setup.py index 710d45465..58e5adb40 100644 --- a/monkey/monkey_island/cc/server_setup.py +++ b/monkey/monkey_island/cc/server_setup.py @@ -23,10 +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 @@ -39,9 +35,9 @@ 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, ) @@ -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..d8caa96d5 100644 --- a/monkey/monkey_island/cc/services/initialize.py +++ b/monkey/monkey_island/cc/services/initialize.py @@ -10,6 +10,10 @@ from common.agent_configuration import ( DEFAULT_RANSOMWARE_AGENT_CONFIGURATION, AgentConfiguration, ) +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 @@ -33,12 +37,12 @@ from monkey_island.cc.repository import ( IUserRepository, JSONFileUserRepository, LocalStorageFileRepository, + MongoAgentEventRepository, MongoAgentRepository, MongoCredentialsRepository, 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) @@ -106,9 +111,7 @@ def _register_repositories(container: DIContainer, data_dir: Path): ICredentialsRepository, container.resolve(MongoCredentialsRepository) ) container.register_instance(IUserRepository, container.resolve(JSONFileUserRepository)) - - # TODO: Replace with MongoEventRepository - container.register_instance(IAgentEventRepository, StubbedEventRepository()) + container.register_instance(IAgentEventRepository, container.resolve(MongoAgentEventRepository)) container.register_instance(INodeRepository, container.resolve(MongoNodeRepository)) container.register_instance(IMachineRepository, container.resolve(MongoMachineRepository)) @@ -130,9 +133,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/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/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_agent_event_encryption.py b/monkey/tests/unit_tests/monkey_island/cc/repository/test_agent_event_encryption.py new file mode 100644 index 000000000..156fe9571 --- /dev/null +++ b/monkey/tests/unit_tests/monkey_island/cc/repository/test_agent_event_encryption.py @@ -0,0 +1,64 @@ +import uuid +from typing import Dict, Sequence + +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 + + +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 +def key_file(tmp_path): + return tmp_path / "test_key.bin" + + +@pytest.fixture +def serializer(): + return PydanticAgentEventSerializer(FakeAgentEvent) + + +def test_agent_event_encryption__encrypts(repository_encryptor, serializer): + data = serializer.serialize(EVENT) + encrypted_data = encrypt_event(repository_encryptor.encrypt, data) + + # 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_list_data"] is not EVENT.list_data + assert encrypted_data["encrypted_dict_data"] is not EVENT.dict_data + + +def test_agent_event_encryption__decrypts(repository_encryptor, serializer): + data = serializer.serialize(EVENT) + encrypted_data = encrypt_event(repository_encryptor.encrypt, 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(repository_encryptor): + data = "Not a dict." + + with pytest.raises(TypeError): + encrypt_event(repository_encryptor.encrypt, data, fields=[]) + + +def test_agent_event_encryption__decryption_throws(repository_encryptor): + data = "Not a dict." + + with pytest.raises(TypeError): + decrypt_event(repository_encryptor.decrypt, data) 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 52% 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 99a780b3d..a97ca8209 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 @@ -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, @@ -13,11 +15,16 @@ from common.agent_event_serializers import ( from common.agent_events import AbstractAgentEvent from monkey_island.cc.repository import ( IAgentEventRepository, - MongoEventRepository, + MongoAgentEventRepository, RemovalError, 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): @@ -54,8 +61,15 @@ 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 mongo_repository( + mongo_client, event_serializer_registry, repository_encryptor +) -> IAgentEventRepository: + return MongoAgentEventRepository(mongo_client, event_serializer_registry, repository_encryptor) @pytest.fixture @@ -75,9 +89,11 @@ 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 MongoEventRepository(error_raising_mongo_client, event_serializer_registry) + return MongoAgentEventRepository( + error_raising_mongo_client, event_serializer_registry, repository_encryptor + ) def assert_same_contents(a, b): @@ -86,7 +102,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() @@ -94,7 +110,16 @@ 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__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, ): event = FakeAgentEvent(source=uuid.uuid4()) @@ -103,48 +128,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) @@ -152,7 +179,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): @@ -160,7 +187,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 @@ -170,8 +197,63 @@ 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): 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]) 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..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 ( @@ -22,20 +21,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() @@ -180,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 @@ -194,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()) diff --git a/vulture_allowlist.py b/vulture_allowlist.py index 53001be59..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 @@ -304,7 +299,6 @@ IAgentEventRepository.save_event IAgentEventRepository.get_events_by_type IAgentEventRepository.get_events_by_tag IAgentEventRepository.get_events_by_source -MongoEventRepository # pydantic base models