diff --git a/monkey/common/event_serializers/__init__.py b/monkey/common/event_serializers/__init__.py index 1adb5e3dd..388edf6f3 100644 --- a/monkey/common/event_serializers/__init__.py +++ b/monkey/common/event_serializers/__init__.py @@ -1,3 +1,4 @@ -from .i_event_serialize import IEventSerializer, JSONSerializable +from .i_event_serializer import IEventSerializer, JSONSerializable from .event_serializer_registry import EventSerializerRegistry from .pydantic_event_serializer import PydanticEventSerializer +from .register import register_common_agent_event_serializers diff --git a/monkey/common/event_serializers/i_event_serialize.py b/monkey/common/event_serializers/i_event_serializer.py similarity index 100% rename from monkey/common/event_serializers/i_event_serialize.py rename to monkey/common/event_serializers/i_event_serializer.py diff --git a/monkey/common/event_serializers/register.py b/monkey/common/event_serializers/register.py new file mode 100644 index 000000000..6c533c9bf --- /dev/null +++ b/monkey/common/event_serializers/register.py @@ -0,0 +1,9 @@ +from common.events import CredentialsStolenEvent + +from . import EventSerializerRegistry, PydanticEventSerializer + + +def register_common_agent_event_serializers(event_serializer_registry: EventSerializerRegistry): + event_serializer_registry[CredentialsStolenEvent] = PydanticEventSerializer( + CredentialsStolenEvent + ) diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index abb0964ca..8a5c4e26b 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -11,6 +11,10 @@ from pubsub.core import Publisher import infection_monkey.tunnel as tunnel from common.event_queue import IAgentEventQueue, PyPubSubAgentEventQueue +from common.event_serializers import ( + EventSerializerRegistry, + register_common_agent_event_serializers, +) from common.events import CredentialsStolenEvent from common.network.network_utils import address_to_ip_port from common.utils.argparse_types import positive_int @@ -172,6 +176,8 @@ class InfectionMonkey: if firewall.is_enabled(): firewall.add_firewall_rule() + _ = self._setup_agent_event_serializers() + control_channel = ControlChannel( self._control_client.server_address, GUID, self._control_client.proxies ) @@ -194,6 +200,13 @@ class InfectionMonkey: register_signal_handlers(self._master) + # TODO: This is just a placeholder for now. We will modify/integrate it with PR #2279. + def _setup_agent_event_serializers(self) -> EventSerializerRegistry: + agent_event_serializer_registry = EventSerializerRegistry() + register_common_agent_event_serializers(agent_event_serializer_registry) + + return agent_event_serializer_registry + def _build_master(self): local_network_interfaces = InfectionMonkey._get_local_network_interfaces() diff --git a/monkey/monkey_island/cc/server_setup.py b/monkey/monkey_island/cc/server_setup.py index 9fc5fc030..95084d703 100644 --- a/monkey/monkey_island/cc/server_setup.py +++ b/monkey/monkey_island/cc/server_setup.py @@ -23,6 +23,10 @@ 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.event_serializers import ( # noqa: E402 + EventSerializerRegistry, + register_common_agent_event_serializers, +) 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 @@ -59,6 +63,7 @@ def run_monkey_island(): _send_analytics(deployment, version) container = _initialize_di_container(ip_addresses, version, config_options.data_dir) setup_island_event_handlers(container) + _setup_agent_event_serializers(container) _initialize_mongodb_connection(config_options.start_mongodb, config_options.data_dir) _start_island_server(ip_addresses, island_args.setup_only, config_options, container) @@ -129,6 +134,13 @@ def _initialize_di_container( return container +def _setup_agent_event_serializers(container: DIContainer): + agent_event_serializer_registry = EventSerializerRegistry() + register_common_agent_event_serializers(agent_event_serializer_registry) + + container.register_instance(EventSerializerRegistry, agent_event_serializer_registry) + + def _initialize_mongodb_connection(start_mongodb: bool, data_dir: Path): mongo_db_process = None if start_mongodb: