Agent: Reimplement HTTPIslandAPIClient.get_pba_file() as a method

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}")

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

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)