1 #!/usr/bin/env python 2 # Copyright 2014 The Chromium Authors. All rights reserved. 3 # Use of this source code is governed by a BSD-style license that can be 4 # found in the LICENSE file. 5 # 6 # Simple client/server script for generating an unlimited TCP stream. 7 # see shadow.sh for how it's intended to be used. 8 9 import socket 10 import sys 11 import thread 12 import time 13 14 sent = 0 15 received = 0 16 17 def Sink(socket): 18 global received 19 while True: 20 tmp = socket.recv(4096) 21 received += len(tmp) 22 if not tmp: 23 break; 24 25 def Spew(socket): 26 global sent 27 data = " " * 4096 28 while True: 29 tmp = socket.send(data) 30 if tmp <= 0: 31 break 32 sent += tmp; 33 34 def PrintStats(): 35 global sent 36 global received 37 last_report = time.time() 38 last_sent = 0 39 last_received = 0 40 while True: 41 time.sleep(5) 42 now = time.time(); 43 sent_now = sent 44 received_now = received 45 delta = now - last_report 46 sent_mbps = ((sent_now - last_sent) * 8.0 / 1000000) / delta 47 received_mbps = ((received_now - last_received) * 8.0 / 1000000) / delta 48 print "Sent: %5.2f mbps Received: %5.2f mbps" % (sent_mbps, received_mbps) 49 last_report = now 50 last_sent = sent_now 51 last_received = received_now 52 53 def Serve(socket, upload=True, download=True): 54 while True: 55 (s, addr) = socket.accept() 56 if upload: 57 thread.start_new_thread(Spew, (s,)) 58 if download: 59 thread.start_new_thread(Sink, (s,)) 60 61 def Receiver(port, upload=True, download=True): 62 s = socket.socket(socket.AF_INET, socket.SOCK_STREAM) 63 s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) 64 s.bind(('', port)) 65 s.listen(5) 66 thread.start_new_thread(Serve, (s, upload, download)) 67 68 69 def Connect(to_hostport, upload=True, download=False): 70 s = socket.socket(socket.AF_INET, socket.SOCK_STREAM) 71 s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) 72 s.connect(to_hostport) 73 if upload: 74 thread.start_new_thread(Spew, (s,)) 75 if download: 76 thread.start_new_thread(Sink, (s,)) 77 78 79 def Usage(): 80 print "One of:" 81 print "%s listen <port>" % sys.arv[0] 82 print "%s upload <host> <port>" % sys.arv[0] 83 print "%s download <host> <port>" % sys.arv[0] 84 print "%s updown <host> <port>" % sys.arv[0] 85 sys.exit(1) 86 87 if len(sys.argv) < 2: 88 Usage() 89 if sys.argv[1] == "listen": 90 Receiver(int(sys.argv[2])) 91 elif sys.argv[1] == "download": 92 Connect( (sys.argv[2], int(sys.argv[3])), upload=False, download=True) 93 elif sys.argv[1] == "upload": 94 Connect( (sys.argv[2], int(sys.argv[3])), upload=True, download=False) 95 elif sys.argv[1] == "updown": 96 Connect( (sys.argv[2], int(sys.argv[3])), upload=True, download=True) 97 else: 98 Usage() 99 100 PrintStats() 101