diff --git a/monkey/infection_monkey/credential_repository/__init__.py b/monkey/infection_monkey/credential_repository/__init__.py index 62eae389b..f52bdb0c0 100644 --- a/monkey/infection_monkey/credential_repository/__init__.py +++ b/monkey/infection_monkey/credential_repository/__init__.py @@ -2,3 +2,6 @@ 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/credential_repository/add_credentials_from_event.py b/monkey/infection_monkey/credential_repository/add_credentials_from_event.py new file mode 100644 index 000000000..123774581 --- /dev/null +++ b/monkey/infection_monkey/credential_repository/add_credentials_from_event.py @@ -0,0 +1,11 @@ +from common.events import CredentialsStolenEvent + +from . import IPropagationCredentialsRepository + + +class add_credentials_from_event_to_propagation_credentials_repository: + def __init__(self, credentials_repository: IPropagationCredentialsRepository): + self._credentials_repository = credentials_repository + + def __call__(self, event: CredentialsStolenEvent): + self._credentials_repository.add_credentials(event.stolen_credentials) diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index 665f97222..258d05a36 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -10,6 +10,7 @@ from pubsub.core import Publisher import infection_monkey.tunnel as tunnel from common.event_queue import PyPubSubEventQueue +from common.events import CredentialsStolenEvent from common.network.network_utils import address_to_ip_port from common.utils.argparse_types import positive_int from common.utils.attack_utils import ScanStatus, UsageEnum @@ -23,6 +24,7 @@ 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 CachingAgentRepository, ExploiterWrapper from infection_monkey.exploit.hadoop import HadoopExploiter @@ -198,6 +200,9 @@ class InfectionMonkey: local_network_interfaces = InfectionMonkey._get_local_network_interfaces() _event_queue = PyPubSubEventQueue(Publisher()) + _event_queue.subscribe_type( + CredentialsStolenEvent, add_credentials_from_event_to_propagation_credentials_repository + ) # TODO control_channel and control_client have same responsibilities, merge them control_channel = ControlChannel( 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 new file mode 100644 index 000000000..e01321fce --- /dev/null +++ b/monkey/tests/unit_tests/infection_monkey/credential_store/test_add_credentials_from_event.py @@ -0,0 +1,30 @@ +from unittest.mock import MagicMock +from uuid import UUID + +from common.credentials import Credentials, Password, Username +from common.events import CredentialsStolenEvent +from infection_monkey.credential_repository import ( + IPropagationCredentialsRepository, + add_credentials_from_event_to_propagation_credentials_repository, +) + +credentials = [Credentials(identity=Username("test_username"), secret=Password("some_password"))] + +credentials_stolen_event = CredentialsStolenEvent( + source=UUID("f811ad00-5a68-4437-bd51-7b5cc1768ad5"), + target=None, + timestamp=0.0, + tags=frozenset({"stolen credentials"}), + stolen_credentials=credentials, +) + + +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( + mock_propagation_credentials_repository + ) + + fn(credentials_stolen_event) + + assert mock_propagation_credentials_repository.add_credentials.called_with(credentials) diff --git a/vulture_allowlist.py b/vulture_allowlist.py index 8c0beb5ef..f547e8e5c 100644 --- a/vulture_allowlist.py +++ b/vulture_allowlist.py @@ -217,7 +217,7 @@ _event_queue # TODO DELETE IN #2176 CredentialsStolenEvent - +add_credentials_from_event_to_propagation_credentials_repository # TODO DELETE AFTER RESOURCE REFACTORING