Moved a function into common, since Monkey doesn't have ring as a dependency

Also renamed it and added UTs
This commit is contained in:
Shay Nehmad 2020-06-03 16:18:19 +03:00
parent ca87ff1330
commit 9ea6718d37
5 changed files with 26 additions and 8 deletions

View File

@ -0,0 +1,12 @@
from urllib.parse import urlparse
def get_host_from_network_location(network_location: str) -> str:
"""
URL structure is "<scheme>://<net_loc>/<path>;<params>?<query>#<fragment>" (https://tools.ietf.org/html/rfc1808.html)
And the net_loc is "<user>:<password>@<host>:<port>" (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)

View File

@ -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"

View File

@ -33,7 +33,7 @@ from infection_monkey.telemetry.attack.t1106_telem import T1106Telem
from common.utils.attack_utils import ScanStatus, UsageEnum from common.utils.attack_utils import ScanStatus, UsageEnum
from common.version import get_version from common.version import get_version
from infection_monkey.exploit.HostExploiter import HostExploiter 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" 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) LOG.debug("default server set to: %s" % self._default_server)
def is_started_on_island(self): 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 return is_running_on_server(island_ip) and WormConfiguration.depth == WormConfiguration.max_depth
def log_arguments(self): def log_arguments(self):

View File

@ -5,7 +5,6 @@ import socket
import struct import struct
import sys import sys
from typing import List from typing import List
from urllib.parse import urlparse
from netifaces import interfaces, ifaddresses, AF_INET from netifaces import interfaces, ifaddresses, AF_INET
from ring import lru from ring import lru
@ -86,8 +85,3 @@ def get_subnets():
] ]
) )
return subnets return subnets
def remove_port_from_ip_string(ip_string: str) -> str:
url = urlparse("http://" + ip_string)
return str(url.hostname)