Agent: Use SocketAddress in notify_disconnect() in network relay utils

This commit is contained in:
Shreya Malviya 2022-09-26 13:44:40 +05:30
parent 105a2b39cf
commit 4c76543a28
2 changed files with 19 additions and 11 deletions

View File

@ -121,6 +121,11 @@ class InfectionMonkey:
# 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._island_address = SocketAddress(
IPv4Address(self._cmd_island_ip), self._cmd_island_port
)
self._control_client = ControlClient(
server_address=server, island_api_client=self._island_api_client
)
@ -231,10 +236,9 @@ class InfectionMonkey:
config = self._control_channel.get_config()
relay_port = get_free_tcp_port()
island_address = SocketAddress(IPv4Address(self._cmd_island_ip), self._cmd_island_port)
self._relay = TCPRelay(
relay_port,
island_address,
self._island_address,
client_disconnect_timeout=config.keep_tunnel_open_time,
)
@ -488,7 +492,7 @@ class InfectionMonkey:
def _close_tunnel(self):
logger.info(f"Quitting tunnel {self._cmd_island_ip}")
notify_disconnect(self._cmd_island_ip, self._cmd_island_port)
notify_disconnect(self._island_address)
def _send_log(self):
monkey_log_path = get_agent_log_path()

View File

@ -92,22 +92,26 @@ def send_remove_from_waitlist_control_message_to_relays(servers: Iterable[str]):
def _send_remove_from_waitlist_control_message_to_relay(server: str):
ip, port = address_to_ip_port(server)
server_address = SocketAddress(IPv4Address(ip), int(port))
notify_disconnect(IPv4Address(ip), int(port))
notify_disconnect(server_address)
def notify_disconnect(server_ip: IPv4Address, server_port: int):
def notify_disconnect(server_address: SocketAddress):
"""
Tell upstream relay that we no longer need the relay.
Tell upstream relay that we no longer need the relay
:param server_ip: The IP address of the server to notify.
:param server_port: The port of the server to notify.
:param server_address: The address of the server to notify
"""
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as d_socket:
d_socket.settimeout(LONG_REQUEST_TIMEOUT)
try:
d_socket.connect((str(server_ip), server_port))
d_socket.connect((str(server_address.ip), server_address.port))
d_socket.sendall(RELAY_CONTROL_MESSAGE_REMOVE_FROM_WAITLIST)
logger.info(f"Control message was sent to the server/relay {server_ip}:{server_port}")
logger.info(
"Control message was sent to the server/relay "
f"{server_address.ip}:{server_address.port}"
)
except OSError as err:
logger.error(f"Error connecting to socket {server_ip}:{server_port}: {err}")
logger.error(
f"Error connecting to socket {server_address.ip}:{server_address.port}: {err}"
)