From 56f07e01882fcb45677b5e26bd538765ac25f12d Mon Sep 17 00:00:00 2001
From: Ilija Lazoroski <ilija.la@live.com>
Date: Tue, 23 Nov 2021 11:13:14 +0100
Subject: [PATCH] Agent: Add control channel server property

---
 monkey/infection_monkey/control_channel.py   | 15 ++++++++-------
 monkey/infection_monkey/i_control_channel.py |  7 +++++++
 2 files changed, 15 insertions(+), 7 deletions(-)

diff --git a/monkey/infection_monkey/control_channel.py b/monkey/infection_monkey/control_channel.py
index d8ec17c6e..90076cbdb 100644
--- a/monkey/infection_monkey/control_channel.py
+++ b/monkey/infection_monkey/control_channel.py
@@ -14,14 +14,15 @@ logger = logging.getLogger(__name__)
 
 
 class ControlChannel(IControlChannel):
+    control_channel_server = WormConfiguration.current_server
+
     def should_agent_stop(self) -> bool:
-        server = WormConfiguration.current_server
-        if not server:
+        if not self.control_channel_server:
             return
 
         try:
             response = requests.get(  # noqa: DUO123
-                f"{server}/api/monkey_control/{GUID}",
+                f"{self.control_channel_server}/api/monkey_control/{GUID}",
                 verify=False,
                 timeout=SHORT_REQUEST_TIMEOUT,
             )
@@ -32,16 +33,16 @@ class ControlChannel(IControlChannel):
             logger.error(f"An error occurred while trying to connect to server. {e}")
 
     def get_config(self) -> dict:
-        return ControlClient.load_control_config()
+        ControlClient.load_control_config()
+        return WormConfiguration.as_dict()
 
     def get_credentials_for_propagation(self) -> dict:
-        server = WormConfiguration.current_server
-        if not server:
+        if not self.control_channel_server:
             return
 
         try:
             response = requests.get(  # noqa: DUO123
-                f"{server}/api/propagationCredentials",
+                f"{self.control_channel_server}/api/propagationCredentials",
                 verify=False,
                 timeout=SHORT_REQUEST_TIMEOUT,
             )
diff --git a/monkey/infection_monkey/i_control_channel.py b/monkey/infection_monkey/i_control_channel.py
index eb1a4d5b2..6f287671d 100644
--- a/monkey/infection_monkey/i_control_channel.py
+++ b/monkey/infection_monkey/i_control_channel.py
@@ -2,6 +2,13 @@ import abc
 
 
 class IControlChannel(metaclass=abc.ABCMeta):
+    @property
+    @abc.abstractmethod
+    def control_channel_server(self):
+        """
+        :return: Worm configuration server
+        """
+
     @abc.abstractmethod
     def should_agent_stop(self) -> bool:
         """