Agent: Disconnect from upstream relay

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
@@ -424,12 +424,8 @@ class InfectionMonkey:
         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