forked from p34709852/monkey
Island: Add MongoEventRepository to DI container
This commit is contained in:
parent
71110c61a7
commit
dc1eeefbc1
|
@ -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:
|
||||
|
|
|
@ -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
|
||||
"""
|
||||
|
|
|
@ -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
|
||||
|
|
Loading…
Reference in New Issue