[PYTHON] openflow nox LAVI<->ENVI json message translate.py (allowing multiple ENVI connection)

I have edited the original translate.py to allow multiple ENVI connections by expanding the LAVI and ENVI connection sockets

#!/usr/bin/python
##
# Utility translate binary command to JSON equivalent for compatibility
# between ENVI and new LAVI in nox-destiny
#
# @author ykk
# @date June 2010
#
import socket
import select
import struct
import time
import simplejson
import getopt
import sys

def usage():
    """Display usage
    """
    print "Usage "+sys.argv[0]+" [lavihost, default=localhost]"
    print "\tConnect binary ENVI to JSON LAVI"
    print  "Options:"
    print "-h/--help\n\tPrint this usage guide"
    print "-e/--enviport\n\tSpecify port number for ENVI (default:2503)"
    print "-l/--laviport\n\tSpecify port number for LAVI (default:2703)"

#Parse options and arguments
try:
    opts, args = getopt.getopt(sys.argv[1:], "he:l:",
                              ["help","laviport=","enviport="])
except getopt.GetoptError:
    print "Option error!"
    usage()
    sys.exit(2)

#Parse options
noxhost = "localhost"
noxport = 2703
serveport = 2503
for opt,arg in opts:
    if (opt in ("-h","--help")):
        usage()
        sys.exit(0)
    elif (opt in ("-e","--enviport")):
        serveport = int(arg)
    elif (opt in ("-l","--laviport")):
        noxport = int(arg)

#Parse host
if (len(args) >= 1):
    noxhost = args[0]

def envi2lavi(binary):
    jsonmsg = {}
    (length, mtype, xid) = struct.unpack_from("!HBL",binary)
    if (mtype == 0):
        jsonmsg["type"] = "disconnect"
    elif (mtype == 1):
        jsonmsg["type"] = "ping"
jsonmsg["type"] = "tcp"
jsonmsg["type"] = "udp"
    elif (mtype == 2):
        jsonmsg["type"] = "echo"
    elif (mtype == 16): #Node
        (reqtype, nodetype) = struct.unpack_from("!BH",binary,7)
        jsonmsg["type"] = "lavi"
        if (reqtype == 1): #Req
            jsonmsg["command"] = "request"
        elif (reqtype == 2): #Sub
            jsonmsg["command"] = "subscribe"
        elif (reqtype == 3): #Unsub
            jsonmsg["command"] = "unsubscribe"
        else:
            return {}
       
        if (nodetype == 0): #Unknown
            jsonmsg["node_type"] = "all"
        elif (nodetype == 1): #OpenFlow
            jsonmsg["node_type"] = "switch"
        else:
            return {}
    elif (mtype == 19): #Link
        (reqtype, linktype) = struct.unpack_from("!BH",binary,7)
        jsonmsg["type"] = "lavi"
        if (reqtype == 1): #Req
            jsonmsg["command"] = "request"
        elif (reqtype == 2): #Sub
            jsonmsg["command"] = "subscribe"
        elif (reqtype == 3): #Unsub
            jsonmsg["command"] = "unsubscribe"
        else:
            return {}
       
        if (linktype == 0): #Unknown
            jsonmsg["link_type"] = "all"
        elif (linktype == 1): #SW2SW
            jsonmsg["link_type"] = "sw2sw"
        else:
            return {}
    elif (mtype == 22): #Flow
        (reqtype, flowtype) = struct.unpack_from("!BH",binary,7)
        jsonmsg["type"] = "lavi"
        if (reqtype == 1): #Req
            jsonmsg["command"] = "request"
        elif (reqtype == 2): #Sub
            jsonmsg["command"] = "subscribe"
        elif (reqtype == 3): #Unsub
            jsonmsg["command"] = "unsubscribe"
        else:
            return {}
       
        if (flowtype == 0): #Unknown
            jsonmsg["flow_type"] = "all"
        elif (flowtype == 1): #OpenFlow
            jsonmsg["flow_type"] = "network"
        else:
            return {}
    else:
        return
    return jsonmsg
   
