Agent: Simplify error handling in ControlChannel

This commit is contained in:
Kekoa Kaaikala 2022-09-20 18:32:20 +00:00
parent 417d167026
commit 022630ddc1
1 changed files with 3 additions and 14 deletions

View File

@ -7,13 +7,7 @@ from urllib3 import disable_warnings
from common.agent_configuration import AgentConfiguration from common.agent_configuration import AgentConfiguration
from common.credentials import Credentials from common.credentials import Credentials
from infection_monkey.i_control_channel import IControlChannel, IslandCommunicationError from infection_monkey.i_control_channel import IControlChannel, IslandCommunicationError
from infection_monkey.island_api_client import ( from infection_monkey.island_api_client import IIslandAPIClient, IslandAPIError
IIslandAPIClient,
IslandAPIConnectionError,
IslandAPIRequestError,
IslandAPIRequestFailedError,
IslandAPITimeoutError,
)
disable_warnings() # noqa: DUO131 disable_warnings() # noqa: DUO131
@ -26,13 +20,8 @@ def handle_island_api_errors(func):
try: try:
print(args) print(args)
func(*args, **kwargs) func(*args, **kwargs)
except ( except IslandAPIError as err:
IslandAPIConnectionError, raise IslandCommunicationError(err)
IslandAPIRequestError,
IslandAPIRequestFailedError,
IslandAPITimeoutError,
) as e:
raise IslandCommunicationError(e)
return wrapper return wrapper