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)