Agent: Handle ControlClient exceptions in AutomatedMaster

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,

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