From 9e6a569393e3e198d9a6579a78069a978689ab1a Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Mon, 15 Aug 2022 19:14:42 +0000 Subject: [PATCH 1/2] Agent: Update credentials repository to cache per-instance --- .../aggregating_propagation_credentials_repository.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/monkey/infection_monkey/credential_repository/aggregating_propagation_credentials_repository.py b/monkey/infection_monkey/credential_repository/aggregating_propagation_credentials_repository.py index 3050e5f85..59ff279e0 100644 --- a/monkey/infection_monkey/credential_repository/aggregating_propagation_credentials_repository.py +++ b/monkey/infection_monkey/credential_repository/aggregating_propagation_credentials_repository.py @@ -1,5 +1,5 @@ import logging -from typing import Any, Iterable, Sequence +from typing import Any, Iterable from common.credentials import CredentialComponentType, Credentials, ICredentialComponent from infection_monkey.custom_types import PropagationCredentials @@ -29,6 +29,11 @@ class AggregatingPropagationCredentialsRepository(IPropagationCredentialsReposit } self._control_channel = control_channel + # Ensure caching happens per-instance instead of being shared across instances + self._get_credentials_from_control_channel = request_cache(CREDENTIALS_POLL_PERIOD_SEC)( + self._control_channel.get_credentials_for_propagation + ) + def add_credentials(self, credentials_to_add: Iterable[Credentials]): for credentials in credentials_to_add: if credentials.identity: @@ -64,10 +69,6 @@ class AggregatingPropagationCredentialsRepository(IPropagationCredentialsReposit self._stored_credentials = {} logger.error(f"Error while attempting to retrieve credentials for propagation: {ex}") - @request_cache(CREDENTIALS_POLL_PERIOD_SEC) - def _get_credentials_from_control_channel(self) -> Sequence[Credentials]: - return self._control_channel.get_credentials_for_propagation() - def _set_attribute(self, attribute_to_be_set: str, credentials_values: Iterable[Any]): if not credentials_values: return From e4f7707b66e9f7b223a091d70c085c2a00b2fd40 Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Mon, 15 Aug 2022 19:17:37 +0000 Subject: [PATCH 2/2] Agent: Return credentials when credentials propagation fails --- ...ggregating_propagation_credentials_repository.py | 5 ++--- ...ggregating_propagation_credentials_repository.py | 13 +++++++++++++ 2 files changed, 15 insertions(+), 3 deletions(-) diff --git a/monkey/infection_monkey/credential_repository/aggregating_propagation_credentials_repository.py b/monkey/infection_monkey/credential_repository/aggregating_propagation_credentials_repository.py index 59ff279e0..ab380fdef 100644 --- a/monkey/infection_monkey/credential_repository/aggregating_propagation_credentials_repository.py +++ b/monkey/infection_monkey/credential_repository/aggregating_propagation_credentials_repository.py @@ -63,12 +63,11 @@ class AggregatingPropagationCredentialsRepository(IPropagationCredentialsReposit try: propagation_credentials = self._get_credentials_from_control_channel() self.add_credentials(propagation_credentials) - - return self._stored_credentials except Exception as ex: - self._stored_credentials = {} logger.error(f"Error while attempting to retrieve credentials for propagation: {ex}") + return self._stored_credentials + def _set_attribute(self, attribute_to_be_set: str, credentials_values: Iterable[Any]): if not credentials_values: return diff --git a/monkey/tests/unit_tests/infection_monkey/credential_store/test_aggregating_propagation_credentials_repository.py b/monkey/tests/unit_tests/infection_monkey/credential_store/test_aggregating_propagation_credentials_repository.py index f2cc0dcfa..c992aaf85 100644 --- a/monkey/tests/unit_tests/infection_monkey/credential_store/test_aggregating_propagation_credentials_repository.py +++ b/monkey/tests/unit_tests/infection_monkey/credential_store/test_aggregating_propagation_credentials_repository.py @@ -122,3 +122,16 @@ def test_all_keys_if_credentials_empty(): assert "exploit_password_list" in actual_stored_credentials assert "exploit_ntlm_hash_list" in actual_stored_credentials assert "exploit_ssh_keys" in actual_stored_credentials + + +def test_credentials_obtained_if_propagation_credentials_fails(): + control_channel = MagicMock() + control_channel.get_credentials_for_propagation.return_value = EMPTY_CHANNEL_CREDENTIALS + control_channel.get_credentials_for_propagation.side_effect = Exception( + "No credentials for you!" + ) + credentials_repository = AggregatingPropagationCredentialsRepository(control_channel) + + credentials = credentials_repository.get_credentials() + + assert credentials is not None