Agent: Use SocketAddress in TCPRelay

This commit is contained in:
Shreya Malviya 2022-09-26 13:07:47 +05:30
parent dbaa56c39d
commit 426647c5b9
2 changed files with 6 additions and 5 deletions

View File

@ -21,6 +21,7 @@ from common.network.network_utils import (
get_my_ip_addresses, get_my_ip_addresses,
get_network_interfaces, get_network_interfaces,
) )
from common.types import SocketAddress
from common.utils.argparse_types import positive_int from common.utils.argparse_types import positive_int
from common.utils.attack_utils import ScanStatus, UsageEnum from common.utils.attack_utils import ScanStatus, UsageEnum
from common.version import get_version from common.version import get_version
@ -230,10 +231,10 @@ class InfectionMonkey:
config = self._control_channel.get_config() config = self._control_channel.get_config()
relay_port = get_free_tcp_port() relay_port = get_free_tcp_port()
island_address = SocketAddress(IPv4Address(self._cmd_island_ip), self._cmd_island_port)
self._relay = TCPRelay( self._relay = TCPRelay(
relay_port, relay_port,
IPv4Address(self._cmd_island_ip), island_address,
self._cmd_island_port,
client_disconnect_timeout=config.keep_tunnel_open_time, client_disconnect_timeout=config.keep_tunnel_open_time,
) )

View File

@ -3,6 +3,7 @@ from logging import getLogger
from threading import Lock, Thread from threading import Lock, Thread
from time import sleep from time import sleep
from common.types import SocketAddress
from infection_monkey.network.relay import ( from infection_monkey.network.relay import (
RelayConnectionHandler, RelayConnectionHandler,
RelayUserHandler, RelayUserHandler,
@ -22,15 +23,14 @@ class TCPRelay(Thread, InterruptableThreadMixin):
def __init__( def __init__(
self, self,
relay_port: int, relay_port: int,
dest_addr: IPv4Address, dest_address: SocketAddress,
dest_port: int,
client_disconnect_timeout: float, client_disconnect_timeout: float,
): ):
self._user_handler = RelayUserHandler( self._user_handler = RelayUserHandler(
new_client_timeout=client_disconnect_timeout, new_client_timeout=client_disconnect_timeout,
client_disconnect_timeout=client_disconnect_timeout, client_disconnect_timeout=client_disconnect_timeout,
) )
self._pipe_spawner = TCPPipeSpawner(dest_addr, dest_port) self._pipe_spawner = TCPPipeSpawner(dest_address.ip, dest_address.port)
relay_filter = RelayConnectionHandler(self._pipe_spawner, self._user_handler) relay_filter = RelayConnectionHandler(self._pipe_spawner, self._user_handler)
self._connection_handler = TCPConnectionHandler( self._connection_handler = TCPConnectionHandler(
bind_host="", bind_host="",