From d188b06980e64a65cdd26447e67b6a8acf310699 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Mon, 19 Sep 2022 18:17:11 +0530 Subject: [PATCH] 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, + )