From 8ec580e19cc2a95c25d11bc2d14d2f9f938d9bfd Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Wed, 15 Dec 2021 15:43:38 +0100 Subject: [PATCH] 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 | 4 ---- monkey/infection_monkey/i_control_channel.py | 4 ++++ .../master/automated_master.py | 7 ++----- .../infection_monkey/master/control_channel.py | 18 +++++++++++------- 4 files changed, 17 insertions(+), 16 deletions(-) diff --git a/monkey/common/utils/exceptions.py b/monkey/common/utils/exceptions.py index 89caae27a..cc70cbc51 100644 --- a/monkey/common/utils/exceptions.py +++ b/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""" diff --git a/monkey/infection_monkey/i_control_channel.py b/monkey/infection_monkey/i_control_channel.py index eb1a4d5b2..33539417c 100644 --- a/monkey/infection_monkey/i_control_channel.py +++ b/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""" diff --git a/monkey/infection_monkey/master/automated_master.py b/monkey/infection_monkey/master/automated_master.py index b3a064a19..8342b4b7c 100644 --- a/monkey/infection_monkey/master/automated_master.py +++ b/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") diff --git a/monkey/infection_monkey/master/control_channel.py b/monkey/infection_monkey/master/control_channel.py index 8f8f30406..52b565d55 100644 --- a/monkey/infection_monkey/master/control_channel.py +++ b/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)