Move remote_port to a designated file and add UT

monkey/common/network/network_utils.py

@@ -1,3 +1,4 @@
+import re
 from urllib.parse import urlparse
 
 
@@ -10,3 +11,10 @@ def get_host_from_network_location(network_location: str) -> str:
     """
     url = urlparse("http://" + network_location)
     return str(url.hostname)
+
+
+def remove_port(url):
+    parsed = urlparse(url)
+    with_port = f'{parsed.scheme}://{parsed.netloc}'
+    without_port = re.sub(':[0-9]+$', '', with_port)
+    return without_port

monkey/common/network/test_network_utils.py

@@ -1,12 +1,17 @@
 from unittest import TestCase
 
-from common.network.network_utils import get_host_from_network_location
+from common.network.network_utils import get_host_from_network_location, remove_port
 
 
 class TestNetworkUtils(TestCase):
-    def test_remove_port_from_ip_string(self):
+    def test_get_host_from_network_location(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"
+    
+    def test_remove_port_from_url(self):
+        assert remove_port('https://google.com:80') == 'https://google.com'
+        assert remove_port('https://8.8.8.8:65336') == 'https://8.8.8.8'
+        assert remove_port('ftp://ftpserver.com:21/hello/world') == 'ftp://ftpserver.com'

monkey/infection_monkey/exploit/drupal.py

@@ -5,23 +5,16 @@ Implementation is based on:
 """
 
 import logging
-import re
 import requests
-from urllib.parse import urljoin, urlparse
+from urllib.parse import urljoin
 from infection_monkey.exploit.web_rce import WebRCE
+from network.network_utils import remove_port
 
 __author__ = 'Ophir Harpaz'
 
 LOG = logging.getLogger(__name__)
 
 
-def remove_port(url):
-    parsed = urlparse(url)
-    with_port = f'{parsed.scheme}://{parsed.netloc}'
-    without_port = re.sub(':[0-9]+$', '', with_port)
-    return without_port
-
-
 def build_url(*args) -> str:
     f = ''
     for x in args: