#!/usr/bin/python

# Tests movement between various configurations of bonding, VLANs and
# physical devices.

# Requires the UUID of a bonded PIF device as a parameter. This PIF
# and all of its slaves must also have at least one VLAN on top. The
# script interface-reconfigure-test-setup take a system with multiple
# non-management interfaces, bond them together and create appropriate
# VLAN devices for you.

import XenAPI
import sys, os
import getopt

#/* Standard interface flags (netdevice->flags). */
IFF_UP = 0x01

session = None

host = None
bond = None
bond_vlans = None
slaves = None
slave_vlans = {}

class Usage(Exception):
    def __init__(self, msg):
        Exception.__init__(self)
        self.msg = msg

class Error(Exception):
    def __init__(self, msg):
        Exception.__init__(self)
        self.msg = msg

class Failure(Exception):
    def __init__(self):
        Exception.__init__(self)
        #self.msg = "FAIL: " + test + "\n"
        #for m in msg:
        #    self.msg = self.msg + "FAIL: " + m + "\n"

def get_current_host():
    f = open("@INVENTORY@")
    hosts = [h.strip() for h in f.readlines() if h.startswith("INSTALLATION_UUID=")]
    f.close()

    hosts = [l[len("INSTALLATION_UUID='"):len(l)-1] for l in hosts]

    return session.xenapi.host.get_by_uuid(hosts[0])

def interface_name(pif):
    """Construct an interface name from the given PIF record."""

    pifrec = session.xenapi.PIF.get_record(pif)

    if pifrec['VLAN'] == '-1':
        return pifrec['device']
    else:
        return "%(device)s.%(VLAN)s" % pifrec

def get_pif_by_uuid(uuid):
    try:
        pif = session.xenapi.PIF.get_by_uuid(uuid)
    except XenAPI.Failure, err:
        raise Usage("Unknown PIF \"%s\": \"%s\"" % (uuid, err))

    return pif

def get_bond_slaves_of_pif(pif):
    pifrec = session.xenapi.PIF.get_record(pif)

    bmo = pifrec['bond_master_of']
    if len(bmo) > 1:
        raise Error("Bond-master-of contains too many elements")
    
    if len(bmo) == 0:
        return []
    else:
        bondrec = session.xenapi.Bond.get_record(bmo[0])
        return bondrec['slaves']

def get_vlans_of_pif(pif):
    """Returns a list of PIFs which are VLANs on top of the given pif."""
    
    pifrec = session.xenapi.PIF.get_record(pif)
    
    return [session.xenapi.VLAN.get_record(v)['untagged_PIF'] for v in pifrec['VLAN_slave_of']]

