From 841183d8e7b9e461a1aca1cb1f4b6d13db41b29b Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Mon, 19 Sep 2022 19:53:54 -0400 Subject: [PATCH] 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)