Merge pull request #2257 from guardicore/2216-notify-relay-when-finished

2216 notify relay when finished

monkey/infection_monkey/monkey.py

@@ -9,7 +9,6 @@ from typing import List
 
 from pubsub.core import Publisher
 
-import infection_monkey.tunnel as tunnel
 from common.event_queue import IAgentEventQueue, PyPubSubAgentEventQueue
 from common.events import CredentialsStolenEvent
 from common.network.network_utils import address_to_ip_port
@@ -45,6 +44,7 @@ from infection_monkey.network.info import get_free_tcp_port, get_network_interfa
 from infection_monkey.network.relay import TCPRelay
 from infection_monkey.network.relay.utils import (
     find_server,
+    notify_disconnect,
     send_remove_from_waitlist_control_message_to_relays,
 )
 from infection_monkey.network_scanning.elasticsearch_fingerprinter import ElasticSearchFingerprinter
@@ -415,21 +415,18 @@ class InfectionMonkey:
             ).send()  # Signal the server (before closing the tunnel)
 
             self._close_tunnel()
-            self._singleton.unlock()
         except Exception as e:
             logger.error(f"An error occurred while cleaning up the monkey agent: {e}")
             if deleted is None:
                 InfectionMonkey._self_delete()
+        finally:
+            self._singleton.unlock()
 
         logger.info("Monkey is shutting down")
 
     def _close_tunnel(self):
-        tunnel_address = (
-            self._control_client.proxies.get("https", "").replace("http://", "").split(":")[0]
-        )
-        if tunnel_address:
-            logger.info("Quitting tunnel %s", tunnel_address)
-            tunnel.quit_tunnel(tunnel_address)
+        logger.info(f"Quitting tunnel {self._cmd_island_ip}")
+        notify_disconnect(self._cmd_island_ip, self._cmd_island_port)
 
     def _send_log(self):
         monkey_log_path = get_agent_log_path()

monkey/infection_monkey/network/relay/__init__.py

@@ -7,3 +7,4 @@ from .sockets_pipe import SocketsPipe
 from .tcp_connection_handler import TCPConnectionHandler
 from .tcp_pipe_spawner import TCPPipeSpawner
 from .tcp_relay import TCPRelay
+from .utils import notify_disconnect

monkey/infection_monkey/network/relay/utils.py

@@ -1,5 +1,6 @@
 import logging
 import socket
+from ipaddress import IPv4Address
 from typing import Iterable, Optional
 
 import requests
@@ -49,14 +50,21 @@ 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)
+    notify_disconnect(IPv4Address(ip), int(port))
+
+
+def notify_disconnect(server_ip: IPv4Address, server_port: int):
+    """
+    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.
+    """
     with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as d_socket:
-        d_socket.settimeout(MEDIUM_REQUEST_TIMEOUT)
-
-        ip, port = address_to_ip_port(server)
-        logger.info(f"Control message was sent to the server/relay {server}")
-
         try:
-            d_socket.connect((ip, int(port)))
-            d_socket.send(RELAY_CONTROL_MESSAGE_REMOVE_FROM_WAITLIST)
+            d_socket.connect((server_ip, server_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}")
         except OSError as err:
-            logger.error(f"Error connecting to socket {server}: {err}")
+            logger.error(f"Error connecting to socket {server_ip}:{server_port}: {err}")