diff --git a/monkey/infection_monkey/exploit/caching_agent_binary_repository.py b/monkey/infection_monkey/exploit/caching_agent_binary_repository.py index 745aae112..72ff61446 100644 --- a/monkey/infection_monkey/exploit/caching_agent_binary_repository.py +++ b/monkey/infection_monkey/exploit/caching_agent_binary_repository.py @@ -6,6 +6,7 @@ import requests from common import OperatingSystem from common.common_consts.timeouts import MEDIUM_REQUEST_TIMEOUT +from infection_monkey.island_api_client import IIslandAPIClient from . import IAgentBinaryRepository @@ -17,9 +18,10 @@ class CachingAgentBinaryRepository(IAgentBinaryRepository): request is actually sent to the island for each requested binary. """ - def __init__(self, island_url: str): + def __init__(self, island_url: str, island_api_client: IIslandAPIClient): self._island_url = island_url self._lock = threading.Lock() + self._island_api_client = island_api_client def get_agent_binary( self, operating_system: OperatingSystem, architecture: str = None diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index 96da82225..6a7e1cfb1 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -111,12 +111,12 @@ class InfectionMonkey: self._opts = self._get_arguments(args) # TODO: Revisit variable names - server, island_api_client = self._connect_to_island_api() + server, self.island_api_client = self._connect_to_island_api() # TODO: `address_to_port()` should return the port as an integer. self._cmd_island_ip, self._cmd_island_port = address_to_ip_port(server) self._cmd_island_port = int(self._cmd_island_port) self._control_client = ControlClient( - server_address=server, island_api_client=island_api_client + server_address=server, island_api_client=self.island_api_client ) # TODO Refactor the telemetry messengers to accept control client @@ -315,7 +315,8 @@ class InfectionMonkey: puppet.load_plugin("ssh", SSHFingerprinter(), PluginType.FINGERPRINTER) agent_binary_repository = CachingAgentBinaryRepository( - f"https://{self._control_client.server_address}" + island_url=f"https://{self._control_client.server_address}", + island_api_client=self.island_api_client, ) exploit_wrapper = ExploiterWrapper( self._telemetry_messenger, event_queue, agent_binary_repository