From 8a609e0871c3768374c047c522692c718a061153 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Mon, 26 Sep 2022 18:13:48 +0530 Subject: [PATCH] Agent: Simplify logic in SocketAddress object creation in monkey.py and network/relay/utils.py --- monkey/infection_monkey/monkey.py | 6 ++---- monkey/infection_monkey/network/relay/utils.py | 4 +--- 2 files changed, 3 insertions(+), 7 deletions(-) diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index 071f742b1..a1cbdc2f2 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -3,7 +3,7 @@ import logging import os import subprocess import sys -from ipaddress import IPv4Address, IPv4Interface +from ipaddress import IPv4Interface from pathlib import Path, WindowsPath from typing import List, Mapping, Optional, Tuple @@ -123,9 +123,7 @@ class InfectionMonkey: self._cmd_island_ip, self._cmd_island_port = address_to_ip_port(server) self._cmd_island_port = int(self._cmd_island_port) - self._island_address = SocketAddress( - IPv4Address(self._cmd_island_ip), self._cmd_island_port - ) + self._island_address = SocketAddress(self._cmd_island_ip, self._cmd_island_port) self._control_client = ControlClient( server_address=server, island_api_client=self._island_api_client diff --git a/monkey/infection_monkey/network/relay/utils.py b/monkey/infection_monkey/network/relay/utils.py index 3eca4ac01..737a07c9a 100644 --- a/monkey/infection_monkey/network/relay/utils.py +++ b/monkey/infection_monkey/network/relay/utils.py @@ -1,7 +1,6 @@ import logging import socket from contextlib import suppress -from ipaddress import IPv4Address from typing import Dict, Iterable, Iterator, Optional from common.common_consts.timeouts import LONG_REQUEST_TIMEOUT @@ -90,8 +89,7 @@ def send_remove_from_waitlist_control_message_to_relays(servers: Iterable[str]): def _send_remove_from_waitlist_control_message_to_relay(server: str): - ip, port = address_to_ip_port(server) - server_address = SocketAddress(IPv4Address(ip), int(port)) + server_address = SocketAddress(*address_to_ip_port(server)) notify_disconnect(server_address)