diff --git a/monkey/infection_monkey/i_control_channel.py b/monkey/infection_monkey/i_control_channel.py index b90308018..03a5bb4a9 100644 --- a/monkey/infection_monkey/i_control_channel.py +++ b/monkey/infection_monkey/i_control_channel.py @@ -1,6 +1,8 @@ import abc +from typing import Sequence from common.configuration import AgentConfiguration +from common.credentials import Credentials class IControlChannel(metaclass=abc.ABCMeta): @@ -21,10 +23,11 @@ class IControlChannel(metaclass=abc.ABCMeta): pass @abc.abstractmethod - def get_credentials_for_propagation(self) -> dict: + def get_credentials_for_propagation(self) -> Sequence[Credentials]: """ - :return: A dictionary containing propagation credentials data - :rtype: dict + Get credentials to use during propagation + + :return: A Sequence containing propagation credentials data """ pass diff --git a/monkey/infection_monkey/master/control_channel.py b/monkey/infection_monkey/master/control_channel.py index 6daf107fc..cdf660842 100644 --- a/monkey/infection_monkey/master/control_channel.py +++ b/monkey/infection_monkey/master/control_channel.py @@ -1,13 +1,13 @@ import json import logging from pprint import pformat -from typing import Mapping +from typing import Mapping, Sequence import requests from common.common_consts.timeouts import SHORT_REQUEST_TIMEOUT from common.configuration import AgentConfiguration -from infection_monkey.custom_types import PropagationCredentials +from common.credentials import Credentials from infection_monkey.i_control_channel import IControlChannel, IslandCommunicationError requests.packages.urllib3.disable_warnings() @@ -71,7 +71,7 @@ class ControlChannel(IControlChannel): ) as e: raise IslandCommunicationError(e) - def get_credentials_for_propagation(self) -> PropagationCredentials: + def get_credentials_for_propagation(self) -> Sequence[Credentials]: propagation_credentials_url = ( f"https://{self._control_channel_server}/api/propagation-credentials" ) @@ -84,9 +84,9 @@ class ControlChannel(IControlChannel): ) response.raise_for_status() - return json.loads(response.content.decode())["propagation_credentials"] + return [Credentials.from_mapping(credentials) for credentials in response.json] except ( - json.JSONDecodeError, + requests.exceptions.JSONDecodeError, requests.exceptions.ConnectionError, requests.exceptions.Timeout, requests.exceptions.TooManyRedirects,