Merge pull request #2016 from guardicore/1996-agent-worm-config-decouple

1996 agent worm config decouple

Author: mssalvatore

monkey/infection_monkey/config.py

@@ -8,7 +8,7 @@
     "exploit_user_list",
     "exploit_ssh_keys",
 ]
-LOCAL_CONFIG_VARS = ["name", "id", "current_server", "max_depth"]
+LOCAL_CONFIG_VARS = ["name", "id", "max_depth"]
 HIDDEN_FIELD_REPLACEMENT_CONTENT = "hidden"
 
 
@@ -62,10 +62,6 @@ def as_dict(self):
     # depth of propagation
     depth = 2
     max_depth = None
-    current_server = ""
-
-    # Configuration servers to try to connect to, in this order.
-    command_servers = []
 
     keep_tunnel_open_time = 30
 

monkey/infection_monkey/control.py

@@ -3,6 +3,7 @@
 import platform
 from pprint import pformat
 from socket import gethostname
+from typing import Mapping, Optional
 
 import requests
 from requests.exceptions import ConnectionError
@@ -23,11 +24,17 @@
 PBA_FILE_DOWNLOAD = "https://%s/api/pba/download/%s"
 
 
-class ControlClient(object):
-    proxies = {}
+class ControlClient:
+    # TODO When we have mechanism that support telemetry messenger
+    # with control clients, then this needs to be removed
+    # https://github.com/guardicore/monkey/blob/133f7f5da131b481561141171827d1f9943f6aec/monkey/infection_monkey/telemetry/base_telem.py
+    control_client_object = None
 
-    @staticmethod
-    def wakeup(parent=None):
+    def __init__(self, server_address: str, proxies: Optional[Mapping[str, str]] = None):
+        self.proxies = {} if not proxies else proxies
+        self.server_address = server_address
+
+    def wakeup(self, parent=None):
         if parent:
             logger.debug("parent: %s" % (parent,))
 
@@ -45,67 +52,51 @@ def wakeup(parent=None):
             "launch_time": agent_process.get_start_time(),
         }
 
-        if ControlClient.proxies:
-            monkey["tunnel"] = ControlClient.proxies.get("https")
+        if self.proxies:
+            monkey["tunnel"] = self.proxies.get("https")
 
         requests.post(  # noqa: DUO123
-            f"https://{WormConfiguration.current_server}/api/agent",
+            f"https://{self.server_address}/api/agent",
             data=json.dumps(monkey),
             headers={"content-type": "application/json"},
             verify=False,
-            proxies=ControlClient.proxies,
+            proxies=self.proxies,
             timeout=MEDIUM_REQUEST_TIMEOUT,
         )
 
-    @staticmethod
-    def find_server(default_tunnel=None):
-        logger.debug(
-            "Trying to wake up with Monkey Island servers list: %r"
-            % WormConfiguration.command_servers
-        )
+    def find_server(self, default_tunnel=None):
+        logger.debug(f"Trying to wake up with Monkey Island server: {self.server_address}")
         if default_tunnel:
             logger.debug("default_tunnel: %s" % (default_tunnel,))
 
-        current_server = ""
-
-        for server in WormConfiguration.command_servers:
-            try:
-                current_server = server
-
-                debug_message = "Trying to connect to server: %s" % server
-                if ControlClient.proxies:
-                    debug_message += " through proxies: %s" % ControlClient.proxies
-                logger.debug(debug_message)
-                requests.get(  # noqa: DUO123
-                    f"https://{server}/api?action=is-up",
-                    verify=False,
-                    proxies=ControlClient.proxies,
-                    timeout=MEDIUM_REQUEST_TIMEOUT,
-                )
-                WormConfiguration.current_server = current_server
-                break
-
-            except ConnectionError as exc:
-                current_server = ""
-                logger.warning("Error connecting to control server %s: %s", server, exc)
-
-        if current_server:
+        try:
+            debug_message = "Trying to connect to server: %s" % self.server_address
+            if self.proxies:
+                debug_message += " through proxies: %s" % self.proxies
+            logger.debug(debug_message)
+            requests.get(  # noqa: DUO123
+                f"https://{self.server_address}/api?action=is-up",
+                verify=False,
+                proxies=self.proxies,
+                timeout=MEDIUM_REQUEST_TIMEOUT,
+            )
             return True
