From 022630ddc119566f062d5af877de16f86843f1cf Mon Sep 17 00:00:00 2001
From: Kekoa Kaaikala <kekoa.kaaikala@gmail.com>
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