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
9import socket
10import sys
11import thread
12import time
13
14sent = 0
15received = 0
16
17def Sink(socket):
18  global received
19  while True:
20    tmp = socket.recv(4096)
21    received += len(tmp)
22    if not tmp:
23      break;
24
25def 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
34def 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
53def 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
61def 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
69def 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
79def 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
87if len(sys.argv) < 2:
88  Usage()
89if sys.argv[1] == "listen":
90  Receiver(int(sys.argv[2]))
91elif sys.argv[1] == "download":
92  Connect( (sys.argv[2], int(sys.argv[3])), upload=False, download=True)
93elif sys.argv[1] == "upload":
94  Connect( (sys.argv[2], int(sys.argv[3])), upload=True, download=False)
95elif sys.argv[1] == "updown":
96  Connect( (sys.argv[2], int(sys.argv[3])), upload=True, download=True)
97else:
98  Usage()
99
100PrintStats()
101