#!/usr/bin/env python

import sys, time, socket
from websocket import create_connection, WebSocketTimeoutException
#websocket._exceptions.WebSocketTimeoutException

def checkws(ws,hostname,timeout,debug):
        # create websocket
        ws = create_connection(ws, subprotocols=["binary", "base64"])
        # set timeout
        ws.settimeout(timeout);
        # terminal initialization
        ws.send("\x0d")
        ws.send("\xff\xfd\x01")
        ws.send("\xff\xfb\x03")
        ws.send("\xff\xfb\x18")
        ws.send("\xff\xfb\x1f")
        ws.send("\xff\xfd\x03")
        ws.send("\xff\xfd\x18")
        ws.send("\xff\xfd\x1f")
        ws.send("\xff\xfe\x00")
        ws.send("\xff\xfc\x00")
        # initial receive
        try:
                result = ""
                result += ws.recv()
        except WebSocketTimeoutException:
                print "timeout"
        except socket.error:
                print "socket closed"
                sys.exit(3)
        if debug > 1: print result
        # send three carriage returns
        ws.send("\r\r\r")
        start = time.clock()
        duration = 0
        # receive result
        try:
                result = ""
                while True:
                        result += ws.recv()
                        if duration == 0 and result.count(hostname) >= 3:
                                duration = time.clock() - start
                                break
        except WebSocketTimeoutException:
                print "timeout"
        except socket.error:
                print "socket closed"
                sys.exit(3)
        if debug > 1: print result
        ws.close()
        # count occurences of hostname
        if result.count(hostname) >= 3:
                print "%f *sec elapsed*" % duration
                return 0
        else:
                print "UNUSABLE"
                return 2

def main():
        if len(sys.argv) <= 4:
                sys.stdout.write(str(sys.argv[0]))
                print ": ws url, hostname, timeout and debug (e.g., \"ws://192.168.76.210:19406/websockify?token=870668d2-8855-4208-9d48-9b1a141271b2\" iosv-1 0.5 0) needed as argument! bailing out"
                return 1
        else:
                ws = str(sys.argv[1]).strip()
                hostname = str(sys.argv[2]).strip()
                timeout = float(sys.argv[3])
                debug = int(sys.argv[4])
                return checkws(ws,hostname,timeout,debug)

if __name__ == '__main__':
        sys.exit(main())