def lavi2envi(jsonmsg):
    envimsg = ""
    msgtype = None
    #Check for basic keys
    if  (not ("type" in jsonmsg and "command" in jsonmsg)):
        print "Invalid message "+str(jsonmsg)
        return ""
   
    if (jsonmsg["type"] == "lavi"):
        if ("node_type" in jsonmsg):
            #Nodes
            if (jsonmsg["node_type"] == "switch" or
                jsonmsg["node_type"] == "host"):
                #Switches
                if (jsonmsg["command"] == "add"):
                    msgtype = 17
                elif (jsonmsg["command"] == "delete"):
                    msgtype = 18
                else:
                    return ""

                for sw in jsonmsg["node_id"]:
                    envimsg += struct.pack("!HQ",1,int(sw,16))
            else:
                return ""
        elif ("flow_type" in jsonmsg):
            #Flows
            enviaggmsg = ""
            if (jsonmsg["flow_type"] == "network"):
                if (jsonmsg["command"] == "add"):
                    msgtype = 23   
                    for flow in jsonmsg["flows"]:
                        envimsg = struct.pack("!LHL", 1, 0, int(jsonmsg["flow_id"], 16)) + \
                                  struct.pack("!HQH", 1, int(flow["src id"],16), int(flow["src port"],16)) + \
                                  struct.pack("!HQH", 1, int(flow["dst id"],16), int(flow["dst port"],16)) + \
                                  struct.pack("!H", 0)
                        enviaggmsg += struct.pack("!HBL",len(envimsg)+7,msgtype,0)+envimsg
                elif (jsonmsg["command"] == "delete"):
                    msgtype = 24
    #modified#######################################
  # for flow in jsonmsg["flows"]:
                    ################################################       
                    envimsg = struct.pack("!LHL", 1, 0, int(jsonmsg["flow_id"], 16)) + \
                              struct.pack("!HQH", 1, 0, 0) + \
                              struct.pack("!HQH", 1, 0, 0) + \
                              struct.pack("!H", 0)
                    enviaggmsg += struct.pack("!HBL",len(envimsg)+7,msgtype,0)+envimsg
            elif (jsonmsg["flow_type"] == "host"):
                if (jsonmsg["command"] == "add"):
                    msgtype = 23   
                    for flow in jsonmsg["flows"]:
                        if ("dst port" in flow):
                            envimsg = struct.pack("!LHL", 1, 0, int(jsonmsg["flow_id"], 16)) + \
                                      struct.pack("!HQH", 1, int(flow["src id"],16), 0) + \
                                      struct.pack("!HQH", 1, int(flow["dst id"],16), int(flow["dst port"],16)) + \
                                      struct.pack("!H", 0)
                            enviaggmsg += struct.pack("!HBL",len(envimsg)+7,msgtype,0)+envimsg
                        else:
                            envimsg = struct.pack("!LHL", 1, 0, int(jsonmsg["flow_id"], 16)) + \
                                      struct.pack("!HQH", 1, int(flow["src id"],16), int(flow["src port"],16)) + \
                                      struct.pack("!HQH", 1, int(flow["dst id"],16), 0) + \
                                      struct.pack("!H", 0)
                            enviaggmsg += struct.pack("!HBL",len(envimsg)+7,msgtype,0)+envimsg
                elif (jsonmsg["command"] == "delete"):
                    msgtype = 24
                    #modified####################################
  # for flow in jsonmag["flows"]:
                    #############################################
                    envimsg = struct.pack("!LHL", 1, 0, int(jsonmsg["flow_id"], 16)) + \
                              struct.pack("!HQH", 1, 0, 0) + \
                              struct.pack("!HQH", 1, 0, 0) + \
                              struct.pack("!H", 0)
                    enviaggmsg += struct.pack("!HBL",len(envimsg)+7,msgtype,0)+envimsg

            return enviaggmsg

        elif ("link_type" in jsonmsg):
            #Links
            if (jsonmsg["link_type"] == "sw2sw"):
                if (jsonmsg["command"] == "add"):
                    msgtype = 20
                    for link in jsonmsg["links"]:
                        envimsg += struct.pack("!HHQHHQHQ",
                                              1,
                                              1,int(link["src id"], 16),int(link["src port"], 16),
                                              1,int(link["dst id"], 16),int(link["dst port"], 16),
                                              1e7)
                elif (jsonmsg["command"] == "delete"):
                    msgtype = 21
                    for link in jsonmsg["links"]:
                        envimsg += struct.pack("!HHQHHQH",
                                              1,
                                              1,int(link["src id"], 16),int(link["src port"], 16),
                                              1,int(link["dst id"], 16),int(link["dst port"], 16))
                else:
                    return ""
            elif (jsonmsg["link_type"] == "host2sw"):
                #Host to switch
                if (jsonmsg["command"] == "add"):
                    msgtype = 20
                    for link in jsonmsg["links"]:
                        envimsg += struct.pack("!HHQHHQHQ",
                                              1,
                                              1, int(link["src id"], 16),0,
                                              1, int(link["dst id"], 16),int(link["dst port"], 16),
                                              1e7)
                elif (jsonmsg["command"] == "delete"):
                    msgtype = 21
                    for link in jsonmsg["links"]:
                        envimsg += struct.pack("!HHQHHQH",
                                              1,
                                              1,int(link["src id"], 16),0,
                                              1, 0,int(link["src port"], 16),
                                              1, int(link["dst id"], 16),int(link["dst port"], 16))
                else:
                    return ""
        else:
            print "Unidentified message type"
            print jsonmsg
            return ""
    else:
        print "Not a LAVI message"
        return ""

    return struct.pack("!HBL",len(envimsg)+7,msgtype,0)+envimsg

