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