From 4c76543a287ddc634b95e7def5c79e0bbb5b012d Mon Sep 17 00:00:00 2001
From: Shreya Malviya <shreya.malviya@gmail.com>
Date: Mon, 26 Sep 2022 13:44:40 +0530
Subject: [PATCH] Agent: Use SocketAddress in notify_disconnect() in network
 relay utils

---
 monkey/infection_monkey/monkey.py             | 10 +++++++---
 .../infection_monkey/network/relay/utils.py   | 20 +++++++++++--------
 2 files changed, 19 insertions(+), 11 deletions(-)

diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py
index 425234760..b3e2c1d4c 100644
--- a/monkey/infection_monkey/monkey.py
+++ b/monkey/infection_monkey/monkey.py
@@ -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()
diff --git a/monkey/infection_monkey/network/relay/utils.py b/monkey/infection_monkey/network/relay/utils.py
index 3f7cc4629..f09a076f3 100644
--- a/monkey/infection_monkey/network/relay/utils.py
+++ b/monkey/infection_monkey/network/relay/utils.py
@@ -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}"
+            )