Merge pull request #2315 from guardicore/2180-encrypt-event-data

2180 encrypt event data

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

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_agent_event_repository import MongoAgentEventRepository
-from .mongo_event_repository import MongoEventRepository

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

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

monkey/monkey_island/cc/repository/stubbed_event_repository.py

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

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:

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))
-
+    container.register_instance(IAgentEventRepository, container.resolve(MongoAgentEventRepository))
-    # TODO: Replace with MongoEventRepository
-    container.register_instance(IAgentEventRepository, StubbedEventRepository())
 
     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
     """

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

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

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)

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:
+def key_file(tmp_path):
-    return MongoEventRepository(mongo_client, event_serializer_registry)
+    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])

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(
+def get_all_credentials_in_mongo(mongo_client: MongoClient) -> Iterable[Mapping[str, Any]]:
-    mongo_client: MongoClient,
-) -> Iterable[Mapping[str, 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())

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 (
+from monkey_island.cc.repository import MongoAgentRepository, MongoMachineRepository
-    MongoAgentRepository,
-    MongoMachineRepository,
-    StubbedEventRepository,
-)
 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