From fa9225370e71db8ff2d75b4cabb00f11ee1abcba Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Mon, 19 Sep 2022 18:37:31 +0530 Subject: [PATCH] Agent: Add handle_island_errors() decorator to http_island_api_client.py --- .../http_island_api_client.py | 19 +++++++------------ 1 file changed, 7 insertions(+), 12 deletions(-) diff --git a/monkey/infection_monkey/island_api_client/http_island_api_client.py b/monkey/infection_monkey/island_api_client/http_island_api_client.py index 5523a10a0..abd4fa9d1 100644 --- a/monkey/infection_monkey/island_api_client/http_island_api_client.py +++ b/monkey/infection_monkey/island_api_client/http_island_api_client.py @@ -1,3 +1,4 @@ +import functools import logging import requests @@ -9,19 +10,11 @@ from . import IIslandAPIClient, IslandAPIConnectionError, IslandAPIError, Island logger = logging.getLogger(__name__) -class HTTPIslandAPIClient(IIslandAPIClient): - """ - A client for the Island's HTTP API - """ - - def __init__(self, island_server: str): +def handle_island_errors(fn): + @functools.wraps(fn) + def decorated(*args, **kwargs): try: - requests.get( # noqa: DUO123 - f"https://{island_server}/api?action=is-up", - verify=False, - timeout=MEDIUM_REQUEST_TIMEOUT, - ) - self._island_server = island_server + return fn(*args, **kwargs) except requests.exceptions.ConnectionError as err: raise IslandAPIConnectionError(err) except TimeoutError as err: @@ -29,6 +22,8 @@ class HTTPIslandAPIClient(IIslandAPIClient): except Exception as err: raise IslandAPIError(err) + return decorated + # TODO: set server address as object property when init is called in find_server and pass # object around? won't need to pass island server and create object in every function def send_log(self, data: str):