Agent: Add get_config to IslandAPIClient

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] = []
 

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
+        """

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)
 

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()