From 763cf578c74f70bcc24c001e950bdc3db90269de Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 29 Mar 2022 11:52:43 -0400 Subject: [PATCH] Agent: Move credentials request caching to AggregatingCredentialsStore The ControlChannel shouldn't be concerned with caching. It's mission should be to service requests. The caching is more appropriately placed in the AggregatingCredentialsStore. --- .../credential_store/aggregating_credentials_store.py | 9 ++++++++- monkey/infection_monkey/master/control_channel.py | 7 ++----- 2 files changed, 10 insertions(+), 6 deletions(-) diff --git a/monkey/infection_monkey/credential_store/aggregating_credentials_store.py b/monkey/infection_monkey/credential_store/aggregating_credentials_store.py index 61883a49c..696b87bd5 100644 --- a/monkey/infection_monkey/credential_store/aggregating_credentials_store.py +++ b/monkey/infection_monkey/credential_store/aggregating_credentials_store.py @@ -5,11 +5,14 @@ from common.common_consts.credential_component_type import CredentialComponentTy from infection_monkey.i_control_channel import IControlChannel from infection_monkey.i_puppet import Credentials from infection_monkey.typing import PropagationCredentials +from infection_monkey.utils.decorators import request_cache from .i_credentials_store import ICredentialsStore logger = logging.getLogger(__name__) +CREDENTIALS_POLL_PERIOD_SEC = 30 + class AggregatingCredentialsStore(ICredentialsStore): def __init__(self, control_channel: IControlChannel): @@ -52,7 +55,7 @@ class AggregatingCredentialsStore(ICredentialsStore): def get_credentials(self) -> PropagationCredentials: try: - propagation_credentials = self._control_channel.get_credentials_for_propagation() + propagation_credentials = self._get_credentials_from_control_channel() # Needs to be reworked when exploiters accepts sequence of Credentials self._aggregate_credentials(propagation_credentials) @@ -62,6 +65,10 @@ class AggregatingCredentialsStore(ICredentialsStore): 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) -> PropagationCredentials: + return self._control_channel.get_credentials_for_propagation() + def _aggregate_credentials(self, credentials_to_aggr: Mapping): for cred_attr, credentials_values in credentials_to_aggr.items(): self._set_attribute(cred_attr, credentials_values) diff --git a/monkey/infection_monkey/master/control_channel.py b/monkey/infection_monkey/master/control_channel.py index d6781af7f..eb6dd12d5 100644 --- a/monkey/infection_monkey/master/control_channel.py +++ b/monkey/infection_monkey/master/control_channel.py @@ -7,14 +7,12 @@ from common.common_consts.timeouts import SHORT_REQUEST_TIMEOUT from infection_monkey.config import WormConfiguration from infection_monkey.control import ControlClient from infection_monkey.i_control_channel import IControlChannel, IslandCommunicationError -from infection_monkey.utils.decorators import request_cache +from infection_monkey.typing import PropagationCredentials requests.packages.urllib3.disable_warnings() logger = logging.getLogger(__name__) -CREDENTIALS_POLL_PERIOD_SEC = 30 - class ControlChannel(IControlChannel): def __init__(self, server: str, agent_id: str): @@ -69,8 +67,7 @@ class ControlChannel(IControlChannel): ) as e: raise IslandCommunicationError(e) - @request_cache(CREDENTIALS_POLL_PERIOD_SEC) - def get_credentials_for_propagation(self) -> dict: + def get_credentials_for_propagation(self) -> PropagationCredentials: propagation_credentials_url = ( f"https://{self._control_channel_server}/api/propagation-credentials/{self._agent_id}" )