Agent: Reduce duplication due to error handling

This commit is contained in:
Kekoa Kaaikala 2022-09-19 20:19:29 +00:00
parent b260dcc5cb
commit 42633c066f
1 changed files with 51 additions and 88 deletions

View File

@ -35,7 +35,7 @@ def handle_island_errors(fn):
def decorated(*args, **kwargs):
try:
return fn(*args, **kwargs)
except requests.exceptions.ConnectionError as err:
except (requests.exceptions.ConnectionError, requests.exceptions.TooManyRedirects) as err:
raise IslandAPIConnectionError(err)
except requests.exceptions.HTTPError as err:
if 400 <= err.response.status_code < 500:
@ -54,6 +54,17 @@ def handle_island_errors(fn):
return decorated
def convert_json_error_to_island_api_error(fn):
@functools.wraps(fn)
def wrapper(*args, **kwargs):
try:
fn(*args, **kwargs)
except json.JSONDecodeError as e:
raise IslandAPIRequestFailedError(e)
return wrapper
class HTTPIslandAPIClient(IIslandAPIClient):
"""
A client for the Island's HTTP API
@ -125,7 +136,6 @@ class HTTPIslandAPIClient(IIslandAPIClient):
response.raise_for_status()
def register_agent(self, agent_registration_data: AgentRegistrationData):
try:
url = f"https://{agent_registration_data.cc_server}/api/agents"
response = requests.post( # noqa: DUO123
url,
@ -134,17 +144,10 @@ class HTTPIslandAPIClient(IIslandAPIClient):
timeout=SHORT_REQUEST_TIMEOUT,
)
response.raise_for_status()
except (
requests.exceptions.ConnectionError,
requests.exceptions.TooManyRedirects,
requests.exceptions.HTTPError,
) as e:
raise IslandAPIConnectionError(e)
except requests.exceptions.Timeout as e:
raise IslandAPITimeoutError(e)
@handle_island_errors
@convert_json_error_to_island_api_error
def should_agent_stop(self, island_server: str, agent_id: str) -> bool:
try:
url = f"https://{island_server}/api/monkey-control" f"/needs-to-stop/{agent_id}"
response = requests.get( # noqa: DUO123
url,
@ -155,23 +158,10 @@ class HTTPIslandAPIClient(IIslandAPIClient):
json_response = json.loads(response.content.decode())
return json_response["stop_agent"]
except (
requests.exceptions.ConnectionError,
requests.exceptions.TooManyRedirects,
) as e:
raise IslandAPIConnectionError(e)
except requests.exceptions.Timeout as e:
raise IslandAPITimeoutError(e)
except requests.exceptions.HTTPError as e:
if e.errno >= 500:
raise IslandAPIRequestFailedError(e)
else:
raise IslandAPIRequestError(e)
except json.JSONDecodeError as e:
raise IslandAPIRequestFailedError(e)
@handle_island_errors
@convert_json_error_to_island_api_error
def get_config(self, island_server: str) -> AgentConfiguration:
try:
response = requests.get( # noqa: DUO123
f"https://{island_server}/api/agent-configuration",
verify=False,
@ -184,23 +174,10 @@ class HTTPIslandAPIClient(IIslandAPIClient):
logger.debug(f"Received configuration:\n{pformat(config_dict)}")
return AgentConfiguration(**config_dict)
except (
requests.exceptions.ConnectionError,
requests.exceptions.TooManyRedirects,
) as e:
raise IslandAPIConnectionError(e)
except requests.exceptions.Timeout as e:
raise IslandAPITimeoutError(e)
except requests.exceptions.HTTPError as e:
if e.errno >= 500:
raise IslandAPIRequestFailedError(e)
else:
raise IslandAPIRequestError(e)
except json.JSONDecodeError as e:
raise IslandAPIRequestFailedError(e)
@handle_island_errors
@convert_json_error_to_island_api_error
def get_credentials_for_propagation(self, island_server: str) -> Sequence[Credentials]:
try:
response = requests.get( # noqa: DUO123
f"https://{island_server}/api/propagation-credentials",
verify=False,
@ -209,20 +186,6 @@ class HTTPIslandAPIClient(IIslandAPIClient):
response.raise_for_status()
return [Credentials(**credentials) for credentials in response.json()]
except (
requests.exceptions.ConnectionError,
requests.exceptions.TooManyRedirects,
) as e:
raise IslandAPIConnectionError(e)
except requests.exceptions.Timeout as e:
raise IslandAPITimeoutError(e)
except requests.exceptions.HTTPError as e:
if e.errno >= 500:
raise IslandAPIRequestFailedError(e)
else:
raise IslandAPIRequestError(e)
except json.JSONDecodeError as e:
raise IslandAPIRequestFailedError(e)
def _serialize_events(self, events: Sequence[AbstractAgentEvent]) -> JSONSerializable:
serialized_events: List[JSONSerializable] = []