class Test(object):
    def __init__(self, test_name):
        self.__test_name = test_name
    def name(self):
        return self.__test_name

    def assert_device_exists(self, pif, exists):
        rec = session.xenapi.PIF.get_record(pif)
        rec['interface-name'] = interface_name(pif)
        path = os.path.join("/sys/class/net", rec['interface-name'])
        if os.path.isdir(path) != exists:
            if exists:
                self.fail(["PIF device %(uuid)s/%(interface-name)s does not exist" % rec])
            else:
                self.fail(["PIF device %(uuid)s/%(interface-name)s exists" % rec])

    def __assert_device_bridge(self, pif, exists):
        rec = session.xenapi.PIF.get_record(pif)
        bridge = session.xenapi.network.get_record(rec['network'])['bridge']
        path = os.path.join("/sys/class/net", bridge)
        if os.path.isdir(path) != exists:
            if exists:
                self.fail([("PIF %(uuid)s/%(device)s bridge " % rec) + bridge + " does not exist"])
            else:
                self.fail([("PIF %(uuid)s/%(device)s bridge " % rec) + bridge + " exists"])

    def __assert_device_attached(self, pif, attached):
        rec = session.xenapi.PIF.get_record(pif)
        rec['interface-name'] = interface_name(pif)
        
        if attached != rec['currently_attached']:
            if attached:
                self.fail(["PIF device %(uuid)s/%(interface-name)s is not currently attached" % rec])
            else:
                self.fail(["PIF device %(uuid)s/%(interface-name)s is currently attached" % rec])

    def __assert_device_up(self, pif, up):
        rec = session.xenapi.PIF.get_record(pif)
        rec['interface-name'] = interface_name(pif)

        path = os.path.join("/sys/class/net", rec['interface-name'], "flags")
        
        if not os.path.exists(path):
            current_state = False
            if not up:
                return
        else:
            f = open(path)
            flags = int(f.readline().strip(),0)
            f.close()

            current_state = (flags & IFF_UP) == IFF_UP
            if current_state == up:
                return

        self.fail([("PIF %(uuid)s/%(interface-name)s " % rec) + ("in state %s not %s" % (current_state, up))])

    def assert_device_state(self, pif, attached=None, up=None, bridge=None):
        """Confirms that the device state is as expected.

attached: regarded as attached by xapi
up:       the interface is up as far as the kernel is concerned
bridge:   the bridge exists

Value my be True, False or None. None means don't care and is mainly
used for bridge where spurious bridges are create on the device
underlying a VLAN."""

        if attached is not None:
            self.__assert_device_attached(pif, attached)

        if up is not None:
            self.__assert_device_up(pif, up)
            
        if bridge is not None:
            self.__assert_device_bridge(pif, bridge)

    def assert_device_is_enslaved(self, pif, master):
        rec = session.xenapi.PIF.get_record(pif)

        path = os.path.join("/sys/class/net", rec['device'], "master")
        if os.path.islink(path):
            m = os.readlink(path)
            m = m[m.rfind("/")+1:]
        else:
            m = None

        if m != master:
            self.fail([("PIF operstate %(uuid)s/%(device)s " % rec) + "enslaved to %s expected %s" % (m, master)])

    def attach_pif(self, pif):
        rec = session.xenapi.PIF.get_record(pif)
        rec['interface-name'] = interface_name(pif)

        self.log("attach %(uuid)s %(interface-name)s" % rec)
        session.xenapi.network.attach(rec['network'], host)

        #self.assert_bridge_exists(pif, True)

    def detach_pif(self, pif):
        rec = session.xenapi.PIF.get_record(pif)
        rec['interface-name'] = interface_name(pif)

        self.log("detach %(uuid)s %(interface-name)s" % rec)
        session.xenapi.PIF.unplug(pif)

        self.assert_device_is_enslaved(pif, None)

        #self.assert_bridge_exists(pif, False)

    def run(self):
        pass
        #print "TEST: " + self.__test_name
    def log(self, s):
        print self.__test_name + ": " + s
        
    def success(self):
        self.log("PASS")
        print ""
    def fail(self, msg):
        for m in msg:
            self.log("FAIL: " + m)
        raise Failure()

class Test_start_of_day(Test):
    """Ensure sane starting environment.

    Checks that all slaves are not enslaved and that bond device has not been created."""
    
    def __init__(self):
        Test.__init__(self, "start-of-day")

    def run(self):
        Test.run(self)
        all = [bond] + bond_vlans + slaves            
        all = [(pif,session.xenapi.PIF.get_record(pif)) for pif in all]

        self.log("interfaces on correct host")
        for rec in [rec for (_,rec) in all if rec['host'] != host]:
            self.fail(["PIF %(uuid)s/%(device)s is not on correct host" % rec])
        
        self.log("no management interface involved in test")
        for (_,rec) in all:
            if rec['management']:
                self.fail(["PIF %(uuid)s/%(device)s is management interface" % rec, \
                           "Cannot run test against management interface"])

        self.log("slaves devices exist, are detached and unenslaved")
        #for (pif,_) in [(pif,session.xenapi.PIF.get_record(pif)) for pif in slaves]:
        for pif in slaves:
            self.assert_device_state(pif, attached=False, up=False, bridge=False)
            self.assert_device_is_enslaved(pif, None)

        self.log("bond is detached and device does not exist")
        self.assert_device_state(bond, attached=False, up=False, bridge=False)
        self.assert_device_exists(bond, False)

        self.log("all vlans are detached and device does not exist")
        all_vlans = bond_vlans
        [all_vlans.extend(v) for v in slave_vlans.values()]
        for vlan in all_vlans:
            self.assert_device_state(vlan, attached=False, up=False, bridge=False)
            self.assert_device_exists(vlan, False)

        self.success()

