forked from p15670423/monkey
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.
This commit is contained in:
parent
b49d9d9b9a
commit
763cf578c7
|
@ -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)
|
||||
|
|
|
@ -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}"
|
||||
)
|
||||
|
|
Loading…
Reference in New Issue