import atexit
import pyroute2

#-------------------------------

def InitializeLibrary(library):
    if library == 'NDB':
        global ndb
        ndb = pyroute2.NDB()
        @atexit.register
        def close_ndb():
            ndb.close()
    elif library == 'IPR':
        global iproute
        iproute = pyroute2.IPRoute()
        @atexit.register
        def close_iproute():
            iproute.close()

#-------------------------------

def GetInterfaces(library):
    interfaces = []
    if library == 'NDB':
        for itf in ndb.interfaces:
            interfaces.append(Interface(ndb.interfaces[itf.ifname]))
    elif library == 'IPR':
        for itf in iproute.get_links():
            interfaces.append(Interface(itf))
    return interfaces

def GetInterface(library, ifname):
    if library == 'NDB':
        return Interface(ndb.interfaces[ifname])
    elif library == 'IPR':
        return Interface(iproute.get_links(ifname=ifname)[0])
    return None

class Interface:
    def __init__(self, interface):
        self.ndb = None
        self.ipr = None

        self._set_interface_(interface)
        self._initialize_attributes_()

    def __getitem__(self, key):
        return getattr(self, key, None)

    def _set_interface_(self, interface):
        if isinstance(interface, pyroute2.netlink.rtnl.ifinfmsg.ifinfmsg):
            self.ipr = interface
        elif isinstance(interface, pyroute2.ndb.objects.interface.Interface):
            self.ndb = interface

    def _initialize_attributes_(self):
        if self.ndb:
            self.index = self.ndb['index']
            self.ifname = self.ndb['ifname']
            self.kind = self.ndb['kind']
            self.address = self.ndb['address']
            if 'netns' in self.ndb:
                 self.netns = self.ndb['netns']
            else:
                 self.netns = None
            self.ipaddr = []
            for address in self.ndb.ipaddr:
                ipaddress = {}
                ipaddress['address'] = address['address']
                ipaddress['prefixlen'] = address['prefixlen']
                self.ipaddr.append(ipaddress)
            self.routes = []
            for r in self.ndb.routes:
                route = {}
                route['type'] = r['type']
                route['family'] = r['family']
                route['dst'] = r['dst']
                route['gateway'] = r['gateway']
                route['dst_len'] = r['dst_len']
                self.routes.append(route)
        elif self.ipr:
            self.index = self.ipr['index']
            self.ifname = self.ipr.get_attr('IFLA_IFNAME')
            self.kind = self.ipr.get_attr('IFLA_LINKINFO').get_attr('IFLA_INFO_KIND') if self.ipr.get_attr('IFLA_LINKINFO') else None
            self.address = self.ipr.get_attr('IFLA_ADDRESS')
            self.netns = None
            self.ipaddr = []
            for address in iproute.get_addr(label=self.ipr.get_attr('IFLA_IFNAME')):
                ipaddress = {}
                ipaddress['address'] = address.get_attr('IFA_ADDRESS')
                ipaddress['prefixlen'] = address['prefixlen']
                self.ipaddr.append(ipaddress)
            self.routes = []
            for r in iproute.get_routes():
                attrs = dict(r['attrs'])
                if attrs['RTA_OIF'] == self['index']:
                    if 'RTA_PRIORITY' in attrs and not attrs['RTA_PRIORITY'] < 255:
                        continue
                    route = {}
                    route['type'] = r['type']
                    route['family'] = r['family']
                    route['dst'] = attrs['RTA_DST'] if 'RTA_DST' in attrs else ''
                    route['gateway'] = attrs['RTA_GATEWAY'] if 'RTA_GATEWAY' in attrs else None
                    route['dst_len'] = r['dst_len']
                    self.routes.append(route)

    def up(self):
        if self.ndb:
             self.ndb.set('state', 'up').commit()
        elif self.ipr:
             iproute.link("set", index=self.index, state="up")

    def down(self):
        if self.ndb:
             self.ndb.set('state', 'down').commit()
        elif self.ipr:
             iproute.link("set", index=self.index, state="down")

def CreateInterface(library, ifname, kind, peer):
    try:
        if library == 'NDB':
            ndb.interfaces.create(ifname=ifname, kind=kind, peer=peer).commit()
            return Interface(ndb.interfaces[ifname])
        elif library == 'IPR':
            iproute.link("add", ifname=ifname, kind=kind, peer=peer)
            return Interface(iproute.get_links(ifname=ifname)[0])
        return None
    except Exception as e:
        return None

def RemoveInterface(library, ifname):
    try:
        if library == 'NDB':
            ndb.interfaces[ifname].remove().commit()
        elif library == 'IPR':
            iproute.link("del", index=iproute.link_lookup(ifname=ifname)[0])
    except Exception as e:
        return

def AddPort(library, bridge, ifname):
    try:
        if library == 'NDB':
            ndb.interfaces[bridge].add_port(ifname).commit()
        elif library == 'IPR':
            iproute.link("set", index=iproute.link_lookup(ifname=ifname)[0], master=iproute.link_lookup(ifname=bridge)[0])
    except Exception as e:
        return

def DelPort(library, bridge, ifname):
    try:
        if library == 'NDB':
            ndb.interfaces[bridge].del_port(ifname).commit()
        elif library == 'IPR':
            iproute.link("set", index=iproute.link_lookup(ifname=ifname)[0], master=0)
    except Exception as e:
        return

#-------------------------------

def test_Variable(v1, v2, f):
    assert v1 == v2, "{} not equal ({}, {})".format(f, v1, v2)

def test_Interface():
    itf1 = GetInterface('NDB', 'br0')
    itf2 = GetInterface('IPR', 'br0')

    test_Variable(itf1.ifname, itf2.ifname, 'ifname')
    test_Variable(itf1.kind, itf2.kind, 'kind')
    test_Variable(itf1.ipaddr(), itf2.ipaddr(), 'ipaddr')
    test_Variable(itf1.routes, itf2.routes, 'routes')

    print("OK")
#test_Interface()

def test_CreateInterface():
    CreateInterface('NDB', 'ndb0', 'veth', 'ndb1')
    CreateInterface('IPR', 'ipr0', 'veth', 'ipr1')

    itf1 = GetInterface('NDB', 'ndb0')
    itf2 = GetInterface('IPR', 'ipr0')

    itf1.up();
    itf2.up();

    itf3 = GetInterface('NDB', 'ndb1')
    itf4 = GetInterface('IPR', 'ipr1')

    itf3.up();
    itf4.up();

    AddPort('NDB', 'br0', 'ndb0')
    AddPort('IPR', 'br0', 'ipr0')
    
#test_CreateInterface()

def test_RemoveInterface():
    DelPort('NDB', 'br0', 'ndb0')
    DelPort('IPR', 'br0', 'ipr0')

    RemoveInterface('NDB', 'ndb0')
    RemoveInterface('IPR', 'ipr0')
#test_RemoveInterface()