#!/usr/bin/env python
# -*- coding: UTF-8 -*-

def_port = 28584    # Default port
def_size = 51200    # Stuff to send/recv (in KB)


######################################################################
######################################################################

import getopt
import socket
import sys
import time

version = "0.0.1"
date = "04-jan-2004"

######################################################################

def showlonghelp():
    print """
Speed test was written to quickly test the real throughput of a network
connection, such as an ethernet link. The more data you use (-s flag),the
more reliable the speed test is.

Since it generates a lot of traffic on the network, it should NOT be used if
the user lacks the authorization of the net admin or any other equivalent
person.
The author cannot be held responsible of *any* kind of problem that the misuse
of this tool might yield.

Speed test was written by Daniele Nicolucci (Jollino), with some suggestions
and moral help by Enrico Franchi (RiK0).
Speed test comes with no warranty or support whatsoever. Use at your own risk.
If you like it, though, send me a postcard of your hometown. Mail me at
<jollino@sogno.net> and I'll give you my address."""

######################################################################

def showusage():
    print "Usage: %s [-v | -h | -H] [-l | -r host] [-p port] [-s size]" % sys.argv[0]
    print
    print "-v\t\tshow the version info"
    print "-h\t\tshow this help screen"
    print "-H\t\tshow the long help"
    print "-l\t\tgo into listen mode"
    print "-r\t\tgo into send mode, using host 'host' as recipient"
    print "-p port\t\tuse the specified port (default: %s)" % def_port
    print "-s size\t\tsend/receive these many kilobytes (default: %s KB)" % def_size

######################################################################

def showversion():
    print "This is version %s of Speed test, completed on %s." % (version, date)
    print "Written by Daniele Nicolucci (Jollino), <jollino@sogno.net>."

######################################################################

def do_send(host, port, size):
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)

    print "Connecting to %s:%s..." % (host, port)
    try:
        s.connect((host, port))
    except:
        print "FATAL ERROR -- unable to connect to %s:%s." % (host, port)
        sys.exit(4)
    print "Connected to %s:%s.\n" % (host, port)


    sent = 0
    sizeb = size*1024   # size in bytes
    start = time.time()
    print "Sending %s KB to %s:%s..." % (size, host, port)
    
    while sent < sizeb:
        try:
            reallysent = s.send("12345678901234567890123456789012"*32) # 1 KB
        except:
            print "FATAL ERROR -- broken pipe (%s KB sent)." % round(sent/1024,3)
            sys.exit(5)
        sent += reallysent
    
    #s.shutdown(1)

    end = time.time()
    elapsed = round(end - start, 5)
    speed = round(size/elapsed, 3)

    print "Finished sending %s KB to %s:%s.\n" % (size, host, port)
    print "Time elapsed: %s seconds." % elapsed
    print "Average speed: %s KB/s.\n" % speed

    data = s.recv(1024)
    print "Acknowledge received."
    s.close()

######################################################################

def do_listen(port, size):
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    
    try:
        s.bind(("", port))
    except:
        print "FATAL ERROR -- unable to bind to port %s." % port
        sys.exit(4)
    s.listen(1)
    print "Waiting for connection..."
    
    conn, addr = s.accept()
    print "Connection estabilshed by %s:%s.\n" % (addr[0], addr[1])
    #s.shutdown(2)
    s.close()

    got = 0
    sizeb = size*1024   # size in bytes
    start = time.time()

    print "Receiving %s KB from %s:%s..." % (size, addr[0], addr[1])
    while got < sizeb:
        data = conn.recv(1024)
        if not data:
            print "FATAL ERROR -- connection lost (%s KB received)." % round(got/1024,3)
            sys.exit(5)
        got += len(data)
        
    end = time.time()
    elapsed = round(end - start, 5)
    speed = round(size/elapsed, 3)

    print "Finished receiving %s KB from %s:%s." % (size, addr[0], addr[1])
    print "Sending acknowledgement.\n"
    conn.close()

    print "Time elapsed: %s seconds." % elapsed
    print "Average speed: %s KB/s." % speed

######################################################################

def main():
    print "+--------------------- --------- ---- --- -- -"
    print "| Welcome to Speed test, version %s" % version
    print "| Written by Daniele Nicolucci (Jollino)"
    print "+--------------------- --------- ---- --- -- -\n"

    print "Checking options..."
    
    if (len(sys.argv) == 1):
        print "No options used.  Try \"%s -h\"." % sys.argv[0]
        sys.exit(1)

    try:
        opts, args = getopt.getopt(sys.argv[1:], "vhHlr:p:s:")
    except:
        print "Wrong options used.  Try \"%s -h\"." % sys.argv[0]
        sys.exit(1)

    mode = None
    size = def_size
    port = 28584
    host = None
    
    for opt, arg in opts:
        if opt == "-v":
            print "Version info requested.\n"
            showversion()
            sys.exit(1)
        if opt == "-h":
            print "Usage help requested.\n"
            showusage()
            sys.exit(1)
        if opt == "-H":
            print "Long help requested.\n"
            showlonghelp()
            sys.exit(1)
        if opt == "-l":
            if mode is not None:
                print "FATAL ERROR -- cannot use both -l and -r."
                sys.exit(2)
            mode = "listen"
        if opt == "-r":
            if mode is not None:
                print "FATAL ERROR -- cannot use both -l and -r."
                sys.exit(2)
            mode = "send"
            host = arg
        if opt == "-s":
            size = int(arg)
        if opt == "-p":
            port = int(arg)
            if port < 1024 or port > 65535:
                print "FATAL ERROR -- port must be within 1024 and 65535."
                sys.exit(2)
    
    if mode is None or (mode is "send" and host is None):
        print "FATAL ERROR -- internal error."
        sys.exit(3)

    print "Options ok.\n"
    print "Summary:"
    print "  Test mode ....... %s" % mode
    if mode == "send":
        print "  Remote host ..... %s" % host
    print "  Port ............ %s" % port
    print "  Size of data .... %s KB\n" % size

    try:
        if mode == "send":
            do_send(host, port, size)
        elif mode == "listen":
            do_listen(port, size)
        else:
            print "FATAL ERROR -- internal error."
    except KeyboardInterrupt:
        print "EXIT -- User pressed CTRL-C or some kind of catastrophe took place."

######################################################################

if __name__ == "__main__":
    main()