From f299e61b208f43798fa79a7dffad0428c462c03d Mon Sep 17 00:00:00 2001
From: Ilija Lazoroski <ilija.la@live.com>
Date: Wed, 15 Dec 2021 12:07:13 +0100
Subject: [PATCH] Agent: Handle ControlClient exceptions in AutomatedMaster

---
 .../master/automated_master.py                | 19 +++++++++++++++----
 .../master/control_channel.py                 | 10 ++--------
 2 files changed, 17 insertions(+), 12 deletions(-)

diff --git a/monkey/infection_monkey/master/automated_master.py b/monkey/infection_monkey/master/automated_master.py
index 75a5a5dbf..cf7e66102 100644
--- a/monkey/infection_monkey/master/automated_master.py
+++ b/monkey/infection_monkey/master/automated_master.py
@@ -93,15 +93,26 @@ class AutomatedMaster(IMaster):
             time.sleep(CHECK_FOR_TERMINATE_INTERVAL_SEC)
 
     def _check_for_stop(self):
-        if self._control_channel.should_agent_stop():
-            logger.debug('Received the "stop" signal from the Island')
-            self._stop.set()
+        try:
+            if self._control_channel.should_agent_stop():
+                logger.debug('Received the "stop" signal from the Island')
+                self._stop.set()
+        except Exception as e:
+            self._failed_stop += 1
+            if self._failed_stop > 5:
+                logger.error(f"An error occurred while trying to check for agent stop: {e}")
+                self._stop.set()
 
     def _master_thread_should_run(self):
         return (not self._stop.is_set()) and self._simulation_thread.is_alive()
 
     def _run_simulation(self):
-        config = self._control_channel.get_config()["config"]
+
+        try:
+            config = self._control_channel.get_config()["config"]
+        except Exception as e:
+            logger.error(f"An error occurred while fetching configuration: {e}")
+            return
 
         system_info_collector_thread = create_daemon_thread(
             target=self._run_plugins,
diff --git a/monkey/infection_monkey/master/control_channel.py b/monkey/infection_monkey/master/control_channel.py
index 5fdd03942..4ad871997 100644
--- a/monkey/infection_monkey/master/control_channel.py
+++ b/monkey/infection_monkey/master/control_channel.py
@@ -37,10 +37,7 @@ class ControlChannel(IControlChannel):
             response = json.loads(response.content.decode())
             return response["stop_agent"]
         except Exception as e:
-            # TODO: Evaluate how this exception is handled; don't just log and ignore it.
-            logger.error(f"An error occurred while trying to connect to server. {e}")
-
-        return True
+            raise Exception(f"An error occurred while trying to connect to server. {e}")
 
     def get_config(self) -> dict:
         try:
@@ -53,13 +50,10 @@ class ControlChannel(IControlChannel):
 
             return json.loads(response.content.decode())
         except Exception as exc:
-            # TODO: Evaluate how this exception is handled; don't just log and ignore it.
-            logger.warning(
+            raise Exception(
                 "Error connecting to control server %s: %s", WormConfiguration.current_server, exc
             )
 
-        return {}
-
     def get_credentials_for_propagation(self) -> dict:
         try:
             response = requests.get(  # noqa: DUO123