Agent: Implement event serializer logic for SendAllAgentEventsToIsland

This commit is contained in:
Shreya Malviya 2022-09-14 16:00:56 +05:30
parent aeda96db06
commit bb8b4742af
2 changed files with 16 additions and 6 deletions

View File

@ -177,7 +177,7 @@ class InfectionMonkey:
if firewall.is_enabled():
firewall.add_firewall_rule()
_ = self._setup_agent_event_serializers()
self._agent_event_serializer_registry = self._setup_agent_event_serializers()
control_channel = ControlChannel(
self._control_client.server_address, GUID, self._control_client.proxies
@ -221,7 +221,10 @@ class InfectionMonkey:
event_queue = PyPubSubAgentEventQueue(Publisher())
InfectionMonkey._subscribe_events(
event_queue, propagation_credentials_repository, self._control_client.server_address
event_queue,
propagation_credentials_repository,
self._control_client.server_address,
self._agent_event_serializer_registry,
)
puppet = self._build_puppet(propagation_credentials_repository, event_queue)
@ -247,6 +250,7 @@ class InfectionMonkey:
event_queue: IAgentEventQueue,
propagation_credentials_repository: IPropagationCredentialsRepository,
server_address: str,
agent_event_serializer_registry: EventSerializerRegistry,
):
event_queue.subscribe_type(
CredentialsStolenEvent,
@ -254,7 +258,9 @@ class InfectionMonkey:
propagation_credentials_repository
),
)
event_queue.subscribe_all_events(SendAllAgentEventsToIsland(server_address).send_event)
event_queue.subscribe_all_events(
SendAllAgentEventsToIsland(server_address, agent_event_serializer_registry).send_event
)
@staticmethod
def _get_local_network_interfaces() -> List[IPv4Interface]:

View File

@ -6,6 +6,7 @@ from time import sleep
import requests
from common.common_consts.timeouts import MEDIUM_REQUEST_TIMEOUT
from common.event_serializers import EventSerializerRegistry
from common.event_serializers.i_event_serializer import JSONSerializable
from common.events import AbstractAgentEvent
@ -21,8 +22,11 @@ class SendAllAgentEventsToIsland:
Sends information about the events carried out by the Agent to the Island in batches
"""
def __init__(self, server_address: str):
def __init__(
self, server_address: str, agent_event_serializer_registry: EventSerializerRegistry
):
self._server_address = server_address
self._agent_event_serializer_registry = agent_event_serializer_registry
self._agent_events_to_island_sender = AgentEventsToIslandSender(self._server_address)
self._agent_events_to_island_sender.start()
@ -38,8 +42,8 @@ class SendAllAgentEventsToIsland:
)
def _serialize_event(self, event: AbstractAgentEvent):
# get serializer from registry and serialize
pass
serializer = self._agent_event_serializer_registry[event.__class__]
return serializer.serialize(event)
class AgentEventsToIslandSender: