From d6795492a4ba9059cd43cdd936c2720b1a2e5e17 Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Mon, 19 Sep 2022 19:53:22 +0000 Subject: [PATCH] Agent: Add get_config to IslandAPIClient --- .../http_island_api_client.py | 31 +++++++++++++++++++ .../island_api_client/i_island_api_client.py | 14 +++++++++ .../master/control_channel.py | 24 +++----------- .../master/test_control_channel.py | 21 +++++++++++++ 4 files changed, 71 insertions(+), 19 deletions(-) diff --git a/monkey/infection_monkey/island_api_client/http_island_api_client.py b/monkey/infection_monkey/island_api_client/http_island_api_client.py index 97c5404d0..f92fcc946 100644 --- a/monkey/infection_monkey/island_api_client/http_island_api_client.py +++ b/monkey/infection_monkey/island_api_client/http_island_api_client.py @@ -1,11 +1,13 @@ import functools import json import logging +from pprint import pformat from typing import List, Sequence import requests from common import AgentRegistrationData, OperatingSystem +from common.agent_configuration import AgentConfiguration from common.agent_event_serializers import AgentEventSerializerRegistry, JSONSerializable from common.agent_events import AbstractAgentEvent from common.common_consts.timeouts import ( @@ -167,6 +169,35 @@ class HTTPIslandAPIClient(IIslandAPIClient): except json.JSONDecodeError as e: raise IslandAPIRequestFailedError(e) + def get_config(self, island_server: str) -> AgentConfiguration: + try: + response = requests.get( # noqa: DUO123 + f"https://{island_server}/api/agent-configuration", + verify=False, + timeout=SHORT_REQUEST_TIMEOUT, + ) + response.raise_for_status() + + config_dict = json.loads(response.text) + + logger.debug(f"Received configuration:\n{pformat(config_dict)}") + + return AgentConfiguration(**config_dict) + except ( + requests.exceptions.ConnectionError, + requests.exceptions.TooManyRedirects, + ) as e: + raise IslandAPIConnectionError(e) + except requests.exceptions.Timeout as e: + raise IslandAPITimeoutError(e) + except requests.exceptions.HTTPError as e: + if e.errno >= 500: + raise IslandAPIRequestFailedError(e) + else: + raise IslandAPIRequestError(e) + except json.JSONDecodeError as e: + raise IslandAPIRequestFailedError(e) + def _serialize_events(self, events: Sequence[AbstractAgentEvent]) -> JSONSerializable: serialized_events: List[JSONSerializable] = [] diff --git a/monkey/infection_monkey/island_api_client/i_island_api_client.py b/monkey/infection_monkey/island_api_client/i_island_api_client.py index cbfbbf306..064511a73 100644 --- a/monkey/infection_monkey/island_api_client/i_island_api_client.py +++ b/monkey/infection_monkey/island_api_client/i_island_api_client.py @@ -5,6 +5,7 @@ from common import OperatingSystem from common.agent_events import AbstractAgentEvent from common import AgentRegistrationData +from common.agent_configuration import AgentConfiguration class IIslandAPIClient(ABC): @@ -118,3 +119,16 @@ class IIslandAPIClient(ABC): :raises IslandAPITimeoutError: If the command timed out :return: True if the agent should stop, otherwise False """ + + @abstractmethod + def get_config(self, island_server: str) -> AgentConfiguration: + """ + Get agent configuration from the island + + :param island_server: The server to query + :raises IslandAPIConnectionError: If the client could not connect to the island + :raises IslandAPIRequestError: If there was a problem with the client request + :raises IslandAPIRequestFailedError: If the server experienced an error + :raises IslandAPITimeoutError: If the command timed out + :return: Agent configuration + """ diff --git a/monkey/infection_monkey/master/control_channel.py b/monkey/infection_monkey/master/control_channel.py index f61091e8f..d01a6a233 100644 --- a/monkey/infection_monkey/master/control_channel.py +++ b/monkey/infection_monkey/master/control_channel.py @@ -1,6 +1,4 @@ -import json import logging -from pprint import pformat from typing import Optional, Sequence from uuid import UUID @@ -68,24 +66,12 @@ class ControlChannel(IControlChannel): def get_config(self) -> AgentConfiguration: try: - response = requests.get( # noqa: DUO123 - f"https://{self._control_channel_server}/api/agent-configuration", - verify=False, - timeout=SHORT_REQUEST_TIMEOUT, - ) - response.raise_for_status() - - config_dict = json.loads(response.text) - - logger.debug(f"Received configuration:\n{pformat(config_dict)}") - - return AgentConfiguration(**config_dict) + return self._island_api_client.get_config(self._control_channel_server) except ( - json.JSONDecodeError, - requests.exceptions.ConnectionError, - requests.exceptions.Timeout, - requests.exceptions.TooManyRedirects, - requests.exceptions.HTTPError, + IslandAPIConnectionError, + IslandAPIRequestError, + IslandAPIRequestFailedError, + IslandAPITimeoutError, ) as e: raise IslandCommunicationError(e) diff --git a/monkey/tests/unit_tests/infection_monkey/master/test_control_channel.py b/monkey/tests/unit_tests/infection_monkey/master/test_control_channel.py index 534ef157b..dbb5ffb75 100644 --- a/monkey/tests/unit_tests/infection_monkey/master/test_control_channel.py +++ b/monkey/tests/unit_tests/infection_monkey/master/test_control_channel.py @@ -89,3 +89,24 @@ def test_control_channel__should_agent_stop_raises_on_request_failed_error( with pytest.raises(IslandCommunicationError): control_channel.should_agent_stop() + + +def test_control_channel__get_config(control_channel, island_api_client): + control_channel.get_config() + assert island_api_client.get_config.called_once() + + +@pytest.mark.parametrize( + "api_error", + [ + IslandAPIConnectionError, + IslandAPIRequestError, + IslandAPIRequestFailedError, + IslandAPITimeoutError, + ], +) +def test_control_channel__get_config_raises_error(control_channel, island_api_client, api_error): + island_api_client.get_config.side_effect = api_error() + + with pytest.raises(IslandCommunicationError): + control_channel.get_config()