p20319658
/
monkey
forked from p34709852/monkey
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Agent: Use SocketAddress in notify_disconnect() in network relay utils

This commit is contained in:
Shreya Malviya 2022-09-26 13:44:40 +05:30
parent 105a2b39cf
commit 4c76543a28
2 changed files with 19 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
monkey/infection_monkey/monkey.py
View File

@ -121,6 +121,11 @@ class InfectionMonkey:
# TODO: `address_to_port()` should return the port as an integer. # TODO: `address_to_port()` should return the port as an integer.
self._cmd_island_ip, self._cmd_island_port = address_to_ip_port(server) self._cmd_island_ip, self._cmd_island_port = address_to_ip_port(server)
self._cmd_island_port = int(self._cmd_island_port) self._cmd_island_port = int(self._cmd_island_port)
self._island_address = SocketAddress(
IPv4Address(self._cmd_island_ip), self._cmd_island_port
)
self._control_client = ControlClient( self._control_client = ControlClient(
server_address=server, island_api_client=self._island_api_client server_address=server, island_api_client=self._island_api_client
) )
@ -231,10 +236,9 @@ class InfectionMonkey:
config = self._control_channel.get_config() config = self._control_channel.get_config()
relay_port = get_free_tcp_port() relay_port = get_free_tcp_port()
island_address = SocketAddress(IPv4Address(self._cmd_island_ip), self._cmd_island_port)
self._relay = TCPRelay( self._relay = TCPRelay(
relay_port, relay_port,
island_address, self._island_address,
client_disconnect_timeout=config.keep_tunnel_open_time, client_disconnect_timeout=config.keep_tunnel_open_time,
) )
@ -488,7 +492,7 @@ class InfectionMonkey:
def _close_tunnel(self): def _close_tunnel(self):
logger.info(f"Quitting tunnel {self._cmd_island_ip}") logger.info(f"Quitting tunnel {self._cmd_island_ip}")
notify_disconnect(self._cmd_island_ip, self._cmd_island_port) notify_disconnect(self._island_address)
def _send_log(self): def _send_log(self):
monkey_log_path = get_agent_log_path() monkey_log_path = get_agent_log_path()

20
monkey/infection_monkey/network/relay/utils.py
View File

@ -92,22 +92,26 @@ def send_remove_from_waitlist_control_message_to_relays(servers: Iterable[str]):
def _send_remove_from_waitlist_control_message_to_relay(server: str): def _send_remove_from_waitlist_control_message_to_relay(server: str):
ip, port = address_to_ip_port(server) ip, port = address_to_ip_port(server)
server_address = SocketAddress(IPv4Address(ip), int(port)) server_address = SocketAddress(IPv4Address(ip), int(port))
notify_disconnect(IPv4Address(ip), int(port)) notify_disconnect(server_address)
def notify_disconnect(server_ip: IPv4Address, server_port: int): def notify_disconnect(server_address: SocketAddress):
""" """
Tell upstream relay that we no longer need the relay. Tell upstream relay that we no longer need the relay
:param server_ip: The IP address of the server to notify. :param server_address: The 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: with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as d_socket:
d_socket.settimeout(LONG_REQUEST_TIMEOUT) d_socket.settimeout(LONG_REQUEST_TIMEOUT)
try: try:
d_socket.connect((str(server_ip), server_port)) d_socket.connect((str(server_address.ip), server_address.port))
d_socket.sendall(RELAY_CONTROL_MESSAGE_REMOVE_FROM_WAITLIST) d_socket.sendall(RELAY_CONTROL_MESSAGE_REMOVE_FROM_WAITLIST)
logger.info(f"Control message was sent to the server/relay {server_ip}:{server_port}") logger.info(
"Control message was sent to the server/relay "
f"{server_address.ip}:{server_address.port}"
)
except OSError as err: except OSError as err:
logger.error(f"Error connecting to socket {server_ip}:{server_port}: {err}") logger.error(
f"Error connecting to socket {server_address.ip}:{server_address.port}: {err}"
)