Agent: Implement HTTPIslandAPIClient.get_pba_file and use in ControlClient

This commit is contained in:
Shreya Malviya 2022-09-19 18:17:11 +05:30 committed by Mike Salvatore
parent 1b92ec78fb
commit d188b06980
2 changed files with 10 additions and 9 deletions

View File

@ -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

View File

@ -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,
)