class Test_device_attach(Test):
    """Ensure we can attach a basic device."""

    def __init__(self):
        Test.__init__(self, "device-attach")

    def run(self):
        Test.run(self)

        dev = slaves[0]
        
        self.attach_pif(dev)
        self.assert_device_state(dev, attached=True, up=True, bridge=True)
        self.detach_pif(dev)        
        self.assert_device_state(dev, attached=False, up=False, bridge=False)
        
        self.success()

class Test_vlan_attach(Test):
    """Ensure we can attach a VLAN device."""

    def __init__(self):
        Test.__init__(self, "vlan-attach")

    def run(self):
        Test.run(self)
        dev = slaves[0]
        vlan = slave_vlans[dev][0]

        self.attach_pif(vlan)
        self.assert_device_exists(vlan, True)
        self.assert_device_state(dev, attached=False, up=True, bridge=True)
        # bridge is there as a side-effect but don't really care.
        self.assert_device_state(vlan, attached=True, up=True, bridge=None)
        
        self.detach_pif(vlan)
        self.assert_device_exists(vlan, False)
        self.assert_device_state(dev, attached=False, up=False, bridge=False)
        self.assert_device_state(vlan, attached=False, up=False, bridge=False)
        
        self.success()

class Test_bond_attach(Test):
    """Ensure we can attach a basic bond."""

    def __init__(self):
        Test.__init__(self, "bond-attach")

    def run(self):
        Test.run(self)

        self.attach_pif(bond)
        self.assert_device_exists(bond, True)
        self.assert_device_state(bond, attached=True, up=True, bridge=True)
        
        # Ensure slaves are enslaved and up
        dev = session.xenapi.PIF.get_record(bond)['device']
        for slave in slaves:
            self.assert_device_state(slave, attached=False, up=True, bridge=False)
            self.assert_device_is_enslaved(slave, dev)

        self.detach_pif(bond)
        self.assert_device_exists(bond, False)
        self.assert_device_state(bond, attached=False, up=False, bridge=False)
        # Ensure slaves are unenslaved and down
        for slave in slaves:
            self.assert_device_state(slave, attached=False, up=False, bridge=False)
            self.assert_device_is_enslaved(slave, None)

        self.success()

class Test_bond_vlan_attach(Test):
    """Ensure we can attach a VLAN bond device."""

    def __init__(self):
        Test.__init__(self, "bond-vlan-attach")

    def run(self):
        Test.run(self)
        
        vlan = bond_vlans[0]

        self.attach_pif(vlan)
        self.assert_device_state(vlan, attached=True, up=True, bridge=True)
        self.assert_device_state(bond, attached=False, up=True, bridge=None)
        
        self.detach_pif(vlan)
        self.assert_device_state(bond, attached=False, up=False, bridge=False)
        self.assert_device_state(vlan, attached=False, up=False, bridge=False)
        self.assert_device_exists(vlan, False)

        self.success()
        
