Agent: Implement should retry task in automated master

Add handling of known requests exceptions in ControlClient.
Moved IslandCommunicationError to IControlChannel

monkey/common/utils/exceptions.py

@@ -56,7 +56,3 @@ class DomainControllerNameFetchError(FailedExploitationError):
 
 class InvalidConfigurationError(Exception):
     """ Raise when configuration is invalid """
-
-
-class IslandCommunicationError(Exception):
-    """Raise when unable to connect to control client"""

monkey/infection_monkey/i_control_channel.py

@@ -25,3 +25,7 @@ class IControlChannel(metaclass=abc.ABCMeta):
         :rtype: dict
         """
         pass
+
+
+class IslandCommunicationError(Exception):
+    """Raise when unable to connect to control client"""

monkey/infection_monkey/master/automated_master.py

@@ -3,8 +3,7 @@ import threading
 import time
 from typing import Any, Callable, Dict, List, Tuple
 
-from common.utils.exceptions import IslandCommunicationError
-from infection_monkey.i_control_channel import IControlChannel
+from infection_monkey.i_control_channel import IControlChannel, IslandCommunicationError
 from infection_monkey.i_master import IMaster
 from infection_monkey.i_puppet import IPuppet
 from infection_monkey.model import VictimHostFactory
@@ -22,7 +21,7 @@ SHUTDOWN_TIMEOUT = 5
 NUM_SCAN_THREADS = 16  # TODO: Adjust this to the optimal number of scan threads
 NUM_EXPLOIT_THREADS = 4  # TODO: Adjust this to the optimal number of exploit threads
 CHECK_FOR_STOP_AGENT_COUNT = 5
-CHECK_FOR_CONFIG_COUNT = 1
+CHECK_FOR_CONFIG_COUNT = 3
 
 logger = logging.getLogger()
 
@@ -48,8 +47,6 @@ class AutomatedMaster(IMaster):
         self._stop = threading.Event()
         self._master_thread = create_daemon_thread(target=self._run_master_thread)
         self._simulation_thread = create_daemon_thread(target=self._run_simulation)
-        self._failed_stop = 0
-        self._failed_config = 0
 
     def start(self):
         logger.info("Starting automated breach and attack simulation")

monkey/infection_monkey/master/control_channel.py

@@ -4,10 +4,9 @@ import logging
 import requests
 
 from common.common_consts.timeouts import SHORT_REQUEST_TIMEOUT
-from common.utils.exceptions import IslandCommunicationError
 from infection_monkey.config import WormConfiguration
 from infection_monkey.control import ControlClient
-from infection_monkey.i_control_channel import IControlChannel
+from infection_monkey.i_control_channel import IControlChannel, IslandCommunicationError
 
 requests.packages.urllib3.disable_warnings()
 
@@ -34,6 +33,7 @@ class ControlChannel(IControlChannel):
                 proxies=ControlClient.proxies,
                 timeout=SHORT_REQUEST_TIMEOUT,
             )
+            response.raise_for_status()
 
             response = json.loads(response.content.decode())
             return response["stop_agent"]
@@ -74,11 +74,15 @@ class ControlChannel(IControlChannel):
                 proxies=ControlClient.proxies,
                 timeout=SHORT_REQUEST_TIMEOUT,
             )
+            response.raise_for_status()
 
             response = json.loads(response.content.decode())["propagation_credentials"]
             return response
-        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 {}
+        except (
+            json.JSONDecodeError,
+            requests.exceptions.ConnectionError,
+            requests.exceptions.Timeout,
+            requests.exceptions.TooManyRedirects,
+            requests.exceptions.HTTPError,
+        ) as e:
+            raise IslandCommunicationError(e)