diff --git a/monkey/infection_monkey/agent_event_handlers/__init__.py b/monkey/infection_monkey/agent_event_handlers/__init__.py index 3e0f9b03c..382131267 100644 --- a/monkey/infection_monkey/agent_event_handlers/__init__.py +++ b/monkey/infection_monkey/agent_event_handlers/__init__.py @@ -1 +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/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/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 e38cbc96e..f7a6eb11c 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -21,8 +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_forwarder import AgentEventForwarder -from infection_monkey.agent_event_handlers import 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 ( @@ -32,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 @@ -300,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 ) 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"