class Test_switch_bond_to_slave(Test):
    """Ensure we can switch from a bond to one of its slaves."""

    def __init__(self):
        Test.__init__(self, "switch-bond-to-slave")

    def run(self):
        Test.run(self)

        # Attach bond
        self.attach_pif(bond)
        self.assert_device_state(bond, attached=True, up=True, bridge=True)
        dev = session.xenapi.PIF.get_record(bond)['device']
        for slave in slaves:
            self.assert_device_state(slave, attached=False, up=True, bridge=False)
            self.assert_device_is_enslaved(slave, dev)

        # Switch to a slave 
        self.attach_pif(slaves[0])
        self.assert_device_state(slaves[0], attached=True, up=True, bridge=True)
        # ... confirm that device is no longer enslaved to the bond
        self.assert_device_is_enslaved(slaves[0], None)
        # ... confirm that bond is detached and device does not exist
        self.assert_device_state(bond, attached=False, up=False, bridge=False)
        self.assert_device_exists(bond, False)
        # ... confirm other slaves are down
        for slave in slaves:
            if slave == slaves[0]:
                continue
            self.assert_device_state(slave, attached=False, up=False)
            self.assert_device_is_enslaved(slave, None)

        # Switch back to bond
        self.attach_pif(bond)
        self.assert_device_state(bond, attached=True, up=True, bridge=True)
        dev = session.xenapi.PIF.get_record(bond)['device']
        for slave in slaves:
            self.assert_device_state(slave, attached=False, up=True, bridge=False)
            self.assert_device_is_enslaved(slave, dev)

        # Clean up
        self.detach_pif(bond)
        self.assert_device_state(bond, attached=False, up=False, bridge=False)
        self.assert_device_exists(bond, False)
        for slave in slaves:
            self.assert_device_state(slave, attached=False, up=False, bridge=False)
            self.assert_device_is_enslaved(slave, None)

        self.success()

class Test_vlan_master_attach(Test):
    """Check VLAN master behaviour on slave detach.

Tests the attach/detach behavior of a non-VLAN PIF when a slave VLAN
is attached and detached.
    """

    def __init__(self):
        Test.__init__(self, "vlan-master-attach")

    def run(self):
        Test.run(self)

        dev = slaves[0]
        vlan = slave_vlans[dev][0]

        self.log("Attach a VLAN device")
        self.attach_pif(vlan)
        self.assert_device_state(dev,  attached=False, up=True, bridge=None)
        self.assert_device_state(vlan, attached=True,  up=True, bridge=True)

        self.log("Attach the VLANs master")
        self.attach_pif(dev)
        self.assert_device_state(dev,  attached=True, up=True, bridge=True)
        self.assert_device_state(vlan, attached=True, up=True, bridge=True)

        self.log("Detach the VLAN, the master should remain")
        self.detach_pif(vlan)
        self.assert_device_state(dev,  attached=True,  up=True,  bridge=True)
        self.assert_device_state(vlan, attached=False, up=False, bridge=False)
        self.assert_device_exists(vlan, False)

        self.log("Detach the master")
        self.detach_pif(dev)
        self.assert_device_state(dev,  attached=False, up=False, bridge=False)
        self.assert_device_state(vlan, attached=False, up=False, bridge=False)
        self.assert_device_exists(vlan, False)

        self.log("Reattach the VLAN")
        self.attach_pif(vlan)
        self.assert_device_state(dev,  attached=False, up=True, bridge=None)
        self.assert_device_state(vlan, attached=True,  up=True, bridge=True)

        self.log("Get rid of the VLAN, master remains detached")
        self.detach_pif(vlan)
        self.assert_device_state(dev,  attached=False, up=False, bridge=False)
        self.assert_device_state(vlan, attached=False, up=False, bridge=False)
        self.assert_device_exists(vlan, False)

        self.log("Attach the master. ")
        self.attach_pif(dev)
        self.assert_device_state(dev,  attached=True,  up=True,  bridge=True)
        self.assert_device_state(vlan, attached=False, up=False, bridge=False)
        self.assert_device_exists(vlan, False)

        self.log("Attach the VLAN")
        self.attach_pif(vlan)
        self.assert_device_state(dev,  attached=True, up=True, bridge=None)
        self.assert_device_state(vlan, attached=True, up=True, bridge=True)

        self.log("Detach the master, VLAN remains attached and master remains up")
        self.detach_pif(dev)
        self.assert_device_state(dev,  attached=False, up=True, bridge=None)
        self.assert_device_state(vlan, attached=True,  up=True, bridge=True)

        self.log("Detach the VLAN, master goes away")
        self.detach_pif(vlan)
        self.assert_device_state(dev,  attached=False, up=False, bridge=False)
        self.assert_device_state(vlan, attached=False, up=False, bridge=False)
        
        self.success()

