Island: Add MongoEventRepository to DI container

This commit is contained in:
Kekoa Kaaikala 2022-09-16 14:36:40 +00:00
parent 71110c61a7
commit dc1eeefbc1
3 changed files with 17 additions and 18 deletions

View File

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

View File

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

View File

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