system/tcptest.py

   1 #!/usr/bin/env python
   2
   3 # Thierry Parmentelat <thierry.parmentelat@inria.fr>
   4 # Copyright (C) 2010 INRIA
   5 #
   6 import sys
   7 import time
   8 import subprocess
   9 import socket
  10 import SocketServer
  11 import threading
  12 from optparse import OptionParser
  13
  14 def myprint(message, id='client'):
  15     now=time.strftime("%H:%M:%S", time.localtime())
  16     print "*",now,'(%s)' % id, '--',message
  17     sys.stdout.flush()
  18
  19 def show_network_status(id):
  20     myprint("ip address show", id=id)
  21     subprocess.call(['ip', 'address', 'show'])
  22     myprint("ip route show", id=id)
  23     subprocess.call(['ip', 'route', 'show'])
  24
  25 class EchoRequestHandler(SocketServer.StreamRequestHandler):
  26     def handle(self):
  27         line = self.rfile.readline()
  28         self.wfile.write(line)
  29
  30 class UppercaseRequestHandler(SocketServer.StreamRequestHandler):
  31     def handle(self):
  32         line = self.rfile.readline()
  33         self.wfile.write(line.upper())
  34
  35 class Server:
  36     """
  37     A TCP server, running for some finite amount of time
  38     """
  39     def main(self):
  40         parser = OptionParser()
  41         parser.add_option("-p", "--port", action="store", dest="port", type="int",
  42                           default=10000, help="port number")
  43         parser.add_option("-a", "--address", action="store", dest="address",
  44                           default=socket.gethostname(), help="address")
  45         parser.add_option("-t", "--timeout", action="store", dest="timeout", type="int",
  46                           default="0")
  47         (options, args) = parser.parse_args()
  48
  49         if len(args) != 0:
  50             parser.print_help()
  51             sys.exit(1)
  52
  53         show_network_status(id='server')
  54         server = SocketServer.TCPServer((options.address, options.port),
  55                                         UppercaseRequestHandler)
  56         try:
  57             if options.timeout:
  58                 t = threading.Thread(target=server.serve_forever)
  59                 t.setDaemon(True) # don't hang on exit
  60                 t.start()
  61                 time.sleep(options.timeout)
  62                 sys.exit(0)
  63             else:
  64                 server.serve_forever()
  65         except KeyboardInterrupt as e:
  66             print 'Bailing out on keyboard interrupt'
  67             sys.exit(1)
  68
  69 class Ready:
  70     """
  71     A utility that does exit(0) iff network as perceived
  72     from the sliver is ready. Designed to be run before Server,
  73     so one can wait for the right conditions.
  74     """
  75     def main(self):
  76         parser = OptionParser()
  77         # by default use another port so we don't run into
  78         # the SO_LINGER kind of trouble
  79         parser.add_option("-p", "--port", action="store", dest="port", type="int",
  80                           default=9999, help="port number")
  81         parser.add_option("-a", "--address", action="store", dest="address",
  82                           default=socket.gethostname(), help="address")
  83         (options, args) = parser.parse_args()
  84
  85         def can_bind ():
  86             s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
  87             try:
  88                 s.bind((options.address, options.port))
  89                 return True
  90             except Exception as e:
  91                 print e
  92                 return False
  93
  94         def eth0_has_ipv4():
  95             command = "ip address show eth0 | grep -q ' inet '"
  96             return subprocess.check_call(command, shell=True) == 0
  97
  98         sys.exit(0 if can_bind() and eth0_has_ipv4() else 1)
  99
 100 class Client:
 101     """
 102     Runs a client against a Server instance
 103     """
 104     def main(self):
 105         parser = OptionParser()
 106         parser.add_option("-p","--port", action="store", dest="port", type="int",
 107                           default=10000, help="port number")
 108         parser.add_option("-a","--address", action="store", dest="address",
 109                           default=socket.gethostname(), help="address")
 110         parser.add_option("-s","--sleep", action="store", dest="sleep", type="int",
 111                           default=1, help="sleep seconds")
 112         parser.add_option("-l","--loops", action="store", dest="loops", type="int",
 113                           default=1, help="iteration loops")
 114
 115         (options, args) = parser.parse_args()
 116         if len(args) != 0:
 117             parser.print_help()
 118             sys.exit(1)
 119
 120         result=True
 121         for i in range(1,options.loops+1):
 122             s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
 123             s.connect((options.address, options.port))
 124             mout=i*'ping ' + '\n'
 125             min=mout.upper()
 126             if s.send(mout) != len(mout):
 127                 myprint("cannot send %s"%mout.strip())
 128                 result=False
 129                 break
 130             line=s.recv(len(min))
 131             if line is not line:
 132                 myprint("unexpected reception\ngot:%s\nexpected: %s",line,min)
 133                 result=False
 134             else:
 135                 myprint("OK:%s"%mout.strip())
 136             # leave the connection open, but the last one (so 1 iter returns fast)
 137             if i != options.loops:
 138                 time.sleep(options.sleep)
 139             myprint("disconnecting")
 140             s.close()
 141         myprint("Done")
 142         exit_return=0
 143         if not result:
 144             exit_return=1
 145         sys.exit(exit_return)
 146
 147 if __name__ == '__main__':
 148     for arg in sys.argv[1:]:
 149         if arg.find("client") >= 0:
 150             sys.argv.remove(arg)
 151             Client().main()
 152         elif arg.find("server") >= 0:
 153             sys.argv.remove(arg)
 154             Server().main()
 155         elif arg.find("ready") >= 0:
 156             sys.argv.remove(arg)
 157             Ready().main()
 158     print 'you must specify either --client or --server'
 159     sys.exit(1)