+        except ConnectionError as exc:
+            logger.warning("Error connecting to control server %s: %s", self.server_address, exc)
+
+        if self.proxies:
+            return False
         else:
-            if ControlClient.proxies:
-                return False
+            logger.info("Starting tunnel lookup...")
+            proxy_find = tunnel.find_tunnel(default=default_tunnel)
+            if proxy_find:
+                self.set_proxies(proxy_find)
+                return self.find_server()
             else:
-                logger.info("Starting tunnel lookup...")
-                proxy_find = tunnel.find_tunnel(default=default_tunnel)
-                if proxy_find:
-                    ControlClient.set_proxies(proxy_find)
-                    return ControlClient.find_server()
-                else:
-                    logger.info("No tunnel found")
-                    return False
-
-    @staticmethod
-    def set_proxies(proxy_find):
+                logger.info("No tunnel found")
+                return False
+
+    def set_proxies(self, proxy_find):
         """
         Note: The proxy schema changes between different versions of requests and urllib3,
         which causes the machine to not open a tunnel back.
@@ -120,13 +111,12 @@ def set_proxies(proxy_find):
         proxy_address, proxy_port = proxy_find
         logger.info("Found tunnel at %s:%s" % (proxy_address, proxy_port))
         if is_windows_os():
-            ControlClient.proxies["https"] = f"http://{proxy_address}:{proxy_port}"
+            self.proxies["https"] = f"http://{proxy_address}:{proxy_port}"
         else:
-            ControlClient.proxies["https"] = f"{proxy_address}:{proxy_port}"
+            self.proxies["https"] = f"{proxy_address}:{proxy_port}"
 
-    @staticmethod
-    def send_telemetry(telem_category, json_data: str):
-        if not WormConfiguration.current_server:
+    def send_telemetry(self, telem_category, json_data: str):
+        if not self.server_address:
             logger.error(
                 "Trying to send %s telemetry before current server is established, aborting."
                 % telem_category
@@ -135,53 +125,45 @@ def send_telemetry(telem_category, json_data: str):
         try:
             telemetry = {"monkey_guid": GUID, "telem_category": telem_category, "data": json_data}
             requests.post(  # noqa: DUO123
-                "https://%s/api/telemetry" % (WormConfiguration.current_server,),
+                "https://%s/api/telemetry" % (self.server_address,),
                 data=json.dumps(telemetry),
                 headers={"content-type": "application/json"},
                 verify=False,
-                proxies=ControlClient.proxies,
+                proxies=self.proxies,
                 timeout=MEDIUM_REQUEST_TIMEOUT,
             )
         except Exception as exc:
-            logger.warning(
-                "Error connecting to control server %s: %s", WormConfiguration.current_server, exc
-            )
+            logger.warning(f"Error connecting to control server {self.server_address}: {exc}")
 
-    @staticmethod
-    def send_log(log):
-        if not WormConfiguration.current_server:
+    def send_log(self, log):
+        if not self.server_address:
             return
         try:
             telemetry = {"monkey_guid": GUID, "log": json.dumps(log)}
             requests.post(  # noqa: DUO123
-                "https://%s/api/log" % (WormConfiguration.current_server,),
+                "https://%s/api/log" % (self.server_address,),
                 data=json.dumps(telemetry),
                 headers={"content-type": "application/json"},
                 verify=False,
-                proxies=ControlClient.proxies,
+                proxies=self.proxies,
                 timeout=MEDIUM_REQUEST_TIMEOUT,
             )
         except Exception as exc:
-            logger.warning(
-                "Error connecting to control server %s: %s", WormConfiguration.current_server, exc
-            )
+            logger.warning(f"Error connecting to control server {self.server_address}: {exc}")
 
-    @staticmethod
-    def load_control_config():
-        if not WormConfiguration.current_server:
+    def load_control_config(self):
+        if not self.server_address:
             return
         try:
             reply = requests.get(  # noqa: DUO123
-                f"https://{WormConfiguration.current_server}/api/agent/",
+                f"https://{self.server_address}/api/agent/",
                 verify=False,
-                proxies=ControlClient.proxies,
+                proxies=self.proxies,
                 timeout=MEDIUM_REQUEST_TIMEOUT,
             )
 
         except Exception as exc:
-            logger.warning(
-                "Error connecting to control server %s: %s", WormConfiguration.current_server, exc
-            )
+            logger.warning(f"Error connecting to control server {self.server_address}: {exc}")
             return
 
         try:
@@ -194,18 +176,17 @@ def load_control_config():
             # we don't continue with default conf here because it might be dangerous
             logger.error(
                 "Error parsing JSON reply from control server %s (%s): %s",
-                WormConfiguration.current_server,
+                self.server_address,
                 reply._content,
                 exc,
             )
             raise Exception("Couldn't load from from server's configuration, aborting. %s" % exc)
 
-    @staticmethod
-    def create_control_tunnel():
-        if not WormConfiguration.current_server:
+    def create_control_tunnel(self):
+        if not self.server_address:
             return None
 
-        my_proxy = ControlClient.proxies.get("https", "").replace("https://", "")
+        my_proxy = self.proxies.get("https", "").replace("https://", "")
         if my_proxy:
             proxy_class = TcpProxy
             try:
@@ -224,13 +205,12 @@ def create_control_tunnel():
             target_port=target_port,
         )
 
-    @staticmethod
-    def get_pba_file(filename):
+    def get_pba_file(self, filename):
         try:
             return requests.get(  # noqa: DUO123
-                PBA_FILE_DOWNLOAD % (WormConfiguration.current_server, filename),
+                PBA_FILE_DOWNLOAD % (self.server_address, filename),
                 verify=False,
-                proxies=ControlClient.proxies,
+                proxies=self.proxies,
                 timeout=LONG_REQUEST_TIMEOUT,
             )
         except requests.exceptions.RequestException:

monkey/infection_monkey/master/control_channel.py

@@ -1,10 +1,10 @@
 import json
 import logging
+from typing import Mapping
 
 import requests
 
 from common.common_consts.timeouts import SHORT_REQUEST_TIMEOUT
-from infection_monkey.control import ControlClient
 from infection_monkey.custom_types import PropagationCredentials
 from infection_monkey.i_control_channel import IControlChannel, IslandCommunicationError
 
@@ -14,9 +14,10 @@
 
 
 class ControlChannel(IControlChannel):
-    def __init__(self, server: str, agent_id: str):
+    def __init__(self, server: str, agent_id: str, proxies: Mapping[str, str]):
         self._agent_id = agent_id
         self._control_channel_server = server
+        self._proxies = proxies
 
     def should_agent_stop(self) -> bool:
         if not self._control_channel_server:
@@ -30,7 +31,7 @@ def should_agent_stop(self) -> bool:
             response = requests.get(  # noqa: DUO123
                 url,
                 verify=False,
-                proxies=ControlClient.proxies,
+                proxies=self._proxies,
                 timeout=SHORT_REQUEST_TIMEOUT,
             )
             response.raise_for_status()
@@ -51,7 +52,7 @@ def get_config(self) -> dict:
             response = requests.get(  # noqa: DUO123
                 f"https://{self._control_channel_server}/api/agent",
                 verify=False,
-                proxies=ControlClient.proxies,
+                proxies=self._proxies,
                 timeout=SHORT_REQUEST_TIMEOUT,
             )
             response.raise_for_status()
@@ -74,7 +75,7 @@ def get_credentials_for_propagation(self) -> PropagationCredentials:
             response = requests.get(  # noqa: DUO123
                 propagation_credentials_url,
                 verify=False,
-                proxies=ControlClient.proxies,
+                proxies=self._proxies,
                 timeout=SHORT_REQUEST_TIMEOUT,
             )
             response.raise_for_status()