Agent: Update credentials repository to cache per-instance

This commit is contained in:
Kekoa Kaaikala 2022-08-15 19:14:42 +00:00
parent 96f794e192
commit 9e6a569393
1 changed files with 6 additions and 5 deletions

View File

@ -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