async def test_socks5_sendto_fail(socks5_server):
"""
Test if sending a UDP packet without a successful association fails.
"""
await socks5_server.start()
client = Socks5Client(('127.0.0.1', socks5_server.port), Mock())
with pytest.raises(Socks5Error):
client.sendto(b'\x00', ('127.0.0.1', 123))
async def test_socks5_tcp_connect(socks5_server):
"""
Test is sending a TCP connect request to the server succeeds.
"""
await socks5_server.start()
client = Socks5Client(('127.0.0.1', socks5_server.port), Mock())
await client.connect_tcp(('127.0.0.1', 123))
assert client.transport is not None
assert client.connection is None
async def test_socks5_udp_associate(socks5_server):
"""
Test is sending a UDP associate request to the server succeeds.
"""
await socks5_server.start()
client = Socks5Client(('127.0.0.1', socks5_server.port), Mock())
await client.associate_udp()
assert client.transport is not None
assert client.connection is not None
assert client.connection.transport is not None
async def test_socks5_write(socks5_server):
"""
Test is sending a TCP data to the server succeeds.
"""
await socks5_server.start()
client = Socks5Client(('127.0.0.1', socks5_server.port), Mock())
await client.connect_tcp(('127.0.0.1', 123))
client.write(b' ')
await sleep(.1)
socks5_server.output_stream.on_socks5_tcp_data.assert_called_once_with(socks5_server.sessions[0],
('127.0.0.1', 123), b' ')
async def _wrap_create_connection(self, protocol_factory, host, port,
**kwargs):
client = Socks5Client(self.proxy_addr, lambda *_: None)
if 'timeout' in kwargs and hasattr(kwargs['timeout'], 'sock_connect'):
await wait_for(client.connect_tcp((host, port)),
timeout=kwargs['timeout'].sock_connect)
else:
await client.connect_tcp((host, port))
proto = protocol_factory()
transport = client.transport
transport._protocol = proto # pylint: disable=W0212
proto.transport = transport
return transport, proto
async def send_request(self, data, tracker_session):
transport = self.transport
proxy = tracker_session.proxy
if proxy:
transport = self.proxy_transports.get(
proxy, Socks5Client(proxy, self.datagram_received))
if not transport.associated:
await transport.associate_udp()
if proxy not in self.proxy_transports:
self.proxy_transports[proxy] = transport
try:
transport.sendto(
data, (tracker_session.ip_address, tracker_session.port))
f = self.tracker_sessions[
tracker_session.transaction_id] = Future()
return await f
except OSError as e:
self._logger.warning("Unable to write data to %s:%d - %s",
tracker_session.ip_address,
tracker_session.port, e)
return RuntimeError("Unable to write to socket - " + str(e))
async def test_socks5_sendto_success(socks5_server):
"""
Test if sending/receiving a UDP packet works correctly.
"""
await socks5_server.start()
data = b'\x00'
target = ('127.0.0.1', 123)
client = Socks5Client(('127.0.0.1', socks5_server.port), Mock())
await client.associate_udp()
client.sendto(data, target)
await sleep(0.1)
socks5_server.output_stream.on_socks5_udp_data.assert_called_once()
connection = socks5_server.output_stream.on_socks5_udp_data.call_args[0][0]
request = socks5_server.output_stream.on_socks5_udp_data.call_args[0][1]
assert request.data == data
assert request.destination == target
packet = socks5_serializer.pack_serializable(UdpPacket(0, 0, target, data))
client.callback.assert_not_called()
connection.send_datagram(packet)
await sleep(0.1)
client.callback.assert_called_once_with(data, target)