diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index f8b523b66..26949ab8e 100644 --- a/monkey/infection_monkey/monkey.py +++ b/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._monkey_inbound_tunnel.start() + self._relay.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( - self._monkey_inbound_tunnel, self._cmd_island_ip, self._cmd_island_port, on_island - ) + return VictimHostFactory(None, 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: - self._monkey_inbound_tunnel.stop() - self._monkey_inbound_tunnel.join() + if self._relay and self._relay.is_alive(): + self._relay.stop() + self._relay.join() if firewall.is_enabled(): firewall.remove_firewall_rule() diff --git a/monkey/infection_monkey/network/relay/__init__.py b/monkey/infection_monkey/network/relay/__init__.py index ad5d325c5..8a4ec4731 100644 --- a/monkey/infection_monkey/network/relay/__init__.py +++ b/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 diff --git a/monkey/infection_monkey/network/relay/utils.py b/monkey/infection_monkey/network/relay/utils.py new file mode 100644 index 000000000..e4321ad43 --- /dev/null +++ b/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 diff --git a/monkey/infection_monkey/telemetry/messengers/exploit_intercepting_telemetry_messenger.py b/monkey/infection_monkey/telemetry/messengers/exploit_intercepting_telemetry_messenger.py index 1e5c6591a..f151591f3 100644 --- a/monkey/infection_monkey/telemetry/messengers/exploit_intercepting_telemetry_messenger.py +++ b/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 = relay + self._relay_user_handler = relay_user_handler 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: TCPRelay, + relay_user_handler: RelayUserHandler, ): telemetry_messenger.send_telemetry(telemetry) @@ -36,12 +33,11 @@ def _send_telemetry( def _( telemetry: ExploitTelem, telemetry_messenger: ITelemetryMessenger, - tunnel: MonkeyTunnel, - relay: TCPRelay, + relay_user_handler: RelayUserHandler, ): if telemetry.propagation_result is True: - tunnel.set_wait_for_exploited_machines() - if relay: - relay.add_potential_user(IPv4Address(telemetry.host["ip_addr"])) + if relay_user_handler: + address = IPv4Address(str(telemetry.host["ip_addr"])) + relay_user_handler.add_potential_user(address) telemetry_messenger.send_telemetry(telemetry) diff --git a/monkey/tests/unit_tests/infection_monkey/telemetry/messengers/test_exploit_intercepting_telemetry_messenger.py b/monkey/tests/unit_tests/infection_monkey/telemetry/messengers/test_exploit_intercepting_telemetry_messenger.py index 0758a5e4d..50533fc75 100644 --- a/monkey/tests/unit_tests/infection_monkey/telemetry/messengers/test_exploit_intercepting_telemetry_messenger.py +++ b/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 = MagicMock() + mock_relay_user_handler = 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.add_potential_user.called + assert not mock_relay_user_handler.add_potential_user.called def test_propagation_successful_exploit_telemetry(): mock_telemetry_messenger = MagicMock() - mock_tunnel = MagicMock() - mock_relay = MagicMock() + mock_relay_user_handler = 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.add_potential_user.called + assert mock_relay_user_handler.add_potential_user.called def test_propagation_failed_exploit_telemetry(): mock_telemetry_messenger = MagicMock() - mock_tunnel = MagicMock() - mock_relay = MagicMock() + mock_relay_user_handler = 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.add_potential_user.called + assert not mock_relay_user_handler.add_potential_user.called