Agent: Change configuration to object in control channel

monkey/infection_monkey/i_control_channel.py

@@ -1,5 +1,7 @@
 import abc
 
+from common.configuration import AgentConfiguration
+
 
 class IControlChannel(metaclass=abc.ABCMeta):
     @abc.abstractmethod
@@ -11,10 +13,10 @@ class IControlChannel(metaclass=abc.ABCMeta):
         """
 
     @abc.abstractmethod
-    def get_config(self) -> dict:
+    def get_config(self) -> AgentConfiguration:
         """
-        :return: A dictionary containing Agent Configuration
-        :rtype: dict
+        :return: An AgentConfiguration object
+        :rtype: AgentConfiguration
         """
         pass
 

monkey/infection_monkey/master/control_channel.py

@@ -5,6 +5,7 @@ from typing import Mapping
 import requests
 
 from common.common_consts.timeouts import SHORT_REQUEST_TIMEOUT
+from common.configuration import AgentConfiguration
 from infection_monkey.custom_types import PropagationCredentials
 from infection_monkey.i_control_channel import IControlChannel, IslandCommunicationError
 
@@ -47,7 +48,7 @@ class ControlChannel(IControlChannel):
         ) as e:
             raise IslandCommunicationError(e)
 
-    def get_config(self) -> dict:
+    def get_config(self) -> AgentConfiguration:
         try:
             response = requests.get(  # noqa: DUO123
                 f"https://{self._control_channel_server}/api/agent",
@@ -57,7 +58,7 @@ class ControlChannel(IControlChannel):
             )
             response.raise_for_status()
 
-            return json.loads(response.content.decode())
+            return AgentConfiguration.from_dict(json.loads(response.text)["config"])
         except (
             json.JSONDecodeError,
             requests.exceptions.ConnectionError,

monkey/infection_monkey/monkey.py

@@ -172,8 +172,9 @@ class InfectionMonkey:
         )
 
         config = control_channel.get_config()
+        keep_tunnel_open_time = control_channel.get_config().keep_tunnel_open_time
         self._monkey_inbound_tunnel = self._control_client.create_control_tunnel(
-            config["config"]["keep_tunnel_open_time"]
+            keep_tunnel_open_time
         )
         if self._monkey_inbound_tunnel and should_propagate(config, self._current_depth):
             self._inbound_tunnel_opened = True