Agent: Replace MonkeyTunnel with TCPRelay

monkey/infection_monkey/monkey.py

@@ -3,7 +3,7 @@ import logging
 import os
 import subprocess
 import sys
-from ipaddress import IPv4Interface
+from ipaddress import IPv4Address, IPv4Interface
 from pathlib import Path, WindowsPath
 from typing import List
 
@@ -41,7 +41,13 @@ from infection_monkey.master import AutomatedMaster
 from infection_monkey.master.control_channel import ControlChannel
 from infection_monkey.model import VictimHostFactory
 from infection_monkey.network.firewall import app as firewall
-from infection_monkey.network.info import get_network_interfaces
+from infection_monkey.network.info import get_free_tcp_port, get_network_interfaces
+from infection_monkey.network.relay import (
+    build_tcprelay_deps,
+    RelayUserHandler,
+    TCPRelay,
+)
+from infection_monkey.network.tools import connect
 from infection_monkey.network_scanning.elasticsearch_fingerprinter import ElasticSearchFingerprinter
 from infection_monkey.network_scanning.http_fingerprinter import HTTPFingerprinter
 from infection_monkey.network_scanning.mssql_fingerprinter import MSSQLFingerprinter
@@ -100,11 +106,11 @@ class InfectionMonkey:
         # TODO Refactor the telemetry messengers to accept control client
         # and remove control_client_object
         ControlClient.control_client_object = self._control_client
-        self._monkey_inbound_tunnel = None
         self._telemetry_messenger = LegacyTelemetryMessengerAdapter()
         self._current_depth = self._opts.depth
         self._master = None
-        self._inbound_tunnel_opened = False
+        self._relay_user_handler: RelayUserHandler
+        self._relay: TCPRelay
 
     @staticmethod
     def _get_arguments(args):
@@ -180,14 +186,27 @@ class InfectionMonkey:
         control_channel.register_agent(self._opts.parent)
 
         config = control_channel.get_config()
