relay-throughput-testing/src/throughput_client.py

#!/usr/bin/python3
#
import throughput_protocols
import basic_protocols
import useful
import os
import argparse
import logging
import socket
import time
#
if __name__ == '__main__':
	logging.basicConfig(level=logging.DEBUG)
	#
	parser = argparse.ArgumentParser(description='Test the network throughput (optionally through a proxy).')
	parser.add_argument('ip', type=str, help='destination ip address')
	parser.add_argument('port', type=int, help='destination port')
	parser.add_argument('num_bytes', type=useful.parse_bytes,
	                    help='number of bytes to send (can also end with \'B\', \'KiB\', \'MiB\', or \'GiB\')', metavar='num-bytes')
	parser.add_argument('--proxy', type=str, help='proxy ip address and port', metavar=('ip','port'), nargs=2)
	parser.add_argument('--fake-proxy', action='store_true', help='connecting to a fake-tor proxy')
	parser.add_argument('--wait', type=int,
	                    help='wait until the given time before pushing data (time in seconds since epoch)', metavar='time')
	parser.add_argument('--buffer-len', type=useful.parse_bytes,
	                    help='size of the send and receive buffers (can also end with \'B\', \'KiB\', \'MiB\', or \'GiB\')', metavar='bytes')
	parser.add_argument('--no-accel', action='store_true', help='don\'t use C acceleration (use pure Python)')
	args = parser.parse_args()
	#
	'''
	endpoint = ('127.0.0.1', 4741)
	#endpoint = ('127.0.0.1', 8627)
	#proxy = ('127.0.0.1', 9003+int(sys.argv[3])-1)
	proxy = ('127.0.0.1', int(sys.argv[3]))
	#proxy = ('127.0.0.1', 9003)
	#proxy = ('127.0.0.1', 12849)
	#proxy = None
	'''
	#
	endpoint = (args.ip, args.port)
	client_socket = socket.socket()
	protocols = []
	#
	if args.proxy is None:
		logging.debug('Socket %d connecting to endpoint %r...', client_socket.fileno(), endpoint)
		client_socket.connect(endpoint)
	elif not args.fake_proxy:
		proxy_username = bytes([x for x in os.urandom(12) if x != 0])
		proxy_endpoint = (args.proxy[0], int(args.proxy[1]))
		#
		logging.debug('Socket %d connecting to proxy %r...', client_socket.fileno(), proxy_endpoint)
		client_socket.connect(proxy_endpoint)
		#
		proxy_protocol = basic_protocols.Socks4Protocol(client_socket, endpoint, username=proxy_username)
		protocols.append(proxy_protocol)
	elif args.fake_proxy:
		proxy_endpoint = (args.proxy[0], int(args.proxy[1]))
		#
		logging.debug('Socket %d connecting to fake proxy %r...', client_socket.fileno(), proxy_endpoint)
		client_socket.connect(proxy_endpoint)
		#
		proxy_protocol = basic_protocols.FakeProxyProtocol(client_socket, endpoint)
		protocols.append(proxy_protocol)
	#
	group_id_bytes = args.wait.to_bytes(8, byteorder='big') if args.wait is not None else b''
	if args.wait is not None:
		push_start_cb = lambda: time.sleep(args.wait-time.time())
	else:
		push_start_cb = None
	#
	throughput_protocol = throughput_protocols.ClientProtocol(client_socket, args.num_bytes,
	                                             custom_data=group_id_bytes,
	                                             send_buffer_len=args.buffer_len,
	                                             use_acceleration=(not args.no_accel),
	                                             push_start_cb=push_start_cb)
	protocols.append(throughput_protocol)
	#
	combined_protocol = basic_protocols.ChainedProtocol(protocols)
	combined_protocol.run()
#