while True:
    #Get ENVI connection
    server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    server.setsockopt( socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    server.bind(('',serveport))
    server.listen(10)
    print "Waiting for ENVI connection"

    socks = [server]
    ec_list = []
    lc_list = []

    #Poll and translate
    lavibuffer = ""
    eb={}
    lb={}
    lc={}
    ec={}
    connected_clients=0
    while True:

        (read, write, err) = select.select(socks,
                                          [],
                                          socks, 1)

        for e in err:
            print "Error:", e

        for r in read:
            if (server == r ):

                (enviconn, saddr) = server.accept()
                socks.append(enviconn)
                ec_list.append(enviconn)
                print "Got connection from "+str(saddr) + " SOCK: ", enviconn, " size: ", len(socks);
                eb[str(enviconn)]=""

                # create corresponding laviconn
                # Get LAVI connection
                laviconn = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
                socks.append(laviconn)
                lc_list.append(laviconn)
   
                try:
                    laviconn.connect((noxhost, noxport))
                    print "Connected to NOX at "+str((noxhost,noxport)), " SOCK:", laviconn
                except socket.error:
                    print "Cannot connect to LAVI @ "+noxhost
   
                lb[str(laviconn)]=""

                # cross reference
                lc[str(enviconn)]=laviconn;
                ec[str(laviconn)]=enviconn;
               

            elif ( r in lc_list):
                try:
                    data = r.recv(1)
                    if (not data):
                        print "NO DATA BREAK"
                        break
                    lb[str(r)] += data
                    try:
                        jsonmsg = simplejson.loads( lb[str(r)] )
                        binarymsg = lavi2envi(jsonmsg)
                        if (len(binarymsg) != 0):
#                            print "SEND MESSAGE TO ENVI to: ", r, " from: ", ec[str(r)]
                            ec[str(r)].send(binarymsg)

                        lb[str(r)] = ""
                    except ValueError:
                        pass
                except socket.error:
                    print "SOCKET ERROR: BREAK"
                    break
            elif ( r in ec_list):
                try:
                    data = None
                    data = r.recv(1)
                except: # this is probably socket error. delete this
                    r.close()
                    lc[str(r)].close()
                    socks.remove(lc[str(r)])
                    lc_list.remove(lc[str(r)])
                    socks.remove(r)
                    ec_list.remove(r)
                   
                    print "A: delete socket: ", r, " size: ", len(socks);
                    continue
                 
                if (not data):
                    r.close()
                    lc[str(r)].close()
                    try:
                        socks.remove(lc[str(r)])
                        lc_list.remove(lc[str(r)])
                        socks.remove(r)
                        ec_list.remove(r)

                    except:
                        pass
                    print "B: delete socket: ", r, " size: ", len(socks);

                else:
                    eb[str(r)] += data;
                    if (len(eb[str(r)]) > 2):
                        (length, ) = struct.unpack_from("!H",eb[str(r)])

                        if (length == len(eb[str(r)])):
                            jsonmsg = envi2lavi(eb[str(r)])
                            if (jsonmsg != None):
                                lc[str(r)].send(simplejson.dumps(jsonmsg))
                            eb[str(r)] = ""

            else:
                print "Error!: Unregistered socket found"

    #Close
    print "Closing connections"
    if (clientok):
        laviconn.shutdown(1)
        laviconn.close()
    time.sleep(1)