Agent, Common, UT: Separate IP and Port in monkey

Instead of splitting IP/port on demand, separate the IP and port from monkey commandline parameter and pass them to VictimHostFactory

monkey/common/network/network_utils.py

@@ -1,4 +1,5 @@
 import re
+from typing import Optional, Tuple
 from urllib.parse import urlparse
 
 
@@ -20,3 +21,11 @@ def remove_port(url):
     with_port = f"{parsed.scheme}://{parsed.netloc}"
     without_port = re.sub(":[0-9]+(?=$|/)", "", with_port)
     return without_port
+
+
+def address_to_ip_port(address: str) -> Tuple[str, Optional[str]]:
+    if ":" in address:
+        ip, port = address.split(":")
+        return ip, port or None
+    else:
+        return address, None

monkey/infection_monkey/model/host.py

@@ -1,3 +1,6 @@
+from typing import Optional
+
+
 class VictimHost(object):
     def __init__(self, ip_addr: str, domain_name: str = ""):
         self.ip_addr = ip_addr
@@ -42,5 +45,5 @@ class VictimHost(object):
         victim += "target monkey: %s" % self.monkey_exe
         return victim
 
-    def set_default_server(self, default_server):
-        self.default_server = default_server
+    def set_island_address(self, ip: str, port: Optional[str]):
+        self.default_server = f"{ip}:{port}" if port else f"{ip}"

monkey/infection_monkey/model/victim_host_factory.py

@@ -1,5 +1,5 @@
 import logging
-from typing import Optional
+from typing import Optional, Tuple
 
 from infection_monkey.model import VictimHost
 from infection_monkey.network import NetworkAddress
@@ -13,13 +13,13 @@ class VictimHostFactory:
     def __init__(
         self,
         tunnel: Optional[MonkeyTunnel],
-        default_server: Optional[str],
-        default_port: Optional[str],
+        island_ip: Optional[str],
+        island_port: Optional[str],
         on_island: bool,
     ):
         self.tunnel = tunnel
-        self.default_server = default_server
-        self.default_port = default_port
+        self.island_ip = island_ip
+        self.island_port = island_port
         self.on_island = on_island
 
     def build_victim_host(self, network_address: NetworkAddress) -> VictimHost:
@@ -29,19 +29,22 @@ class VictimHostFactory:
         if self.tunnel:
             victim_host.default_tunnel = self.tunnel.get_tunnel_for_ip(victim_host.ip_addr)
 
-        if self.default_server:
-            victim_host.set_default_server(self._get_formatted_default_server(victim_host.ip_addr))
+        if self.island_ip:
+            ip, port = self._choose_island_address(victim_host.ip_addr)
+            victim_host.set_island_address(ip, port)
 
         logger.debug(f"Default tunnel for {victim_host} set to {victim_host.default_tunnel}")
         logger.debug(f"Default server for {victim_host} set to {victim_host.default_server}")
 
         return victim_host
 
-    def _get_formatted_default_server(self, ip: str):
+    def _choose_island_address(self, victim_ip: str) -> Tuple[str, Optional[str]]:
+        # Victims need to connect back to the interface they can reach
+        # On island, choose the right interface to pass to children monkeys
         if self.on_island:
-            default_server_port = f":{self.default_port}" if self.default_port else ""
-            interface = get_interface_to_target(ip)
+            default_server_port = self.island_port if self.island_port else None
+            interface = get_interface_to_target(victim_ip)
 
-            return f"{interface}{default_server_port}"
+            return interface, default_server_port
         else:
-            return self.default_server
+            return self.island_ip, self.island_port

monkey/infection_monkey/monkey.py

@@ -7,6 +7,7 @@ import time
 from typing import List
 
 import infection_monkey.tunnel as tunnel
+from common.network.network_utils import address_to_ip_port
 from common.utils.attack_utils import ScanStatus, UsageEnum
 from common.version import get_version
 from infection_monkey.config import GUID, WormConfiguration
@@ -40,8 +41,8 @@ class InfectionMonkey:
         logger.info("Monkey is initializing...")
         self._singleton = SystemSingleton()
         self._opts = self._get_arguments(args)
+        self._cmd_island_ip, self._cmd_island_port = address_to_ip_port(self._opts.server)
         self._default_server = self._opts.server
-        self._default_server_port = None
         # TODO used in propogation phase
         self._monkey_inbound_tunnel = None
 
@@ -119,8 +120,6 @@ class InfectionMonkey:
                 "Monkey couldn't find server with {} default tunnel.".format(self._opts.tunnel)
             )
 
-        self._set_default_port()
-
         ControlClient.wakeup(parent=self._opts.parent)
         ControlClient.load_control_config()
 
@@ -185,7 +184,7 @@ class InfectionMonkey:
         logger.debug(f"This agent is running on the island: {on_island}")
 
         return VictimHostFactory(
-            self._monkey_inbound_tunnel, self._default_server, self._default_server_port, on_island
+            self._monkey_inbound_tunnel, self._cmd_island_ip, self._cmd_island_port, on_island
         )
 
     def _running_on_island(self, local_network_interfaces: List[NetworkInterface]) -> bool:
