diff --git a/monkey/common/network/network_utils.py b/monkey/common/network/network_utils.py index 2b01d1974..3c87d5737 100644 --- a/monkey/common/network/network_utils.py +++ b/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 diff --git a/monkey/infection_monkey/model/host.py b/monkey/infection_monkey/model/host.py index 4331bcf7e..3bbd1dfb8 100644 --- a/monkey/infection_monkey/model/host.py +++ b/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}" diff --git a/monkey/infection_monkey/model/victim_host_factory.py b/monkey/infection_monkey/model/victim_host_factory.py index 5e49e5ffd..a6b56532e 100644 --- a/monkey/infection_monkey/model/victim_host_factory.py +++ b/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 diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index e2e3b4253..f26c92b3b 100644 --- a/monkey/infection_monkey/monkey.py +++ b/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() diff --git a/monkey/tests/unit_tests/common/network/test_network_utils.py b/monkey/tests/unit_tests/common/network/test_network_utils.py index 0376cd6d5..e7d82e649 100644 --- a/monkey/tests/unit_tests/common/network/test_network_utils.py +++ b/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 diff --git a/monkey/tests/unit_tests/infection_monkey/model/test_victim_host_factory.py b/monkey/tests/unit_tests/infection_monkey/model/test_victim_host_factory.py index 2b5250c8c..2b7c10864 100644 --- a/monkey/tests/unit_tests/infection_monkey/model/test_victim_host_factory.py +++ b/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) diff --git a/monkey/tests/unit_tests/infection_monkey/utils/test_commands.py b/monkey/tests/unit_tests/infection_monkey/utils/test_commands.py index 5d33cb8ae..db9ddbbe7 100644 --- a/monkey/tests/unit_tests/infection_monkey/utils/test_commands.py +++ b/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