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

View File

@ -6,6 +6,7 @@ from time import sleep
import requests import requests
from common.common_consts.timeouts import MEDIUM_REQUEST_TIMEOUT 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.event_serializers.i_event_serializer import JSONSerializable
from common.events import AbstractAgentEvent 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 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._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 = AgentEventsToIslandSender(self._server_address)
self._agent_events_to_island_sender.start() self._agent_events_to_island_sender.start()
@ -38,8 +42,8 @@ class SendAllAgentEventsToIsland:
) )
def _serialize_event(self, event: AbstractAgentEvent): def _serialize_event(self, event: AbstractAgentEvent):
# get serializer from registry and serialize serializer = self._agent_event_serializer_registry[event.__class__]
pass return serializer.serialize(event)
class AgentEventsToIslandSender: class AgentEventsToIslandSender: