#!/usr/bin/python3
#
import throughput_protocols
import basic_protocols
import useful
import os
import argparse
import logging
import socket
#
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('--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 = (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)
	else:
		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)
	#
	throughput_protocol = throughput_protocols.ClientProtocol(client_socket, args.num_bytes,
	                                             wait_until=args.wait,
	                                             send_buffer_len=args.buffer_len,
	                                             use_acceleration=(not args.no_accel))
	protocols.append(throughput_protocol)
	#
	combined_protocol = basic_protocols.ChainedProtocol(protocols)
	combined_protocol.run()
#