From 94dbd9a8e2be24f4178f749d8cc80fb6a60480df Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Mon, 13 Jun 2022 13:02:17 +0200 Subject: [PATCH] Agent: Add proxies to the initialization of ControlChannel --- monkey/infection_monkey/master/control_channel.py | 10 ++++++---- monkey/infection_monkey/monkey.py | 8 ++++++-- 2 files changed, 12 insertions(+), 6 deletions(-) diff --git a/monkey/infection_monkey/master/control_channel.py b/monkey/infection_monkey/master/control_channel.py index 028437faf..0834f28f7 100644 --- a/monkey/infection_monkey/master/control_channel.py +++ b/monkey/infection_monkey/master/control_channel.py @@ -1,5 +1,6 @@ import json import logging +from typing import Mapping import requests @@ -14,9 +15,10 @@ logger = logging.getLogger(__name__) class ControlChannel(IControlChannel): - def __init__(self, server: str, agent_id: str): + def __init__(self, server: str, agent_id: str, proxies: Mapping[str, str]): self._agent_id = agent_id self._control_channel_server = server + self._proxies = proxies def should_agent_stop(self) -> bool: if not self._control_channel_server: @@ -30,7 +32,7 @@ class ControlChannel(IControlChannel): response = requests.get( # noqa: DUO123 url, verify=False, - proxies=ControlClient.proxies, + proxies=self._proxies, timeout=SHORT_REQUEST_TIMEOUT, ) response.raise_for_status() @@ -51,7 +53,7 @@ class ControlChannel(IControlChannel): response = requests.get( # noqa: DUO123 f"https://{self._control_channel_server}/api/agent", verify=False, - proxies=ControlClient.proxies, + proxies=self._proxies, timeout=SHORT_REQUEST_TIMEOUT, ) response.raise_for_status() @@ -74,7 +76,7 @@ class ControlChannel(IControlChannel): response = requests.get( # noqa: DUO123 propagation_credentials_url, verify=False, - proxies=ControlClient.proxies, + proxies=self._proxies, timeout=SHORT_REQUEST_TIMEOUT, ) response.raise_for_status() diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index e8be1f05d..eda19632a 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -128,7 +128,9 @@ class InfectionMonkey: run_aws_environment_check(self._telemetry_messenger) - should_stop = ControlChannel(WormConfiguration.current_server, GUID).should_agent_stop() + should_stop = ControlChannel( + WormConfiguration.current_server, GUID, self.cc_client.proxies + ).should_agent_stop() if should_stop: logger.info("The Monkey Island has instructed this agent to stop") return @@ -177,7 +179,9 @@ class InfectionMonkey: local_network_interfaces = InfectionMonkey._get_local_network_interfaces() # TODO control_channel and control_client have same responsibilities, merge them - control_channel = ControlChannel(self.cc_client.server_address, GUID) + control_channel = ControlChannel( + self.cc_client.server_address, GUID, self.cc_client.proxies + ) control_client = self.cc_client credentials_store = AggregatingCredentialsStore(control_channel)