Home | History | Annotate | Download | only in utility
      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