p59342861
/
monkey
forked from p15670423/monkey
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

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
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

139
monkey/infection_monkey/island_api_client/http_island_api_client.py
View File

@ -35,7 +35,7 @@ def handle_island_errors(fn):
def decorated(*args, **kwargs): def decorated(*args, **kwargs):
try: try:
return fn(*args, **kwargs) return fn(*args, **kwargs)
except requests.exceptions.ConnectionError as err: except (requests.exceptions.ConnectionError, requests.exceptions.TooManyRedirects) as err:
raise IslandAPIConnectionError(err) raise IslandAPIConnectionError(err)
except requests.exceptions.HTTPError as err: except requests.exceptions.HTTPError as err:
if 400 <= err.response.status_code < 500: if 400 <= err.response.status_code < 500:
@ -54,6 +54,17 @@ def handle_island_errors(fn):
return decorated 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): class HTTPIslandAPIClient(IIslandAPIClient):
""" """
A client for the Island's HTTP API A client for the Island's HTTP API
@ -125,104 +136,56 @@ class HTTPIslandAPIClient(IIslandAPIClient):
response.raise_for_status() response.raise_for_status()
def register_agent(self, agent_registration_data: AgentRegistrationData): def register_agent(self, agent_registration_data: AgentRegistrationData):
try: url = f"https://{agent_registration_data.cc_server}/api/agents"
url = f"https://{agent_registration_data.cc_server}/api/agents" response = requests.post( # noqa: DUO123
response = requests.post( # noqa: DUO123 url,
url, json=agent_registration_data.dict(simplify=True),
json=agent_registration_data.dict(simplify=True), verify=False,
verify=False, timeout=SHORT_REQUEST_TIMEOUT,
timeout=SHORT_REQUEST_TIMEOUT, )
) response.raise_for_status()
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: 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}"
url = f"https://{island_server}/api/monkey-control" f"/needs-to-stop/{agent_id}" response = requests.get( # noqa: DUO123
response = requests.get( # noqa: DUO123 url,
url, verify=False,
verify=False, timeout=SHORT_REQUEST_TIMEOUT,
timeout=SHORT_REQUEST_TIMEOUT, )
) response.raise_for_status()
response.raise_for_status()
json_response = json.loads(response.content.decode()) json_response = json.loads(response.content.decode())
return json_response["stop_agent"] 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: def get_config(self, island_server: str) -> AgentConfiguration:
try: response = requests.get( # noqa: DUO123
response = requests.get( # noqa: DUO123 f"https://{island_server}/api/agent-configuration",
f"https://{island_server}/api/agent-configuration", verify=False,
verify=False, timeout=SHORT_REQUEST_TIMEOUT,
timeout=SHORT_REQUEST_TIMEOUT, )
) response.raise_for_status()
response.raise_for_status()
config_dict = json.loads(response.text) config_dict = json.loads(response.text)
logger.debug(f"Received configuration:\n{pformat(config_dict)}") logger.debug(f"Received configuration:\n{pformat(config_dict)}")
return AgentConfiguration(**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]: def get_credentials_for_propagation(self, island_server: str) -> Sequence[Credentials]:
try: response = requests.get( # noqa: DUO123
response = requests.get( # noqa: DUO123 f"https://{island_server}/api/propagation-credentials",
f"https://{island_server}/api/propagation-credentials", verify=False,
verify=False, timeout=SHORT_REQUEST_TIMEOUT,
timeout=SHORT_REQUEST_TIMEOUT, )
) response.raise_for_status()
response.raise_for_status()
return [Credentials(**credentials) for credentials in response.json()] 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: def _serialize_events(self, events: Sequence[AbstractAgentEvent]) -> JSONSerializable:
serialized_events: List[JSONSerializable] = [] serialized_events: List[JSONSerializable] = []