p41275903
/
monkey
forked from p15670423/monkey
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

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
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
monkey/infection_monkey/credential_repository/aggregating_propagation_credentials_repository.py
View File

@ -1,5 +1,5 @@
import logging import logging
from typing import Any, Iterable, Sequence from typing import Any, Iterable
from common.credentials import CredentialComponentType, Credentials, ICredentialComponent from common.credentials import CredentialComponentType, Credentials, ICredentialComponent
from infection_monkey.custom_types import PropagationCredentials from infection_monkey.custom_types import PropagationCredentials
@ -29,6 +29,11 @@ class AggregatingPropagationCredentialsRepository(IPropagationCredentialsReposit
} }
self._control_channel = control_channel 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]): def add_credentials(self, credentials_to_add: Iterable[Credentials]):
for credentials in credentials_to_add: for credentials in credentials_to_add:
if credentials.identity: if credentials.identity:
@ -64,10 +69,6 @@ class AggregatingPropagationCredentialsRepository(IPropagationCredentialsReposit
self._stored_credentials = {} self._stored_credentials = {}
logger.error(f"Error while attempting to retrieve credentials for propagation: {ex}") 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]): def _set_attribute(self, attribute_to_be_set: str, credentials_values: Iterable[Any]):
if not credentials_values: if not credentials_values:
return return