Agent: Pass and accept island_api_client in CachingAgentBinaryRepository's constructor

This commit is contained in:
Shreya Malviya 2022-09-20 14:54:25 +05:30
parent 3100e6c010
commit 9da1ad4c46
2 changed files with 7 additions and 4 deletions

View File

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

View File

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