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

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:
Mike Salvatore 2022-03-29 11:52:43 -04:00
parent b49d9d9b9a
commit 763cf578c7
2 changed files with 10 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
monkey/infection_monkey/credential_store/aggregating_credentials_store.py
View File

@ -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_control_channel import IControlChannel
from infection_monkey.i_puppet import Credentials from infection_monkey.i_puppet import Credentials
from infection_monkey.typing import PropagationCredentials from infection_monkey.typing import PropagationCredentials
from infection_monkey.utils.decorators import request_cache
from .i_credentials_store import ICredentialsStore from .i_credentials_store import ICredentialsStore
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
CREDENTIALS_POLL_PERIOD_SEC = 30
class AggregatingCredentialsStore(ICredentialsStore): class AggregatingCredentialsStore(ICredentialsStore):
def __init__(self, control_channel: IControlChannel): def __init__(self, control_channel: IControlChannel):
@ -52,7 +55,7 @@ class AggregatingCredentialsStore(ICredentialsStore):
def get_credentials(self) -> PropagationCredentials: def get_credentials(self) -> PropagationCredentials:
try: 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 # Needs to be reworked when exploiters accepts sequence of Credentials
self._aggregate_credentials(propagation_credentials) self._aggregate_credentials(propagation_credentials)
@ -62,6 +65,10 @@ class AggregatingCredentialsStore(ICredentialsStore):
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) -> PropagationCredentials:
return self._control_channel.get_credentials_for_propagation()
def _aggregate_credentials(self, credentials_to_aggr: Mapping): def _aggregate_credentials(self, credentials_to_aggr: Mapping):
for cred_attr, credentials_values in credentials_to_aggr.items(): for cred_attr, credentials_values in credentials_to_aggr.items():
self._set_attribute(cred_attr, credentials_values) self._set_attribute(cred_attr, credentials_values)

7
monkey/infection_monkey/master/control_channel.py
View File

@ -7,14 +7,12 @@ from common.common_consts.timeouts import SHORT_REQUEST_TIMEOUT
from infection_monkey.config import WormConfiguration from infection_monkey.config import WormConfiguration
from infection_monkey.control import ControlClient from infection_monkey.control import ControlClient
from infection_monkey.i_control_channel import IControlChannel, IslandCommunicationError 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() requests.packages.urllib3.disable_warnings()
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
CREDENTIALS_POLL_PERIOD_SEC = 30
class ControlChannel(IControlChannel): class ControlChannel(IControlChannel):
def __init__(self, server: str, agent_id: str): def __init__(self, server: str, agent_id: str):
@ -69,8 +67,7 @@ class ControlChannel(IControlChannel):
) as e: ) as e:
raise IslandCommunicationError(e) raise IslandCommunicationError(e)
@request_cache(CREDENTIALS_POLL_PERIOD_SEC) def get_credentials_for_propagation(self) -> PropagationCredentials:
def get_credentials_for_propagation(self) -> dict:
propagation_credentials_url = ( propagation_credentials_url = (
f"https://{self._control_channel_server}/api/propagation-credentials/{self._agent_id}" f"https://{self._control_channel_server}/api/propagation-credentials/{self._agent_id}"
) )