From 9ea6718d37b5ea00ea46ab132fe681b552a21818 Mon Sep 17 00:00:00 2001 From: Shay Nehmad Date: Wed, 3 Jun 2020 16:18:19 +0300 Subject: [PATCH] Moved a function into common, since Monkey doesn't have ring as a dependency Also renamed it and added UTs --- monkey/common/network/network_utils.py | 12 ++++++++++++ monkey/common/network/test_network_utils.py | 12 ++++++++++++ ...tion_utils_test.py => test_segmentation_utils.py} | 0 monkey/infection_monkey/monkey.py | 4 ++-- monkey/monkey_island/cc/network_utils.py | 6 ------ 5 files changed, 26 insertions(+), 8 deletions(-) create mode 100644 monkey/common/network/network_utils.py create mode 100644 monkey/common/network/test_network_utils.py rename monkey/common/network/{segmentation_utils_test.py => test_segmentation_utils.py} (100%) diff --git a/monkey/common/network/network_utils.py b/monkey/common/network/network_utils.py new file mode 100644 index 000000000..230e494bf --- /dev/null +++ b/monkey/common/network/network_utils.py @@ -0,0 +1,12 @@ +from urllib.parse import urlparse + + +def get_host_from_network_location(network_location: str) -> str: + """ + URL structure is ":///;?#" (https://tools.ietf.org/html/rfc1808.html) + And the net_loc is ":@:" (https://tools.ietf.org/html/rfc1738#section-3.1) + :param network_location: server network location + :return: host part of the network location + """ + url = urlparse("http://" + network_location) + return str(url.hostname) diff --git a/monkey/common/network/test_network_utils.py b/monkey/common/network/test_network_utils.py new file mode 100644 index 000000000..b4194aa27 --- /dev/null +++ b/monkey/common/network/test_network_utils.py @@ -0,0 +1,12 @@ +from unittest import TestCase + +from common.network.network_utils import get_host_from_network_location + + +class TestNetworkUtils(TestCase): + def test_remove_port_from_ip_string(self): + assert get_host_from_network_location("127.0.0.1:12345") == "127.0.0.1" + assert get_host_from_network_location("127.0.0.1:12345") == "127.0.0.1" + assert get_host_from_network_location("127.0.0.1") == "127.0.0.1" + assert get_host_from_network_location("www.google.com:8080") == "www.google.com" + assert get_host_from_network_location("user:password@host:8080") == "host" diff --git a/monkey/common/network/segmentation_utils_test.py b/monkey/common/network/test_segmentation_utils.py similarity index 100% rename from monkey/common/network/segmentation_utils_test.py rename to monkey/common/network/test_segmentation_utils.py diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index a31ea4d47..3495fd7cc 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -33,7 +33,7 @@ from infection_monkey.telemetry.attack.t1106_telem import T1106Telem from common.utils.attack_utils import ScanStatus, UsageEnum from common.version import get_version from infection_monkey.exploit.HostExploiter import HostExploiter -from monkey_island.cc.network_utils import remove_port_from_ip_string +from common.network.network_utils import get_host_from_network_location MAX_DEPTH_REACHED_MESSAGE = "Reached max depth, shutting down" @@ -386,7 +386,7 @@ class InfectionMonkey(object): LOG.debug("default server set to: %s" % self._default_server) def is_started_on_island(self): - island_ip = remove_port_from_ip_string(self._default_server) + island_ip = get_host_from_network_location(self._default_server) return is_running_on_server(island_ip) and WormConfiguration.depth == WormConfiguration.max_depth def log_arguments(self): diff --git a/monkey/monkey_island/cc/network_utils.py b/monkey/monkey_island/cc/network_utils.py index 903485723..d399d4255 100644 --- a/monkey/monkey_island/cc/network_utils.py +++ b/monkey/monkey_island/cc/network_utils.py @@ -5,7 +5,6 @@ import socket import struct import sys from typing import List -from urllib.parse import urlparse from netifaces import interfaces, ifaddresses, AF_INET from ring import lru @@ -86,8 +85,3 @@ def get_subnets(): ] ) return subnets - - -def remove_port_from_ip_string(ip_string: str) -> str: - url = urlparse("http://" + ip_string) - return str(url.hostname)