def cleanup():
    if bond == None or slaves == None:
        return
    
    all = [(pif, session.xenapi.PIF.get_record(pif)) for pif in [bond] + slaves]

    for (pif, rec) in [(pif,rec) for (pif,rec) in all if rec['currently_attached']]:
        print "cleanup: unplug %(uuid)s %(device)s" % rec
        session.xenapi.PIF.unplug(pif)
        
def main(argv=None):
    global host
    global bond
    global bond_vlans
    global slaves
    global slave_vlans
    All_Tests = []
    selected_tests = []
    
    All_Tests.append(Test_device_attach())
    All_Tests.append(Test_vlan_attach())
    All_Tests.append(Test_bond_attach())
    All_Tests.append(Test_bond_vlan_attach())
    All_Tests.append(Test_switch_bond_to_slave())
    All_Tests.append(Test_vlan_master_attach())

    try:
        if argv is None:
            argv = sys.argv

        try:
            shortops = "ht:"
            longops = [ "help" ]
            arglist, args = getopt.gnu_getopt(argv[1:], shortops, longops)
        except getopt.GetoptError, msg:
            raise Usage(msg)

        for o, a in arglist:
            if o == "-t":
                if a == "list":
                    print "Available tests:"
                    for t in All_Tests:
                        print "*** %s:" % t.name()
                        print t.__doc__
                        print
                    return
                for tname in a.split(","):
                    ts = [x for x in All_Tests if x.name() == tname]
                    if len(ts) == 0:
                        raise Error("No test(s) named %s" % tname)
                    selected_tests.extend(ts)
            elif o == "-h" or o == "--help":
                print __doc__ % {'command-name': os.path.basename(argv[0])}
                return 0

        if len(args) < 1:
            raise Usage("Required option <bond-uuid> not present")

        host = get_current_host()
        print "Host is %(uuid)s %(hostname)s" % session.xenapi.host.get_record(host)
        print

        bond = get_pif_by_uuid(args[0])
        print "Bond interface is %(uuid)s %(device)s" % session.xenapi.PIF.get_record(bond)

        bond_vlans = get_vlans_of_pif(bond)
        for vlan in bond_vlans:
            print " vlan%(VLAN)s %(uuid)s %(device)s" % session.xenapi.PIF.get_record(vlan)
        if len(bond_vlans) == 0:
            raise Error("bond has no vlans")
        
        slaves = get_bond_slaves_of_pif(bond)
        for slave in slaves:
            print " slave %(uuid)s %(device)s" % session.xenapi.PIF.get_record(slave)
            slave_vlans[slave] = get_vlans_of_pif(slave)
            for vlan in slave_vlans[slave]:
                print "   vlan%(VLAN)s %(uuid)s %(device)s" % session.xenapi.PIF.get_record(vlan)
                #if len(bond_vlans) == 0:
                #    raise Error("bond has no vlans")
        if len(slaves) == 0:
            raise Error("bond has no slaves")
        elif len(slaves) < 2:
            raise Error("require at least two slaves")

        if len(slave_vlans[slaves[0]]) == 0:
            raise Error("first slave has no vlans")
        print ""

        Test_start_of_day().run()
        
        if selected_tests == []:
            selected_tests = All_Tests

        for test in selected_tests:
            test.run()
            
    except Failure, err:
        #print "FAILED"
        cleanup()
        return 1
    except Usage, err:
        print >>sys.stderr, err.msg
        #print >>sys.stderr, "For help use --help."
        cleanup()
        return 2

    cleanup()
    return 0

if __name__ == "__main__":
    try:
        session = XenAPI.xapi_local()
        session.xenapi.login_with_password("root", "")
        rc = main()
    finally:
        session.xenapi.session.logout()
        
    sys.exit(rc)
