Agent: Use new credentials format in ControlChannel

This commit is contained in:
Mike Salvatore 2022-07-18 09:55:08 -04:00
parent 19a7bfd8e6
commit ef4fbb30cc
2 changed files with 11 additions and 8 deletions

View File

@ -1,6 +1,8 @@
import abc import abc
from typing import Sequence
from common.configuration import AgentConfiguration from common.configuration import AgentConfiguration
from common.credentials import Credentials
class IControlChannel(metaclass=abc.ABCMeta): class IControlChannel(metaclass=abc.ABCMeta):
@ -21,10 +23,11 @@ class IControlChannel(metaclass=abc.ABCMeta):
pass pass
@abc.abstractmethod @abc.abstractmethod
def get_credentials_for_propagation(self) -> dict: def get_credentials_for_propagation(self) -> Sequence[Credentials]:
""" """
:return: A dictionary containing propagation credentials data Get credentials to use during propagation
:rtype: dict
:return: A Sequence containing propagation credentials data
""" """
pass pass

View File

@ -1,13 +1,13 @@
import json import json
import logging import logging
from pprint import pformat from pprint import pformat
from typing import Mapping from typing import Mapping, Sequence
import requests import requests
from common.common_consts.timeouts import SHORT_REQUEST_TIMEOUT from common.common_consts.timeouts import SHORT_REQUEST_TIMEOUT
from common.configuration import AgentConfiguration 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 from infection_monkey.i_control_channel import IControlChannel, IslandCommunicationError
requests.packages.urllib3.disable_warnings() requests.packages.urllib3.disable_warnings()
@ -71,7 +71,7 @@ class ControlChannel(IControlChannel):
) as e: ) as e:
raise IslandCommunicationError(e) raise IslandCommunicationError(e)
def get_credentials_for_propagation(self) -> PropagationCredentials: def get_credentials_for_propagation(self) -> Sequence[Credentials]:
propagation_credentials_url = ( propagation_credentials_url = (
f"https://{self._control_channel_server}/api/propagation-credentials" f"https://{self._control_channel_server}/api/propagation-credentials"
) )
@ -84,9 +84,9 @@ class ControlChannel(IControlChannel):
) )
response.raise_for_status() response.raise_for_status()
return json.loads(response.content.decode())["propagation_credentials"] return [Credentials.from_mapping(credentials) for credentials in response.json]
except ( except (
json.JSONDecodeError, requests.exceptions.JSONDecodeError,
requests.exceptions.ConnectionError, requests.exceptions.ConnectionError,
requests.exceptions.Timeout, requests.exceptions.Timeout,
requests.exceptions.TooManyRedirects, requests.exceptions.TooManyRedirects,