From 1b92ec78fb043879141543dc7fa9314a8fdf249c Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Mon, 19 Sep 2022 18:14:24 +0530 Subject: [PATCH 01/19] Agent: Implement HTTPIslandAPIClient.send_log and use in ControlClient --- monkey/infection_monkey/control.py | 9 ++------- .../island_api_client/http_island_api_client.py | 11 +++++++++++ 2 files changed, 13 insertions(+), 7 deletions(-) diff --git a/monkey/infection_monkey/control.py b/monkey/infection_monkey/control.py index 58ab43fa6..d1b935899 100644 --- a/monkey/infection_monkey/control.py +++ b/monkey/infection_monkey/control.py @@ -9,6 +9,7 @@ from urllib3 import disable_warnings from common.common_consts.timeouts import LONG_REQUEST_TIMEOUT, MEDIUM_REQUEST_TIMEOUT from common.network.network_utils import get_my_ip_addresses from infection_monkey.config import GUID +from infection_monkey.island_api_client import HTTPIslandAPIClient from infection_monkey.network.info import get_host_subnets from infection_monkey.utils import agent_process @@ -78,13 +79,7 @@ class ControlClient: return try: telemetry = {"monkey_guid": GUID, "log": json.dumps(log)} - requests.post( # noqa: DUO123 - "https://%s/api/log" % (self.server_address,), - data=json.dumps(telemetry), - headers={"content-type": "application/json"}, - verify=False, - timeout=MEDIUM_REQUEST_TIMEOUT, - ) + HTTPIslandAPIClient(self.server_address).send_log(json.dumps(telemetry)) except Exception as exc: logger.warning(f"Error connecting to control server {self.server_address}: {exc}") 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 d9f9b1d9e..ce5e4338b 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 @@ -21,9 +21,20 @@ class HTTPIslandAPIClient(IIslandAPIClient): verify=False, timeout=MEDIUM_REQUEST_TIMEOUT, ) + self._island_server = island_server except requests.exceptions.ConnectionError as err: raise IslandAPIConnectionError(err) except TimeoutError as err: raise IslandAPITimeoutError(err) except Exception as err: raise IslandAPIError(err) + + # 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): + requests.post( # noqa: DUO123 + "https://%s/api/log" % (self._island_server,), + json=data, + verify=False, + timeout=MEDIUM_REQUEST_TIMEOUT, + ) From d188b06980e64a65cdd26447e67b6a8acf310699 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Mon, 19 Sep 2022 18:17:11 +0530 Subject: [PATCH 02/19] Agent: Implement HTTPIslandAPIClient.get_pba_file and use in ControlClient --- monkey/infection_monkey/control.py | 10 ++-------- .../island_api_client/http_island_api_client.py | 9 ++++++++- 2 files changed, 10 insertions(+), 9 deletions(-) diff --git a/monkey/infection_monkey/control.py b/monkey/infection_monkey/control.py index d1b935899..646da5239 100644 --- a/monkey/infection_monkey/control.py +++ b/monkey/infection_monkey/control.py @@ -6,7 +6,7 @@ from socket import gethostname import requests from urllib3 import disable_warnings -from common.common_consts.timeouts import LONG_REQUEST_TIMEOUT, MEDIUM_REQUEST_TIMEOUT +from common.common_consts.timeouts import MEDIUM_REQUEST_TIMEOUT from common.network.network_utils import get_my_ip_addresses from infection_monkey.config import GUID from infection_monkey.island_api_client import HTTPIslandAPIClient @@ -17,8 +17,6 @@ disable_warnings() # noqa DUO131 logger = logging.getLogger(__name__) -PBA_FILE_DOWNLOAD = "https://%s/api/pba/download/%s" - class ControlClient: # TODO When we have mechanism that support telemetry messenger @@ -85,10 +83,6 @@ class ControlClient: def get_pba_file(self, filename): try: - return requests.get( # noqa: DUO123 - PBA_FILE_DOWNLOAD % (self.server_address, filename), - verify=False, - timeout=LONG_REQUEST_TIMEOUT, - ) + HTTPIslandAPIClient(self.server_address).get_pba_file(filename) except requests.exceptions.RequestException: return False 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 ce5e4338b..5523a10a0 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 @@ -2,7 +2,7 @@ import logging import requests -from common.common_consts.timeouts import MEDIUM_REQUEST_TIMEOUT +from common.common_consts.timeouts import LONG_REQUEST_TIMEOUT, MEDIUM_REQUEST_TIMEOUT from . import IIslandAPIClient, IslandAPIConnectionError, IslandAPIError, IslandAPITimeoutError @@ -38,3 +38,10 @@ class HTTPIslandAPIClient(IIslandAPIClient): verify=False, timeout=MEDIUM_REQUEST_TIMEOUT, ) + + def get_pba_file(self, filename: str): + return requests.get( # noqa: DUO123 + "https://%s/api/pba/download/%s" % (self.server_address, filename), + verify=False, + timeout=LONG_REQUEST_TIMEOUT, + ) From fa9225370e71db8ff2d75b4cabb00f11ee1abcba Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Mon, 19 Sep 2022 18:37:31 +0530 Subject: [PATCH 03/19] 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): From 8ab17a96e3375e1e70ed6636cdba22e678221138 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Mon, 19 Sep 2022 18:38:37 +0530 Subject: [PATCH 04/19] Agent: Fix functions defined in HTTPIslandAPIClient and use the handle_island_errors() decorator on them --- .../http_island_api_client.py | 47 +++++++++++++------ 1 file changed, 32 insertions(+), 15 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 abd4fa9d1..df1c2814d 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 @@ -24,19 +24,36 @@ def handle_island_errors(fn): 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): - requests.post( # noqa: DUO123 - "https://%s/api/log" % (self._island_server,), - json=data, - verify=False, - timeout=MEDIUM_REQUEST_TIMEOUT, - ) - def get_pba_file(self, filename: str): - return requests.get( # noqa: DUO123 - "https://%s/api/pba/download/%s" % (self.server_address, filename), - verify=False, - timeout=LONG_REQUEST_TIMEOUT, - ) +class HTTPIslandAPIClient(IIslandAPIClient): + """ + A client for the Island's HTTP API + """ + + @handle_island_errors + def __init__(self, island_server: str): + requests.get( # noqa: DUO123 + f"https://{island_server}/api?action=is-up", + verify=False, + timeout=MEDIUM_REQUEST_TIMEOUT, + ) + self._island_server = island_server + + # 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 + @handle_island_errors + def send_log(self, data: str): + requests.post( # noqa: DUO123 + "https://%s/api/log" % (self._island_server,), + json=data, + verify=False, + timeout=MEDIUM_REQUEST_TIMEOUT, + ) + + @handle_island_errors + def get_pba_file(self, filename: str): + return requests.get( # noqa: DUO123 + "https://%s/api/pba/download/%s" % (self.server_address, filename), + verify=False, + timeout=LONG_REQUEST_TIMEOUT, + ) From d07760fe60c72c5838e93154fee36686d0315555 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Mon, 19 Sep 2022 18:44:31 +0530 Subject: [PATCH 05/19] Agent: Make HTTPIslandAPIClient.send_log() and HTTPIslandAPIClient.get_pba_file() static --- monkey/infection_monkey/control.py | 4 ++-- .../island_api_client/http_island_api_client.py | 11 ++++++----- 2 files changed, 8 insertions(+), 7 deletions(-) diff --git a/monkey/infection_monkey/control.py b/monkey/infection_monkey/control.py index 646da5239..2bb5e324f 100644 --- a/monkey/infection_monkey/control.py +++ b/monkey/infection_monkey/control.py @@ -77,12 +77,12 @@ class ControlClient: return try: telemetry = {"monkey_guid": GUID, "log": json.dumps(log)} - HTTPIslandAPIClient(self.server_address).send_log(json.dumps(telemetry)) + HTTPIslandAPIClient.send_log(self.server_address, json.dumps(telemetry)) except Exception as exc: logger.warning(f"Error connecting to control server {self.server_address}: {exc}") def get_pba_file(self, filename): try: - HTTPIslandAPIClient(self.server_address).get_pba_file(filename) + HTTPIslandAPIClient.get_pba_file(self.server_address, filename) except requests.exceptions.RequestException: return False 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 df1c2814d..b22eabad7 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 @@ -37,23 +37,24 @@ class HTTPIslandAPIClient(IIslandAPIClient): verify=False, timeout=MEDIUM_REQUEST_TIMEOUT, ) - self._island_server = island_server # 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 + @staticmethod @handle_island_errors - def send_log(self, data: str): + def send_log(server_address: str, data: str): requests.post( # noqa: DUO123 - "https://%s/api/log" % (self._island_server,), + "https://%s/api/log" % (server_address,), json=data, verify=False, timeout=MEDIUM_REQUEST_TIMEOUT, ) + @staticmethod @handle_island_errors - def get_pba_file(self, filename: str): + def get_pba_file(server_address: str, filename: str): return requests.get( # noqa: DUO123 - "https://%s/api/pba/download/%s" % (self.server_address, filename), + "https://%s/api/pba/download/%s" % (server_address, filename), verify=False, timeout=LONG_REQUEST_TIMEOUT, ) From 365376a19013ce8157a4e38ca619c4b962581ab9 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Mon, 19 Sep 2022 18:46:34 +0530 Subject: [PATCH 06/19] Agent: Change exception handling and log message in ControlClient.get_pba_file() --- monkey/infection_monkey/control.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/monkey/infection_monkey/control.py b/monkey/infection_monkey/control.py index 2bb5e324f..56db8a652 100644 --- a/monkey/infection_monkey/control.py +++ b/monkey/infection_monkey/control.py @@ -84,5 +84,5 @@ class ControlClient: def get_pba_file(self, filename): try: HTTPIslandAPIClient.get_pba_file(self.server_address, filename) - except requests.exceptions.RequestException: - return False + except Exception as exc: + logger.warning(f"Error connecting to control server {self.server_address}: {exc}") From 393bec29e729af723009e5c7e608ef3a325e658f Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Mon, 19 Sep 2022 18:47:30 +0530 Subject: [PATCH 07/19] Agent: Remove comment from HTTPIslandAPIClient --- .../island_api_client/http_island_api_client.py | 2 -- 1 file changed, 2 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 b22eabad7..ceddc7de2 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 @@ -38,8 +38,6 @@ class HTTPIslandAPIClient(IIslandAPIClient): timeout=MEDIUM_REQUEST_TIMEOUT, ) - # 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 @staticmethod @handle_island_errors def send_log(server_address: str, data: str): From 99366052421be4074cfcc3c540f399e642c22fde Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Mon, 19 Sep 2022 19:21:03 -0400 Subject: [PATCH 08/19] Agent: Add IIslandAPIClient.send_log() --- .../island_api_client/i_island_api_client.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/monkey/infection_monkey/island_api_client/i_island_api_client.py b/monkey/infection_monkey/island_api_client/i_island_api_client.py index 4a168b3d5..2db9ab239 100644 --- a/monkey/infection_monkey/island_api_client/i_island_api_client.py +++ b/monkey/infection_monkey/island_api_client/i_island_api_client.py @@ -17,3 +17,16 @@ class IIslandAPIClient(ABC): :raises IslandAPIError: If an unexpected error occurs while attempting to connect to the island """ + + @abstractmethod + def send_log(self, log_contents: str): + """ + Send the contents of the agent's log to the island + + :param log_contents: The contents of the agent's log + :raises IslandAPIConnectionError: If a connection cannot be made to the island + :raises IslandAPITimeoutError: If a timeout occurs while attempting to connect to the island + or send the log contents + :raises IslandAPIError: If an unexpected error occurred while attempting to send the + contents of the agent's log to the island + """ From aba342b3f221c93d7f910ccadc8aee225fa42f17 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Mon, 19 Sep 2022 19:41:14 -0400 Subject: [PATCH 09/19] Agent: Pass IIslandAPIClient to ControlClient.__init__() --- monkey/infection_monkey/control.py | 5 +++-- monkey/infection_monkey/monkey.py | 4 +++- 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/monkey/infection_monkey/control.py b/monkey/infection_monkey/control.py index 56db8a652..82a7250c4 100644 --- a/monkey/infection_monkey/control.py +++ b/monkey/infection_monkey/control.py @@ -9,7 +9,7 @@ from urllib3 import disable_warnings from common.common_consts.timeouts import MEDIUM_REQUEST_TIMEOUT from common.network.network_utils import get_my_ip_addresses from infection_monkey.config import GUID -from infection_monkey.island_api_client import HTTPIslandAPIClient +from infection_monkey.island_api_client import HTTPIslandAPIClient, IIslandAPIClient from infection_monkey.network.info import get_host_subnets from infection_monkey.utils import agent_process @@ -24,8 +24,9 @@ class ControlClient: # https://github.com/guardicore/monkey/blob/133f7f5da131b481561141171827d1f9943f6aec/monkey/infection_monkey/telemetry/base_telem.py control_client_object = None - def __init__(self, server_address: str): + def __init__(self, server_address: str, island_api_client: IIslandAPIClient): self.server_address = server_address + self._island_api_client = island_api_client def wakeup(self, parent=None): if parent: diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index 680db9f6f..96da82225 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -115,7 +115,9 @@ class InfectionMonkey: # TODO: `address_to_port()` should return the port as an integer. self._cmd_island_ip, self._cmd_island_port = address_to_ip_port(server) self._cmd_island_port = int(self._cmd_island_port) - self._control_client = ControlClient(server_address=server) + self._control_client = ControlClient( + server_address=server, island_api_client=island_api_client + ) # TODO Refactor the telemetry messengers to accept control client # and remove control_client_object From 92da3b78db12d3abe9ab6ca361e25224890ca45d Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Mon, 19 Sep 2022 19:44:16 -0400 Subject: [PATCH 10/19] Agent: Reimplement HTTPIslandAPIClient.send_log() as member, not static --- monkey/infection_monkey/control.py | 2 +- .../island_api_client/http_island_api_client.py | 9 +++++---- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/monkey/infection_monkey/control.py b/monkey/infection_monkey/control.py index 82a7250c4..9d7daec6e 100644 --- a/monkey/infection_monkey/control.py +++ b/monkey/infection_monkey/control.py @@ -78,7 +78,7 @@ class ControlClient: return try: telemetry = {"monkey_guid": GUID, "log": json.dumps(log)} - HTTPIslandAPIClient.send_log(self.server_address, json.dumps(telemetry)) + self._island_api_client.send_log(json.dumps(telemetry)) except Exception as exc: logger.warning(f"Error connecting to control server {self.server_address}: {exc}") 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 ceddc7de2..98bca7562 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 @@ -38,12 +38,13 @@ class HTTPIslandAPIClient(IIslandAPIClient): timeout=MEDIUM_REQUEST_TIMEOUT, ) - @staticmethod + self._island_server = island_server + @handle_island_errors - def send_log(server_address: str, data: str): + def send_log(self, log_contents: str): requests.post( # noqa: DUO123 - "https://%s/api/log" % (server_address,), - json=data, + f"https://{self._island_server}/api/log", + json=log_contents, verify=False, timeout=MEDIUM_REQUEST_TIMEOUT, ) From aa3c6c2f4dfa10cce4e24d366657cf4b430cab50 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Mon, 19 Sep 2022 19:47:14 -0400 Subject: [PATCH 11/19] Agent: Add IIslandAPIClient.get_pba_file() --- .../island_api_client/i_island_api_client.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/monkey/infection_monkey/island_api_client/i_island_api_client.py b/monkey/infection_monkey/island_api_client/i_island_api_client.py index 2db9ab239..8479f39ec 100644 --- a/monkey/infection_monkey/island_api_client/i_island_api_client.py +++ b/monkey/infection_monkey/island_api_client/i_island_api_client.py @@ -30,3 +30,17 @@ class IIslandAPIClient(ABC): :raises IslandAPIError: If an unexpected error occurred while attempting to send the contents of the agent's log to the island """ + + @abstractmethod + def get_pba_file(self, filename: str) -> bytes: + """ + Get a custom PBA file from the island + + :param filename: The name of the custom PBA file + :return: The contents of the custom PBA file in bytes + :raises IslandAPIConnectionError: If a connection cannot be made to the island + :raises IslandAPITimeoutError: If a timeout occurs while attempting to connect to the island + or retrieve the custom PBA file + :raises IslandAPIError: If an unexpected error occurred while attempting to retrieve the + custom PBA file + """ From 841183d8e7b9e461a1aca1cb1f4b6d13db41b29b Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Mon, 19 Sep 2022 19:53:54 -0400 Subject: [PATCH 12/19] Agent: Reimplement HTTPIslandAPIClient.get_pba_file() as a method --- monkey/infection_monkey/control.py | 4 ++-- .../island_api_client/http_island_api_client.py | 9 +++++---- .../post_breach/custom_pba/custom_pba.py | 4 ++-- 3 files changed, 9 insertions(+), 8 deletions(-) diff --git a/monkey/infection_monkey/control.py b/monkey/infection_monkey/control.py index 9d7daec6e..440f58712 100644 --- a/monkey/infection_monkey/control.py +++ b/monkey/infection_monkey/control.py @@ -9,7 +9,7 @@ from urllib3 import disable_warnings from common.common_consts.timeouts import MEDIUM_REQUEST_TIMEOUT from common.network.network_utils import get_my_ip_addresses from infection_monkey.config import GUID -from infection_monkey.island_api_client import HTTPIslandAPIClient, IIslandAPIClient +from infection_monkey.island_api_client import IIslandAPIClient from infection_monkey.network.info import get_host_subnets from infection_monkey.utils import agent_process @@ -84,6 +84,6 @@ class ControlClient: def get_pba_file(self, filename): try: - HTTPIslandAPIClient.get_pba_file(self.server_address, filename) + return self._island_api_client.get_pba_file(filename) except Exception as exc: logger.warning(f"Error connecting to control server {self.server_address}: {exc}") 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 98bca7562..bf6d085db 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 @@ -49,11 +49,12 @@ class HTTPIslandAPIClient(IIslandAPIClient): timeout=MEDIUM_REQUEST_TIMEOUT, ) - @staticmethod @handle_island_errors - def get_pba_file(server_address: str, filename: str): - return requests.get( # noqa: DUO123 - "https://%s/api/pba/download/%s" % (server_address, filename), + def get_pba_file(self, filename: str): + response = requests.get( # noqa: DUO123 + f"https://{self._island_server}/api/pba/download/{filename}", verify=False, timeout=LONG_REQUEST_TIMEOUT, ) + + return response.content diff --git a/monkey/infection_monkey/post_breach/custom_pba/custom_pba.py b/monkey/infection_monkey/post_breach/custom_pba/custom_pba.py index f5087a91a..34fb73147 100644 --- a/monkey/infection_monkey/post_breach/custom_pba/custom_pba.py +++ b/monkey/infection_monkey/post_breach/custom_pba/custom_pba.py @@ -76,7 +76,7 @@ class CustomPBA(PBA): pba_file_contents = self.control_client.get_pba_file(filename) status = None - if not pba_file_contents or not pba_file_contents.content: + if not pba_file_contents: logger.error("Island didn't respond with post breach file.") status = ScanStatus.SCANNED @@ -97,7 +97,7 @@ class CustomPBA(PBA): try: with open(os.path.join(dst_dir, filename), "wb") as written_PBA_file: - written_PBA_file.write(pba_file_contents.content) + written_PBA_file.write(pba_file_contents) return True except IOError as e: logger.error("Can not upload post breach file to target machine: %s" % e) From cb8fda0beccf4bb2140a35f8ba902df0a97022f5 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Mon, 19 Sep 2022 19:57:28 -0400 Subject: [PATCH 13/19] Agent: Add HTTPIslandAPIClient._api_url attribute --- .../island_api_client/http_island_api_client.py | 5 +++-- 1 file changed, 3 insertions(+), 2 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 bf6d085db..9d85447ea 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 @@ -39,11 +39,12 @@ class HTTPIslandAPIClient(IIslandAPIClient): ) self._island_server = island_server + self._api_url = f"https://{self._island_server}/api" @handle_island_errors def send_log(self, log_contents: str): requests.post( # noqa: DUO123 - f"https://{self._island_server}/api/log", + f"{self._api_url}/log", json=log_contents, verify=False, timeout=MEDIUM_REQUEST_TIMEOUT, @@ -52,7 +53,7 @@ class HTTPIslandAPIClient(IIslandAPIClient): @handle_island_errors def get_pba_file(self, filename: str): response = requests.get( # noqa: DUO123 - f"https://{self._island_server}/api/pba/download/{filename}", + f"{self._api_url}/pba/download/{filename}", verify=False, timeout=LONG_REQUEST_TIMEOUT, ) From a724758caa0b0a9183ad472d19c3014cffe239b8 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Mon, 19 Sep 2022 20:04:52 -0400 Subject: [PATCH 14/19] Agent: Handle HTTPErrors in HTTPIslandAPIClient --- .../island_api_client/http_island_api_client.py | 16 +++++++++++++++- 1 file changed, 15 insertions(+), 1 deletion(-) 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 9d85447ea..15afd5e44 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 @@ -5,7 +5,14 @@ import requests from common.common_consts.timeouts import LONG_REQUEST_TIMEOUT, MEDIUM_REQUEST_TIMEOUT -from . import IIslandAPIClient, IslandAPIConnectionError, IslandAPIError, IslandAPITimeoutError +from . import ( + IIslandAPIClient, + IslandAPIConnectionError, + IslandAPIError, + IslandAPIRequestError, + IslandAPIRequestFailedError, + IslandAPITimeoutError, +) logger = logging.getLogger(__name__) @@ -17,6 +24,13 @@ def handle_island_errors(fn): return fn(*args, **kwargs) except requests.exceptions.ConnectionError as err: raise IslandAPIConnectionError(err) + except requests.exceptions.HTTPError as err: + if 400 <= err.response.status_code < 500: + raise IslandAPIRequestError(err) + elif 500 <= err.response.status_code < 600: + raise IslandAPIRequestFailedError(err) + else: + raise IslandAPIError(err) except TimeoutError as err: raise IslandAPITimeoutError(err) except Exception as err: From 107a15b5f00aeab79051b78767b2aff12d0302f2 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Mon, 19 Sep 2022 20:05:19 -0400 Subject: [PATCH 15/19] Agent: Call raise_for_status() in HTTPIslandAPIClient --- .../island_api_client/http_island_api_client.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 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 15afd5e44..37feb8942 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 @@ -46,23 +46,25 @@ class HTTPIslandAPIClient(IIslandAPIClient): @handle_island_errors def __init__(self, island_server: str): - requests.get( # noqa: DUO123 + response = requests.get( # noqa: DUO123 f"https://{island_server}/api?action=is-up", verify=False, timeout=MEDIUM_REQUEST_TIMEOUT, ) + response.raise_for_status() self._island_server = island_server self._api_url = f"https://{self._island_server}/api" @handle_island_errors def send_log(self, log_contents: str): - requests.post( # noqa: DUO123 + response = requests.post( # noqa: DUO123 f"{self._api_url}/log", json=log_contents, verify=False, timeout=MEDIUM_REQUEST_TIMEOUT, ) + response.raise_for_status() @handle_island_errors def get_pba_file(self, filename: str): @@ -71,5 +73,6 @@ class HTTPIslandAPIClient(IIslandAPIClient): verify=False, timeout=LONG_REQUEST_TIMEOUT, ) + response.raise_for_status() return response.content From 0c13298bbbdffb2f89b1e8c7c22ee4956683c501 Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Tue, 20 Sep 2022 07:49:17 +0200 Subject: [PATCH 16/19] UT: Add error handling tests for HTTPIslandAPIClient send_log and get_pba_file --- .../test_http_island_api_client.py | 39 +++++++++++++++++++ 1 file changed, 39 insertions(+) diff --git a/monkey/tests/unit_tests/infection_monkey/island_api_client/test_http_island_api_client.py b/monkey/tests/unit_tests/infection_monkey/island_api_client/test_http_island_api_client.py index f213b0569..ffeada1c5 100644 --- a/monkey/tests/unit_tests/infection_monkey/island_api_client/test_http_island_api_client.py +++ b/monkey/tests/unit_tests/infection_monkey/island_api_client/test_http_island_api_client.py @@ -10,8 +10,11 @@ from infection_monkey.island_api_client import ( ) SERVER = "1.1.1.1:9999" +PBA_FILE = "dummy.pba" ISLAND_URI = f"https://{SERVER}/api?action=is-up" +ISLAND_SEND_LOG_URI = f"https://{SERVER}/api/log" +ISLAND_GET_PBA_FILE_URI = f"https://{SERVER}/api/pba/download/{PBA_FILE}" @pytest.mark.parametrize( @@ -28,3 +31,39 @@ def test_island_api_client(actual_error, expected_error): with pytest.raises(expected_error): HTTPIslandAPIClient(SERVER) + + +@pytest.mark.parametrize( + "actual_error, expected_error", + [ + (requests.exceptions.ConnectionError, IslandAPIConnectionError), + (TimeoutError, IslandAPITimeoutError), + (Exception, IslandAPIError), + ], +) +def test_island_api_client__send_log(actual_error, expected_error): + with requests_mock.Mocker() as m: + m.get(ISLAND_URI) + island_api_client = HTTPIslandAPIClient(SERVER) + + with pytest.raises(expected_error): + m.post(ISLAND_SEND_LOG_URI, exc=actual_error) + island_api_client.send_log(log_contents="some_data") + + +@pytest.mark.parametrize( + "actual_error, expected_error", + [ + (requests.exceptions.ConnectionError, IslandAPIConnectionError), + (TimeoutError, IslandAPITimeoutError), + (Exception, IslandAPIError), + ], +) +def test_island_api_client__get_pba_file(actual_error, expected_error): + with requests_mock.Mocker() as m: + m.get(ISLAND_URI) + island_api_client = HTTPIslandAPIClient(SERVER) + + with pytest.raises(expected_error): + m.get(ISLAND_GET_PBA_FILE_URI, exc=actual_error) + island_api_client.get_pba_file(filename=PBA_FILE) From 1b4f834f4638486b88c67cf07935efe4fb42c07e Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Tue, 20 Sep 2022 08:50:55 +0200 Subject: [PATCH 17/19] UT: Add status code tests for HTTIslandAPIClient --- .../test_http_island_api_client.py | 51 +++++++++++++++++++ 1 file changed, 51 insertions(+) diff --git a/monkey/tests/unit_tests/infection_monkey/island_api_client/test_http_island_api_client.py b/monkey/tests/unit_tests/infection_monkey/island_api_client/test_http_island_api_client.py index ffeada1c5..bd6bfcb41 100644 --- a/monkey/tests/unit_tests/infection_monkey/island_api_client/test_http_island_api_client.py +++ b/monkey/tests/unit_tests/infection_monkey/island_api_client/test_http_island_api_client.py @@ -6,6 +6,8 @@ from infection_monkey.island_api_client import ( HTTPIslandAPIClient, IslandAPIConnectionError, IslandAPIError, + IslandAPIRequestError, + IslandAPIRequestFailedError, IslandAPITimeoutError, ) @@ -33,6 +35,21 @@ def test_island_api_client(actual_error, expected_error): HTTPIslandAPIClient(SERVER) +@pytest.mark.parametrize( + "status_code, expected_error", + [ + (401, IslandAPIRequestError), + (501, IslandAPIRequestFailedError), + ], +) +def test_island_api_client__status_code(status_code, expected_error): + with requests_mock.Mocker() as m: + m.get(ISLAND_URI, status_code=status_code) + + with pytest.raises(expected_error): + HTTPIslandAPIClient(SERVER) + + @pytest.mark.parametrize( "actual_error, expected_error", [ @@ -51,6 +68,23 @@ def test_island_api_client__send_log(actual_error, expected_error): island_api_client.send_log(log_contents="some_data") +@pytest.mark.parametrize( + "status_code, expected_error", + [ + (401, IslandAPIRequestError), + (501, IslandAPIRequestFailedError), + ], +) +def test_island_api_client_send_log__status_code(status_code, expected_error): + with requests_mock.Mocker() as m: + m.get(ISLAND_URI) + island_api_client = HTTPIslandAPIClient(SERVER) + + with pytest.raises(expected_error): + m.post(ISLAND_SEND_LOG_URI, status_code=status_code) + island_api_client.send_log(log_contents="some_data") + + @pytest.mark.parametrize( "actual_error, expected_error", [ @@ -67,3 +101,20 @@ def test_island_api_client__get_pba_file(actual_error, expected_error): with pytest.raises(expected_error): m.get(ISLAND_GET_PBA_FILE_URI, exc=actual_error) island_api_client.get_pba_file(filename=PBA_FILE) + + +@pytest.mark.parametrize( + "status_code, expected_error", + [ + (401, IslandAPIRequestError), + (501, IslandAPIRequestFailedError), + ], +) +def test_island_api_client_get_pba_file__status_code(status_code, expected_error): + with requests_mock.Mocker() as m: + m.get(ISLAND_URI) + island_api_client = HTTPIslandAPIClient(SERVER) + + with pytest.raises(expected_error): + m.get(ISLAND_GET_PBA_FILE_URI, status_code=status_code) + island_api_client.get_pba_file(filename=PBA_FILE) From 1480203627c5458a0e1d9551652bf8a7267b3deb Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Tue, 20 Sep 2022 12:35:21 +0530 Subject: [PATCH 18/19] Agent: Modify docstrings in IIslandAPIClient --- .../island_api_client/i_island_api_client.py | 24 +++++++++++++------ 1 file changed, 17 insertions(+), 7 deletions(-) diff --git a/monkey/infection_monkey/island_api_client/i_island_api_client.py b/monkey/infection_monkey/island_api_client/i_island_api_client.py index 8479f39ec..fefba9973 100644 --- a/monkey/infection_monkey/island_api_client/i_island_api_client.py +++ b/monkey/infection_monkey/island_api_client/i_island_api_client.py @@ -9,10 +9,14 @@ class IIslandAPIClient(ABC): @abstractmethod def __init__(self, island_server: str): """ - Construct and island API client and connect it to the island + Construct an island API client and connect it to the island :param island_server: The socket address of the API :raises IslandAPIConnectionError: If the client cannot successfully connect to the island + :raises IslandAPIRequestError: If an error occurs while attempting to connect to the + island due to an issue in the request sent from the client + :raises IslandAPIRequestFailedError: If an error occurs while attempting to connect to the + island due to an error on the server :raises IslandAPITimeoutError: If a timeout occurs while attempting to connect to the island :raises IslandAPIError: If an unexpected error occurs while attempting to connect to the island @@ -24,10 +28,13 @@ class IIslandAPIClient(ABC): Send the contents of the agent's log to the island :param log_contents: The contents of the agent's log - :raises IslandAPIConnectionError: If a connection cannot be made to the island + :raises IslandAPIConnectionError: If the client cannot successfully connect to the island + :raises IslandAPIRequestError: If an error occurs while attempting to connect to the + island due to an issue in the request sent from the client + :raises IslandAPIRequestFailedError: If an error occurs while attempting to connect to the + island due to an error on the server :raises IslandAPITimeoutError: If a timeout occurs while attempting to connect to the island - or send the log contents - :raises IslandAPIError: If an unexpected error occurred while attempting to send the + :raises IslandAPIError: If an unexpected error occurs while attempting to send the contents of the agent's log to the island """ @@ -38,9 +45,12 @@ class IIslandAPIClient(ABC): :param filename: The name of the custom PBA file :return: The contents of the custom PBA file in bytes - :raises IslandAPIConnectionError: If a connection cannot be made to the island + :raises IslandAPIConnectionError: If the client cannot successfully connect to the island + :raises IslandAPIRequestError: If an error occurs while attempting to connect to the + island due to an issue in the request sent from the client + :raises IslandAPIRequestFailedError: If an error occurs while attempting to connect to the + island due to an error on the server :raises IslandAPITimeoutError: If a timeout occurs while attempting to connect to the island - or retrieve the custom PBA file - :raises IslandAPIError: If an unexpected error occurred while attempting to retrieve the + :raises IslandAPIError: If an unexpected error occurs while attempting to retrieve the custom PBA file """ From 3100e6c01008d9122e800e2a87827b30150a639e Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Tue, 20 Sep 2022 12:37:05 +0530 Subject: [PATCH 19/19] Project: Remove outdated entried from Vulture allowlist --- vulture_allowlist.py | 12 ------------ 1 file changed, 12 deletions(-) diff --git a/vulture_allowlist.py b/vulture_allowlist.py index 83fb82005..53001be59 100644 --- a/vulture_allowlist.py +++ b/vulture_allowlist.py @@ -9,12 +9,6 @@ from common.agent_configuration.agent_sub_configurations import ( ) from common.credentials import Credentials, LMHash, NTHash from infection_monkey.exploit.log4shell_utils.ldap_server import LDAPServerFactory -from infection_monkey.island_api_client import ( - HTTPIslandAPIClient, - IIslandAPIClient, - IslandAPIRequestError, - IslandAPIRequestFailedError, -) from monkey_island.cc.event_queue import IslandEventTopic, PyPubSubIslandEventQueue from monkey_island.cc.models import Report from monkey_island.cc.models.networkmap import Arc, NetworkMap @@ -334,9 +328,3 @@ CC_TUNNEL IslandEventTopic.AGENT_CONNECTED IslandEventTopic.CLEAR_SIMULATION_DATA IslandEventTopic.RESET_AGENT_CONFIGURATION - -# TODO: Remove after #2292 is closed -IIslandAPIClient -HTTPIslandAPIClient -IslandAPIRequestFailedError -IslandAPIRequestError