@@ -195,12 +194,6 @@ class InfectionMonkey:
     def _is_another_monkey_running(self):
         return not self._singleton.try_lock()
 
-    def _set_default_port(self):
-        try:
-            self._default_server_port = self._default_server.split(":")[1]
-        except KeyError:
-            self._default_server_port = ""
-
     def cleanup(self):
         logger.info("Monkey cleanup started")
         self._wait_for_exploited_machine_connection()

monkey/tests/unit_tests/common/network/test_network_utils.py

@@ -1,6 +1,10 @@
 from unittest import TestCase
 
-from common.network.network_utils import get_host_from_network_location, remove_port
+from common.network.network_utils import (
+    address_to_ip_port,
+    get_host_from_network_location,
+    remove_port,
+)
 
 
 class TestNetworkUtils(TestCase):
@@ -15,3 +19,17 @@ class TestNetworkUtils(TestCase):
         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"
+
+
+def test_address_to_ip_port():
+    ip, port = address_to_ip_port("192.168.65.1:5000")
+    assert ip == "192.168.65.1"
+    assert port == "5000"
+
+
+def test_address_to_ip_port_no_port():
+    ip, port = address_to_ip_port("192.168.65.1")
+    assert port is None
+
+    ip, port = address_to_ip_port("192.168.65.1:")
+    assert port is None

monkey/tests/unit_tests/infection_monkey/model/test_victim_host_factory.py

@@ -22,13 +22,13 @@ def mock_get_interface_to_target(monkeypatch):
 
 def test_factory_no_tunnel():
     factory = VictimHostFactory(
-        tunnel=None, default_server="192.168.56.1", default_port="5000", on_island=False
+        tunnel=None, island_ip="192.168.56.1", island_port="5000", on_island=False
     )
     network_address = NetworkAddress("192.168.56.2", None)
 
     victim = factory.build_victim_host(network_address)
 
-    assert victim.default_server == "192.168.56.1"
+    assert victim.default_server == "192.168.56.1:5000"
     assert victim.ip_addr == "192.168.56.2"
     assert victim.default_tunnel is None
     assert victim.domain_name == ""
@@ -36,13 +36,13 @@ def test_factory_no_tunnel():
 
 def test_factory_with_tunnel(mock_tunnel):
     factory = VictimHostFactory(
-        tunnel=mock_tunnel, default_server="192.168.56.1", default_port="5000", on_island=False
+        tunnel=mock_tunnel, island_ip="192.168.56.1", island_port="5000", on_island=False
     )
     network_address = NetworkAddress("192.168.56.2", None)
 
     victim = factory.build_victim_host(network_address)
 
-    assert victim.default_server == "192.168.56.1"
+    assert victim.default_server == "192.168.56.1:5000"
     assert victim.ip_addr == "192.168.56.2"
     assert victim.default_tunnel == "1.2.3.4:1234"
     assert victim.domain_name == ""
@@ -50,7 +50,7 @@ def test_factory_with_tunnel(mock_tunnel):
 
 def test_factory_on_island(mock_tunnel):
     factory = VictimHostFactory(
-        tunnel=mock_tunnel, default_server="192.168.56.1", default_port="99", on_island=True
+        tunnel=mock_tunnel, island_ip="192.168.56.1", island_port="99", on_island=True
     )
     network_address = NetworkAddress("192.168.56.2", "www.bogus.monkey")
 
@@ -65,7 +65,7 @@ def test_factory_on_island(mock_tunnel):
 @pytest.mark.parametrize("default_port", ["", None])
 def test_factory_no_port(mock_tunnel, default_port):
     factory = VictimHostFactory(
-        tunnel=mock_tunnel, default_server="192.168.56.1", default_port=default_port, on_island=True
+        tunnel=mock_tunnel, island_ip="192.168.56.1", island_port=default_port, on_island=True
     )
     network_address = NetworkAddress("192.168.56.2", "www.bogus.monkey")
 
@@ -75,9 +75,7 @@ def test_factory_no_port(mock_tunnel, default_port):
 
 
 def test_factory_no_default_server(mock_tunnel):
-    factory = VictimHostFactory(
-        tunnel=mock_tunnel, default_server=None, default_port="", on_island=True
-    )
+    factory = VictimHostFactory(tunnel=mock_tunnel, island_ip=None, island_port="", on_island=True)
     network_address = NetworkAddress("192.168.56.2", "www.bogus.monkey")
 
     victim = factory.build_victim_host(network_address)

monkey/tests/unit_tests/infection_monkey/utils/test_commands.py

@@ -96,9 +96,9 @@ def test_get_monkey_commandline_linux():
 
 def test_build_monkey_commandline():
     example_host = VictimHost(ip_addr="bla")
-    example_host.set_default_server("101010")
+    example_host.set_island_address("101010", "5000")
 
-    expected = f" -p {GUID} -s 101010 -d 0 -l /home/bla"
+    expected = f" -p {GUID} -s 101010:5000 -d 0 -l /home/bla"
     actual = build_monkey_commandline(target_host=example_host, depth=0, location="/home/bla")
 
     assert expected == actual