-        self._monkey_inbound_tunnel = self._control_client.create_control_tunnel(
+
-            config.keep_tunnel_open_time
+        local_port = get_free_tcp_port()
+        sock, ip_str, port = connect([self._opts.server])
+        sock.close()
+        user_handler, connection_handler, pipe_spawner = build_tcprelay_deps(
+            local_port,
+            IPv4Address(ip_str),
+            port,
+            client_disconnect_timeout=config.keep_tunnel_open_time,
         )
-        if self._monkey_inbound_tunnel and maximum_depth_reached(
+        self._relay_user_handler = user_handler
+        self._relay = TCPRelay(
+            self._relay_user_handler,
+            connection_handler,
+            pipe_spawner,
+        )
+
+        if self._relay and maximum_depth_reached(
             config.propagation.maximum_depth, self._current_depth
         ):
-            self._inbound_tunnel_opened = True
+            self._relay.start()
-            self._monkey_inbound_tunnel.start()
 
         StateTelem(is_done=False, version=get_version()).send()
         TunnelTelem(self._control_client.proxies).send()
@@ -215,7 +234,7 @@ class InfectionMonkey:
         victim_host_factory = self._build_victim_host_factory(local_network_interfaces)
 
         telemetry_messenger = ExploitInterceptingTelemetryMessenger(
-            self._telemetry_messenger, self._monkey_inbound_tunnel
+            self._telemetry_messenger, self._relay_user_handler
         )
 
         self._master = AutomatedMaster(
@@ -374,9 +393,7 @@ class InfectionMonkey:
         on_island = self._running_on_island(local_network_interfaces)
         logger.debug(f"This agent is running on the island: {on_island}")
 
-        return VictimHostFactory(
+        return VictimHostFactory(None, self._cmd_island_ip, self._cmd_island_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[IPv4Interface]) -> bool:
         server_ip, _ = address_to_ip_port(self._control_client.server_address)
@@ -394,9 +411,9 @@ class InfectionMonkey:
 
             reset_signal_handlers()
 
-            if self._inbound_tunnel_opened:
+            if self._relay and self._relay.is_alive():
-                self._monkey_inbound_tunnel.stop()
+                self._relay.stop()
-                self._monkey_inbound_tunnel.join()
+                self._relay.join()
 
             if firewall.is_enabled():
                 firewall.remove_firewall_rule()

monkey/infection_monkey/network/relay/__init__.py

@@ -3,3 +3,5 @@ from .relay_user_handler import RelayUser, RelayUserHandler
 from .sockets_pipe import SocketsPipe
 from .tcp_connection_handler import TCPConnectionHandler
 from .tcp_pipe_spawner import TCPPipeSpawner
+from .tcp_relay import TCPRelay
+from .utils import build_tcprelay_deps

monkey/infection_monkey/network/relay/utils.py

@@ -0,0 +1,22 @@
+from . import TCPConnectionHandler, TCPPipeSpawner, RelayUserHandler, RelayConnectionHandler
+from ipaddress import IPv4Address
+from typing import Tuple
+
+
+def build_tcprelay_deps(
+    local_port: int, dest_addr: IPv4Address, dest_port: int, client_disconnect_timeout: float
+) -> Tuple[RelayUserHandler, TCPPipeSpawner, TCPConnectionHandler]:
+
+    # TODO: Add the timeouts
+    relay_user_handler = RelayUserHandler()
+    pipe_spawner = TCPPipeSpawner(dest_addr, dest_port)
+    relay_filter = RelayConnectionHandler(pipe_spawner, relay_user_handler)
+    connection_handler = TCPConnectionHandler(
+        bind_host="",
+        bind_port=local_port,
+        client_connected=[
+            relay_filter.handle_new_connection,
+        ],
+    )
+
+    return relay_user_handler, pipe_spawner, connection_handler

monkey/infection_monkey/telemetry/messengers/exploit_intercepting_telemetry_messenger.py

@@ -1,23 +1,21 @@
 from functools import singledispatch
 from ipaddress import IPv4Address
 
-from infection_monkey.network.relay.tcp_relay import TCPRelay
+from infection_monkey.network.relay import RelayUserHandler
 from infection_monkey.telemetry.exploit_telem import ExploitTelem
 from infection_monkey.telemetry.i_telem import ITelem
 from infection_monkey.telemetry.messengers.i_telemetry_messenger import ITelemetryMessenger
-from infection_monkey.tunnel import MonkeyTunnel
 
 
 class ExploitInterceptingTelemetryMessenger(ITelemetryMessenger):
     def __init__(
-        self, telemetry_messenger: ITelemetryMessenger, tunnel: MonkeyTunnel, relay: TCPRelay
+        self, telemetry_messenger: ITelemetryMessenger, relay_user_handler: RelayUserHandler
     ):
         self._telemetry_messenger = telemetry_messenger
-        self._tunnel = tunnel
+        self._relay_user_handler = relay_user_handler
-        self._relay = relay
 
     def send_telemetry(self, telemetry: ITelem):
-        _send_telemetry(telemetry, self._telemetry_messenger, self._tunnel, self._relay)
+        _send_telemetry(telemetry, self._telemetry_messenger, self._relay_user_handler)
 
 
 # Note: We can use @singledispatchmethod instead of @singledispatch if we migrate to Python 3.8 or
@@ -26,8 +24,7 @@ class ExploitInterceptingTelemetryMessenger(ITelemetryMessenger):
 def _send_telemetry(
     telemetry: ITelem,
     telemetry_messenger: ITelemetryMessenger,
-    tunnel: MonkeyTunnel,
+    relay_user_handler: RelayUserHandler,
-    relay: TCPRelay,
 ):
     telemetry_messenger.send_telemetry(telemetry)
 
@@ -36,12 +33,11 @@ def _send_telemetry(
 def _(
     telemetry: ExploitTelem,
     telemetry_messenger: ITelemetryMessenger,
-    tunnel: MonkeyTunnel,
+    relay_user_handler: RelayUserHandler,
-    relay: TCPRelay,
 ):
     if telemetry.propagation_result is True:
-        tunnel.set_wait_for_exploited_machines()
+        if relay_user_handler:
-        if relay:
+            address = IPv4Address(str(telemetry.host["ip_addr"]))
-            relay.add_potential_user(IPv4Address(telemetry.host["ip_addr"]))
+            relay_user_handler.add_potential_user(address)
 
     telemetry_messenger.send_telemetry(telemetry)

monkey/tests/unit_tests/infection_monkey/telemetry/messengers/test_exploit_intercepting_telemetry_messenger.py

@@ -20,49 +20,43 @@ class MockExploitTelem(ExploitTelem):
 
 def test_generic_telemetry(TestTelem):
     mock_telemetry_messenger = MagicMock()
-    mock_tunnel = MagicMock()
+    mock_relay_user_handler = MagicMock()
-    mock_relay = MagicMock()
 
     telemetry_messenger = ExploitInterceptingTelemetryMessenger(
-        mock_telemetry_messenger, mock_tunnel, mock_relay
+        mock_telemetry_messenger, mock_relay_user_handler
     )
 
     telemetry_messenger.send_telemetry(TestTelem())
 
     assert mock_telemetry_messenger.send_telemetry.called
-    assert not mock_tunnel.set_wait_for_exploited_machines.called
+    assert not mock_relay_user_handler.add_potential_user.called
-    assert not mock_relay.add_potential_user.called
 
 
 def test_propagation_successful_exploit_telemetry():
     mock_telemetry_messenger = MagicMock()
-    mock_tunnel = MagicMock()
+    mock_relay_user_handler = MagicMock()
-    mock_relay = MagicMock()
     mock_exploit_telem = MockExploitTelem(True)
 
     telemetry_messenger = ExploitInterceptingTelemetryMessenger(
-        mock_telemetry_messenger, mock_tunnel, mock_relay
+        mock_telemetry_messenger, mock_relay_user_handler
     )
 
     telemetry_messenger.send_telemetry(mock_exploit_telem)
 
     assert mock_telemetry_messenger.send_telemetry.called
-    assert mock_tunnel.set_wait_for_exploited_machines.called
+    assert mock_relay_user_handler.add_potential_user.called
-    assert mock_relay.add_potential_user.called
 
 
 def test_propagation_failed_exploit_telemetry():
     mock_telemetry_messenger = MagicMock()
-    mock_tunnel = MagicMock()
+    mock_relay_user_handler = MagicMock()
-    mock_relay = MagicMock()
     mock_exploit_telem = MockExploitTelem(False)
 
     telemetry_messenger = ExploitInterceptingTelemetryMessenger(
-        mock_telemetry_messenger, mock_tunnel, mock_relay
+        mock_telemetry_messenger, mock_relay_user_handler
     )
 
     telemetry_messenger.send_telemetry(mock_exploit_telem)
 
     assert mock_telemetry_messenger.send_telemetry.called
-    assert not mock_tunnel.set_wait_for_exploited_machines.called
+    assert not mock_relay_user_handler.add_potential_user.called
-    assert not mock_relay.add_potential_user.called