Agent: Pass IIslandAPIClient to ControlClient.__init__()

This commit is contained in:
Mike Salvatore 2022-09-19 19:41:14 -04:00
parent 9936605242
commit aba342b3f2
2 changed files with 6 additions and 3 deletions

View File

@ -9,7 +9,7 @@ from urllib3 import disable_warnings
from common.common_consts.timeouts import MEDIUM_REQUEST_TIMEOUT from common.common_consts.timeouts import MEDIUM_REQUEST_TIMEOUT
from common.network.network_utils import get_my_ip_addresses from common.network.network_utils import get_my_ip_addresses
from infection_monkey.config import GUID from infection_monkey.config import GUID
from infection_monkey.island_api_client import HTTPIslandAPIClient from infection_monkey.island_api_client import HTTPIslandAPIClient, IIslandAPIClient
from infection_monkey.network.info import get_host_subnets from infection_monkey.network.info import get_host_subnets
from infection_monkey.utils import agent_process from infection_monkey.utils import agent_process
@ -24,8 +24,9 @@ class ControlClient:
# https://github.com/guardicore/monkey/blob/133f7f5da131b481561141171827d1f9943f6aec/monkey/infection_monkey/telemetry/base_telem.py # https://github.com/guardicore/monkey/blob/133f7f5da131b481561141171827d1f9943f6aec/monkey/infection_monkey/telemetry/base_telem.py
control_client_object = None control_client_object = None
def __init__(self, server_address: str): def __init__(self, server_address: str, island_api_client: IIslandAPIClient):
self.server_address = server_address self.server_address = server_address
self._island_api_client = island_api_client
def wakeup(self, parent=None): def wakeup(self, parent=None):
if parent: if parent:

View File

@ -115,7 +115,9 @@ class InfectionMonkey:
# TODO: `address_to_port()` should return the port as an integer. # 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_ip, self._cmd_island_port = address_to_ip_port(server)
self._cmd_island_port = int(self._cmd_island_port) self._cmd_island_port = int(self._cmd_island_port)
self._control_client = ControlClient(server_address=server) self._control_client = ControlClient(
server_address=server, island_api_client=island_api_client
)
# TODO Refactor the telemetry messengers to accept control client # TODO Refactor the telemetry messengers to accept control client
# and remove control_client_object # and remove control_client_object