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