Agent: Convert _cmd_island_port to int in one place

This commit is contained in:
Mike Salvatore 2022-09-09 07:52:12 -04:00
parent 3516fa1fec
commit 0246f6b981
2 changed files with 5 additions and 2 deletions

View File

@ -1,6 +1,7 @@
from typing import Optional, Tuple from typing import Optional, Tuple
# TODO: `address_to_port()` should return the port as an integer.
def address_to_ip_port(address: str) -> Tuple[str, Optional[str]]: def address_to_ip_port(address: str) -> Tuple[str, Optional[str]]:
""" """
Split a string containing an IP address (and optionally a port) into IP and Port components. Split a string containing an IP address (and optionally a port) into IP and Port components.

View File

@ -103,7 +103,9 @@ class InfectionMonkey:
# TODO: Revisit variable names # TODO: Revisit variable names
server = self._get_server() server = self._get_server()
# 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_ip = int(self._cmd_island_ip)
self._control_client = ControlClient(server_address=server) self._control_client = ControlClient(server_address=server)
# TODO Refactor the telemetry messengers to accept control client # TODO Refactor the telemetry messengers to accept control client
@ -189,7 +191,7 @@ class InfectionMonkey:
self._relay = TCPRelay( self._relay = TCPRelay(
relay_port, relay_port,
IPv4Address(self._cmd_island_ip), IPv4Address(self._cmd_island_ip),
int(self._cmd_island_port), self._cmd_island_port,
client_disconnect_timeout=config.keep_tunnel_open_time, client_disconnect_timeout=config.keep_tunnel_open_time,
) )
@ -427,7 +429,7 @@ class InfectionMonkey:
def _close_tunnel(self): def _close_tunnel(self):
logger.info(f"Quitting tunnel {self._cmd_island_ip}") logger.info(f"Quitting tunnel {self._cmd_island_ip}")
notify_disconnect(self._cmd_island_ip, int(self._cmd_island_port)) notify_disconnect(self._cmd_island_ip, self._cmd_island_port)
def _send_log(self): def _send_log(self):
monkey_log_path = get_agent_log_path() monkey_log_path = get_agent_log_path()