From 37b884a5b8b552858af3a3a5d735dc7deea8b9cf Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Mon, 3 Oct 2022 14:02:48 -0400 Subject: [PATCH 1/2] Agent: Move agent_event_forwarder.py to agent_event_handlers package --- monkey/infection_monkey/agent_event_handlers/__init__.py | 1 + .../{ => agent_event_handlers}/agent_event_forwarder.py | 0 monkey/infection_monkey/monkey.py | 3 +-- .../unit_tests/infection_monkey/test_agent_event_forwarder.py | 2 +- 4 files changed, 3 insertions(+), 3 deletions(-) rename monkey/infection_monkey/{ => agent_event_handlers}/agent_event_forwarder.py (100%) diff --git a/monkey/infection_monkey/agent_event_handlers/__init__.py b/monkey/infection_monkey/agent_event_handlers/__init__.py index 3e0f9b03c..925bf7f63 100644 --- a/monkey/infection_monkey/agent_event_handlers/__init__.py +++ b/monkey/infection_monkey/agent_event_handlers/__init__.py @@ -1 +1,2 @@ from .notify_relay_on_propagation import notify_relay_on_propagation +from .agent_event_forwarder import AgentEventForwarder diff --git a/monkey/infection_monkey/agent_event_forwarder.py b/monkey/infection_monkey/agent_event_handlers/agent_event_forwarder.py similarity index 100% rename from monkey/infection_monkey/agent_event_forwarder.py rename to monkey/infection_monkey/agent_event_handlers/agent_event_forwarder.py diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index e38cbc96e..7ca2162f5 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -21,8 +21,7 @@ from common.types import SocketAddress from common.utils.argparse_types import positive_int from common.utils.attack_utils import ScanStatus, UsageEnum from common.version import get_version -from infection_monkey.agent_event_forwarder import AgentEventForwarder -from infection_monkey.agent_event_handlers import notify_relay_on_propagation +from infection_monkey.agent_event_handlers import AgentEventForwarder, notify_relay_on_propagation from infection_monkey.config import GUID from infection_monkey.control import ControlClient from infection_monkey.credential_collectors import ( diff --git a/monkey/tests/unit_tests/infection_monkey/test_agent_event_forwarder.py b/monkey/tests/unit_tests/infection_monkey/test_agent_event_forwarder.py index aa2d1381a..07c7b5d52 100644 --- a/monkey/tests/unit_tests/infection_monkey/test_agent_event_forwarder.py +++ b/monkey/tests/unit_tests/infection_monkey/test_agent_event_forwarder.py @@ -3,7 +3,7 @@ from unittest.mock import MagicMock import pytest -from infection_monkey.agent_event_forwarder import BatchingAgentEventForwarder +from infection_monkey.agent_event_handlers.agent_event_forwarder import BatchingAgentEventForwarder from infection_monkey.island_api_client import IIslandAPIClient SERVER = "1.1.1.1:9999" From adee0b40631437384db2bc838d79ccf62d6b8561 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Mon, 3 Oct 2022 14:11:20 -0400 Subject: [PATCH 2/2] Agent: Move add_credentials_from_event to agent_event_handlers package --- monkey/infection_monkey/agent_event_handlers/__init__.py | 3 +++ .../add_stolen_credentials_to_repository.py} | 5 ++--- .../infection_monkey/credential_repository/__init__.py | 3 --- monkey/infection_monkey/monkey.py | 9 ++++++--- .../credential_store/test_add_credentials_from_event.py | 8 ++++---- 5 files changed, 15 insertions(+), 13 deletions(-) rename monkey/infection_monkey/{credential_repository/add_credentials_from_event.py => agent_event_handlers/add_stolen_credentials_to_repository.py} (76%) diff --git a/monkey/infection_monkey/agent_event_handlers/__init__.py b/monkey/infection_monkey/agent_event_handlers/__init__.py index 925bf7f63..382131267 100644 --- a/monkey/infection_monkey/agent_event_handlers/__init__.py +++ b/monkey/infection_monkey/agent_event_handlers/__init__.py @@ -1,2 +1,5 @@ from .notify_relay_on_propagation import notify_relay_on_propagation from .agent_event_forwarder import AgentEventForwarder +from .add_stolen_credentials_to_repository import ( + add_stolen_credentials_to_propagation_credentials_repository, +) diff --git a/monkey/infection_monkey/credential_repository/add_credentials_from_event.py b/monkey/infection_monkey/agent_event_handlers/add_stolen_credentials_to_repository.py similarity index 76% rename from monkey/infection_monkey/credential_repository/add_credentials_from_event.py rename to monkey/infection_monkey/agent_event_handlers/add_stolen_credentials_to_repository.py index 639a26bcf..16ead8aa5 100644 --- a/monkey/infection_monkey/credential_repository/add_credentials_from_event.py +++ b/monkey/infection_monkey/agent_event_handlers/add_stolen_credentials_to_repository.py @@ -1,13 +1,12 @@ import logging from common.agent_events import CredentialsStolenEvent - -from . import IPropagationCredentialsRepository +from infection_monkey.credential_repository import IPropagationCredentialsRepository logger = logging.getLogger(__name__) -class add_credentials_from_event_to_propagation_credentials_repository: +class add_stolen_credentials_to_propagation_credentials_repository: def __init__(self, credentials_repository: IPropagationCredentialsRepository): self._credentials_repository = credentials_repository diff --git a/monkey/infection_monkey/credential_repository/__init__.py b/monkey/infection_monkey/credential_repository/__init__.py index f52bdb0c0..62eae389b 100644 --- a/monkey/infection_monkey/credential_repository/__init__.py +++ b/monkey/infection_monkey/credential_repository/__init__.py @@ -2,6 +2,3 @@ from .i_propagation_credentials_repository import IPropagationCredentialsReposit from .aggregating_propagation_credentials_repository import ( AggregatingPropagationCredentialsRepository, ) -from .add_credentials_from_event import ( - add_credentials_from_event_to_propagation_credentials_repository, -) diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index 7ca2162f5..f7a6eb11c 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -21,7 +21,11 @@ from common.types import SocketAddress from common.utils.argparse_types import positive_int from common.utils.attack_utils import ScanStatus, UsageEnum from common.version import get_version -from infection_monkey.agent_event_handlers import AgentEventForwarder, notify_relay_on_propagation +from infection_monkey.agent_event_handlers import ( + AgentEventForwarder, + add_stolen_credentials_to_propagation_credentials_repository, + notify_relay_on_propagation, +) from infection_monkey.config import GUID from infection_monkey.control import ControlClient from infection_monkey.credential_collectors import ( @@ -31,7 +35,6 @@ from infection_monkey.credential_collectors import ( from infection_monkey.credential_repository import ( AggregatingPropagationCredentialsRepository, IPropagationCredentialsRepository, - add_credentials_from_event_to_propagation_credentials_repository, ) from infection_monkey.exploit import CachingAgentBinaryRepository, ExploiterWrapper from infection_monkey.exploit.hadoop import HadoopExploiter @@ -299,7 +302,7 @@ class InfectionMonkey: ): agent_event_queue.subscribe_type( CredentialsStolenEvent, - add_credentials_from_event_to_propagation_credentials_repository( + add_stolen_credentials_to_propagation_credentials_repository( propagation_credentials_repository ), ) diff --git a/monkey/tests/unit_tests/infection_monkey/credential_store/test_add_credentials_from_event.py b/monkey/tests/unit_tests/infection_monkey/credential_store/test_add_credentials_from_event.py index 5327603fe..df22c54ca 100644 --- a/monkey/tests/unit_tests/infection_monkey/credential_store/test_add_credentials_from_event.py +++ b/monkey/tests/unit_tests/infection_monkey/credential_store/test_add_credentials_from_event.py @@ -3,10 +3,10 @@ from uuid import UUID from common.agent_events import CredentialsStolenEvent from common.credentials import Credentials, Password, Username -from infection_monkey.credential_repository import ( - IPropagationCredentialsRepository, - add_credentials_from_event_to_propagation_credentials_repository, +from infection_monkey.agent_event_handlers import ( + add_stolen_credentials_to_propagation_credentials_repository, ) +from infection_monkey.credential_repository import IPropagationCredentialsRepository credentials = [ Credentials( @@ -25,7 +25,7 @@ credentials_stolen_event = CredentialsStolenEvent( def test_add_credentials_from_event_to_propagation_credentials_repository(): mock_propagation_credentials_repository = MagicMock(spec=IPropagationCredentialsRepository) - fn = add_credentials_from_event_to_propagation_credentials_repository( + fn = add_stolen_credentials_to_propagation_credentials_repository( mock_propagation_credentials_repository )