diff --git a/monkey/infection_monkey/i_control_channel.py b/monkey/infection_monkey/i_control_channel.py index 33539417c..b90308018 100644 --- a/monkey/infection_monkey/i_control_channel.py +++ b/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 diff --git a/monkey/infection_monkey/master/control_channel.py b/monkey/infection_monkey/master/control_channel.py index c93af43e0..8567a301f 100644 --- a/monkey/infection_monkey/master/control_channel.py +++ b/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, diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index 466de6664..afcf460b2 100644 --- a/monkey/infection_monkey/monkey.py +++ b/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