From f8855d290d466131b181fe6ac98f86110e8ca025 Mon Sep 17 00:00:00 2001
From: vakarisz <vakarisz@yahoo.com>
Date: Fri, 17 Jun 2022 16:38:52 +0300
Subject: [PATCH] Agent: Remove WormConfiguration and cleanup related
 infrastructure

---
 monkey/infection_monkey/config.py  |  3 ---
 monkey/infection_monkey/control.py | 38 +++---------------------------
 monkey/infection_monkey/main.py    |  5 ----
 monkey/infection_monkey/monkey.py  |  9 +++++--
 4 files changed, 10 insertions(+), 45 deletions(-)

diff --git a/monkey/infection_monkey/config.py b/monkey/infection_monkey/config.py
index 8e783dbf5..1b7b7f60f 100644
--- a/monkey/infection_monkey/config.py
+++ b/monkey/infection_monkey/config.py
@@ -64,6 +64,3 @@ class Configuration(object):
     max_depth = None
 
     keep_tunnel_open_time = 30
-
-
-WormConfiguration = Configuration()
diff --git a/monkey/infection_monkey/control.py b/monkey/infection_monkey/control.py
index 985778b7c..8d1e48a22 100644
--- a/monkey/infection_monkey/control.py
+++ b/monkey/infection_monkey/control.py
@@ -1,7 +1,6 @@
 import json
 import logging
 import platform
-from pprint import pformat
 from socket import gethostname
 from typing import Mapping, Optional
 
@@ -10,7 +9,7 @@ from requests.exceptions import ConnectionError
 
 import infection_monkey.tunnel as tunnel
 from common.common_consts.timeouts import LONG_REQUEST_TIMEOUT, MEDIUM_REQUEST_TIMEOUT
-from infection_monkey.config import GUID, WormConfiguration
+from infection_monkey.config import GUID
 from infection_monkey.network.info import get_host_subnets, local_ips
 from infection_monkey.transport.http import HTTPConnectProxy
 from infection_monkey.transport.tcp import TcpProxy
@@ -151,38 +150,7 @@ class ControlClient:
         except Exception as exc:
             logger.warning(f"Error connecting to control server {self.server_address}: {exc}")
 
-    def load_control_config(self):
-        if not self.server_address:
-            return
-        try:
-            reply = requests.get(  # noqa: DUO123
-                f"https://{self.server_address}/api/agent/",
-                verify=False,
-                proxies=self.proxies,
-                timeout=MEDIUM_REQUEST_TIMEOUT,
-            )
-
-        except Exception as exc:
-            logger.warning(f"Error connecting to control server {self.server_address}: {exc}")
-            return
-
-        try:
-            WormConfiguration.from_kv(reply.json().get("config"))
-            formatted_config = pformat(
-                WormConfiguration.hide_sensitive_info(WormConfiguration.as_dict())
-            )
-            logger.info(f"New configuration was loaded from server:\n{formatted_config}")
-        except Exception as exc:
-            # 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",
-                self.server_address,
-                reply._content,
-                exc,
-            )
-            raise Exception("Couldn't load from from server's configuration, aborting. %s" % exc)
-
-    def create_control_tunnel(self):
+    def create_control_tunnel(self, keep_tunnel_open_time: int):
         if not self.server_address:
             return None
 
@@ -200,7 +168,7 @@ class ControlClient:
 
         return tunnel.MonkeyTunnel(
             proxy_class,
-            keep_tunnel_open_time=WormConfiguration.keep_tunnel_open_time,
+            keep_tunnel_open_time=keep_tunnel_open_time,
             target_addr=target_addr,
             target_port=target_port,
         )
diff --git a/monkey/infection_monkey/main.py b/monkey/infection_monkey/main.py
index 3c71c9721..0102503ca 100644
--- a/monkey/infection_monkey/main.py
+++ b/monkey/infection_monkey/main.py
@@ -5,13 +5,11 @@ import os
 import sys
 import traceback
 from multiprocessing import freeze_support
-from pprint import pformat
 
 # dummy import for pyinstaller
 # noinspection PyUnresolvedReferences
 import infection_monkey.post_breach  # noqa: F401
 from common.version import get_version
-from infection_monkey.config import WormConfiguration
 from infection_monkey.dropper import MonkeyDrops
 from infection_monkey.model import DROPPER_ARG, MONKEY_ARG
 from infection_monkey.monkey import InfectionMonkey
@@ -57,9 +55,6 @@ def main():
     mode_args, mode_specific_args = arg_parser.parse_known_args()
     mode = mode_args.mode
 
-    formatted_config = pformat(WormConfiguration.hide_sensitive_info(WormConfiguration.as_dict()))
-    print(f"Loaded Configuration:\n{formatted_config}")
-
     try:
         if MONKEY_ARG == mode:
             log_path = get_agent_log_path()
diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py
index d7a051193..a0fe5f009 100644
--- a/monkey/infection_monkey/monkey.py
+++ b/monkey/infection_monkey/monkey.py
@@ -149,7 +149,6 @@ class InfectionMonkey:
             raise Exception(f"Monkey couldn't find server with {self._opts.tunnel} default tunnel.")
 
         self._control_client.wakeup(parent=self._opts.parent)
-        self._control_client.load_control_config()
 
     def _current_server_is_set(self) -> bool:
         if self._control_client.find_server(default_tunnel=self._opts.tunnel):
@@ -165,7 +164,13 @@ class InfectionMonkey:
         if firewall.is_enabled():
             firewall.add_firewall_rule()
 
-        self._monkey_inbound_tunnel = self._control_client.create_control_tunnel()
+        control_channel = ControlChannel(
+            self._control_client.server_address, GUID, self._control_client.proxies
+        )
+        keep_tunnel_open_time = control_channel.get_config()["config"]["keep_tunnel_open_time"]
+        self._monkey_inbound_tunnel = self._control_client.create_control_tunnel(
+            keep_tunnel_open_time
+        )
         if self._monkey_inbound_tunnel and self._propagation_enabled():
             self._monkey_inbound_tunnel.start()