From 76cf34b5f0c138bcd00c6b9b17ec7ab7a49d97d8 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Mon, 5 Sep 2022 19:00:20 +0530 Subject: [PATCH 01/56] Common: Rename AbstractEvent -> AbstractAgentEvent --- monkey/common/event_queue/i_event_queue.py | 6 +++--- .../common/event_queue/pypubsub_event_queue.py | 16 ++++++++-------- monkey/common/event_queue/types.py | 4 ++-- .../event_serializer_registry.py | 16 +++++++++------- .../event_serializers/i_event_serialize.py | 6 +++--- monkey/common/events/__init__.py | 2 +- monkey/common/events/abstract_event.py | 6 +++--- .../common/events/credentials_stolen_events.py | 4 ++-- 8 files changed, 31 insertions(+), 29 deletions(-) diff --git a/monkey/common/event_queue/i_event_queue.py b/monkey/common/event_queue/i_event_queue.py index b205b8a98..128d0517b 100644 --- a/monkey/common/event_queue/i_event_queue.py +++ b/monkey/common/event_queue/i_event_queue.py @@ -1,7 +1,7 @@ from abc import ABC, abstractmethod from typing import Type -from common.events import AbstractEvent +from common.events import AbstractAgentEvent from . import EventSubscriber @@ -22,7 +22,7 @@ class IEventQueue(ABC): pass @abstractmethod - def subscribe_type(self, event_type: Type[AbstractEvent], subscriber: EventSubscriber): + def subscribe_type(self, event_type: Type[AbstractAgentEvent], subscriber: EventSubscriber): """ Subscribes a subscriber to the specified event type @@ -44,7 +44,7 @@ class IEventQueue(ABC): pass @abstractmethod - def publish(self, event: AbstractEvent): + def publish(self, event: AbstractAgentEvent): """ Publishes an event with the given data diff --git a/monkey/common/event_queue/pypubsub_event_queue.py b/monkey/common/event_queue/pypubsub_event_queue.py index 6dde760a8..9c3e4dc6d 100644 --- a/monkey/common/event_queue/pypubsub_event_queue.py +++ b/monkey/common/event_queue/pypubsub_event_queue.py @@ -3,7 +3,7 @@ from typing import Type from pubsub.core import Publisher -from common.events import AbstractEvent +from common.events import AbstractAgentEvent from . import EventSubscriber, IEventQueue @@ -20,7 +20,7 @@ class PyPubSubEventQueue(IEventQueue): def subscribe_all_events(self, subscriber: EventSubscriber): self._subscribe(_ALL_EVENTS_TOPIC, subscriber) - def subscribe_type(self, event_type: Type[AbstractEvent], subscriber: EventSubscriber): + def subscribe_type(self, event_type: Type[AbstractAgentEvent], subscriber: EventSubscriber): # pypubsub.pub.subscribe needs a string as the topic/event name event_type_topic = PyPubSubEventQueue._get_type_topic(event_type) self._subscribe(event_type_topic, subscriber) @@ -60,31 +60,31 @@ class PyPubSubEventQueue(IEventQueue): # scope. Adding subscribers to self._refs prevents them from ever going out of scope. self._refs.append(subscriber) - def publish(self, event: AbstractEvent): + def publish(self, event: AbstractAgentEvent): self._publish_to_all_events_topic(event) self._publish_to_type_topic(event) self._publish_to_tags_topics(event) - def _publish_to_all_events_topic(self, event: AbstractEvent): + def _publish_to_all_events_topic(self, event: AbstractAgentEvent): self._publish_event(_ALL_EVENTS_TOPIC, event) - def _publish_to_type_topic(self, event: AbstractEvent): + def _publish_to_type_topic(self, event: AbstractAgentEvent): event_type_topic = PyPubSubEventQueue._get_type_topic(event.__class__) self._publish_event(event_type_topic, event) - def _publish_to_tags_topics(self, event: AbstractEvent): + def _publish_to_tags_topics(self, event: AbstractAgentEvent): for tag in event.tags: tag_topic = PyPubSubEventQueue._get_tag_topic(tag) self._publish_event(tag_topic, event) - def _publish_event(self, topic: str, event: AbstractEvent): + def _publish_event(self, topic: str, event: AbstractAgentEvent): logger.debug(f"Publishing a {event.__class__.__name__} event to {topic}") self._pypubsub_publisher.sendMessage(topic, event=event) # Appending a unique string to the topics for type and tags prevents bugs caused by collisions # between type names and tag names. @staticmethod - def _get_type_topic(event_type: Type[AbstractEvent]) -> str: + def _get_type_topic(event_type: Type[AbstractAgentEvent]) -> str: return f"{event_type.__name__}-type" @staticmethod diff --git a/monkey/common/event_queue/types.py b/monkey/common/event_queue/types.py index f78651306..f3f5f4ef8 100644 --- a/monkey/common/event_queue/types.py +++ b/monkey/common/event_queue/types.py @@ -1,5 +1,5 @@ from typing import Callable -from common.events import AbstractEvent +from common.events import AbstractAgentEvent -EventSubscriber = Callable[[AbstractEvent], None] +EventSubscriber = Callable[[AbstractAgentEvent], None] diff --git a/monkey/common/event_serializers/event_serializer_registry.py b/monkey/common/event_serializers/event_serializer_registry.py index fd02045d1..b4f43f7b0 100644 --- a/monkey/common/event_serializers/event_serializer_registry.py +++ b/monkey/common/event_serializers/event_serializer_registry.py @@ -1,7 +1,7 @@ from typing import Type, Union from common.event_serializers import IEventSerializer -from common.events import AbstractEvent +from common.events import AbstractAgentEvent class EventSerializerRegistry: @@ -21,9 +21,11 @@ class EventSerializerRegistry: def __init__(self): self._registry = {} - def __setitem__(self, event_class: Type[AbstractEvent], event_serializer: IEventSerializer): - if not issubclass(event_class, AbstractEvent): - raise TypeError(f"Event class must be of type: {AbstractEvent.__name__}") + def __setitem__( + self, event_class: Type[AbstractAgentEvent], event_serializer: IEventSerializer + ): + if not issubclass(event_class, AbstractAgentEvent): + raise TypeError(f"Event class must be of type: {AbstractAgentEvent.__name__}") if not isinstance(event_serializer, IEventSerializer): raise TypeError(f"Event serializer must be of type: {IEventSerializer.__name__}") @@ -31,10 +33,10 @@ class EventSerializerRegistry: self._registry[event_class] = event_serializer self._registry[event_class.__name__] = event_serializer - def __getitem__(self, event_class: Union[str, Type[AbstractEvent]]) -> IEventSerializer: - if not (isinstance(event_class, str) or issubclass(event_class, AbstractEvent)): + def __getitem__(self, event_class: Union[str, Type[AbstractAgentEvent]]) -> IEventSerializer: + if not (isinstance(event_class, str) or issubclass(event_class, AbstractAgentEvent)): raise TypeError( - f"Registry get key {event_class} must be of type: {AbstractEvent.__name__} or " + f"Registry get key {event_class} must be of type: {AbstractAgentEvent.__name__} or " f"{str.__name__}" ) diff --git a/monkey/common/event_serializers/i_event_serialize.py b/monkey/common/event_serializers/i_event_serialize.py index 5544e6ede..75d95be3c 100644 --- a/monkey/common/event_serializers/i_event_serialize.py +++ b/monkey/common/event_serializers/i_event_serialize.py @@ -1,7 +1,7 @@ from abc import ABC, abstractmethod from typing import Dict, List, Union -from common.events import AbstractEvent +from common.events import AbstractAgentEvent JSONSerializable = Union[ Dict[str, "JSONSerializable"], List["JSONSerializable"], int, str, float, bool, None @@ -14,7 +14,7 @@ class IEventSerializer(ABC): """ @abstractmethod - def serialize(self, event: AbstractEvent) -> JSONSerializable: + def serialize(self, event: AbstractAgentEvent) -> JSONSerializable: """ Serializes an event @@ -24,7 +24,7 @@ class IEventSerializer(ABC): pass @abstractmethod - def deserialize(self, serialized_event: JSONSerializable) -> AbstractEvent: + def deserialize(self, serialized_event: JSONSerializable) -> AbstractAgentEvent: """ Deserializes an event diff --git a/monkey/common/events/__init__.py b/monkey/common/events/__init__.py index da5b88234..b996452f7 100644 --- a/monkey/common/events/__init__.py +++ b/monkey/common/events/__init__.py @@ -1,2 +1,2 @@ -from .abstract_event import AbstractEvent +from .abstract_event import AbstractAgentEvent from .credentials_stolen_events import CredentialsStolenEvent diff --git a/monkey/common/events/abstract_event.py b/monkey/common/events/abstract_event.py index 4b4edbf99..89f515187 100644 --- a/monkey/common/events/abstract_event.py +++ b/monkey/common/events/abstract_event.py @@ -7,13 +7,13 @@ from uuid import UUID, getnode @dataclass(frozen=True) -class AbstractEvent(ABC): +class AbstractAgentEvent(ABC): """ An event that was initiated or observed by an agent Agents perform actions and collect data. These actions and data are represented as "events". - Subtypes of `AbstractEvent` will have additional properties that provide context and information - about the event. + Subtypes of `AbstractAgentEvent` will have additional properties that provide context and + information about the event. Attributes: :param source: The UUID of the agent that observed the event diff --git a/monkey/common/events/credentials_stolen_events.py b/monkey/common/events/credentials_stolen_events.py index ffd7707e7..184ee81c3 100644 --- a/monkey/common/events/credentials_stolen_events.py +++ b/monkey/common/events/credentials_stolen_events.py @@ -3,11 +3,11 @@ from typing import Sequence from common.credentials import Credentials -from . import AbstractEvent +from . import AbstractAgentEvent @dataclass(frozen=True) -class CredentialsStolenEvent(AbstractEvent): +class CredentialsStolenEvent(AbstractAgentEvent): """ An event that occurs when an agent collects credentials from the victim From ce43a46a9c0aa5d9dc3d22af22afc0353a4c36f0 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Mon, 5 Sep 2022 19:01:33 +0530 Subject: [PATCH 02/56] Common: Rename abstract_event.py -> abstract_agent_event.py --- monkey/common/events/__init__.py | 2 +- .../events/{abstract_event.py => abstract_agent_event.py} | 0 2 files changed, 1 insertion(+), 1 deletion(-) rename monkey/common/events/{abstract_event.py => abstract_agent_event.py} (100%) diff --git a/monkey/common/events/__init__.py b/monkey/common/events/__init__.py index b996452f7..c89f41f0e 100644 --- a/monkey/common/events/__init__.py +++ b/monkey/common/events/__init__.py @@ -1,2 +1,2 @@ -from .abstract_event import AbstractAgentEvent +from .abstract_agent_event import AbstractAgentEvent from .credentials_stolen_events import CredentialsStolenEvent diff --git a/monkey/common/events/abstract_event.py b/monkey/common/events/abstract_agent_event.py similarity index 100% rename from monkey/common/events/abstract_event.py rename to monkey/common/events/abstract_agent_event.py From f577e48d725aa9d8931fdf29538bdabc9a479b79 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Mon, 5 Sep 2022 19:03:23 +0530 Subject: [PATCH 03/56] UT: Use AbstractAgentEvent in place of AbstractEvent --- .../common/event_queue/test_pypubsub_event_queue.py | 12 ++++++------ .../test_event_serializer_registry.py | 8 ++++---- 2 files changed, 10 insertions(+), 10 deletions(-) diff --git a/monkey/tests/unit_tests/common/event_queue/test_pypubsub_event_queue.py b/monkey/tests/unit_tests/common/event_queue/test_pypubsub_event_queue.py index f1648d2f1..0da3c9f01 100644 --- a/monkey/tests/unit_tests/common/event_queue/test_pypubsub_event_queue.py +++ b/monkey/tests/unit_tests/common/event_queue/test_pypubsub_event_queue.py @@ -7,14 +7,14 @@ import pytest from pubsub.core import Publisher from common.event_queue import EventSubscriber, IEventQueue, PyPubSubEventQueue -from common.events import AbstractEvent +from common.events import AbstractAgentEvent EVENT_TAG_1 = "event tag 1" EVENT_TAG_2 = "event tag 2" @dataclass(frozen=True) -class TestEvent1(AbstractEvent): +class TestEvent1(AbstractAgentEvent): __test__ = False source: UUID = UUID("f811ad00-5a68-4437-bd51-7b5cc1768ad5") target: Union[UUID, IPv4Address, None] = None @@ -23,7 +23,7 @@ class TestEvent1(AbstractEvent): @dataclass(frozen=True) -class TestEvent2(AbstractEvent): +class TestEvent2(AbstractAgentEvent): __test__ = False source: UUID = UUID("e810ad01-6b67-9446-fc58-9b8d717653f7") target: Union[UUID, IPv4Address, None] = None @@ -37,8 +37,8 @@ def event_queue() -> IEventQueue: @pytest.fixture -def event_queue_subscriber() -> Callable[[AbstractEvent], None]: - def fn(event: AbstractEvent): +def event_queue_subscriber() -> Callable[[AbstractAgentEvent], None]: + def fn(event: AbstractAgentEvent): fn.call_count += 1 fn.call_types.add(event.__class__) fn.call_tags |= event.tags @@ -117,7 +117,7 @@ def test_keep_subscriber_in_scope(event_queue: IEventQueue): class MyCallable: called = False - def __call__(self, event: AbstractEvent): + def __call__(self, event: AbstractAgentEvent): MyCallable.called = True def subscribe(): diff --git a/monkey/tests/unit_tests/common/event_serializers/test_event_serializer_registry.py b/monkey/tests/unit_tests/common/event_serializers/test_event_serializer_registry.py index 43f909bec..b4ddd6612 100644 --- a/monkey/tests/unit_tests/common/event_serializers/test_event_serializer_registry.py +++ b/monkey/tests/unit_tests/common/event_serializers/test_event_serializer_registry.py @@ -4,21 +4,21 @@ from unittest.mock import MagicMock import pytest from common.event_serializers import EventSerializerRegistry, IEventSerializer -from common.events import AbstractEvent +from common.events import AbstractAgentEvent @dataclass(frozen=True) -class SomeEvent(AbstractEvent): +class SomeEvent(AbstractAgentEvent): some_param: int = field(default=435) @dataclass(frozen=True) -class OtherEvent(AbstractEvent): +class OtherEvent(AbstractAgentEvent): other_param: float = field(default=123.456) @dataclass(frozen=True) -class NoneEvent(AbstractEvent): +class NoneEvent(AbstractAgentEvent): none_param: float = field(default=1.0) From 903d43fe266af97f8ec530b74592a447c69dfa96 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Mon, 5 Sep 2022 19:09:47 +0530 Subject: [PATCH 04/56] Common: Rename IEventQueue -> IAgentEventQueue --- monkey/common/event_queue/__init__.py | 2 +- monkey/common/event_queue/i_event_queue.py | 2 +- monkey/common/event_queue/pypubsub_event_queue.py | 4 ++-- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/monkey/common/event_queue/__init__.py b/monkey/common/event_queue/__init__.py index f88abe300..1d2c0e5fc 100644 --- a/monkey/common/event_queue/__init__.py +++ b/monkey/common/event_queue/__init__.py @@ -1,3 +1,3 @@ from .types import EventSubscriber -from .i_event_queue import IEventQueue +from .i_event_queue import IAgentEventQueue from .pypubsub_event_queue import PyPubSubEventQueue diff --git a/monkey/common/event_queue/i_event_queue.py b/monkey/common/event_queue/i_event_queue.py index 128d0517b..15e35e75a 100644 --- a/monkey/common/event_queue/i_event_queue.py +++ b/monkey/common/event_queue/i_event_queue.py @@ -6,7 +6,7 @@ from common.events import AbstractAgentEvent from . import EventSubscriber -class IEventQueue(ABC): +class IAgentEventQueue(ABC): """ Manages subscription and publishing of events """ diff --git a/monkey/common/event_queue/pypubsub_event_queue.py b/monkey/common/event_queue/pypubsub_event_queue.py index 9c3e4dc6d..b189d8fdf 100644 --- a/monkey/common/event_queue/pypubsub_event_queue.py +++ b/monkey/common/event_queue/pypubsub_event_queue.py @@ -5,14 +5,14 @@ from pubsub.core import Publisher from common.events import AbstractAgentEvent -from . import EventSubscriber, IEventQueue +from . import EventSubscriber, IAgentEventQueue _ALL_EVENTS_TOPIC = "all_events_topic" logger = logging.getLogger(__name__) -class PyPubSubEventQueue(IEventQueue): +class PyPubSubEventQueue(IAgentEventQueue): def __init__(self, pypubsub_publisher: Publisher): self._pypubsub_publisher = pypubsub_publisher self._refs = [] From 74e17ef00605d1e1bc02af99c7275e6c9ef03a26 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Mon, 5 Sep 2022 19:10:54 +0530 Subject: [PATCH 05/56] Agent: Rename IEventQueue -> IAgentEventQueue --- .../mimikatz_credential_collector.py | 4 ++-- .../ssh_collector/ssh_credential_collector.py | 4 ++-- .../ssh_collector/ssh_handler.py | 12 ++++++++---- monkey/infection_monkey/exploit/HostExploiter.py | 4 ++-- monkey/infection_monkey/exploit/exploiter_wrapper.py | 6 +++--- monkey/infection_monkey/monkey.py | 6 +++--- 6 files changed, 20 insertions(+), 16 deletions(-) diff --git a/monkey/infection_monkey/credential_collectors/mimikatz_collector/mimikatz_credential_collector.py b/monkey/infection_monkey/credential_collectors/mimikatz_collector/mimikatz_credential_collector.py index a39023e0e..882ced8ca 100644 --- a/monkey/infection_monkey/credential_collectors/mimikatz_collector/mimikatz_credential_collector.py +++ b/monkey/infection_monkey/credential_collectors/mimikatz_collector/mimikatz_credential_collector.py @@ -2,7 +2,7 @@ import logging from typing import Sequence from common.credentials import Credentials, LMHash, NTHash, Password, Username -from common.event_queue import IEventQueue +from common.event_queue import IAgentEventQueue from common.events import CredentialsStolenEvent from infection_monkey.i_puppet import ICredentialCollector from infection_monkey.model import USERNAME_PREFIX @@ -27,7 +27,7 @@ MIMIKATZ_EVENT_TAGS = frozenset( class MimikatzCredentialCollector(ICredentialCollector): - def __init__(self, event_queue: IEventQueue): + def __init__(self, event_queue: IAgentEventQueue): self._event_queue = event_queue def collect_credentials(self, options=None) -> Sequence[Credentials]: diff --git a/monkey/infection_monkey/credential_collectors/ssh_collector/ssh_credential_collector.py b/monkey/infection_monkey/credential_collectors/ssh_collector/ssh_credential_collector.py index b696adf40..ed0fc1a8e 100644 --- a/monkey/infection_monkey/credential_collectors/ssh_collector/ssh_credential_collector.py +++ b/monkey/infection_monkey/credential_collectors/ssh_collector/ssh_credential_collector.py @@ -2,7 +2,7 @@ import logging from typing import Sequence from common.credentials import Credentials -from common.event_queue import IEventQueue +from common.event_queue import IAgentEventQueue from infection_monkey.credential_collectors.ssh_collector import ssh_handler from infection_monkey.i_puppet import ICredentialCollector from infection_monkey.telemetry.messengers.i_telemetry_messenger import ITelemetryMessenger @@ -15,7 +15,7 @@ class SSHCredentialCollector(ICredentialCollector): SSH keys credential collector """ - def __init__(self, telemetry_messenger: ITelemetryMessenger, event_queue: IEventQueue): + def __init__(self, telemetry_messenger: ITelemetryMessenger, event_queue: IAgentEventQueue): self._telemetry_messenger = telemetry_messenger self._event_queue = event_queue diff --git a/monkey/infection_monkey/credential_collectors/ssh_collector/ssh_handler.py b/monkey/infection_monkey/credential_collectors/ssh_collector/ssh_handler.py index 7d8a046f4..ca271a5d8 100644 --- a/monkey/infection_monkey/credential_collectors/ssh_collector/ssh_handler.py +++ b/monkey/infection_monkey/credential_collectors/ssh_collector/ssh_handler.py @@ -4,7 +4,7 @@ import os from typing import Dict, Iterable, Sequence from common.credentials import Credentials, SSHKeypair, Username -from common.event_queue import IEventQueue +from common.event_queue import IAgentEventQueue from common.events import CredentialsStolenEvent from common.utils.attack_utils import ScanStatus from infection_monkey.telemetry.attack.t1005_telem import T1005Telem @@ -31,7 +31,7 @@ SSH_COLLECTOR_EVENT_TAGS = frozenset( def get_ssh_info( - telemetry_messenger: ITelemetryMessenger, event_queue: IEventQueue + telemetry_messenger: ITelemetryMessenger, event_queue: IAgentEventQueue ) -> Iterable[Dict]: # TODO: Remove this check when this is turned into a plugin. if is_windows_os(): @@ -80,7 +80,9 @@ def _get_ssh_struct(name: str, home_dir: str) -> Dict: def _get_ssh_files( - user_info: Iterable[Dict], telemetry_messenger: ITelemetryMessenger, event_queue: IEventQueue + user_info: Iterable[Dict], + telemetry_messenger: ITelemetryMessenger, + event_queue: IAgentEventQueue, ) -> Iterable[Dict]: for info in user_info: path = info["home_dir"] @@ -165,7 +167,9 @@ def to_credentials(ssh_info: Iterable[Dict]) -> Sequence[Credentials]: return ssh_credentials -def _publish_credentials_stolen_event(collected_credentials: Credentials, event_queue: IEventQueue): +def _publish_credentials_stolen_event( + collected_credentials: Credentials, event_queue: IAgentEventQueue +): credentials_stolen_event = CredentialsStolenEvent( tags=SSH_COLLECTOR_EVENT_TAGS, stolen_credentials=[collected_credentials], diff --git a/monkey/infection_monkey/exploit/HostExploiter.py b/monkey/infection_monkey/exploit/HostExploiter.py index 8ca63c3ba..bb5d7f87d 100644 --- a/monkey/infection_monkey/exploit/HostExploiter.py +++ b/monkey/infection_monkey/exploit/HostExploiter.py @@ -4,7 +4,7 @@ from abc import abstractmethod from datetime import datetime from typing import Dict -from common.event_queue import IEventQueue +from common.event_queue import IAgentEventQueue from common.utils.exceptions import FailedExploitationError from infection_monkey.i_puppet import ExploiterResultData from infection_monkey.telemetry.messengers.i_telemetry_messenger import ITelemetryMessenger @@ -60,7 +60,7 @@ class HostExploiter: host, current_depth: int, telemetry_messenger: ITelemetryMessenger, - event_queue: IEventQueue, + event_queue: IAgentEventQueue, agent_binary_repository: IAgentBinaryRepository, options: Dict, interrupt: threading.Event, diff --git a/monkey/infection_monkey/exploit/exploiter_wrapper.py b/monkey/infection_monkey/exploit/exploiter_wrapper.py index 3e369134f..8b69ce9d5 100644 --- a/monkey/infection_monkey/exploit/exploiter_wrapper.py +++ b/monkey/infection_monkey/exploit/exploiter_wrapper.py @@ -1,7 +1,7 @@ import threading from typing import Dict, Type -from common.event_queue import IEventQueue +from common.event_queue import IAgentEventQueue from infection_monkey.model import VictimHost from infection_monkey.telemetry.messengers.i_telemetry_messenger import ITelemetryMessenger @@ -22,7 +22,7 @@ class ExploiterWrapper: self, exploit_class: Type[HostExploiter], telemetry_messenger: ITelemetryMessenger, - event_queue: IEventQueue, + event_queue: IAgentEventQueue, agent_binary_repository: IAgentBinaryRepository, ): self._exploit_class = exploit_class @@ -47,7 +47,7 @@ class ExploiterWrapper: def __init__( self, telemetry_messenger: ITelemetryMessenger, - event_queue: IEventQueue, + event_queue: IAgentEventQueue, agent_binary_repository: IAgentBinaryRepository, ): self._telemetry_messenger = telemetry_messenger diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index 02e6b57ad..9e8e689ed 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -10,7 +10,7 @@ from typing import List from pubsub.core import Publisher import infection_monkey.tunnel as tunnel -from common.event_queue import IEventQueue, PyPubSubEventQueue +from common.event_queue import IAgentEventQueue, PyPubSubEventQueue from common.events import CredentialsStolenEvent from common.network.network_utils import address_to_ip_port from common.utils.argparse_types import positive_int @@ -228,7 +228,7 @@ class InfectionMonkey: @staticmethod def _subscribe_events( - event_queue: IEventQueue, + event_queue: IAgentEventQueue, propagation_credentials_repository: IPropagationCredentialsRepository, ): event_queue.subscribe_type( @@ -249,7 +249,7 @@ class InfectionMonkey: def _build_puppet( self, propagation_credentials_repository: IPropagationCredentialsRepository, - event_queue: IEventQueue, + event_queue: IAgentEventQueue, ) -> IPuppet: puppet = Puppet() From acb049f1f1cbe5a336b8a0565e1b9d2d17492a3b Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Mon, 5 Sep 2022 19:11:29 +0530 Subject: [PATCH 06/56] Island: Rename IEventQueue -> IAgentEventQueue --- monkey/monkey_island/cc/services/initialize.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/monkey/monkey_island/cc/services/initialize.py b/monkey/monkey_island/cc/services/initialize.py index 3ceb489c2..434f40720 100644 --- a/monkey/monkey_island/cc/services/initialize.py +++ b/monkey/monkey_island/cc/services/initialize.py @@ -12,7 +12,7 @@ from common.agent_configuration import ( ) from common.aws import AWSInstance from common.common_consts.telem_categories import TelemCategoryEnum -from common.event_queue import IEventQueue, PyPubSubEventQueue +from common.event_queue import IAgentEventQueue, PyPubSubEventQueue from common.utils.file_utils import get_binary_io_sha256_hash from monkey_island.cc.repository import ( AgentBinaryRepository, @@ -63,7 +63,7 @@ def initialize_services(container: DIContainer, data_dir: Path): ILockableEncryptor, RepositoryEncryptor(data_dir / REPOSITORY_KEY_FILE_NAME) ) container.register(Publisher, Publisher) - container.register_instance(IEventQueue, container.resolve(PyPubSubEventQueue)) + container.register_instance(IAgentEventQueue, container.resolve(PyPubSubEventQueue)) _register_repositories(container, data_dir) _register_services(container) From 28cf860aeb581196a1686507c5a88f7822957428 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Mon, 5 Sep 2022 19:13:58 +0530 Subject: [PATCH 07/56] UT: Use IAgentEventQueue in place of IEventQueue --- .../event_queue/test_pypubsub_event_queue.py | 16 ++++++++-------- .../test_mimikatz_collector.py | 10 +++++----- .../test_ssh_credentials_collector.py | 6 +++--- 3 files changed, 16 insertions(+), 16 deletions(-) diff --git a/monkey/tests/unit_tests/common/event_queue/test_pypubsub_event_queue.py b/monkey/tests/unit_tests/common/event_queue/test_pypubsub_event_queue.py index 0da3c9f01..8e5bfe64e 100644 --- a/monkey/tests/unit_tests/common/event_queue/test_pypubsub_event_queue.py +++ b/monkey/tests/unit_tests/common/event_queue/test_pypubsub_event_queue.py @@ -6,7 +6,7 @@ from uuid import UUID import pytest from pubsub.core import Publisher -from common.event_queue import EventSubscriber, IEventQueue, PyPubSubEventQueue +from common.event_queue import EventSubscriber, IAgentEventQueue, PyPubSubEventQueue from common.events import AbstractAgentEvent EVENT_TAG_1 = "event tag 1" @@ -32,7 +32,7 @@ class TestEvent2(AbstractAgentEvent): @pytest.fixture -def event_queue() -> IEventQueue: +def event_queue() -> IAgentEventQueue: return PyPubSubEventQueue(Publisher()) @@ -50,7 +50,7 @@ def event_queue_subscriber() -> Callable[[AbstractAgentEvent], None]: return fn -def test_subscribe_all(event_queue: IEventQueue, event_queue_subscriber: EventSubscriber): +def test_subscribe_all(event_queue: IAgentEventQueue, event_queue_subscriber: EventSubscriber): event_queue.subscribe_all_events(event_queue_subscriber) event_queue.publish(TestEvent1(tags=frozenset({EVENT_TAG_1, EVENT_TAG_2}))) @@ -65,7 +65,7 @@ def test_subscribe_all(event_queue: IEventQueue, event_queue_subscriber: EventSu @pytest.mark.parametrize("type_to_subscribe", [TestEvent1, TestEvent2]) def test_subscribe_types( - event_queue: IEventQueue, event_queue_subscriber: EventSubscriber, type_to_subscribe + event_queue: IAgentEventQueue, event_queue_subscriber: EventSubscriber, type_to_subscribe ): event_queue.subscribe_type(type_to_subscribe, event_queue_subscriber) @@ -77,7 +77,7 @@ def test_subscribe_types( def test_subscribe_tags_single_type( - event_queue: IEventQueue, event_queue_subscriber: EventSubscriber + event_queue: IAgentEventQueue, event_queue_subscriber: EventSubscriber ): event_queue.subscribe_tag(EVENT_TAG_1, event_queue_subscriber) @@ -91,7 +91,7 @@ def test_subscribe_tags_single_type( def test_subscribe_tags_multiple_types( - event_queue: IEventQueue, event_queue_subscriber: EventSubscriber + event_queue: IAgentEventQueue, event_queue_subscriber: EventSubscriber ): event_queue.subscribe_tag(EVENT_TAG_2, event_queue_subscriber) @@ -105,7 +105,7 @@ def test_subscribe_tags_multiple_types( assert {EVENT_TAG_1, EVENT_TAG_2}.issubset(event_queue_subscriber.call_tags) -def test_type_tag_collision(event_queue: IEventQueue, event_queue_subscriber: EventSubscriber): +def test_type_tag_collision(event_queue: IAgentEventQueue, event_queue_subscriber: EventSubscriber): event_queue.subscribe_type(TestEvent1, event_queue_subscriber) event_queue.publish(TestEvent2(tags=frozenset({TestEvent1.__name__}))) @@ -113,7 +113,7 @@ def test_type_tag_collision(event_queue: IEventQueue, event_queue_subscriber: Ev assert event_queue_subscriber.call_count == 0 -def test_keep_subscriber_in_scope(event_queue: IEventQueue): +def test_keep_subscriber_in_scope(event_queue: IAgentEventQueue): class MyCallable: called = False diff --git a/monkey/tests/unit_tests/infection_monkey/credential_collectors/test_mimikatz_collector.py b/monkey/tests/unit_tests/infection_monkey/credential_collectors/test_mimikatz_collector.py index 12ca9d594..0476d8a13 100644 --- a/monkey/tests/unit_tests/infection_monkey/credential_collectors/test_mimikatz_collector.py +++ b/monkey/tests/unit_tests/infection_monkey/credential_collectors/test_mimikatz_collector.py @@ -4,7 +4,7 @@ from unittest.mock import MagicMock import pytest from common.credentials import Credentials, LMHash, NTHash, Password, Username -from common.event_queue import IEventQueue +from common.event_queue import IAgentEventQueue from common.events import CredentialsStolenEvent from infection_monkey.credential_collectors import MimikatzCredentialCollector from infection_monkey.credential_collectors.mimikatz_collector.mimikatz_credential_collector import ( # noqa: E501 @@ -24,7 +24,7 @@ def patch_pypykatz(win_creds: [WindowsCredentials], monkeypatch): def collect_credentials() -> Sequence[Credentials]: - mock_event_queue = MagicMock(spec=IEventQueue) + mock_event_queue = MagicMock(spec=IAgentEventQueue) return MimikatzCredentialCollector(mock_event_queue).collect_credentials() @@ -120,7 +120,7 @@ def test_pypykatz_result_parsing_no_secrets(monkeypatch): def test_mimikatz_credentials_stolen_event_published(monkeypatch): - mock_event_queue = MagicMock(spec=IEventQueue) + mock_event_queue = MagicMock(spec=IAgentEventQueue) patch_pypykatz([], monkeypatch) mimikatz_credential_collector = MimikatzCredentialCollector(mock_event_queue) @@ -134,7 +134,7 @@ def test_mimikatz_credentials_stolen_event_published(monkeypatch): def test_mimikatz_credentials_stolen_event_tags(monkeypatch): - mock_event_queue = MagicMock(spec=IEventQueue) + mock_event_queue = MagicMock(spec=IAgentEventQueue) patch_pypykatz([], monkeypatch) mimikatz_credential_collector = MimikatzCredentialCollector(mock_event_queue) @@ -146,7 +146,7 @@ def test_mimikatz_credentials_stolen_event_tags(monkeypatch): def test_mimikatz_credentials_stolen_event_stolen_credentials(monkeypatch): - mock_event_queue = MagicMock(spec=IEventQueue) + mock_event_queue = MagicMock(spec=IAgentEventQueue) win_creds = [ WindowsCredentials( username="user2", password="secret2", lm_hash="0182BD0BD4444BF8FC83B5D9042EED2E" diff --git a/monkey/tests/unit_tests/infection_monkey/credential_collectors/test_ssh_credentials_collector.py b/monkey/tests/unit_tests/infection_monkey/credential_collectors/test_ssh_credentials_collector.py index ba12e416f..1c7fbe1b4 100644 --- a/monkey/tests/unit_tests/infection_monkey/credential_collectors/test_ssh_credentials_collector.py +++ b/monkey/tests/unit_tests/infection_monkey/credential_collectors/test_ssh_credentials_collector.py @@ -3,7 +3,7 @@ from unittest.mock import MagicMock import pytest from common.credentials import Credentials, SSHKeypair, Username -from common.event_queue import IEventQueue +from common.event_queue import IAgentEventQueue from infection_monkey.credential_collectors import SSHCredentialCollector @@ -25,7 +25,7 @@ def patch_ssh_handler(ssh_creds, monkeypatch): def test_ssh_credentials_empty_results(monkeypatch, ssh_creds, patch_telemetry_messenger): patch_ssh_handler(ssh_creds, monkeypatch) collected = SSHCredentialCollector( - patch_telemetry_messenger, MagicMock(spec=IEventQueue) + patch_telemetry_messenger, MagicMock(spec=IAgentEventQueue) ).collect_credentials() assert not collected @@ -71,6 +71,6 @@ def test_ssh_info_result_parsing(monkeypatch, patch_telemetry_messenger): Credentials(identity=None, secret=ssh_keypair3), ] collected = SSHCredentialCollector( - patch_telemetry_messenger, MagicMock(spec=IEventQueue) + patch_telemetry_messenger, MagicMock(spec=IAgentEventQueue) ).collect_credentials() assert expected == collected From e14c1ea5f39afb9f32ae3b956365550a53fc79f6 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Mon, 5 Sep 2022 19:16:31 +0530 Subject: [PATCH 08/56] Project: Use correct file path in Vulture allowlist --- vulture_allowlist.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/vulture_allowlist.py b/vulture_allowlist.py index ea46d5b2b..9a9360c66 100644 --- a/vulture_allowlist.py +++ b/vulture_allowlist.py @@ -277,10 +277,10 @@ key_list simulation netmap validate_windows_filename_not_reserved -subscribe_all # common\event_queue\i_event_queue.py -subscribe_type # common\event_queue\i_event_queue.py -subscribe_tag # common\event_queue\i_event_queue.py -publish # common\event_queue\i_event_queue.py +subscribe_all # common\event_queue\i_agent_event_queue.py +subscribe_type # common\event_queue\i_agent_event_queue.py +subscribe_tag # common\event_queue\i_agent_event_queue.py +publish # common\event_queue\i_agent_event_queue.py subscribe_all # common\event_queue\pypubsub_event_queue.py subscribe_type # common\event_queue\pypubsub_event_queue.py subscribe_tag # common\event_queue\pypubsub_event_queue.py From cb9200b3c061eec86295ac33d17c87efb58498c0 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Mon, 5 Sep 2022 19:17:31 +0530 Subject: [PATCH 09/56] Common: Rename i_event_queue.py -> i_agent_event_queue.py --- monkey/common/event_queue/__init__.py | 2 +- .../event_queue/{i_event_queue.py => i_agent_event_queue.py} | 0 2 files changed, 1 insertion(+), 1 deletion(-) rename monkey/common/event_queue/{i_event_queue.py => i_agent_event_queue.py} (100%) diff --git a/monkey/common/event_queue/__init__.py b/monkey/common/event_queue/__init__.py index 1d2c0e5fc..3fd177f5d 100644 --- a/monkey/common/event_queue/__init__.py +++ b/monkey/common/event_queue/__init__.py @@ -1,3 +1,3 @@ from .types import EventSubscriber -from .i_event_queue import IAgentEventQueue +from .i_agent_event_queue import IAgentEventQueue from .pypubsub_event_queue import PyPubSubEventQueue diff --git a/monkey/common/event_queue/i_event_queue.py b/monkey/common/event_queue/i_agent_event_queue.py similarity index 100% rename from monkey/common/event_queue/i_event_queue.py rename to monkey/common/event_queue/i_agent_event_queue.py From a83446e0f2feb9f714fde7a14893fa3b5d679164 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Mon, 5 Sep 2022 19:20:48 +0530 Subject: [PATCH 10/56] Common: Rename PyPubSubEventQueue -> PyPubSubAgentEventQueue --- monkey/common/event_queue/__init__.py | 2 +- monkey/common/event_queue/pypubsub_event_queue.py | 10 +++++----- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/monkey/common/event_queue/__init__.py b/monkey/common/event_queue/__init__.py index 3fd177f5d..5d6a8c564 100644 --- a/monkey/common/event_queue/__init__.py +++ b/monkey/common/event_queue/__init__.py @@ -1,3 +1,3 @@ from .types import EventSubscriber from .i_agent_event_queue import IAgentEventQueue -from .pypubsub_event_queue import PyPubSubEventQueue +from .pypubsub_event_queue import PyPubSubAgentEventQueue diff --git a/monkey/common/event_queue/pypubsub_event_queue.py b/monkey/common/event_queue/pypubsub_event_queue.py index b189d8fdf..7decce957 100644 --- a/monkey/common/event_queue/pypubsub_event_queue.py +++ b/monkey/common/event_queue/pypubsub_event_queue.py @@ -12,7 +12,7 @@ _ALL_EVENTS_TOPIC = "all_events_topic" logger = logging.getLogger(__name__) -class PyPubSubEventQueue(IAgentEventQueue): +class PyPubSubAgentEventQueue(IAgentEventQueue): def __init__(self, pypubsub_publisher: Publisher): self._pypubsub_publisher = pypubsub_publisher self._refs = [] @@ -22,11 +22,11 @@ class PyPubSubEventQueue(IAgentEventQueue): def subscribe_type(self, event_type: Type[AbstractAgentEvent], subscriber: EventSubscriber): # pypubsub.pub.subscribe needs a string as the topic/event name - event_type_topic = PyPubSubEventQueue._get_type_topic(event_type) + event_type_topic = PyPubSubAgentEventQueue._get_type_topic(event_type) self._subscribe(event_type_topic, subscriber) def subscribe_tag(self, tag: str, subscriber: EventSubscriber): - tag_topic = PyPubSubEventQueue._get_tag_topic(tag) + tag_topic = PyPubSubAgentEventQueue._get_tag_topic(tag) self._subscribe(tag_topic, subscriber) def _subscribe(self, topic: str, subscriber: EventSubscriber): @@ -69,12 +69,12 @@ class PyPubSubEventQueue(IAgentEventQueue): self._publish_event(_ALL_EVENTS_TOPIC, event) def _publish_to_type_topic(self, event: AbstractAgentEvent): - event_type_topic = PyPubSubEventQueue._get_type_topic(event.__class__) + event_type_topic = PyPubSubAgentEventQueue._get_type_topic(event.__class__) self._publish_event(event_type_topic, event) def _publish_to_tags_topics(self, event: AbstractAgentEvent): for tag in event.tags: - tag_topic = PyPubSubEventQueue._get_tag_topic(tag) + tag_topic = PyPubSubAgentEventQueue._get_tag_topic(tag) self._publish_event(tag_topic, event) def _publish_event(self, topic: str, event: AbstractAgentEvent): From 5b7d3bfb70fa13153a2cb06627d138576002ba77 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Mon, 5 Sep 2022 19:25:05 +0530 Subject: [PATCH 11/56] Agent: Rename PyPubSubEventQueue -> PyPubSubAgentEventQueue --- monkey/infection_monkey/monkey.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index 9e8e689ed..abb0964ca 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -10,7 +10,7 @@ from typing import List from pubsub.core import Publisher import infection_monkey.tunnel as tunnel -from common.event_queue import IAgentEventQueue, PyPubSubEventQueue +from common.event_queue import IAgentEventQueue, PyPubSubAgentEventQueue from common.events import CredentialsStolenEvent from common.network.network_utils import address_to_ip_port from common.utils.argparse_types import positive_int @@ -205,7 +205,7 @@ class InfectionMonkey: control_channel ) - event_queue = PyPubSubEventQueue(Publisher()) + event_queue = PyPubSubAgentEventQueue(Publisher()) InfectionMonkey._subscribe_events(event_queue, propagation_credentials_repository) puppet = self._build_puppet(propagation_credentials_repository, event_queue) From de74b866bbfae22a0dd1710deaf03d5bfd159be0 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Mon, 5 Sep 2022 19:25:37 +0530 Subject: [PATCH 12/56] Island: Rename PyPubSubEventQueue -> PyPubSubAgentEventQueue --- monkey/monkey_island/cc/services/initialize.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/monkey/monkey_island/cc/services/initialize.py b/monkey/monkey_island/cc/services/initialize.py index 434f40720..b1caab657 100644 --- a/monkey/monkey_island/cc/services/initialize.py +++ b/monkey/monkey_island/cc/services/initialize.py @@ -12,7 +12,7 @@ from common.agent_configuration import ( ) from common.aws import AWSInstance from common.common_consts.telem_categories import TelemCategoryEnum -from common.event_queue import IAgentEventQueue, PyPubSubEventQueue +from common.event_queue import IAgentEventQueue, PyPubSubAgentEventQueue from common.utils.file_utils import get_binary_io_sha256_hash from monkey_island.cc.repository import ( AgentBinaryRepository, @@ -63,7 +63,7 @@ def initialize_services(container: DIContainer, data_dir: Path): ILockableEncryptor, RepositoryEncryptor(data_dir / REPOSITORY_KEY_FILE_NAME) ) container.register(Publisher, Publisher) - container.register_instance(IAgentEventQueue, container.resolve(PyPubSubEventQueue)) + container.register_instance(IAgentEventQueue, container.resolve(PyPubSubAgentEventQueue)) _register_repositories(container, data_dir) _register_services(container) From 331c11751f3c32fbf6c4c5015cabd2e257072efb Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Mon, 5 Sep 2022 19:26:09 +0530 Subject: [PATCH 13/56] UT: Use PyPubSubAgentEventQueue in place of PyPubSubEventQueue --- .../common/event_queue/test_pypubsub_event_queue.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/monkey/tests/unit_tests/common/event_queue/test_pypubsub_event_queue.py b/monkey/tests/unit_tests/common/event_queue/test_pypubsub_event_queue.py index 8e5bfe64e..372930da3 100644 --- a/monkey/tests/unit_tests/common/event_queue/test_pypubsub_event_queue.py +++ b/monkey/tests/unit_tests/common/event_queue/test_pypubsub_event_queue.py @@ -6,7 +6,7 @@ from uuid import UUID import pytest from pubsub.core import Publisher -from common.event_queue import EventSubscriber, IAgentEventQueue, PyPubSubEventQueue +from common.event_queue import EventSubscriber, IAgentEventQueue, PyPubSubAgentEventQueue from common.events import AbstractAgentEvent EVENT_TAG_1 = "event tag 1" @@ -33,7 +33,7 @@ class TestEvent2(AbstractAgentEvent): @pytest.fixture def event_queue() -> IAgentEventQueue: - return PyPubSubEventQueue(Publisher()) + return PyPubSubAgentEventQueue(Publisher()) @pytest.fixture From 21815291a09a9976fb3a2479861f48c8075be4ef Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Mon, 5 Sep 2022 19:26:30 +0530 Subject: [PATCH 14/56] Project: Use correct variable in Vulture allowlist --- vulture_allowlist.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vulture_allowlist.py b/vulture_allowlist.py index 9a9360c66..49986ea85 100644 --- a/vulture_allowlist.py +++ b/vulture_allowlist.py @@ -285,7 +285,7 @@ subscribe_all # common\event_queue\pypubsub_event_queue.py subscribe_type # common\event_queue\pypubsub_event_queue.py subscribe_tag # common\event_queue\pypubsub_event_queue.py publish # common\event_queue\pypubsub_event_queue.py -PyPubSubEventQueue # common\event_queue\pypubsub_event_queue.py +PyPubSubAgentEventQueue # common\event_queue\pypubsub_event_queue.py subscribe_all_events # common\event_queue\pypubsub_event_queue.py From 612132a906c117737bdaf38e8c9de4be18c04d66 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Mon, 5 Sep 2022 19:27:20 +0530 Subject: [PATCH 15/56] Common: Rename pypubsub_event_queue.py -> pypubsub_agent_event_queue.py --- monkey/common/event_queue/__init__.py | 2 +- .../{pypubsub_event_queue.py => pypubsub_agent_event_queue.py} | 0 2 files changed, 1 insertion(+), 1 deletion(-) rename monkey/common/event_queue/{pypubsub_event_queue.py => pypubsub_agent_event_queue.py} (100%) diff --git a/monkey/common/event_queue/__init__.py b/monkey/common/event_queue/__init__.py index 5d6a8c564..34567919c 100644 --- a/monkey/common/event_queue/__init__.py +++ b/monkey/common/event_queue/__init__.py @@ -1,3 +1,3 @@ from .types import EventSubscriber from .i_agent_event_queue import IAgentEventQueue -from .pypubsub_event_queue import PyPubSubAgentEventQueue +from .pypubsub_agent_event_queue import PyPubSubAgentEventQueue diff --git a/monkey/common/event_queue/pypubsub_event_queue.py b/monkey/common/event_queue/pypubsub_agent_event_queue.py similarity index 100% rename from monkey/common/event_queue/pypubsub_event_queue.py rename to monkey/common/event_queue/pypubsub_agent_event_queue.py From e1e4b054cfad5b55c7b3e8ada85af05d0d740cee Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Mon, 5 Sep 2022 19:27:44 +0530 Subject: [PATCH 16/56] Project: Use correct file path in Vulture allowlist --- vulture_allowlist.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/vulture_allowlist.py b/vulture_allowlist.py index 49986ea85..4dc74436c 100644 --- a/vulture_allowlist.py +++ b/vulture_allowlist.py @@ -281,12 +281,12 @@ subscribe_all # common\event_queue\i_agent_event_queue.py subscribe_type # common\event_queue\i_agent_event_queue.py subscribe_tag # common\event_queue\i_agent_event_queue.py publish # common\event_queue\i_agent_event_queue.py -subscribe_all # common\event_queue\pypubsub_event_queue.py -subscribe_type # common\event_queue\pypubsub_event_queue.py -subscribe_tag # common\event_queue\pypubsub_event_queue.py -publish # common\event_queue\pypubsub_event_queue.py -PyPubSubAgentEventQueue # common\event_queue\pypubsub_event_queue.py -subscribe_all_events # common\event_queue\pypubsub_event_queue.py +subscribe_all # common\event_queue\pypubsub_agent_event_queue.py +subscribe_type # common\event_queue\pypubsub_agent_event_queue.py +subscribe_tag # common\event_queue\pypubsub_agent_event_queue.py +publish # common\event_queue\pypubsub_agent_event_queue.py +PyPubSubAgentEventQueue # common\event_queue\pypubsub_agent_event_queue.py +subscribe_all_events # common\event_queue\pypubsub_agent_event_queue.py # TODO: Remove once #2179 is closed From d6df50e32392b7239f07ce114d69203fc4fa16d0 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Mon, 5 Sep 2022 19:36:39 +0530 Subject: [PATCH 17/56] Common: Make docstring in IAgentEventQueue more specific --- monkey/common/event_queue/i_agent_event_queue.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/monkey/common/event_queue/i_agent_event_queue.py b/monkey/common/event_queue/i_agent_event_queue.py index 15e35e75a..b36a3dea9 100644 --- a/monkey/common/event_queue/i_agent_event_queue.py +++ b/monkey/common/event_queue/i_agent_event_queue.py @@ -8,7 +8,7 @@ from . import EventSubscriber class IAgentEventQueue(ABC): """ - Manages subscription and publishing of events + Manages subscription and publishing of events in the Agent """ @abstractmethod From 9bc4d8c241817c67bf95a006cc9e2a4fe8d8dcb1 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Tue, 6 Sep 2022 11:44:10 +0530 Subject: [PATCH 18/56] Common: Fix docstring in IAgentEventQueue.publish --- monkey/common/event_queue/i_agent_event_queue.py | 1 - 1 file changed, 1 deletion(-) diff --git a/monkey/common/event_queue/i_agent_event_queue.py b/monkey/common/event_queue/i_agent_event_queue.py index b36a3dea9..d1ebf25ca 100644 --- a/monkey/common/event_queue/i_agent_event_queue.py +++ b/monkey/common/event_queue/i_agent_event_queue.py @@ -49,7 +49,6 @@ class IAgentEventQueue(ABC): Publishes an event with the given data :param event: Event to publish - :param data: Data to pass to subscribers with the event publish """ pass From f4a0b89e873a08c052d7206dbcd6daa5db41485f Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Tue, 6 Sep 2022 11:44:34 +0530 Subject: [PATCH 19/56] Common: Add IIslandEventQueue --- .../event_queue/i_island_event_queue.py | 37 +++++++++++++++++++ 1 file changed, 37 insertions(+) create mode 100644 monkey/common/event_queue/i_island_event_queue.py diff --git a/monkey/common/event_queue/i_island_event_queue.py b/monkey/common/event_queue/i_island_event_queue.py new file mode 100644 index 000000000..f7bc60c5f --- /dev/null +++ b/monkey/common/event_queue/i_island_event_queue.py @@ -0,0 +1,37 @@ +from abc import ABC, abstractmethod +from enum import Enum, auto +from typing import Any, Callable + + +class IslandEventTopics(Enum): + AGENT_CONNECTED = auto() + CLEAR_SIMULATION_DATA = auto() + RESET_AGENT_CONFIGURATION = auto() + + +class IIslandEventQueue(ABC): + """ + Manages subscription and publishing of events in the Island + """ + + @abstractmethod + def subscribe(self, topic: IslandEventTopics, subscriber: Callable[..., None]): + """ + Subscribes a subscriber to the specified event topic + + :param topic: Event topic to which the subscriber should subscribe + :param subscriber: A subscriber that will receive events + """ + + pass + + @abstractmethod + def publish(self, topic: IslandEventTopics, event_data: Any = None): + """ + Publishes an event topic with the given data + + :param topic: Event topic to publish + :param event_data: Event data to pass to subscribers with the event publish + """ + + pass From 122e09426e23d48587e3d4037ce907383e298d16 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Tue, 6 Sep 2022 12:07:59 +0530 Subject: [PATCH 20/56] Project: Add IIslandEventQueue entries to Vulture allowlist --- vulture_allowlist.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/vulture_allowlist.py b/vulture_allowlist.py index 4dc74436c..976810c9d 100644 --- a/vulture_allowlist.py +++ b/vulture_allowlist.py @@ -316,3 +316,10 @@ CC_TUNNEL Credentials.from_json IJSONSerializable.from_json + +# revisit when implementing PyPubSubIslandEventQueue +AGENT_CONNECTED +CLEAR_SIMULATION_DATA +RESET_AGENT_CONFIGURATION +IIslandEventQueue +event_data From 86018be736d27884b77255c301581100cfe470d9 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Tue, 6 Sep 2022 12:09:48 +0530 Subject: [PATCH 21/56] Common: Rename EventSubscriber -> AgentEventSubscriber --- monkey/common/event_queue/__init__.py | 2 +- monkey/common/event_queue/i_agent_event_queue.py | 10 ++++++---- .../event_queue/pypubsub_agent_event_queue.py | 14 ++++++++------ monkey/common/event_queue/types.py | 2 +- 4 files changed, 16 insertions(+), 12 deletions(-) diff --git a/monkey/common/event_queue/__init__.py b/monkey/common/event_queue/__init__.py index 34567919c..26d29cb49 100644 --- a/monkey/common/event_queue/__init__.py +++ b/monkey/common/event_queue/__init__.py @@ -1,3 +1,3 @@ -from .types import EventSubscriber +from .types import AgentEventSubscriber from .i_agent_event_queue import IAgentEventQueue from .pypubsub_agent_event_queue import PyPubSubAgentEventQueue diff --git a/monkey/common/event_queue/i_agent_event_queue.py b/monkey/common/event_queue/i_agent_event_queue.py index d1ebf25ca..19ef65e2b 100644 --- a/monkey/common/event_queue/i_agent_event_queue.py +++ b/monkey/common/event_queue/i_agent_event_queue.py @@ -3,7 +3,7 @@ from typing import Type from common.events import AbstractAgentEvent -from . import EventSubscriber +from . import AgentEventSubscriber class IAgentEventQueue(ABC): @@ -12,7 +12,7 @@ class IAgentEventQueue(ABC): """ @abstractmethod - def subscribe_all_events(self, subscriber: EventSubscriber): + def subscribe_all_events(self, subscriber: AgentEventSubscriber): """ Subscribes a subscriber to all events @@ -22,7 +22,9 @@ class IAgentEventQueue(ABC): pass @abstractmethod - def subscribe_type(self, event_type: Type[AbstractAgentEvent], subscriber: EventSubscriber): + def subscribe_type( + self, event_type: Type[AbstractAgentEvent], subscriber: AgentEventSubscriber + ): """ Subscribes a subscriber to the specified event type @@ -33,7 +35,7 @@ class IAgentEventQueue(ABC): pass @abstractmethod - def subscribe_tag(self, tag: str, subscriber: EventSubscriber): + def subscribe_tag(self, tag: str, subscriber: AgentEventSubscriber): """ Subscribes a subscriber to the specified event tag diff --git a/monkey/common/event_queue/pypubsub_agent_event_queue.py b/monkey/common/event_queue/pypubsub_agent_event_queue.py index 7decce957..86dbea4af 100644 --- a/monkey/common/event_queue/pypubsub_agent_event_queue.py +++ b/monkey/common/event_queue/pypubsub_agent_event_queue.py @@ -5,7 +5,7 @@ from pubsub.core import Publisher from common.events import AbstractAgentEvent -from . import EventSubscriber, IAgentEventQueue +from . import AgentEventSubscriber, IAgentEventQueue _ALL_EVENTS_TOPIC = "all_events_topic" @@ -17,19 +17,21 @@ class PyPubSubAgentEventQueue(IAgentEventQueue): self._pypubsub_publisher = pypubsub_publisher self._refs = [] - def subscribe_all_events(self, subscriber: EventSubscriber): + def subscribe_all_events(self, subscriber: AgentEventSubscriber): self._subscribe(_ALL_EVENTS_TOPIC, subscriber) - def subscribe_type(self, event_type: Type[AbstractAgentEvent], subscriber: EventSubscriber): + def subscribe_type( + self, event_type: Type[AbstractAgentEvent], subscriber: AgentEventSubscriber + ): # pypubsub.pub.subscribe needs a string as the topic/event name event_type_topic = PyPubSubAgentEventQueue._get_type_topic(event_type) self._subscribe(event_type_topic, subscriber) - def subscribe_tag(self, tag: str, subscriber: EventSubscriber): + def subscribe_tag(self, tag: str, subscriber: AgentEventSubscriber): tag_topic = PyPubSubAgentEventQueue._get_tag_topic(tag) self._subscribe(tag_topic, subscriber) - def _subscribe(self, topic: str, subscriber: EventSubscriber): + def _subscribe(self, topic: str, subscriber: AgentEventSubscriber): try: subscriber_name = subscriber.__name__ except AttributeError: @@ -39,7 +41,7 @@ class PyPubSubAgentEventQueue(IAgentEventQueue): self._pypubsub_publisher.subscribe(topicName=topic, listener=subscriber) self._keep_subscriber_strongref(subscriber) - def _keep_subscriber_strongref(self, subscriber: EventSubscriber): + def _keep_subscriber_strongref(self, subscriber: AgentEventSubscriber): # NOTE: PyPubSub stores subscribers by weak reference. From the documentation: # > PyPubSub holds listeners by weak reference so that the lifetime of the # > callable is not affected by PyPubSub: once the application no longer diff --git a/monkey/common/event_queue/types.py b/monkey/common/event_queue/types.py index f3f5f4ef8..6b57ae712 100644 --- a/monkey/common/event_queue/types.py +++ b/monkey/common/event_queue/types.py @@ -2,4 +2,4 @@ from typing import Callable from common.events import AbstractAgentEvent -EventSubscriber = Callable[[AbstractAgentEvent], None] +AgentEventSubscriber = Callable[[AbstractAgentEvent], None] From eef72b245ffc624afed949ab0e8bcc88758b8078 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Tue, 6 Sep 2022 12:10:29 +0530 Subject: [PATCH 22/56] UT: Use AgentEventSubscriber in place of EventSubscriber --- .../event_queue/test_pypubsub_event_queue.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/monkey/tests/unit_tests/common/event_queue/test_pypubsub_event_queue.py b/monkey/tests/unit_tests/common/event_queue/test_pypubsub_event_queue.py index 372930da3..3ee7ba482 100644 --- a/monkey/tests/unit_tests/common/event_queue/test_pypubsub_event_queue.py +++ b/monkey/tests/unit_tests/common/event_queue/test_pypubsub_event_queue.py @@ -6,7 +6,7 @@ from uuid import UUID import pytest from pubsub.core import Publisher -from common.event_queue import EventSubscriber, IAgentEventQueue, PyPubSubAgentEventQueue +from common.event_queue import AgentEventSubscriber, IAgentEventQueue, PyPubSubAgentEventQueue from common.events import AbstractAgentEvent EVENT_TAG_1 = "event tag 1" @@ -50,7 +50,7 @@ def event_queue_subscriber() -> Callable[[AbstractAgentEvent], None]: return fn -def test_subscribe_all(event_queue: IAgentEventQueue, event_queue_subscriber: EventSubscriber): +def test_subscribe_all(event_queue: IAgentEventQueue, event_queue_subscriber: AgentEventSubscriber): event_queue.subscribe_all_events(event_queue_subscriber) event_queue.publish(TestEvent1(tags=frozenset({EVENT_TAG_1, EVENT_TAG_2}))) @@ -65,7 +65,7 @@ def test_subscribe_all(event_queue: IAgentEventQueue, event_queue_subscriber: Ev @pytest.mark.parametrize("type_to_subscribe", [TestEvent1, TestEvent2]) def test_subscribe_types( - event_queue: IAgentEventQueue, event_queue_subscriber: EventSubscriber, type_to_subscribe + event_queue: IAgentEventQueue, event_queue_subscriber: AgentEventSubscriber, type_to_subscribe ): event_queue.subscribe_type(type_to_subscribe, event_queue_subscriber) @@ -77,7 +77,7 @@ def test_subscribe_types( def test_subscribe_tags_single_type( - event_queue: IAgentEventQueue, event_queue_subscriber: EventSubscriber + event_queue: IAgentEventQueue, event_queue_subscriber: AgentEventSubscriber ): event_queue.subscribe_tag(EVENT_TAG_1, event_queue_subscriber) @@ -91,7 +91,7 @@ def test_subscribe_tags_single_type( def test_subscribe_tags_multiple_types( - event_queue: IAgentEventQueue, event_queue_subscriber: EventSubscriber + event_queue: IAgentEventQueue, event_queue_subscriber: AgentEventSubscriber ): event_queue.subscribe_tag(EVENT_TAG_2, event_queue_subscriber) @@ -105,7 +105,9 @@ def test_subscribe_tags_multiple_types( assert {EVENT_TAG_1, EVENT_TAG_2}.issubset(event_queue_subscriber.call_tags) -def test_type_tag_collision(event_queue: IAgentEventQueue, event_queue_subscriber: EventSubscriber): +def test_type_tag_collision( + event_queue: IAgentEventQueue, event_queue_subscriber: AgentEventSubscriber +): event_queue.subscribe_type(TestEvent1, event_queue_subscriber) event_queue.publish(TestEvent2(tags=frozenset({TestEvent1.__name__}))) From 3dba1bc7d5bacad6b1e90cfbe631fd2996f05827 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Tue, 6 Sep 2022 13:04:51 +0530 Subject: [PATCH 23/56] Common: Import IIslandEventQueue, IslandEventTopics in common/event_queue/__init__.py --- monkey/common/event_queue/__init__.py | 1 + 1 file changed, 1 insertion(+) diff --git a/monkey/common/event_queue/__init__.py b/monkey/common/event_queue/__init__.py index 26d29cb49..18412d2a7 100644 --- a/monkey/common/event_queue/__init__.py +++ b/monkey/common/event_queue/__init__.py @@ -1,3 +1,4 @@ from .types import AgentEventSubscriber from .i_agent_event_queue import IAgentEventQueue +from .i_island_event_queue import IIslandEventQueue, IslandEventTopics from .pypubsub_agent_event_queue import PyPubSubAgentEventQueue From 71e9f68fe6102fbeabe265c6b7faa3c8117d0d4f Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Tue, 6 Sep 2022 13:05:35 +0530 Subject: [PATCH 24/56] Common: Fix IslandEventTopics enum values --- monkey/common/event_queue/i_island_event_queue.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/monkey/common/event_queue/i_island_event_queue.py b/monkey/common/event_queue/i_island_event_queue.py index f7bc60c5f..8066b1385 100644 --- a/monkey/common/event_queue/i_island_event_queue.py +++ b/monkey/common/event_queue/i_island_event_queue.py @@ -1,12 +1,12 @@ from abc import ABC, abstractmethod -from enum import Enum, auto +from enum import Enum from typing import Any, Callable class IslandEventTopics(Enum): - AGENT_CONNECTED = auto() - CLEAR_SIMULATION_DATA = auto() - RESET_AGENT_CONFIGURATION = auto() + AGENT_CONNECTED = "agent_connected" + CLEAR_SIMULATION_DATA = "clear_simulation_data" + RESET_AGENT_CONFIGURATION = "reset_agent_configuration" class IIslandEventQueue(ABC): From c9500cd04f13df54bd42b15ffda16d7511df70c9 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Tue, 6 Sep 2022 13:17:26 +0530 Subject: [PATCH 25/56] Common: Add PyPubSubIslandEventQueue --- .../pypubsub_island_event_queue.py | 64 +++++++++++++++++++ 1 file changed, 64 insertions(+) create mode 100644 monkey/common/event_queue/pypubsub_island_event_queue.py diff --git a/monkey/common/event_queue/pypubsub_island_event_queue.py b/monkey/common/event_queue/pypubsub_island_event_queue.py new file mode 100644 index 000000000..ba6ef3af4 --- /dev/null +++ b/monkey/common/event_queue/pypubsub_island_event_queue.py @@ -0,0 +1,64 @@ +import logging +from typing import Any, Callable + +from pubsub.core import Publisher + +from . import IIslandEventQueue, IslandEventTopics + +logger = logging.getLogger(__name__) + + +class PyPubSubIslandEventQueue(IIslandEventQueue): + def __init__(self, pypubsub_publisher: Publisher): + self._pypubsub_publisher = pypubsub_publisher + self._refs = [] + + def subscribe(self, topic: IslandEventTopics, subscriber: Callable[..., None]): + topic_value = topic.value # needs to be a string for pypubsub + try: + subscriber_name = subscriber.__name__ + except AttributeError: + subscriber_name = subscriber.__class__.__name__ + + logging.debug(f"Subscriber {subscriber_name} subscribed to {topic_value}") + + # NOTE: The subscriber's signature needs to match the MDS (message data specification) of + # the topic, otherwise, errors will arise. The MDS of a topic is set when the topic + # is created, which in our case is when a subscriber subscribes to a topic which + # is new (hasn't been subscribed to before). If the topic is being subscribed to by + # a subscriber for the first time, the topic's MDS will automatically be set + # according to that subscriber's signature. + self._pypubsub_publisher.subscribe(topicName=topic_value, listener=subscriber) + self._keep_subscriber_strongref(subscriber) + + def _keep_subscriber_strongref(self, subscriber: Callable[..., None]): + # NOTE: PyPubSub stores subscribers by weak reference. From the documentation: + # > PyPubSub holds listeners by weak reference so that the lifetime of the + # > callable is not affected by PyPubSub: once the application no longer + # > references the callable, it can be garbage collected and PyPubSub can clean up + # > so it is no longer registered (this happens thanks to the weakref module). + # > Without this, it would be imperative to remember to unsubscribe certain + # > listeners, which is error prone; they would end up living until the + # > application exited. + # + # https://pypubsub.readthedocs.io/en/v4.0.3/usage/usage_basic_tasks.html?#callable + # + # In our case, we're OK with subscribers living until the application exits. We don't + # provide an unsubscribe method (at this time) as subscriptions are expected to last + # for the life of the process. + # + # Specifically, if an instance object of a callable class is created and subscribed, + # we don't want that subscription to end if the callable instance object goes out of + # scope. Adding subscribers to self._refs prevents them from ever going out of scope. + self._refs.append(subscriber) + + def publish(self, topic: IslandEventTopics, event_data: Any = None): + topic_value = topic.value # needs to be a string for pypubsub + + logger.debug(f"Publishing {topic_value} event") + + # NOTE: `event_data` needs to match the MDS (message data specification) of the topic, + # otherwise, errors will arise. The MDS of a topic is set when the topic is created, + # which in our case is when a subscriber subscribes to a topic (in `subscribe()`) + # which is new (hasn't been subscribed to before). + self._pypubsub_publisher.sendMessage(topicName=topic_value, event=event_data) From 4219b6cbd4b3447427090dfe135abf6d11cd6c38 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Tue, 6 Sep 2022 13:19:53 +0530 Subject: [PATCH 26/56] Common: Rename IslandEventTopics -> IslandEventTopic --- monkey/common/event_queue/__init__.py | 2 +- monkey/common/event_queue/i_island_event_queue.py | 6 +++--- monkey/common/event_queue/pypubsub_island_event_queue.py | 6 +++--- 3 files changed, 7 insertions(+), 7 deletions(-) diff --git a/monkey/common/event_queue/__init__.py b/monkey/common/event_queue/__init__.py index 18412d2a7..dabc9527e 100644 --- a/monkey/common/event_queue/__init__.py +++ b/monkey/common/event_queue/__init__.py @@ -1,4 +1,4 @@ from .types import AgentEventSubscriber from .i_agent_event_queue import IAgentEventQueue -from .i_island_event_queue import IIslandEventQueue, IslandEventTopics +from .i_island_event_queue import IIslandEventQueue, IslandEventTopic from .pypubsub_agent_event_queue import PyPubSubAgentEventQueue diff --git a/monkey/common/event_queue/i_island_event_queue.py b/monkey/common/event_queue/i_island_event_queue.py index 8066b1385..a5614c37f 100644 --- a/monkey/common/event_queue/i_island_event_queue.py +++ b/monkey/common/event_queue/i_island_event_queue.py @@ -3,7 +3,7 @@ from enum import Enum from typing import Any, Callable -class IslandEventTopics(Enum): +class IslandEventTopic(Enum): AGENT_CONNECTED = "agent_connected" CLEAR_SIMULATION_DATA = "clear_simulation_data" RESET_AGENT_CONFIGURATION = "reset_agent_configuration" @@ -15,7 +15,7 @@ class IIslandEventQueue(ABC): """ @abstractmethod - def subscribe(self, topic: IslandEventTopics, subscriber: Callable[..., None]): + def subscribe(self, topic: IslandEventTopic, subscriber: Callable[..., None]): """ Subscribes a subscriber to the specified event topic @@ -26,7 +26,7 @@ class IIslandEventQueue(ABC): pass @abstractmethod - def publish(self, topic: IslandEventTopics, event_data: Any = None): + def publish(self, topic: IslandEventTopic, event_data: Any = None): """ Publishes an event topic with the given data diff --git a/monkey/common/event_queue/pypubsub_island_event_queue.py b/monkey/common/event_queue/pypubsub_island_event_queue.py index ba6ef3af4..25be20633 100644 --- a/monkey/common/event_queue/pypubsub_island_event_queue.py +++ b/monkey/common/event_queue/pypubsub_island_event_queue.py @@ -3,7 +3,7 @@ from typing import Any, Callable from pubsub.core import Publisher -from . import IIslandEventQueue, IslandEventTopics +from . import IIslandEventQueue, IslandEventTopic logger = logging.getLogger(__name__) @@ -13,7 +13,7 @@ class PyPubSubIslandEventQueue(IIslandEventQueue): self._pypubsub_publisher = pypubsub_publisher self._refs = [] - def subscribe(self, topic: IslandEventTopics, subscriber: Callable[..., None]): + def subscribe(self, topic: IslandEventTopic, subscriber: Callable[..., None]): topic_value = topic.value # needs to be a string for pypubsub try: subscriber_name = subscriber.__name__ @@ -52,7 +52,7 @@ class PyPubSubIslandEventQueue(IIslandEventQueue): # scope. Adding subscribers to self._refs prevents them from ever going out of scope. self._refs.append(subscriber) - def publish(self, topic: IslandEventTopics, event_data: Any = None): + def publish(self, topic: IslandEventTopic, event_data: Any = None): topic_value = topic.value # needs to be a string for pypubsub logger.debug(f"Publishing {topic_value} event") From 342a4959b3f2cb5fe51c50be5a1c9312d7d7510c Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Tue, 6 Sep 2022 13:20:53 +0530 Subject: [PATCH 27/56] Common: Import PyPubSubIslandEventQueue in common/event_queue/__init__.py --- monkey/common/event_queue/__init__.py | 1 + 1 file changed, 1 insertion(+) diff --git a/monkey/common/event_queue/__init__.py b/monkey/common/event_queue/__init__.py index dabc9527e..04b67e625 100644 --- a/monkey/common/event_queue/__init__.py +++ b/monkey/common/event_queue/__init__.py @@ -2,3 +2,4 @@ from .types import AgentEventSubscriber from .i_agent_event_queue import IAgentEventQueue from .i_island_event_queue import IIslandEventQueue, IslandEventTopic from .pypubsub_agent_event_queue import PyPubSubAgentEventQueue +from .pypubsub_island_event_queue import PyPubSubIslandEventQueue From fb4bfb7be14b31b7437ac99edc6a664527862ceb Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Tue, 6 Sep 2022 13:25:59 +0530 Subject: [PATCH 28/56] Project: Fix PyPubSubIslandEventQueue entries to Vulture allowlist --- vulture_allowlist.py | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/vulture_allowlist.py b/vulture_allowlist.py index 976810c9d..fc05f4c59 100644 --- a/vulture_allowlist.py +++ b/vulture_allowlist.py @@ -8,6 +8,7 @@ from common.agent_configuration.agent_sub_configurations import ( ScanTargetConfiguration, ) from common.credentials import Credentials +from common.event_queue import IslandEventTopic, PyPubSubIslandEventQueue from common.utils import IJSONSerializable from infection_monkey.exploit.log4shell_utils.ldap_server import LDAPServerFactory from monkey_island.cc.models import Report @@ -317,9 +318,7 @@ CC_TUNNEL Credentials.from_json IJSONSerializable.from_json -# revisit when implementing PyPubSubIslandEventQueue -AGENT_CONNECTED -CLEAR_SIMULATION_DATA -RESET_AGENT_CONFIGURATION -IIslandEventQueue -event_data +PyPubSubIslandEventQueue +IslandEventTopic.AGENT_CONNECTED +IslandEventTopic.CLEAR_SIMULATION_DATA +IslandEventTopic.RESET_AGENT_CONFIGURATION From ac2217ce8a8ce4901e38f2ed736f0e2c9b45fe53 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Tue, 6 Sep 2022 13:30:45 +0530 Subject: [PATCH 29/56] Common: Add IslandEventSubecriber type --- monkey/common/event_queue/__init__.py | 2 +- monkey/common/event_queue/i_island_event_queue.py | 6 ++++-- monkey/common/event_queue/pypubsub_island_event_queue.py | 8 ++++---- monkey/common/event_queue/types.py | 1 + 4 files changed, 10 insertions(+), 7 deletions(-) diff --git a/monkey/common/event_queue/__init__.py b/monkey/common/event_queue/__init__.py index 04b67e625..3f1e149b6 100644 --- a/monkey/common/event_queue/__init__.py +++ b/monkey/common/event_queue/__init__.py @@ -1,4 +1,4 @@ -from .types import AgentEventSubscriber +from .types import AgentEventSubscriber, IslandEventSubscriber from .i_agent_event_queue import IAgentEventQueue from .i_island_event_queue import IIslandEventQueue, IslandEventTopic from .pypubsub_agent_event_queue import PyPubSubAgentEventQueue diff --git a/monkey/common/event_queue/i_island_event_queue.py b/monkey/common/event_queue/i_island_event_queue.py index a5614c37f..65eb94ee3 100644 --- a/monkey/common/event_queue/i_island_event_queue.py +++ b/monkey/common/event_queue/i_island_event_queue.py @@ -1,6 +1,8 @@ from abc import ABC, abstractmethod from enum import Enum -from typing import Any, Callable +from typing import Any + +from . import IslandEventSubscriber class IslandEventTopic(Enum): @@ -15,7 +17,7 @@ class IIslandEventQueue(ABC): """ @abstractmethod - def subscribe(self, topic: IslandEventTopic, subscriber: Callable[..., None]): + def subscribe(self, topic: IslandEventTopic, subscriber: IslandEventSubscriber): """ Subscribes a subscriber to the specified event topic diff --git a/monkey/common/event_queue/pypubsub_island_event_queue.py b/monkey/common/event_queue/pypubsub_island_event_queue.py index 25be20633..d5d1df7cd 100644 --- a/monkey/common/event_queue/pypubsub_island_event_queue.py +++ b/monkey/common/event_queue/pypubsub_island_event_queue.py @@ -1,9 +1,9 @@ import logging -from typing import Any, Callable +from typing import Any from pubsub.core import Publisher -from . import IIslandEventQueue, IslandEventTopic +from . import IIslandEventQueue, IslandEventSubscriber, IslandEventTopic logger = logging.getLogger(__name__) @@ -13,7 +13,7 @@ class PyPubSubIslandEventQueue(IIslandEventQueue): self._pypubsub_publisher = pypubsub_publisher self._refs = [] - def subscribe(self, topic: IslandEventTopic, subscriber: Callable[..., None]): + def subscribe(self, topic: IslandEventTopic, subscriber: IslandEventSubscriber): topic_value = topic.value # needs to be a string for pypubsub try: subscriber_name = subscriber.__name__ @@ -31,7 +31,7 @@ class PyPubSubIslandEventQueue(IIslandEventQueue): self._pypubsub_publisher.subscribe(topicName=topic_value, listener=subscriber) self._keep_subscriber_strongref(subscriber) - def _keep_subscriber_strongref(self, subscriber: Callable[..., None]): + def _keep_subscriber_strongref(self, subscriber: IslandEventSubscriber): # NOTE: PyPubSub stores subscribers by weak reference. From the documentation: # > PyPubSub holds listeners by weak reference so that the lifetime of the # > callable is not affected by PyPubSub: once the application no longer diff --git a/monkey/common/event_queue/types.py b/monkey/common/event_queue/types.py index 6b57ae712..bd0b23dff 100644 --- a/monkey/common/event_queue/types.py +++ b/monkey/common/event_queue/types.py @@ -3,3 +3,4 @@ from typing import Callable from common.events import AbstractAgentEvent AgentEventSubscriber = Callable[[AbstractAgentEvent], None] +IslandEventSubscriber = Callable[..., None] From 265e083571c3392e908d5c6c0818016f9035567d Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Tue, 6 Sep 2022 13:33:36 +0530 Subject: [PATCH 30/56] UT: Rename test_pypubsub_event_queue.py -> test_pypubsub_agent_event_queue.py --- ...pypubsub_event_queue.py => test_pypubsub_agent_event_queue.py} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename monkey/tests/unit_tests/common/event_queue/{test_pypubsub_event_queue.py => test_pypubsub_agent_event_queue.py} (100%) diff --git a/monkey/tests/unit_tests/common/event_queue/test_pypubsub_event_queue.py b/monkey/tests/unit_tests/common/event_queue/test_pypubsub_agent_event_queue.py similarity index 100% rename from monkey/tests/unit_tests/common/event_queue/test_pypubsub_event_queue.py rename to monkey/tests/unit_tests/common/event_queue/test_pypubsub_agent_event_queue.py From 71c7a9a5336c7c94e0ed4baa11e27c6e47a7d4cd Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Tue, 6 Sep 2022 17:47:01 +0530 Subject: [PATCH 31/56] Common: Change parameter name event_data -> event in Island event queue --- monkey/common/event_queue/i_island_event_queue.py | 4 ++-- monkey/common/event_queue/pypubsub_island_event_queue.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/monkey/common/event_queue/i_island_event_queue.py b/monkey/common/event_queue/i_island_event_queue.py index 65eb94ee3..c8e5115a9 100644 --- a/monkey/common/event_queue/i_island_event_queue.py +++ b/monkey/common/event_queue/i_island_event_queue.py @@ -28,12 +28,12 @@ class IIslandEventQueue(ABC): pass @abstractmethod - def publish(self, topic: IslandEventTopic, event_data: Any = None): + def publish(self, topic: IslandEventTopic, event: Any = None): """ Publishes an event topic with the given data :param topic: Event topic to publish - :param event_data: Event data to pass to subscribers with the event publish + :param event: Event to pass to subscribers with the event publish """ pass diff --git a/monkey/common/event_queue/pypubsub_island_event_queue.py b/monkey/common/event_queue/pypubsub_island_event_queue.py index d5d1df7cd..8898b343a 100644 --- a/monkey/common/event_queue/pypubsub_island_event_queue.py +++ b/monkey/common/event_queue/pypubsub_island_event_queue.py @@ -52,7 +52,7 @@ class PyPubSubIslandEventQueue(IIslandEventQueue): # scope. Adding subscribers to self._refs prevents them from ever going out of scope. self._refs.append(subscriber) - def publish(self, topic: IslandEventTopic, event_data: Any = None): + def publish(self, topic: IslandEventTopic, event: Any = None): topic_value = topic.value # needs to be a string for pypubsub logger.debug(f"Publishing {topic_value} event") @@ -61,4 +61,4 @@ class PyPubSubIslandEventQueue(IIslandEventQueue): # otherwise, errors will arise. The MDS of a topic is set when the topic is created, # which in our case is when a subscriber subscribes to a topic (in `subscribe()`) # which is new (hasn't been subscribed to before). - self._pypubsub_publisher.sendMessage(topicName=topic_value, event=event_data) + self._pypubsub_publisher.sendMessage(topicName=topic_value, event=event) From 237f6d01b6581a446bc3124fa0382430a4fcd201 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Tue, 6 Sep 2022 17:53:36 +0530 Subject: [PATCH 32/56] UT: Add tests for PyPubSubIslandEventQueue --- .../test_pypubsub_island_event_queue.py | 67 +++++++++++++++++++ 1 file changed, 67 insertions(+) create mode 100644 monkey/tests/unit_tests/common/event_queue/test_pypubsub_island_event_queue.py diff --git a/monkey/tests/unit_tests/common/event_queue/test_pypubsub_island_event_queue.py b/monkey/tests/unit_tests/common/event_queue/test_pypubsub_island_event_queue.py new file mode 100644 index 000000000..cf51780f7 --- /dev/null +++ b/monkey/tests/unit_tests/common/event_queue/test_pypubsub_island_event_queue.py @@ -0,0 +1,67 @@ +from typing import Any, Callable + +import pytest +from pubsub import pub +from pubsub.core import Publisher + +from common.event_queue import ( + IIslandEventQueue, + IslandEventSubscriber, + IslandEventTopic, + PyPubSubIslandEventQueue, +) + + +@pytest.fixture +def event_queue() -> IIslandEventQueue: + return PyPubSubIslandEventQueue(Publisher()) + + +@pytest.fixture +def event_queue_subscriber() -> Callable[..., None]: + def fn(event, topic=pub.AUTO_TOPIC): + fn.call_count += 1 + fn.call_topics |= {topic.getName()} + + fn.call_count = 0 + fn.call_topics = set() + + return fn + + +def test_subscribe_publish( + event_queue: IIslandEventQueue, event_queue_subscriber: IslandEventSubscriber +): + event_queue.subscribe(topic=IslandEventTopic.AGENT_CONNECTED, subscriber=event_queue_subscriber) + event_queue.subscribe( + topic=IslandEventTopic.CLEAR_SIMULATION_DATA, subscriber=event_queue_subscriber + ) + + event_queue.publish(topic=IslandEventTopic.AGENT_CONNECTED) + event_queue.publish(topic=IslandEventTopic.CLEAR_SIMULATION_DATA) + event_queue.publish(topic=IslandEventTopic.RESET_AGENT_CONFIGURATION) + + assert event_queue_subscriber.call_count == 2 + assert event_queue_subscriber.call_topics == { + IslandEventTopic.AGENT_CONNECTED.value, + IslandEventTopic.CLEAR_SIMULATION_DATA.value, + } + + +def test_keep_subscriber_in_scope(event_queue: IIslandEventQueue): + class MyCallable: + called = False + + def __call__(self, event: Any): + MyCallable.called = True + + def subscribe(): + # fn will go out of scope after subscribe() returns. + fn = MyCallable() + event_queue.subscribe(topic=IslandEventTopic.AGENT_CONNECTED, subscriber=fn) + + subscribe() + + event_queue.publish(topic=IslandEventTopic.AGENT_CONNECTED) + + assert MyCallable.called From 27c8a1019b95143f5be345d330d82b4007781910 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Tue, 6 Sep 2022 17:59:17 +0530 Subject: [PATCH 33/56] Island: Register IIslandEventQueue instance in DI container --- monkey/monkey_island/cc/services/initialize.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/monkey/monkey_island/cc/services/initialize.py b/monkey/monkey_island/cc/services/initialize.py index b1caab657..5f5fc76c1 100644 --- a/monkey/monkey_island/cc/services/initialize.py +++ b/monkey/monkey_island/cc/services/initialize.py @@ -14,6 +14,7 @@ from common.aws import AWSInstance from common.common_consts.telem_categories import TelemCategoryEnum from common.event_queue import IAgentEventQueue, PyPubSubAgentEventQueue from common.utils.file_utils import get_binary_io_sha256_hash +from monkey.common.event_queue import IIslandEventQueue, PyPubSubIslandEventQueue from monkey_island.cc.repository import ( AgentBinaryRepository, FileAgentConfigurationRepository, @@ -64,6 +65,7 @@ def initialize_services(container: DIContainer, data_dir: Path): ) container.register(Publisher, Publisher) container.register_instance(IAgentEventQueue, container.resolve(PyPubSubAgentEventQueue)) + container.register_instance(IIslandEventQueue, container.resolve(PyPubSubIslandEventQueue)) _register_repositories(container, data_dir) _register_services(container) From 1a09f26fd9d16f86258802821fde319d1690ac6b Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Tue, 6 Sep 2022 18:07:25 +0530 Subject: [PATCH 34/56] Common: Modify IslandEventTopic enum to not have values --- monkey/common/event_queue/i_island_event_queue.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/monkey/common/event_queue/i_island_event_queue.py b/monkey/common/event_queue/i_island_event_queue.py index c8e5115a9..f5e8ea5c3 100644 --- a/monkey/common/event_queue/i_island_event_queue.py +++ b/monkey/common/event_queue/i_island_event_queue.py @@ -4,11 +4,9 @@ from typing import Any from . import IslandEventSubscriber - -class IslandEventTopic(Enum): - AGENT_CONNECTED = "agent_connected" - CLEAR_SIMULATION_DATA = "clear_simulation_data" - RESET_AGENT_CONFIGURATION = "reset_agent_configuration" +IslandEventTopic = Enum( + "IslandEventTopic", ["AGENT_CONNECTED", "CLEAR_SIMULATION_DATA", "RESET_AGENT_CONFIGURATION"] +) class IIslandEventQueue(ABC): From 004337583a9880fb559727f508f41841c44cb735 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Tue, 6 Sep 2022 18:11:08 +0530 Subject: [PATCH 35/56] Common: Use IslandEventTopic enum's names for pypubsub topics --- .../event_queue/pypubsub_island_event_queue.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/monkey/common/event_queue/pypubsub_island_event_queue.py b/monkey/common/event_queue/pypubsub_island_event_queue.py index 8898b343a..64c1b25ba 100644 --- a/monkey/common/event_queue/pypubsub_island_event_queue.py +++ b/monkey/common/event_queue/pypubsub_island_event_queue.py @@ -14,13 +14,13 @@ class PyPubSubIslandEventQueue(IIslandEventQueue): self._refs = [] def subscribe(self, topic: IslandEventTopic, subscriber: IslandEventSubscriber): - topic_value = topic.value # needs to be a string for pypubsub + topic_name = topic.name # needs to be a string for pypubsub try: subscriber_name = subscriber.__name__ except AttributeError: subscriber_name = subscriber.__class__.__name__ - logging.debug(f"Subscriber {subscriber_name} subscribed to {topic_value}") + logging.debug(f"Subscriber {subscriber_name} subscribed to {topic_name}") # NOTE: The subscriber's signature needs to match the MDS (message data specification) of # the topic, otherwise, errors will arise. The MDS of a topic is set when the topic @@ -28,7 +28,7 @@ class PyPubSubIslandEventQueue(IIslandEventQueue): # is new (hasn't been subscribed to before). If the topic is being subscribed to by # a subscriber for the first time, the topic's MDS will automatically be set # according to that subscriber's signature. - self._pypubsub_publisher.subscribe(topicName=topic_value, listener=subscriber) + self._pypubsub_publisher.subscribe(topicName=topic_name, listener=subscriber) self._keep_subscriber_strongref(subscriber) def _keep_subscriber_strongref(self, subscriber: IslandEventSubscriber): @@ -53,12 +53,12 @@ class PyPubSubIslandEventQueue(IIslandEventQueue): self._refs.append(subscriber) def publish(self, topic: IslandEventTopic, event: Any = None): - topic_value = topic.value # needs to be a string for pypubsub + topic_name = topic.name # needs to be a string for pypubsub - logger.debug(f"Publishing {topic_value} event") + logger.debug(f"Publishing {topic_name} event") # NOTE: `event_data` needs to match the MDS (message data specification) of the topic, # otherwise, errors will arise. The MDS of a topic is set when the topic is created, # which in our case is when a subscriber subscribes to a topic (in `subscribe()`) # which is new (hasn't been subscribed to before). - self._pypubsub_publisher.sendMessage(topicName=topic_value, event=event) + self._pypubsub_publisher.sendMessage(topicName=topic_name, event=event) From 5da8b424b57fe3df0f968c7fa5a45eceefafe1da Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Tue, 6 Sep 2022 18:13:25 +0530 Subject: [PATCH 36/56] UT: Use IslandEventTopic enum's names for pypubsub topics --- .../common/event_queue/test_pypubsub_island_event_queue.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/monkey/tests/unit_tests/common/event_queue/test_pypubsub_island_event_queue.py b/monkey/tests/unit_tests/common/event_queue/test_pypubsub_island_event_queue.py index cf51780f7..d026a288a 100644 --- a/monkey/tests/unit_tests/common/event_queue/test_pypubsub_island_event_queue.py +++ b/monkey/tests/unit_tests/common/event_queue/test_pypubsub_island_event_queue.py @@ -43,8 +43,8 @@ def test_subscribe_publish( assert event_queue_subscriber.call_count == 2 assert event_queue_subscriber.call_topics == { - IslandEventTopic.AGENT_CONNECTED.value, - IslandEventTopic.CLEAR_SIMULATION_DATA.value, + IslandEventTopic.AGENT_CONNECTED.name, + IslandEventTopic.CLEAR_SIMULATION_DATA.name, } From 502a875fdd6d22f350d0cd70cc02c54567989e12 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Tue, 6 Sep 2022 18:57:56 +0530 Subject: [PATCH 37/56] Common: Move Island event queue stuff out of common/ --- monkey/common/event_queue/__init__.py | 4 +- .../event_queue/i_island_event_queue.py | 37 ----------- .../pypubsub_island_event_queue.py | 64 ------------------- monkey/common/event_queue/types.py | 1 - .../monkey_island/cc/event_queue/__init__.py | 3 + 5 files changed, 4 insertions(+), 105 deletions(-) delete mode 100644 monkey/common/event_queue/i_island_event_queue.py delete mode 100644 monkey/common/event_queue/pypubsub_island_event_queue.py create mode 100644 monkey/monkey_island/cc/event_queue/__init__.py diff --git a/monkey/common/event_queue/__init__.py b/monkey/common/event_queue/__init__.py index 3f1e149b6..26d29cb49 100644 --- a/monkey/common/event_queue/__init__.py +++ b/monkey/common/event_queue/__init__.py @@ -1,5 +1,3 @@ -from .types import AgentEventSubscriber, IslandEventSubscriber +from .types import AgentEventSubscriber from .i_agent_event_queue import IAgentEventQueue -from .i_island_event_queue import IIslandEventQueue, IslandEventTopic from .pypubsub_agent_event_queue import PyPubSubAgentEventQueue -from .pypubsub_island_event_queue import PyPubSubIslandEventQueue diff --git a/monkey/common/event_queue/i_island_event_queue.py b/monkey/common/event_queue/i_island_event_queue.py deleted file mode 100644 index f5e8ea5c3..000000000 --- a/monkey/common/event_queue/i_island_event_queue.py +++ /dev/null @@ -1,37 +0,0 @@ -from abc import ABC, abstractmethod -from enum import Enum -from typing import Any - -from . import IslandEventSubscriber - -IslandEventTopic = Enum( - "IslandEventTopic", ["AGENT_CONNECTED", "CLEAR_SIMULATION_DATA", "RESET_AGENT_CONFIGURATION"] -) - - -class IIslandEventQueue(ABC): - """ - Manages subscription and publishing of events in the Island - """ - - @abstractmethod - def subscribe(self, topic: IslandEventTopic, subscriber: IslandEventSubscriber): - """ - Subscribes a subscriber to the specified event topic - - :param topic: Event topic to which the subscriber should subscribe - :param subscriber: A subscriber that will receive events - """ - - pass - - @abstractmethod - def publish(self, topic: IslandEventTopic, event: Any = None): - """ - Publishes an event topic with the given data - - :param topic: Event topic to publish - :param event: Event to pass to subscribers with the event publish - """ - - pass diff --git a/monkey/common/event_queue/pypubsub_island_event_queue.py b/monkey/common/event_queue/pypubsub_island_event_queue.py deleted file mode 100644 index 64c1b25ba..000000000 --- a/monkey/common/event_queue/pypubsub_island_event_queue.py +++ /dev/null @@ -1,64 +0,0 @@ -import logging -from typing import Any - -from pubsub.core import Publisher - -from . import IIslandEventQueue, IslandEventSubscriber, IslandEventTopic - -logger = logging.getLogger(__name__) - - -class PyPubSubIslandEventQueue(IIslandEventQueue): - def __init__(self, pypubsub_publisher: Publisher): - self._pypubsub_publisher = pypubsub_publisher - self._refs = [] - - def subscribe(self, topic: IslandEventTopic, subscriber: IslandEventSubscriber): - topic_name = topic.name # needs to be a string for pypubsub - try: - subscriber_name = subscriber.__name__ - except AttributeError: - subscriber_name = subscriber.__class__.__name__ - - logging.debug(f"Subscriber {subscriber_name} subscribed to {topic_name}") - - # NOTE: The subscriber's signature needs to match the MDS (message data specification) of - # the topic, otherwise, errors will arise. The MDS of a topic is set when the topic - # is created, which in our case is when a subscriber subscribes to a topic which - # is new (hasn't been subscribed to before). If the topic is being subscribed to by - # a subscriber for the first time, the topic's MDS will automatically be set - # according to that subscriber's signature. - self._pypubsub_publisher.subscribe(topicName=topic_name, listener=subscriber) - self._keep_subscriber_strongref(subscriber) - - def _keep_subscriber_strongref(self, subscriber: IslandEventSubscriber): - # NOTE: PyPubSub stores subscribers by weak reference. From the documentation: - # > PyPubSub holds listeners by weak reference so that the lifetime of the - # > callable is not affected by PyPubSub: once the application no longer - # > references the callable, it can be garbage collected and PyPubSub can clean up - # > so it is no longer registered (this happens thanks to the weakref module). - # > Without this, it would be imperative to remember to unsubscribe certain - # > listeners, which is error prone; they would end up living until the - # > application exited. - # - # https://pypubsub.readthedocs.io/en/v4.0.3/usage/usage_basic_tasks.html?#callable - # - # In our case, we're OK with subscribers living until the application exits. We don't - # provide an unsubscribe method (at this time) as subscriptions are expected to last - # for the life of the process. - # - # Specifically, if an instance object of a callable class is created and subscribed, - # we don't want that subscription to end if the callable instance object goes out of - # scope. Adding subscribers to self._refs prevents them from ever going out of scope. - self._refs.append(subscriber) - - def publish(self, topic: IslandEventTopic, event: Any = None): - topic_name = topic.name # needs to be a string for pypubsub - - logger.debug(f"Publishing {topic_name} event") - - # NOTE: `event_data` needs to match the MDS (message data specification) of the topic, - # otherwise, errors will arise. The MDS of a topic is set when the topic is created, - # which in our case is when a subscriber subscribes to a topic (in `subscribe()`) - # which is new (hasn't been subscribed to before). - self._pypubsub_publisher.sendMessage(topicName=topic_name, event=event) diff --git a/monkey/common/event_queue/types.py b/monkey/common/event_queue/types.py index bd0b23dff..6b57ae712 100644 --- a/monkey/common/event_queue/types.py +++ b/monkey/common/event_queue/types.py @@ -3,4 +3,3 @@ from typing import Callable from common.events import AbstractAgentEvent AgentEventSubscriber = Callable[[AbstractAgentEvent], None] -IslandEventSubscriber = Callable[..., None] diff --git a/monkey/monkey_island/cc/event_queue/__init__.py b/monkey/monkey_island/cc/event_queue/__init__.py new file mode 100644 index 000000000..6eab2b6ac --- /dev/null +++ b/monkey/monkey_island/cc/event_queue/__init__.py @@ -0,0 +1,3 @@ +from .types import IslandEventSubscriber +from .i_island_event_queue import IIslandEventQueue, IslandEventTopic +from .pypubsub_island_event_queue import PyPubSubIslandEventQueue From f2e7a3d66f12c285efa52bd0599110b835c62fc3 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Tue, 6 Sep 2022 18:58:33 +0530 Subject: [PATCH 38/56] Island: Add Island event queue stuff to monkey_island/ --- .../cc/event_queue/i_island_event_queue.py | 37 +++++++++++ .../pypubsub_island_event_queue.py | 64 +++++++++++++++++++ monkey/monkey_island/cc/event_queue/types.py | 3 + 3 files changed, 104 insertions(+) create mode 100644 monkey/monkey_island/cc/event_queue/i_island_event_queue.py create mode 100644 monkey/monkey_island/cc/event_queue/pypubsub_island_event_queue.py create mode 100644 monkey/monkey_island/cc/event_queue/types.py diff --git a/monkey/monkey_island/cc/event_queue/i_island_event_queue.py b/monkey/monkey_island/cc/event_queue/i_island_event_queue.py new file mode 100644 index 000000000..f5e8ea5c3 --- /dev/null +++ b/monkey/monkey_island/cc/event_queue/i_island_event_queue.py @@ -0,0 +1,37 @@ +from abc import ABC, abstractmethod +from enum import Enum +from typing import Any + +from . import IslandEventSubscriber + +IslandEventTopic = Enum( + "IslandEventTopic", ["AGENT_CONNECTED", "CLEAR_SIMULATION_DATA", "RESET_AGENT_CONFIGURATION"] +) + + +class IIslandEventQueue(ABC): + """ + Manages subscription and publishing of events in the Island + """ + + @abstractmethod + def subscribe(self, topic: IslandEventTopic, subscriber: IslandEventSubscriber): + """ + Subscribes a subscriber to the specified event topic + + :param topic: Event topic to which the subscriber should subscribe + :param subscriber: A subscriber that will receive events + """ + + pass + + @abstractmethod + def publish(self, topic: IslandEventTopic, event: Any = None): + """ + Publishes an event topic with the given data + + :param topic: Event topic to publish + :param event: Event to pass to subscribers with the event publish + """ + + pass diff --git a/monkey/monkey_island/cc/event_queue/pypubsub_island_event_queue.py b/monkey/monkey_island/cc/event_queue/pypubsub_island_event_queue.py new file mode 100644 index 000000000..64c1b25ba --- /dev/null +++ b/monkey/monkey_island/cc/event_queue/pypubsub_island_event_queue.py @@ -0,0 +1,64 @@ +import logging +from typing import Any + +from pubsub.core import Publisher + +from . import IIslandEventQueue, IslandEventSubscriber, IslandEventTopic + +logger = logging.getLogger(__name__) + + +class PyPubSubIslandEventQueue(IIslandEventQueue): + def __init__(self, pypubsub_publisher: Publisher): + self._pypubsub_publisher = pypubsub_publisher + self._refs = [] + + def subscribe(self, topic: IslandEventTopic, subscriber: IslandEventSubscriber): + topic_name = topic.name # needs to be a string for pypubsub + try: + subscriber_name = subscriber.__name__ + except AttributeError: + subscriber_name = subscriber.__class__.__name__ + + logging.debug(f"Subscriber {subscriber_name} subscribed to {topic_name}") + + # NOTE: The subscriber's signature needs to match the MDS (message data specification) of + # the topic, otherwise, errors will arise. The MDS of a topic is set when the topic + # is created, which in our case is when a subscriber subscribes to a topic which + # is new (hasn't been subscribed to before). If the topic is being subscribed to by + # a subscriber for the first time, the topic's MDS will automatically be set + # according to that subscriber's signature. + self._pypubsub_publisher.subscribe(topicName=topic_name, listener=subscriber) + self._keep_subscriber_strongref(subscriber) + + def _keep_subscriber_strongref(self, subscriber: IslandEventSubscriber): + # NOTE: PyPubSub stores subscribers by weak reference. From the documentation: + # > PyPubSub holds listeners by weak reference so that the lifetime of the + # > callable is not affected by PyPubSub: once the application no longer + # > references the callable, it can be garbage collected and PyPubSub can clean up + # > so it is no longer registered (this happens thanks to the weakref module). + # > Without this, it would be imperative to remember to unsubscribe certain + # > listeners, which is error prone; they would end up living until the + # > application exited. + # + # https://pypubsub.readthedocs.io/en/v4.0.3/usage/usage_basic_tasks.html?#callable + # + # In our case, we're OK with subscribers living until the application exits. We don't + # provide an unsubscribe method (at this time) as subscriptions are expected to last + # for the life of the process. + # + # Specifically, if an instance object of a callable class is created and subscribed, + # we don't want that subscription to end if the callable instance object goes out of + # scope. Adding subscribers to self._refs prevents them from ever going out of scope. + self._refs.append(subscriber) + + def publish(self, topic: IslandEventTopic, event: Any = None): + topic_name = topic.name # needs to be a string for pypubsub + + logger.debug(f"Publishing {topic_name} event") + + # NOTE: `event_data` needs to match the MDS (message data specification) of the topic, + # otherwise, errors will arise. The MDS of a topic is set when the topic is created, + # which in our case is when a subscriber subscribes to a topic (in `subscribe()`) + # which is new (hasn't been subscribed to before). + self._pypubsub_publisher.sendMessage(topicName=topic_name, event=event) diff --git a/monkey/monkey_island/cc/event_queue/types.py b/monkey/monkey_island/cc/event_queue/types.py new file mode 100644 index 000000000..3a1588dbb --- /dev/null +++ b/monkey/monkey_island/cc/event_queue/types.py @@ -0,0 +1,3 @@ +from typing import Callable + +IslandEventSubscriber = Callable[..., None] From 8f35a435915418e207185b24de43be01d61f09f8 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Tue, 6 Sep 2022 19:00:29 +0530 Subject: [PATCH 39/56] Project: Fix import path in Vulture allowlist --- vulture_allowlist.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vulture_allowlist.py b/vulture_allowlist.py index fc05f4c59..a672f319e 100644 --- a/vulture_allowlist.py +++ b/vulture_allowlist.py @@ -8,9 +8,9 @@ from common.agent_configuration.agent_sub_configurations import ( ScanTargetConfiguration, ) from common.credentials import Credentials -from common.event_queue import IslandEventTopic, PyPubSubIslandEventQueue from common.utils import IJSONSerializable from infection_monkey.exploit.log4shell_utils.ldap_server import LDAPServerFactory +from monkey_island.cc.event_queue import IslandEventTopic, PyPubSubIslandEventQueue from monkey_island.cc.models import Report from monkey_island.cc.models.networkmap import Arc, NetworkMap from monkey_island.cc.repository.attack.IMitigationsRepository import IMitigationsRepository From 38c6d53cc545ebd77ee2a5b19eeb6c79f90a92fc Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Tue, 6 Sep 2022 19:05:50 +0530 Subject: [PATCH 40/56] UT: Move test_pypubsub_island_event_queue.py out of common/ and in to monkey_island/ --- .../cc}/event_queue/test_pypubsub_island_event_queue.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) rename monkey/tests/unit_tests/{common => monkey_island/cc}/event_queue/test_pypubsub_island_event_queue.py (97%) diff --git a/monkey/tests/unit_tests/common/event_queue/test_pypubsub_island_event_queue.py b/monkey/tests/unit_tests/monkey_island/cc/event_queue/test_pypubsub_island_event_queue.py similarity index 97% rename from monkey/tests/unit_tests/common/event_queue/test_pypubsub_island_event_queue.py rename to monkey/tests/unit_tests/monkey_island/cc/event_queue/test_pypubsub_island_event_queue.py index d026a288a..d2a2fa84f 100644 --- a/monkey/tests/unit_tests/common/event_queue/test_pypubsub_island_event_queue.py +++ b/monkey/tests/unit_tests/monkey_island/cc/event_queue/test_pypubsub_island_event_queue.py @@ -4,7 +4,7 @@ import pytest from pubsub import pub from pubsub.core import Publisher -from common.event_queue import ( +from monkey_island.cc.event_queue import ( IIslandEventQueue, IslandEventSubscriber, IslandEventTopic, From 3cf332a0790cfc1855423ce67c722d78bd04fe1a Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Tue, 6 Sep 2022 19:37:35 +0530 Subject: [PATCH 41/56] Common: Add PyPubSubPublisherWrapper --- monkey/common/event_queue/__init__.py | 1 + .../event_queue/pypubsub_publisher_wrapper.py | 47 +++++++++++++++++++ 2 files changed, 48 insertions(+) create mode 100644 monkey/common/event_queue/pypubsub_publisher_wrapper.py diff --git a/monkey/common/event_queue/__init__.py b/monkey/common/event_queue/__init__.py index 26d29cb49..1c2e287d7 100644 --- a/monkey/common/event_queue/__init__.py +++ b/monkey/common/event_queue/__init__.py @@ -1,3 +1,4 @@ from .types import AgentEventSubscriber +from .pypubsub_publisher_wrapper import PyPubSubPublisherWrapper from .i_agent_event_queue import IAgentEventQueue from .pypubsub_agent_event_queue import PyPubSubAgentEventQueue diff --git a/monkey/common/event_queue/pypubsub_publisher_wrapper.py b/monkey/common/event_queue/pypubsub_publisher_wrapper.py new file mode 100644 index 000000000..3f9b75d79 --- /dev/null +++ b/monkey/common/event_queue/pypubsub_publisher_wrapper.py @@ -0,0 +1,47 @@ +import logging +from typing import Any, Callable + +from pubsub.core import Publisher + +logger = logging.getLogger(__name__) + + +class PyPubSubPublisherWrapper: + def __init__(self, pypubsub_publisher: Publisher): + self._pypubsub_publisher = pypubsub_publisher + self._refs = [] + + def subscribe(self, topic_name: str, subscriber: Callable): + try: + subscriber_name = subscriber.__name__ + except AttributeError: + subscriber_name = subscriber.__class__.__name__ + + logging.debug(f"Subscriber {subscriber_name} subscribed to {topic_name}") + + self._pypubsub_publisher.subscribe(topicName=topic_name, listener=subscriber) + self._keep_subscriber_strongref(subscriber) + + def _keep_subscriber_strongref(self, subscriber: Callable): + # NOTE: PyPubSub stores subscribers by weak reference. From the documentation: + # > PyPubSub holds listeners by weak reference so that the lifetime of the + # > callable is not affected by PyPubSub: once the application no longer + # > references the callable, it can be garbage collected and PyPubSub can clean up + # > so it is no longer registered (this happens thanks to the weakref module). + # > Without this, it would be imperative to remember to unsubscribe certain + # > listeners, which is error prone; they would end up living until the + # > application exited. + # + # https://pypubsub.readthedocs.io/en/v4.0.3/usage/usage_basic_tasks.html?#callable + # + # In our case, we're OK with subscribers living until the application exits. We don't + # provide an unsubscribe method (at this time) as subscriptions are expected to last + # for the life of the process. + # + # Specifically, if an instance object of a callable class is created and subscribed, + # we don't want that subscription to end if the callable instance object goes out of + # scope. Adding subscribers to self._refs prevents them from ever going out of scope. + self._refs.append(subscriber) + + def publish(self, topic_name: str, event: Any = None): + self._pypubsub_publisher.sendMessage(topicName=topic_name, event=event) From 3c71211b799797faa980bf2aa188eff007b3b55b Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Tue, 6 Sep 2022 19:38:16 +0530 Subject: [PATCH 42/56] Common: Use PyPubSubPublisherWrapper in PyPubSubAgentEventQueue --- .../event_queue/pypubsub_agent_event_queue.py | 36 +++---------------- 1 file changed, 4 insertions(+), 32 deletions(-) diff --git a/monkey/common/event_queue/pypubsub_agent_event_queue.py b/monkey/common/event_queue/pypubsub_agent_event_queue.py index 86dbea4af..d5cebf39a 100644 --- a/monkey/common/event_queue/pypubsub_agent_event_queue.py +++ b/monkey/common/event_queue/pypubsub_agent_event_queue.py @@ -3,6 +3,7 @@ from typing import Type from pubsub.core import Publisher +from common.event_queue import PyPubSubPublisherWrapper from common.events import AbstractAgentEvent from . import AgentEventSubscriber, IAgentEventQueue @@ -14,8 +15,7 @@ logger = logging.getLogger(__name__) class PyPubSubAgentEventQueue(IAgentEventQueue): def __init__(self, pypubsub_publisher: Publisher): - self._pypubsub_publisher = pypubsub_publisher - self._refs = [] + self._pypubsub_publisher_wrapped = PyPubSubPublisherWrapper(pypubsub_publisher) def subscribe_all_events(self, subscriber: AgentEventSubscriber): self._subscribe(_ALL_EVENTS_TOPIC, subscriber) @@ -32,35 +32,7 @@ class PyPubSubAgentEventQueue(IAgentEventQueue): self._subscribe(tag_topic, subscriber) def _subscribe(self, topic: str, subscriber: AgentEventSubscriber): - try: - subscriber_name = subscriber.__name__ - except AttributeError: - subscriber_name = subscriber.__class__.__name__ - - logging.debug(f"Subscriber {subscriber_name} subscribed to {topic}") - self._pypubsub_publisher.subscribe(topicName=topic, listener=subscriber) - self._keep_subscriber_strongref(subscriber) - - def _keep_subscriber_strongref(self, subscriber: AgentEventSubscriber): - # NOTE: PyPubSub stores subscribers by weak reference. From the documentation: - # > PyPubSub holds listeners by weak reference so that the lifetime of the - # > callable is not affected by PyPubSub: once the application no longer - # > references the callable, it can be garbage collected and PyPubSub can clean up - # > so it is no longer registered (this happens thanks to the weakref module). - # > Without this, it would be imperative to remember to unsubscribe certain - # > listeners, which is error prone; they would end up living until the - # > application exited. - # - # https://pypubsub.readthedocs.io/en/v4.0.3/usage/usage_basic_tasks.html?#callable - # - # In our case, we're OK with subscribers living until the application exits. We don't - # provide an unsubscribe method (at this time) as subscriptions are expected to last - # for the life of the process. - # - # Specifically, if an instance object of a callable class is created and subscribed, - # we don't want that subscription to end if the callable instance object goes out of - # scope. Adding subscribers to self._refs prevents them from ever going out of scope. - self._refs.append(subscriber) + self._pypubsub_publisher_wrapped.subscribe(topic, subscriber) def publish(self, event: AbstractAgentEvent): self._publish_to_all_events_topic(event) @@ -81,7 +53,7 @@ class PyPubSubAgentEventQueue(IAgentEventQueue): def _publish_event(self, topic: str, event: AbstractAgentEvent): logger.debug(f"Publishing a {event.__class__.__name__} event to {topic}") - self._pypubsub_publisher.sendMessage(topic, event=event) + self._pypubsub_publisher_wrapped.publish(topic, event) # Appending a unique string to the topics for type and tags prevents bugs caused by collisions # between type names and tag names. From 70468c37fb34a91aee9dbed870fc41ab46884595 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Tue, 6 Sep 2022 19:38:48 +0530 Subject: [PATCH 43/56] Island: Use PyPubSubPublisherWrapper in PyPubSubIslandEventQueue --- .../pypubsub_island_event_queue.py | 37 +++---------------- 1 file changed, 5 insertions(+), 32 deletions(-) diff --git a/monkey/monkey_island/cc/event_queue/pypubsub_island_event_queue.py b/monkey/monkey_island/cc/event_queue/pypubsub_island_event_queue.py index 64c1b25ba..478228d72 100644 --- a/monkey/monkey_island/cc/event_queue/pypubsub_island_event_queue.py +++ b/monkey/monkey_island/cc/event_queue/pypubsub_island_event_queue.py @@ -3,6 +3,8 @@ from typing import Any from pubsub.core import Publisher +from common.event_queue import PyPubSubPublisherWrapper + from . import IIslandEventQueue, IslandEventSubscriber, IslandEventTopic logger = logging.getLogger(__name__) @@ -10,17 +12,10 @@ logger = logging.getLogger(__name__) class PyPubSubIslandEventQueue(IIslandEventQueue): def __init__(self, pypubsub_publisher: Publisher): - self._pypubsub_publisher = pypubsub_publisher - self._refs = [] + self._pypubsub_publisher_wrapped = PyPubSubPublisherWrapper(pypubsub_publisher) def subscribe(self, topic: IslandEventTopic, subscriber: IslandEventSubscriber): topic_name = topic.name # needs to be a string for pypubsub - try: - subscriber_name = subscriber.__name__ - except AttributeError: - subscriber_name = subscriber.__class__.__name__ - - logging.debug(f"Subscriber {subscriber_name} subscribed to {topic_name}") # NOTE: The subscriber's signature needs to match the MDS (message data specification) of # the topic, otherwise, errors will arise. The MDS of a topic is set when the topic @@ -28,29 +23,7 @@ class PyPubSubIslandEventQueue(IIslandEventQueue): # is new (hasn't been subscribed to before). If the topic is being subscribed to by # a subscriber for the first time, the topic's MDS will automatically be set # according to that subscriber's signature. - self._pypubsub_publisher.subscribe(topicName=topic_name, listener=subscriber) - self._keep_subscriber_strongref(subscriber) - - def _keep_subscriber_strongref(self, subscriber: IslandEventSubscriber): - # NOTE: PyPubSub stores subscribers by weak reference. From the documentation: - # > PyPubSub holds listeners by weak reference so that the lifetime of the - # > callable is not affected by PyPubSub: once the application no longer - # > references the callable, it can be garbage collected and PyPubSub can clean up - # > so it is no longer registered (this happens thanks to the weakref module). - # > Without this, it would be imperative to remember to unsubscribe certain - # > listeners, which is error prone; they would end up living until the - # > application exited. - # - # https://pypubsub.readthedocs.io/en/v4.0.3/usage/usage_basic_tasks.html?#callable - # - # In our case, we're OK with subscribers living until the application exits. We don't - # provide an unsubscribe method (at this time) as subscriptions are expected to last - # for the life of the process. - # - # Specifically, if an instance object of a callable class is created and subscribed, - # we don't want that subscription to end if the callable instance object goes out of - # scope. Adding subscribers to self._refs prevents them from ever going out of scope. - self._refs.append(subscriber) + self._pypubsub_publisher_wrapped.subscribe(topic_name, subscriber) def publish(self, topic: IslandEventTopic, event: Any = None): topic_name = topic.name # needs to be a string for pypubsub @@ -61,4 +34,4 @@ class PyPubSubIslandEventQueue(IIslandEventQueue): # otherwise, errors will arise. The MDS of a topic is set when the topic is created, # which in our case is when a subscriber subscribes to a topic (in `subscribe()`) # which is new (hasn't been subscribed to before). - self._pypubsub_publisher.sendMessage(topicName=topic_name, event=event) + self._pypubsub_publisher_wrapped.publish(topic_name, event) From 69813f8cd4727bb364c16af094622e64c8eff1a6 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Tue, 6 Sep 2022 19:41:42 +0530 Subject: [PATCH 44/56] Common: Add explanatory comments about pypubsub's internal working in PyPubSubPublisherWrapper --- .../common/event_queue/pypubsub_publisher_wrapper.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/monkey/common/event_queue/pypubsub_publisher_wrapper.py b/monkey/common/event_queue/pypubsub_publisher_wrapper.py index 3f9b75d79..510b929a8 100644 --- a/monkey/common/event_queue/pypubsub_publisher_wrapper.py +++ b/monkey/common/event_queue/pypubsub_publisher_wrapper.py @@ -19,6 +19,12 @@ class PyPubSubPublisherWrapper: logging.debug(f"Subscriber {subscriber_name} subscribed to {topic_name}") + # NOTE: The subscriber's signature needs to match the MDS (message data specification) of + # the topic, otherwise, errors will arise. The MDS of a topic is set when the topic + # is created, which in our case is when a subscriber subscribes to a topic which + # is new (hasn't been subscribed to before). If the topic is being subscribed to by + # a subscriber for the first time, the topic's MDS will automatically be set + # according to that subscriber's signature. self._pypubsub_publisher.subscribe(topicName=topic_name, listener=subscriber) self._keep_subscriber_strongref(subscriber) @@ -44,4 +50,8 @@ class PyPubSubPublisherWrapper: self._refs.append(subscriber) def publish(self, topic_name: str, event: Any = None): + # NOTE: `event` needs to match the MDS (message data specification) of the topic, + # otherwise, errors will arise. The MDS of a topic is set when the topic is created, + # which in our case is when a subscriber subscribes to a topic (in `subscribe()`) + # which is new (hasn't been subscribed to before). self._pypubsub_publisher.sendMessage(topicName=topic_name, event=event) From c16c6456aa33dd6c4e15f9ff1b89de36f70cb354 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Tue, 6 Sep 2022 19:41:59 +0530 Subject: [PATCH 45/56] Island: Remove unneeded comments from PyPubSubIslandEventQueue --- .../cc/event_queue/pypubsub_island_event_queue.py | 13 ------------- 1 file changed, 13 deletions(-) diff --git a/monkey/monkey_island/cc/event_queue/pypubsub_island_event_queue.py b/monkey/monkey_island/cc/event_queue/pypubsub_island_event_queue.py index 478228d72..e18c12f2a 100644 --- a/monkey/monkey_island/cc/event_queue/pypubsub_island_event_queue.py +++ b/monkey/monkey_island/cc/event_queue/pypubsub_island_event_queue.py @@ -16,22 +16,9 @@ class PyPubSubIslandEventQueue(IIslandEventQueue): def subscribe(self, topic: IslandEventTopic, subscriber: IslandEventSubscriber): topic_name = topic.name # needs to be a string for pypubsub - - # NOTE: The subscriber's signature needs to match the MDS (message data specification) of - # the topic, otherwise, errors will arise. The MDS of a topic is set when the topic - # is created, which in our case is when a subscriber subscribes to a topic which - # is new (hasn't been subscribed to before). If the topic is being subscribed to by - # a subscriber for the first time, the topic's MDS will automatically be set - # according to that subscriber's signature. self._pypubsub_publisher_wrapped.subscribe(topic_name, subscriber) def publish(self, topic: IslandEventTopic, event: Any = None): topic_name = topic.name # needs to be a string for pypubsub - logger.debug(f"Publishing {topic_name} event") - - # NOTE: `event_data` needs to match the MDS (message data specification) of the topic, - # otherwise, errors will arise. The MDS of a topic is set when the topic is created, - # which in our case is when a subscriber subscribes to a topic (in `subscribe()`) - # which is new (hasn't been subscribed to before). self._pypubsub_publisher_wrapped.publish(topic_name, event) From ba52eae8ed9fbeeed46efc99c71361929faf2438 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 6 Sep 2022 12:56:49 -0400 Subject: [PATCH 46/56] Common: Decouple PyPubSubPublisherWrapper from events --- .../event_queue/pypubsub_agent_event_queue.py | 2 +- .../event_queue/pypubsub_publisher_wrapper.py | 6 +-- .../pypubsub_island_event_queue.py | 6 ++- .../test_pypubsub_island_event_queue.py | 37 ++++++++++++++++--- 4 files changed, 40 insertions(+), 11 deletions(-) diff --git a/monkey/common/event_queue/pypubsub_agent_event_queue.py b/monkey/common/event_queue/pypubsub_agent_event_queue.py index d5cebf39a..cf75c5b6a 100644 --- a/monkey/common/event_queue/pypubsub_agent_event_queue.py +++ b/monkey/common/event_queue/pypubsub_agent_event_queue.py @@ -53,7 +53,7 @@ class PyPubSubAgentEventQueue(IAgentEventQueue): def _publish_event(self, topic: str, event: AbstractAgentEvent): logger.debug(f"Publishing a {event.__class__.__name__} event to {topic}") - self._pypubsub_publisher_wrapped.publish(topic, event) + self._pypubsub_publisher_wrapped.publish(topic, event=event) # Appending a unique string to the topics for type and tags prevents bugs caused by collisions # between type names and tag names. diff --git a/monkey/common/event_queue/pypubsub_publisher_wrapper.py b/monkey/common/event_queue/pypubsub_publisher_wrapper.py index 510b929a8..71c7dda48 100644 --- a/monkey/common/event_queue/pypubsub_publisher_wrapper.py +++ b/monkey/common/event_queue/pypubsub_publisher_wrapper.py @@ -1,5 +1,5 @@ import logging -from typing import Any, Callable +from typing import Callable from pubsub.core import Publisher @@ -49,9 +49,9 @@ class PyPubSubPublisherWrapper: # scope. Adding subscribers to self._refs prevents them from ever going out of scope. self._refs.append(subscriber) - def publish(self, topic_name: str, event: Any = None): + def publish(self, topic_name: str, **kwargs): # NOTE: `event` needs to match the MDS (message data specification) of the topic, # otherwise, errors will arise. The MDS of a topic is set when the topic is created, # which in our case is when a subscriber subscribes to a topic (in `subscribe()`) # which is new (hasn't been subscribed to before). - self._pypubsub_publisher.sendMessage(topicName=topic_name, event=event) + self._pypubsub_publisher.sendMessage(topicName=topic_name, **kwargs) diff --git a/monkey/monkey_island/cc/event_queue/pypubsub_island_event_queue.py b/monkey/monkey_island/cc/event_queue/pypubsub_island_event_queue.py index e18c12f2a..55d4f9846 100644 --- a/monkey/monkey_island/cc/event_queue/pypubsub_island_event_queue.py +++ b/monkey/monkey_island/cc/event_queue/pypubsub_island_event_queue.py @@ -21,4 +21,8 @@ class PyPubSubIslandEventQueue(IIslandEventQueue): def publish(self, topic: IslandEventTopic, event: Any = None): topic_name = topic.name # needs to be a string for pypubsub logger.debug(f"Publishing {topic_name} event") - self._pypubsub_publisher_wrapped.publish(topic_name, event) + + if event is None: + self._pypubsub_publisher_wrapped.publish(topic_name) + else: + self._pypubsub_publisher_wrapped.publish(topic_name, event=event) diff --git a/monkey/tests/unit_tests/monkey_island/cc/event_queue/test_pypubsub_island_event_queue.py b/monkey/tests/unit_tests/monkey_island/cc/event_queue/test_pypubsub_island_event_queue.py index d2a2fa84f..2546359fc 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/event_queue/test_pypubsub_island_event_queue.py +++ b/monkey/tests/unit_tests/monkey_island/cc/event_queue/test_pypubsub_island_event_queue.py @@ -1,4 +1,4 @@ -from typing import Any, Callable +from typing import Callable import pytest from pubsub import pub @@ -19,7 +19,7 @@ def event_queue() -> IIslandEventQueue: @pytest.fixture def event_queue_subscriber() -> Callable[..., None]: - def fn(event, topic=pub.AUTO_TOPIC): + def fn(topic=pub.AUTO_TOPIC): fn.call_count += 1 fn.call_topics |= {topic.getName()} @@ -29,10 +29,12 @@ def event_queue_subscriber() -> Callable[..., None]: return fn -def test_subscribe_publish( +def test_subscribe_publish__no_event_body( event_queue: IIslandEventQueue, event_queue_subscriber: IslandEventSubscriber ): - event_queue.subscribe(topic=IslandEventTopic.AGENT_CONNECTED, subscriber=event_queue_subscriber) + event_queue.subscribe( + topic=IslandEventTopic.RESET_AGENT_CONFIGURATION, subscriber=event_queue_subscriber + ) event_queue.subscribe( topic=IslandEventTopic.CLEAR_SIMULATION_DATA, subscriber=event_queue_subscriber ) @@ -43,16 +45,39 @@ def test_subscribe_publish( assert event_queue_subscriber.call_count == 2 assert event_queue_subscriber.call_topics == { - IslandEventTopic.AGENT_CONNECTED.name, + IslandEventTopic.RESET_AGENT_CONFIGURATION.name, IslandEventTopic.CLEAR_SIMULATION_DATA.name, } +def test_subscribe_publish__with_event_body( + event_queue: IIslandEventQueue, event_queue_subscriber: IslandEventSubscriber +): + class MyCallable: + call_count = 0 + event = None + + def __call__(self, event): + MyCallable.call_count += 1 + MyCallable.event = event + + event = "my event!" + my_callable = MyCallable() + event_queue.subscribe(topic=IslandEventTopic.AGENT_CONNECTED, subscriber=my_callable) + + event_queue.publish(topic=IslandEventTopic.AGENT_CONNECTED, event=event) + event_queue.publish(topic=IslandEventTopic.CLEAR_SIMULATION_DATA) + event_queue.publish(topic=IslandEventTopic.RESET_AGENT_CONFIGURATION) + + assert my_callable.call_count == 1 + assert my_callable.event == event + + def test_keep_subscriber_in_scope(event_queue: IIslandEventQueue): class MyCallable: called = False - def __call__(self, event: Any): + def __call__(self): MyCallable.called = True def subscribe(): From e1e119c27aaba2f15d5762ebedf9c916f96d4d58 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 6 Sep 2022 12:57:24 -0400 Subject: [PATCH 47/56] Common: Add missing type hint for PyPubSubPublisherWrapper._refs --- monkey/common/event_queue/pypubsub_publisher_wrapper.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/monkey/common/event_queue/pypubsub_publisher_wrapper.py b/monkey/common/event_queue/pypubsub_publisher_wrapper.py index 71c7dda48..74b4b9a0c 100644 --- a/monkey/common/event_queue/pypubsub_publisher_wrapper.py +++ b/monkey/common/event_queue/pypubsub_publisher_wrapper.py @@ -1,5 +1,5 @@ import logging -from typing import Callable +from typing import Callable, List from pubsub.core import Publisher @@ -9,7 +9,7 @@ logger = logging.getLogger(__name__) class PyPubSubPublisherWrapper: def __init__(self, pypubsub_publisher: Publisher): self._pypubsub_publisher = pypubsub_publisher - self._refs = [] + self._refs: List[Callable] = [] def subscribe(self, topic_name: str, subscriber: Callable): try: From 59c58b3115346764125ff2ca8d21581c3b110ddb Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 6 Sep 2022 12:59:22 -0400 Subject: [PATCH 48/56] UT: Replace fn() with SubscriberSpy callable --- .../event_queue/test_pypubsub_island_event_queue.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/monkey/tests/unit_tests/monkey_island/cc/event_queue/test_pypubsub_island_event_queue.py b/monkey/tests/unit_tests/monkey_island/cc/event_queue/test_pypubsub_island_event_queue.py index 2546359fc..e3cce926f 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/event_queue/test_pypubsub_island_event_queue.py +++ b/monkey/tests/unit_tests/monkey_island/cc/event_queue/test_pypubsub_island_event_queue.py @@ -19,14 +19,15 @@ def event_queue() -> IIslandEventQueue: @pytest.fixture def event_queue_subscriber() -> Callable[..., None]: - def fn(topic=pub.AUTO_TOPIC): - fn.call_count += 1 - fn.call_topics |= {topic.getName()} + class SubscriberSpy: + call_count = 0 + call_topics = set() - fn.call_count = 0 - fn.call_topics = set() + def __call__(self, topic=pub.AUTO_TOPIC): + self.call_count += 1 + self.call_topics |= {topic.getName()} - return fn + return SubscriberSpy() def test_subscribe_publish__no_event_body( From 377bb293fe49ac6df91433a33b277707c24af2dd Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 6 Sep 2022 13:00:01 -0400 Subject: [PATCH 49/56] UT: Use `self` instead of class name --- .../cc/event_queue/test_pypubsub_island_event_queue.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/monkey/tests/unit_tests/monkey_island/cc/event_queue/test_pypubsub_island_event_queue.py b/monkey/tests/unit_tests/monkey_island/cc/event_queue/test_pypubsub_island_event_queue.py index e3cce926f..206009727 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/event_queue/test_pypubsub_island_event_queue.py +++ b/monkey/tests/unit_tests/monkey_island/cc/event_queue/test_pypubsub_island_event_queue.py @@ -59,8 +59,8 @@ def test_subscribe_publish__with_event_body( event = None def __call__(self, event): - MyCallable.call_count += 1 - MyCallable.event = event + self.call_count += 1 + self.event = event event = "my event!" my_callable = MyCallable() From b16d19e0ed856e204ee72e36f0c7a402bddc6a68 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 6 Sep 2022 15:00:24 -0400 Subject: [PATCH 50/56] Common: Rename _pypubsub_publisher_wrappe{d,r} --- monkey/common/event_queue/pypubsub_agent_event_queue.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/monkey/common/event_queue/pypubsub_agent_event_queue.py b/monkey/common/event_queue/pypubsub_agent_event_queue.py index cf75c5b6a..246b743d1 100644 --- a/monkey/common/event_queue/pypubsub_agent_event_queue.py +++ b/monkey/common/event_queue/pypubsub_agent_event_queue.py @@ -15,7 +15,7 @@ logger = logging.getLogger(__name__) class PyPubSubAgentEventQueue(IAgentEventQueue): def __init__(self, pypubsub_publisher: Publisher): - self._pypubsub_publisher_wrapped = PyPubSubPublisherWrapper(pypubsub_publisher) + self._pypubsub_publisher_wrapper = PyPubSubPublisherWrapper(pypubsub_publisher) def subscribe_all_events(self, subscriber: AgentEventSubscriber): self._subscribe(_ALL_EVENTS_TOPIC, subscriber) @@ -32,7 +32,7 @@ class PyPubSubAgentEventQueue(IAgentEventQueue): self._subscribe(tag_topic, subscriber) def _subscribe(self, topic: str, subscriber: AgentEventSubscriber): - self._pypubsub_publisher_wrapped.subscribe(topic, subscriber) + self._pypubsub_publisher_wrapper.subscribe(topic, subscriber) def publish(self, event: AbstractAgentEvent): self._publish_to_all_events_topic(event) @@ -53,7 +53,7 @@ class PyPubSubAgentEventQueue(IAgentEventQueue): def _publish_event(self, topic: str, event: AbstractAgentEvent): logger.debug(f"Publishing a {event.__class__.__name__} event to {topic}") - self._pypubsub_publisher_wrapped.publish(topic, event=event) + self._pypubsub_publisher_wrapper.publish(topic, event=event) # Appending a unique string to the topics for type and tags prevents bugs caused by collisions # between type names and tag names. From 1036189fcc26b17b874a7ace08b4d472ec163bc0 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 6 Sep 2022 15:01:10 -0400 Subject: [PATCH 51/56] Island: Rename _pypubsub_publisher_wrappe{d,r} --- .../cc/event_queue/pypubsub_island_event_queue.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/monkey/monkey_island/cc/event_queue/pypubsub_island_event_queue.py b/monkey/monkey_island/cc/event_queue/pypubsub_island_event_queue.py index 55d4f9846..3cf1399c9 100644 --- a/monkey/monkey_island/cc/event_queue/pypubsub_island_event_queue.py +++ b/monkey/monkey_island/cc/event_queue/pypubsub_island_event_queue.py @@ -12,17 +12,17 @@ logger = logging.getLogger(__name__) class PyPubSubIslandEventQueue(IIslandEventQueue): def __init__(self, pypubsub_publisher: Publisher): - self._pypubsub_publisher_wrapped = PyPubSubPublisherWrapper(pypubsub_publisher) + self._pypubsub_publisher_wrapper = PyPubSubPublisherWrapper(pypubsub_publisher) def subscribe(self, topic: IslandEventTopic, subscriber: IslandEventSubscriber): topic_name = topic.name # needs to be a string for pypubsub - self._pypubsub_publisher_wrapped.subscribe(topic_name, subscriber) + self._pypubsub_publisher_wrapper.subscribe(topic_name, subscriber) def publish(self, topic: IslandEventTopic, event: Any = None): topic_name = topic.name # needs to be a string for pypubsub logger.debug(f"Publishing {topic_name} event") if event is None: - self._pypubsub_publisher_wrapped.publish(topic_name) + self._pypubsub_publisher_wrapper.publish(topic_name) else: - self._pypubsub_publisher_wrapped.publish(topic_name, event=event) + self._pypubsub_publisher_wrapper.publish(topic_name, event=event) From 4e4331c5c32fab9b714fad652b4064ef4ddad0de Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 6 Sep 2022 15:05:29 -0400 Subject: [PATCH 52/56] Common: s/event/kwargs in PyPubSubPublisherWrapper comment --- monkey/common/event_queue/pypubsub_publisher_wrapper.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/monkey/common/event_queue/pypubsub_publisher_wrapper.py b/monkey/common/event_queue/pypubsub_publisher_wrapper.py index 74b4b9a0c..ee2377edd 100644 --- a/monkey/common/event_queue/pypubsub_publisher_wrapper.py +++ b/monkey/common/event_queue/pypubsub_publisher_wrapper.py @@ -50,7 +50,7 @@ class PyPubSubPublisherWrapper: self._refs.append(subscriber) def publish(self, topic_name: str, **kwargs): - # NOTE: `event` needs to match the MDS (message data specification) of the topic, + # NOTE: `kwargs` needs to match the MDS (message data specification) of the topic, # otherwise, errors will arise. The MDS of a topic is set when the topic is created, # which in our case is when a subscriber subscribes to a topic (in `subscribe()`) # which is new (hasn't been subscribed to before). From 84aa993a8ba750f51bf067f8c5186d24772c8d67 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 6 Sep 2022 15:08:50 -0400 Subject: [PATCH 53/56] Island: Improve `event` docstring in IIslandEventQueue --- monkey/monkey_island/cc/event_queue/i_island_event_queue.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/monkey/monkey_island/cc/event_queue/i_island_event_queue.py b/monkey/monkey_island/cc/event_queue/i_island_event_queue.py index f5e8ea5c3..cfee63461 100644 --- a/monkey/monkey_island/cc/event_queue/i_island_event_queue.py +++ b/monkey/monkey_island/cc/event_queue/i_island_event_queue.py @@ -31,7 +31,7 @@ class IIslandEventQueue(ABC): Publishes an event topic with the given data :param topic: Event topic to publish - :param event: Event to pass to subscribers with the event publish + :param event: Event data to publish """ pass From aee4887b646b99aaf04a926b0fa6c9eb93b2110a Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 6 Sep 2022 15:11:02 -0400 Subject: [PATCH 54/56] Island: Fix island event queue import --- monkey/monkey_island/cc/services/initialize.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/monkey/monkey_island/cc/services/initialize.py b/monkey/monkey_island/cc/services/initialize.py index 5f5fc76c1..888bf75c1 100644 --- a/monkey/monkey_island/cc/services/initialize.py +++ b/monkey/monkey_island/cc/services/initialize.py @@ -14,7 +14,7 @@ from common.aws import AWSInstance from common.common_consts.telem_categories import TelemCategoryEnum from common.event_queue import IAgentEventQueue, PyPubSubAgentEventQueue from common.utils.file_utils import get_binary_io_sha256_hash -from monkey.common.event_queue import IIslandEventQueue, PyPubSubIslandEventQueue +from monkey_island.cc.event_queue import IIslandEventQueue, PyPubSubIslandEventQueue from monkey_island.cc.repository import ( AgentBinaryRepository, FileAgentConfigurationRepository, From 166588d00d6fcbb66d03603f29ff74c8a99f0b98 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 6 Sep 2022 15:12:37 -0400 Subject: [PATCH 55/56] Project: Remove PyPubSubAgentEventQueue from vulture_allowlist.py --- vulture_allowlist.py | 1 - 1 file changed, 1 deletion(-) diff --git a/vulture_allowlist.py b/vulture_allowlist.py index a672f319e..619fd8abc 100644 --- a/vulture_allowlist.py +++ b/vulture_allowlist.py @@ -286,7 +286,6 @@ subscribe_all # common\event_queue\pypubsub_agent_event_queue.py subscribe_type # common\event_queue\pypubsub_agent_event_queue.py subscribe_tag # common\event_queue\pypubsub_agent_event_queue.py publish # common\event_queue\pypubsub_agent_event_queue.py -PyPubSubAgentEventQueue # common\event_queue\pypubsub_agent_event_queue.py subscribe_all_events # common\event_queue\pypubsub_agent_event_queue.py From 205848f2a5154eeb9f8d5d0e4a1b64842e184652 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 6 Sep 2022 15:13:35 -0400 Subject: [PATCH 56/56] Project: Remove PyPubSubIslandEventQueue from vulture_allowlist.py --- vulture_allowlist.py | 1 - 1 file changed, 1 deletion(-) diff --git a/vulture_allowlist.py b/vulture_allowlist.py index 619fd8abc..42d46716f 100644 --- a/vulture_allowlist.py +++ b/vulture_allowlist.py @@ -317,7 +317,6 @@ CC_TUNNEL Credentials.from_json IJSONSerializable.from_json -PyPubSubIslandEventQueue IslandEventTopic.AGENT_CONNECTED IslandEventTopic.CLEAR_SIMULATION_DATA IslandEventTopic.RESET_AGENT_CONFIGURATION