From 6c2cab99505cbf01109e12a79854be6d9d79f47a Mon Sep 17 00:00:00 2001
From: Kekoa Kaaikala <kekoa.kaaikala@gmail.com>
Date: Wed, 7 Sep 2022 18:46:17 +0000
Subject: [PATCH] Agent: Disconnect from upstream relay

---
 monkey/infection_monkey/monkey.py                 | 10 +++-------
 monkey/infection_monkey/network/relay/__init__.py |  1 +
 2 files changed, 4 insertions(+), 7 deletions(-)

diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py
index 918f5240c..311ccd9c0 100644
--- a/monkey/infection_monkey/monkey.py
+++ b/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()
diff --git a/monkey/infection_monkey/network/relay/__init__.py b/monkey/infection_monkey/network/relay/__init__.py
index b9eb8a009..50ee96438 100644
--- a/monkey/infection_monkey/network/relay/__init__.py
+++ b/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