Agent: Add proxies to the initialization of ControlChannel

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()

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)