diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index 73468a3c419..5a0ba093857 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -9,7 +9,6 @@ from pubsub.core import Publisher -import infection_monkey.tunnel as tunnel from common.event_queue import IEventQueue, PyPubSubEventQueue from common.events import CredentialsStolenEvent from common.network.network_utils import address_to_ip_port @@ -45,6 +44,7 @@ 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 @@ def cleanup(self): ).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() diff --git a/monkey/infection_monkey/network/relay/__init__.py b/monkey/infection_monkey/network/relay/__init__.py index b9eb8a009d1..50ee9643899 100644 --- a/monkey/infection_monkey/network/relay/__init__.py +++ b/monkey/infection_monkey/network/relay/__init__.py @@ -7,3 +7,4 @@ from .tcp_connection_handler import TCPConnectionHandler from .tcp_pipe_spawner import TCPPipeSpawner from .tcp_relay import TCPRelay +from .utils import notify_disconnect diff --git a/monkey/infection_monkey/network/relay/utils.py b/monkey/infection_monkey/network/relay/utils.py index 096500c5bdb..7ab38581b31 100644 --- a/monkey/infection_monkey/network/relay/utils.py +++ b/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): - 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) + notify_disconnect(IPv4Address(ip), int(port)) + - ip, port = address_to_ip_port(server) - logger.info(f"Control message was sent to the server/relay {server}") +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: 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}")