Agent: Reduce duplication due to error handling

monkey/infection_monkey/island_api_client/http_island_api_client.py

@@ -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] = []