From 022630ddc119566f062d5af877de16f86843f1cf Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Tue, 20 Sep 2022 18:32:20 +0000 Subject: [PATCH] Agent: Simplify error handling in ControlChannel --- .../infection_monkey/master/control_channel.py | 17 +++-------------- 1 file changed, 3 insertions(+), 14 deletions(-) diff --git a/monkey/infection_monkey/master/control_channel.py b/monkey/infection_monkey/master/control_channel.py index eddeb8090..3b8d64740 100644 --- a/monkey/infection_monkey/master/control_channel.py +++ b/monkey/infection_monkey/master/control_channel.py @@ -7,13 +7,7 @@ from urllib3 import disable_warnings from common.agent_configuration import AgentConfiguration from common.credentials import Credentials from infection_monkey.i_control_channel import IControlChannel, IslandCommunicationError -from infection_monkey.island_api_client import ( - IIslandAPIClient, - IslandAPIConnectionError, - IslandAPIRequestError, - IslandAPIRequestFailedError, - IslandAPITimeoutError, -) +from infection_monkey.island_api_client import IIslandAPIClient, IslandAPIError disable_warnings() # noqa: DUO131 @@ -26,13 +20,8 @@ def handle_island_api_errors(func): try: print(args) func(*args, **kwargs) - except ( - IslandAPIConnectionError, - IslandAPIRequestError, - IslandAPIRequestFailedError, - IslandAPITimeoutError, - ) as e: - raise IslandCommunicationError(e) + except IslandAPIError as err: + raise IslandCommunicationError(err) return wrapper