throughput_client.py 3.3 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879
  1. #!/usr/bin/python3
  2. #
  3. import throughput_protocols
  4. import basic_protocols
  5. import useful
  6. import os
  7. import argparse
  8. import logging
  9. import socket
  10. import time
  11. #
  12. if __name__ == '__main__':
  13. logging.basicConfig(level=logging.DEBUG)
  14. #
  15. parser = argparse.ArgumentParser(description='Test the network throughput (optionally through a proxy).')
  16. parser.add_argument('ip', type=str, help='destination ip address')
  17. parser.add_argument('port', type=int, help='destination port')
  18. parser.add_argument('num_bytes', type=useful.parse_bytes,
  19. help='number of bytes to send (can also end with \'B\', \'KiB\', \'MiB\', or \'GiB\')', metavar='num-bytes')
  20. parser.add_argument('--proxy', type=str, help='proxy ip address and port', metavar=('ip','port'), nargs=2)
  21. parser.add_argument('--fake-proxy', action='store_true', help='connecting to a fake-tor proxy')
  22. parser.add_argument('--wait', type=int,
  23. help='wait until the given time before pushing data (time in seconds since epoch)', metavar='time')
  24. parser.add_argument('--buffer-len', type=useful.parse_bytes,
  25. help='size of the send and receive buffers (can also end with \'B\', \'KiB\', \'MiB\', or \'GiB\')', metavar='bytes')
  26. parser.add_argument('--no-accel', action='store_true', help='don\'t use C acceleration (use pure Python)')
  27. args = parser.parse_args()
  28. #
  29. '''
  30. endpoint = ('127.0.0.1', 4741)
  31. #endpoint = ('127.0.0.1', 8627)
  32. #proxy = ('127.0.0.1', 9003+int(sys.argv[3])-1)
  33. proxy = ('127.0.0.1', int(sys.argv[3]))
  34. #proxy = ('127.0.0.1', 9003)
  35. #proxy = ('127.0.0.1', 12849)
  36. #proxy = None
  37. '''
  38. #
  39. endpoint = (args.ip, args.port)
  40. client_socket = socket.socket()
  41. protocols = []
  42. #
  43. if args.proxy is None:
  44. logging.debug('Socket %d connecting to endpoint %r...', client_socket.fileno(), endpoint)
  45. client_socket.connect(endpoint)
  46. elif not args.fake_proxy:
  47. proxy_username = bytes([x for x in os.urandom(12) if x != 0])
  48. proxy_endpoint = (args.proxy[0], int(args.proxy[1]))
  49. #
  50. logging.debug('Socket %d connecting to proxy %r...', client_socket.fileno(), proxy_endpoint)
  51. client_socket.connect(proxy_endpoint)
  52. #
  53. proxy_protocol = basic_protocols.Socks4Protocol(client_socket, endpoint, username=proxy_username)
  54. protocols.append(proxy_protocol)
  55. elif args.fake_proxy:
  56. proxy_endpoint = (args.proxy[0], int(args.proxy[1]))
  57. #
  58. logging.debug('Socket %d connecting to fake proxy %r...', client_socket.fileno(), proxy_endpoint)
  59. client_socket.connect(proxy_endpoint)
  60. #
  61. proxy_protocol = basic_protocols.FakeProxyProtocol(client_socket, endpoint)
  62. protocols.append(proxy_protocol)
  63. #
  64. group_id_bytes = args.wait.to_bytes(8, byteorder='big') if args.wait is not None else b''
  65. if args.wait is not None:
  66. push_start_cb = lambda: time.sleep(args.wait-time.time())
  67. else:
  68. push_start_cb = None
  69. #
  70. throughput_protocol = throughput_protocols.ClientProtocol(client_socket, args.num_bytes,
  71. custom_data=group_id_bytes,
  72. send_buffer_len=args.buffer_len,
  73. use_acceleration=(not args.no_accel),
  74. push_start_cb=push_start_cb)
  75. protocols.append(throughput_protocol)
  76. #
  77. combined_protocol = basic_protocols.ChainedProtocol(protocols)
  78. combined_protocol.run()
  79. #