#!/usr/bin/env python
#
# THIS IS AN EXPERIMENTAL, UNSUPPORTED PLUGIN FOR DEMONSTRATION AND TEST ONLY.
# PLEASE DO NOT USE IT FOR PRODUCTION ENVIRONMENTS.
# THIS PLUGIN CAN BE REMOVED OR MODIFIED AT ANY TIME WITHOUT PREVIOUS WARNING.
#
# A plugin for managing IO virtualization including SR-IOV
#
# Each VM has pass-through VFs configured using the "pci" and "sriovmacs"
# other-config keys. Each of these is a comma separated list of entries
# of the form "<index>/<pciid>" for pci and "<index>/<mac>" for sriovmacs


import XenAPI, inventory
import XenAPIPlugin
import os
import os.path
import stat
import glob
import re
import random
import subprocess
import xml.dom.minidom
import syslog

ipcmd = "/sbin/ip"
hookscripts = ["/etc/xapi.d/vm-pre-start/vm-pre-start-iovirt",
               "/etc/xapi.d/vm-pre-reboot/vm-pre-reboot-iovirt"
]

re_virtfn = re.compile("virtfn(\d+)$")
re_netdev = re.compile("(eth\d+)$")
re_hex16 = re.compile("^[0-9a-f]{4}$")
re_hex16x = re.compile("^0x([0-9a-f]{4})$")

def _install_hook_script():
    # Ensure that the hook script(s) used to configure SR-IOV VF MAC and
    # VLANs is present. If the script isn't present create it. This
    # function will be called whenever a VF is assigned. The intention
    # is that the hook is not used for users not using this plugin because
    # the additional API calls will add a small overhead to the VM.start
    # time which may hurt some use cases.
    for hookscript in hookscripts:
        if os.path.exists(hookscript):
            continue
        syslog.syslog("Creating iovirt hook script at %s" % (hookscript))
        hookdir = os.path.dirname(hookscript)
        if not os.path.exists(hookdir):
            os.makedirs(hookdir)
        f = file(hookscript, "w")
        f.write("""#!/bin/bash
#
# Call the iovirt plugin to set up SR-IOV VF MAC and VLAN config
# if required for this VF.

PLUGIN=iovirt
FN=prep_for_vm

for i
do
  case "$i" in
        -vmuuid) shift; VMUUID=$1; shift;;
  esac
done

if [ -z "$VMUUID" ]; then
    logger -t $(basename $0) "VM UUID not found in args"
fi

. @INVENTORY@

if [ -z "$INSTALLATION_UUID" ]; then
    logger -t $(basename $0) "Could not determine host UUID"
fi

xe host-call-plugin plugin=$PLUGIN fn=$FN host-uuid=$INSTALLATION_UUID args:uuid=$VMUUID
"""
                )
        f.close()
        os.chmod(hookscript, stat.S_IRWXU | stat.S_IRGRP | stat.S_IXGRP | stat.S_IROTH | stat.S_IXOTH)
        
def _get_vfs():
    # Find PCI devices with virtfns, find all virtfns associated with
    # network devices. Return dictionary of pciid => (vf#, ethdev)
    vfnodes = glob.glob("/sys/bus/pci/devices/*/virtfn*")
    vflist = {}
    for vfnode in vfnodes:
        vfpciid = os.path.basename(os.readlink(vfnode))
        r = re_virtfn.search(vfnode)
        if not r:
            raise "Failed to parse VF number for %s" % (vfnode)
        vfnum = r.group(1)
        netdevs = glob.glob("%s/physfn/net:eth*" % (vfnode))
        if len(netdevs) != 1:
            raise "Unexpected length of netdev list for %s: %s" % (vfnode, str(netdevs))
        r = re_netdev.search(netdevs[0])
        if not r:
            raise "Error parsing %s for net device" % (netdevs[0])
        netdev = r.group(1)
        vflist[vfpciid] = (vfnum, netdev)
    return vflist

def _get_devices(vendorid, deviceid):
    # Return device instances with the specified vendorid:deviceid
    # Returns a dictionary of pciid => (vendorid, deviceid)
    # Used for managing the assignment of non-SRIOV pools of devices.
    if not re_hex16.match(vendorid):
        raise "PCI vendor ID '%s' not in expected format" % (vendorid)
    if not re_hex16.match(deviceid):
        raise "PCI device ID '%s' not in expected format" % (deviceid)
    devices = {}
    devnodes = glob.glob("/sys/bus/pci/devices/*")
    for devnode in devnodes:
        pciid = os.path.basename(devnode)
        vendor, device = _get_vendor_and_device_ids(pciid)
        if vendor == vendorid and device == deviceid:
            devices[pciid] = (vendor, device)
    return devices

def _get_vendor_and_device_ids(pciid):
    devnode = ("/sys/bus/pci/devices/%s" % (pciid))
    vendor = None
    device = None
    if os.path.exists("%s/vendor" % (devnode)):
        f = file("%s/vendor" % (devnode))
        d = f.read()
        f.close()
        vendorid = re_hex16x.search(d)
        if vendorid:
            vendor = vendorid.group(1)
    if os.path.exists("%s/device" % (devnode)):
        f = file("%s/device" % (devnode))
        d = f.read()
        f.close()
        deviceid = re_hex16x.search(d)
        if deviceid:
            device = deviceid.group(1)
    return vendor, device

def _get_assignments(session):
    # Return a dict of PCI devices assigned to VMs. This currently assumes
    # a pool of one host. In a pool this will be overly restrictive (i.e.
    # a VM started on another host will be reported as having one of our
    # host's devices) but it'll do for now.
    # Returns a dictionary PCIID => (VMUUID, index, vendorid, deviceid, mac, vlan)
    # where mac and vlan are only (optionally) available for SR-IOV VFs.
    expr = 'field "is_a_template" = "false" and field "is_a_snapshot" = "false"'
    vms = session.xenapi.VM.get_all_records_where(expr)
    devices = {}
    for vm in vms.values():
        # Ignore VMs with no PCI passthrough config or an empty config
        if not vm['other_config'].has_key('pci'):
            continue
        if not vm['other_config']['pci']:
            continue
        vmuuid = vm['uuid']
        assignments = _get_vm_assignments(session, vmuuid)
        for index in assignments.keys():
            devices[assignments[index][0]] = (vmuuid,
                                              index,
                                              assignments[index][1],
                                              assignments[index][2],
                                              assignments[index][3],
                                              assignments[index][4])
    return devices

def enable_iommu(session, args):
    rc = os.system("@OPTDIR@/libexec/xen-cmdline --set-xen iommu=1")
    return str(rc == 0)

#def list_assigned_vfs(session, args):
#    """List VFs on this host that are assigned to VMs that can run here."""
#    x = _get_assignments(session)
#    return `x`

def _get_vm_assignments(session, vmuuid):
    # Return a dictionary of index => (PCIID, vendorid, deviceid, MAC, VLAN)
    # for assignments for the specified VM.
    # This includes SR-IOV NIC VFs and other PCI pass through
    # devices. (MAC and VLAN only includes for the former)
    re_pci = re.compile("^(\d+)/([0-9a-f:\.]+)$")
    re_vlan = re.compile("^(\d+)/([0-9]+)$")
    re_mac = re.compile("^(\d+)/([0-9a-fA-F:]+)$")

    vmref = session.xenapi.VM.get_by_uuid(vmuuid)
    vm = session.xenapi.VM.get_record(vmref)
    assignments = {}

    # Parse the SR-IOV MAC config string
    macs = {}
    if vm['other_config'].has_key('sriovmacs') and vm['other_config']['sriovmacs']:
        for x in vm['other_config']['sriovmacs'].split(","):
            r = re_mac.search(x)
            if not r:
                raise "Failed to parse MAC config '%s' for VM %s" % x, vmuuid
            macs[int(r.group(1))] = r.group(2)
    
    # Parse the SR-IOV VLAN config string
    vlans = {}
    if vm['other_config'].has_key('sriovvlans') and vm['other_config']['sriovvlans']:
        for x in vm['other_config']['sriovvlans'].split(","):
            r = re_vlan.search(x)
            if not r:
                raise "Failed to parse VLAN config '%s' for VM %s" % x, vmuuid
            vlans[int(r.group(1))] = r.group(2)
    
    # Parse the PCI config string
    if vm['other_config'].has_key('pci') and vm['other_config']['pci']:
        for x in vm['other_config']['pci'].split(","):
            r = re_pci.search(x)
            if not r:
                raise "Failed to parse PCI config '%s' for VM %s" % x, vmuuid
            vendorid, deviceid = _get_vendor_and_device_ids(r.group(2))
            if macs.has_key(int(r.group(1))):
                mac = macs[int(r.group(1))]
            else:
                mac = None
            if vlans.has_key(int(r.group(1))):
                vlan = vlans[int(r.group(1))]
            else:
                vlan = None
            assignments[int(r.group(1))] = (r.group(2), vendorid, deviceid, mac, vlan)

    return assignments

def _randomMAC():
    """Return a random MAC in the locally administered range"""
    o1 = (random.randint(0, 63) << 2) | 2
    o2 = random.randint(0, 255)
    o3 = random.randint(0, 255)
    o4 = random.randint(0, 255)
    o5 = random.randint(0, 255)
    o6 = random.randint(0, 255)
    return "%02x:%02x:%02x:%02x:%02x:%02x" % (o1, o2, o3, o4, o5, o6)

def _set_vm_assignments(session, vmuuid, assignments):
    # Set the PCI (and MAC and VLAN for SR-IOV) assignments for the
    # specified VM removing any existing config first. assignments is a
    # dictionary index => (pciid, vendorid, deviceid, mac, vlan) where
    # vendorid and deviceid need not be set.
    indexlist = assignments.keys()
    indexlist.sort()
    pcilist = []
    maclist = []
    vlanlist = []
    for i in indexlist:
        pcilist.append("%u/%s" % (i, assignments[i][0]))
        if assignments[i][3]:
            maclist.append("%u/%s" % (i, assignments[i][3]))
        if assignments[i][4]:
            vlanlist.append("%u/%s" % (i, assignments[i][4]))
    pci = ",".join(pcilist)
    mac = ",".join(maclist)
    vlan = ",".join(vlanlist)
    vmref = session.xenapi.VM.get_by_uuid(vmuuid)
    try:
        session.xenapi.VM.remove_from_other_config(vmref, "pci")
    except:
        pass
    try:
        session.xenapi.VM.remove_from_other_config(vmref, "sriovmacs")
    except:
        pass
    try:
        session.xenapi.VM.remove_from_other_config(vmref, "sriovvlans") 
    except:
        pass
    if len(pci) > 0:
        session.xenapi.VM.add_to_other_config(vmref, "pci", pci)
    if len(mac) > 0:
        session.xenapi.VM.add_to_other_config(vmref, "sriovmacs", mac)
    if len(vlan) > 0:
        session.xenapi.VM.add_to_other_config(vmref, "sriovvlans", vlan)
    
def assign_free_vf(session, args):
    """Assign a free VF on this host to the specified VM."""

    # Ensure the hook script exists to configure VFs before VM.start
    _install_hook_script()
    
    assigned = _get_assignments(session)
    vfs = _get_vfs()
    for vfid in assigned.keys():
        if vfid in vfs:
            del vfs[vfid]
    if not args.has_key("uuid"):
        raise "No VM UUID specified, please use the 'uuid' argument"
    vmuuid = args["uuid"]
    ethdev = None
    vlan = None
    pifref = None
    
    # Allow the caller to specify a the eth device from which the VF should
    # be allocated
    if args.has_key("ethdev"):
        ethdev = args["ethdev"]

    # Specify by network UUID. If this is a VLAN network then this forces the
    # specification of a VLAN tag for the VF.
    if args.has_key("nwuuid"):
        nwuuid = args["nwuuid"]
        nwref = session.xenapi.network.get_by_uuid(nwuuid)
        pifrefs = session.xenapi.network.get_PIFs(nwref)
        # Find PIF for this host
        hostuuid = inventory.get_localhost_uuid()
        hostref = session.xenapi.host.get_by_uuid(hostuuid)
        for pref in pifrefs:
            if session.xenapi.PIF.get_host(pref) == hostref:
                pifref = pref
                break
        if not pifref:
            raise "Could not find PIF record for network %s on this host" % (nwuuid)

    # Specify by PIF UUID. If this is a VLAN PIF then this forces the
    # specification of a VLAN tag for the VF.
    if args.has_key("pifuuid"):
        pifuuid = args["pifuuid"]
        pifref = session.xenapi.PIF.get_by_uuid(pifuuid)

    if pifref:
        v = str(session.xenapi.PIF.get_VLAN(pifref))
        if v and v != "-1":
            vlan = v
        ethdev = session.xenapi.PIF.get_device(pifref)

    # If no ethdev, PIF or network is specified then reject
    if not ethdev:
        raise "Must specify eth device by device, PIF or network."
    
    # If the caller specified a pass through index use that otherwise
    # find the first one not currently configured.
    ourassignments = _get_vm_assignments(session, vmuuid)
    if args.has_key("index"):
        index = int(args["index"])
    else:
        i = 0
        while True:
            if not i in ourassignments.keys():
                index = i
                break
            i = i + 1

    # Use the user specified MAC or create a local adminstered one
    if args.has_key("mac"):
        mac = args["mac"]
    else:
        mac = _randomMAC()

    if args.has_key("vlan"):
        if vlan and args["vlan"] != vlan:
            raise "Cannot override PIF VLAN %s" % (vlan)
        vlan = args["vlan"]
    
    # Choose a suitable VF. Preference is for lower numbered VFs
    revdict = {}
    for vfid in vfs.keys():
        rvf, rdev = vfs[vfid]
        rkey = "%08u%s" % (int(rvf), rdev)
        revdict[rkey] = vfid
    revdictkeys = revdict.keys()
    revdictkeys.sort()
    vfidlist = [revdict[x] for x in revdictkeys]
    myid = None
    for vfid in vfidlist:
        if not ethdev:
            myid = vfid
            break
        if ethdev == vfs[vfid][1]:
            myid = vfid
            break
    if not myid:
        if ethdev:
            raise "No spare VF on %s" % (ethdev)
        raise "No spare VF"
    
    # Set up the config for the VM. No need to fill out the vendor and
    # device id fields for this.
    ourassignments[index] = (myid, None, None, mac, vlan)
    _set_vm_assignments(session, vmuuid, ourassignments)

    vfnum, vfeth = vfs[myid]

    dom = xml.dom.minidom.Document()
    element = dom.createElement("iovirt")
    dom.appendChild(element)
    
    entry = dom.createElement("vf")
    element.appendChild(entry)

    subentry = dom.createElement("pciid")
    entry.appendChild(subentry)
    subentry.appendChild(dom.createTextNode(myid))

    subentry = dom.createElement("device")
    entry.appendChild(subentry)
    subentry.appendChild(dom.createTextNode(vfeth))
    
    subentry = dom.createElement("vfnum")
    entry.appendChild(subentry)
    subentry.appendChild(dom.createTextNode(vfnum))

    if mac:
        subentry = dom.createElement("mac")
        entry.appendChild(subentry)
        subentry.appendChild(dom.createTextNode(mac))
    if vlan:
        subentry = dom.createElement("vlan")
        entry.appendChild(subentry)
        subentry.appendChild(dom.createTextNode(vlan))

    aentry = dom.createElement("assigned")
    entry.appendChild(aentry)
    
    subentry = dom.createElement("vm")
    aentry.appendChild(subentry)
    subentry.appendChild(dom.createTextNode(vmuuid))
    
    subentry = dom.createElement("index")
    aentry.appendChild(subentry)
    subentry.appendChild(dom.createTextNode(str(index)))
    
    return dom.toprettyxml()

def unassign_vf(session, args, genericpci=False):
    """Unassign a VF/device from a VM."""
    if not args.has_key("uuid"):
        raise "No VM UUID specified, please use the 'uuid' argument"
    vmuuid = args["uuid"]

    # Specify the VF/device by either PCI ID, index, or ethX+VFn
    pciid = None
    index = None
    vfdisplay = ""
    if args.has_key("index"):
        index = int(args["index"])
    if args.has_key("pciid"):
        pciid = args["pciid"]
        vfdisplay = pciid
    if args.has_key("ethdev") and args.has_key("vf") and not genericpci:
        ethdev = args["ethdev"]
        vfnum = args["vf"]
        vfdisplay = "%s VF %s" % (ethdev, vfnum)
        vfs = _get_vfs()
        for vfpciid in vfs.keys():
            if (vfnum, ethdev) == vfs[vfpciid]:
                pciid = vfpciid
                break
        if not pciid:
            raise "Unable to find PCI ID for " + vfdisplay

    if not pciid and index == None:
        if genericpci:
            raise "Need to specify either a pciid or index"
        else:
            raise "Need to specify either a pciid, index or ethdev and vf"

    # Current assignments
    current = _get_vm_assignments(session, vmuuid)

    # Remove the specified ID
    if pciid:
        if not pciid in [x[0] for x in current.values()]:
            raise "VM %s does not have %s assigned" % (vmuuid, vfdisplay)
        for i in current.keys():
            if current[i][0] == pciid:
                del current[i]
    else:
        if not current.has_key(index):
            raise "VM %s does not have PCI passthrough index %u" % (index)
        del current[index]

    # Update the config
    _set_vm_assignments(session, vmuuid, current)
    return ""

def assign_free_pci_device(session, args):
    """Assign a free PCI device on this host to the specified VM."""
    assigned = _get_assignments(session)
    if not args.has_key("vendorid"):
        raise "No vendor ID specified, please use the 'vendorid' argument"
    if not args.has_key("deviceid"):
        raise "No device ID specified, please use the 'deviceid' argument"
    # _get_devices will check syntax of the args
    vendorid = args["vendorid"]
    deviceid = args["deviceid"]
    devices = _get_devices(vendorid, deviceid)

    for pciid in assigned.keys():
        if pciid in devices:
            del devices[pciid]
    if not args.has_key("uuid"):
        raise "No VM UUID specified, please use the 'uuid' argument"
    vmuuid = args["uuid"]

    # If the caller specified a pass through index use that otherwise
    # find the first one not currently configured.
    ourassignments = _get_vm_assignments(session, vmuuid)
    if args.has_key("index"):
        index = int(args["index"])
    else:
        i = 0
        while True:
            if not i in ourassignments.keys():
                index = i
                break
            i = i + 1

    # Choose a suitable device. Preference is for lower numbered PCIIDs
    pciids = devices.keys()
    pciids.sort()
    myid = None
    if len(pciids) > 0:
        myid = pciids[0]
    if not myid:
        raise "No spare %s:%s device" % (vendorid, deviceid)
    
    # Set up the config for the VM. No need to fill out the vendor and
    # device id fields for this.
    ourassignments[index] = (myid, None, None, None, None)
    _set_vm_assignments(session, vmuuid, ourassignments)

    dom = xml.dom.minidom.Document()
    element = dom.createElement("iovirt")
    dom.appendChild(element)
    
    entry = dom.createElement("pcidevice")
    element.appendChild(entry)
    
    subentry = dom.createElement("pciid")
    entry.appendChild(subentry)
    subentry.appendChild(dom.createTextNode(myid))
    
    subentry = dom.createElement("vendorid")
    entry.appendChild(subentry)
    subentry.appendChild(dom.createTextNode(vendorid))
    
    subentry = dom.createElement("deviceid")
    entry.appendChild(subentry)
    subentry.appendChild(dom.createTextNode(deviceid))
    
    aentry = dom.createElement("assigned")
    entry.appendChild(aentry)
    
    subentry = dom.createElement("vm")
    aentry.appendChild(subentry)
    subentry.appendChild(dom.createTextNode(vmuuid))
    
    subentry = dom.createElement("index")
    aentry.appendChild(subentry)
    subentry.appendChild(dom.createTextNode(str(index)))
    
    return dom.toprettyxml()

def unassign_pci_device(session, args):
    return unassign_vf(session, args, genericpci=True)
    
def show_summary(session, args):
    """Return a textual summary of SR-IOV configuarion of this host."""
    assigned = _get_assignments(session)
    vfs = _get_vfs()
    vfids = vfs.keys()
    vfids.sort()
    dom = xml.dom.minidom.Document()
    element = dom.createElement("iovirt")
    dom.appendChild(element)

    for vfid in vfids:
        entry = dom.createElement("vf")
        element.appendChild(entry)
        
        vfnum, vfeth = vfs[vfid]

        subentry = dom.createElement("pciid")
        entry.appendChild(subentry)
        subentry.appendChild(dom.createTextNode(vfid))

        subentry = dom.createElement("device")
        entry.appendChild(subentry)
        subentry.appendChild(dom.createTextNode(vfeth))

        subentry = dom.createElement("vfnum")
        entry.appendChild(subentry)
        subentry.appendChild(dom.createTextNode(vfnum))

        if assigned.has_key(vfid):
            vmuuid, vmnum, vendorid, deviceid, mac, vlan = assigned[vfid]
            aentry = dom.createElement("assigned")
            entry.appendChild(aentry)

            subentry = dom.createElement("vm")
            aentry.appendChild(subentry)
            subentry.appendChild(dom.createTextNode(vmuuid))

            subentry = dom.createElement("index")
            aentry.appendChild(subentry)
            subentry.appendChild(dom.createTextNode(str(vmnum)))

            # MAC and VLAN go in the parent node
            if mac:
                subentry = dom.createElement("mac")
                entry.appendChild(subentry)
                subentry.appendChild(dom.createTextNode(mac))
            if vlan:
                subentry = dom.createElement("vlan")
                entry.appendChild(subentry)
                subentry.appendChild(dom.createTextNode(vlan))

    return dom.toprettyxml()

def list_pci_devices(session, args):
    """Return a list of PCI devices of the specified vendorid:deviceid and their assignments."""
    if not args.has_key("vendorid"):
        raise "No vendor ID specified, please use the 'vendorid' argument"
    if not args.has_key("deviceid"):
        raise "No device ID specified, please use the 'deviceid' argument"
    # _get_devices will check syntax of the args
    vendorid = args["vendorid"]
    deviceid = args["deviceid"]
    devices = _get_devices(vendorid, deviceid)
    assigned = _get_assignments(session)
    pciids = devices.keys()
    pciids.sort()
    dom = xml.dom.minidom.Document()
    element = dom.createElement("iovirt")
    dom.appendChild(element)
    for pciid in pciids:
        entry = dom.createElement("pcidevice")
        element.appendChild(entry)
        
        subentry = dom.createElement("pciid")
        entry.appendChild(subentry)
        subentry.appendChild(dom.createTextNode(pciid))

        subentry = dom.createElement("vendorid")
        entry.appendChild(subentry)
        subentry.appendChild(dom.createTextNode(vendorid))

        subentry = dom.createElement("deviceid")
        entry.appendChild(subentry)
        subentry.appendChild(dom.createTextNode(deviceid))

        if assigned.has_key(pciid):
            vmuuid, vmnum, vendorid, deviceid, mac, vlan = assigned[pciid]
            aentry = dom.createElement("assigned")
            entry.appendChild(aentry)

            subentry = dom.createElement("vm")
            aentry.appendChild(subentry)
            subentry.appendChild(dom.createTextNode(vmuuid))

            subentry = dom.createElement("index")
            aentry.appendChild(subentry)
            subentry.appendChild(dom.createTextNode(str(vmnum)))

            if mac:
                subentry = dom.createElement("mac")
                aentry.appendChild(subentry)
                subentry.appendChild(dom.createTextNode(mac))
            if vlan:
                subentry = dom.createElement("vlan")
                aentry.appendChild(subentry)
                subentry.appendChild(dom.createTextNode(vlan))

    return dom.toprettyxml()

def get_vm(session, args):
    """Return a description of the SR-IOV config for the specified VM."""
    if not args.has_key("uuid"):
        raise "No VM UUID specified, please use the 'uuid' argument"
    vmuuid = args["uuid"]
    vfs = _get_vfs()
    current = _get_vm_assignments(session, vmuuid)
    indexlist = current.keys()
    indexlist.sort()
    dom = xml.dom.minidom.Document()
    element = dom.createElement("iovirt")
    dom.appendChild(element)
    vmentry = dom.createElement("vm")
    element.appendChild(vmentry)
    subentry = dom.createElement("uuid")
    vmentry.appendChild(subentry)
    subentry.appendChild(dom.createTextNode(vmuuid))
    for i in indexlist:
        entry = dom.createElement("passthrough")
        vmentry.appendChild(entry)
        subentry = dom.createElement("index")
        entry.appendChild(subentry)
        subentry.appendChild(dom.createTextNode(str(i)))
        
        pciid, vendorid, deviceid, mac, vlan = current[i]
        
        subentry = dom.createElement("pciid")
        entry.appendChild(subentry)
        subentry.appendChild(dom.createTextNode(pciid))
        subentry = dom.createElement("vendorid")
        entry.appendChild(subentry)
        subentry.appendChild(dom.createTextNode(vendorid))
        subentry = dom.createElement("deviceid")
        entry.appendChild(subentry)
        subentry.appendChild(dom.createTextNode(deviceid))
        if mac:
            subentry = dom.createElement("mac")
            entry.appendChild(subentry)
            subentry.appendChild(dom.createTextNode(mac))
        if vlan:
            subentry = dom.createElement("vlan")
            entry.appendChild(subentry)
            subentry.appendChild(dom.createTextNode(vlan))
        if pciid in vfs:
            vfnum, ethdev = vfs[pciid]
            subentry = dom.createElement("vfnum")
            entry.appendChild(subentry)
            subentry.appendChild(dom.createTextNode(vfnum))
            subentry = dom.createElement("device")
            entry.appendChild(subentry)
            subentry.appendChild(dom.createTextNode(ethdev))
            subentry = dom.createElement("pttype")
            entry.appendChild(subentry)
            subentry.appendChild(dom.createTextNode("vf"))
        else:
            subentry = dom.createElement("pttype")
            entry.appendChild(subentry)
            subentry.appendChild(dom.createTextNode("pcidevice"))
    return dom.toprettyxml()

def prep_for_vm(session, args):
    if not args.has_key("uuid"):
        raise "No VM UUID specified, please use the 'uuid' argument"
    vmuuid = args["uuid"]
    current = _get_vm_assignments(session, vmuuid)
    vfs = _get_vfs()
    reply = []
    for i in current.keys():
        pciid, vendorid, deviceid, mac, vlan = current[i]
        if not pciid in vfs.keys():
            # This is probably a non SR-IOV PCI device being passed
            # through. Log that we cannot find details and carry on
            # processing the device list.
            syslog.syslog("Could not find VF details for %s for %s, assume it is a non SR-IOV passthrough" % (pciid, vmuuid))
            continue
        vfnum, ethdev = vfs[pciid]
        if mac:
            # Check MAC really is a MAC
            if not re.match("^([0-9a-fA-F:]+)$", mac):
                raise "Unexpected MAC text '%s' for VM %s" % (mac, vmuuid)
            cmd = "%s link set %s vf %s mac %s" % (ipcmd, ethdev, vfnum, mac)
            reply.append(cmd)
            syslog.syslog("Setting VF MAC with '%s' for VM %s" % (cmd, vmuuid))
            os.system(cmd)
        if not vlan:
            # No VLAN may need explicit clearly of previously assigned VLAN
            vlan = "0"
        # Check VLAN really is an integer
        if not re.match("^\d+$", vlan):
            raise "Unexpected VLAN text '%s' for VM %s" % (vlan, vmuuid)
        cmd2 = "%s link set %s vf %s vlan %s" % (ipcmd, ethdev, vfnum, vlan)
        reply.append(cmd2)
        syslog.syslog("Setting VF VLAN with '%s' for VM %s" % (cmd, vmuuid))
        os.system(cmd2)
    return "\n".join(reply)

def unassign_all(session, args):
    """Clear all SR-IOV and PCI passthrough devices from a VM."""
    if not args.has_key("uuid"):
        raise "No VM UUID specified, please use the 'uuid' argument"
    vmuuid = args["uuid"]
    
    # Update the config
    _set_vm_assignments(session, vmuuid, {})
    return ""

def change_vf_mac(session, args):
    """Change the MAC address for a VF already assigned to a VM."""

    if not args.has_key("mac"):
        raise "Need to specify a new MAC address"
    mac = args["mac"]

    if not args.has_key("uuid"):
        raise "No VM UUID specified, please use the 'uuid' argument"
    vmuuid = args["uuid"]

    # Specify the VF by index
    if not args.has_key("index"):
        raise "Need to specify the VF index"
    index = int(args["index"])

    # Current assignments
    current = _get_vm_assignments(session, vmuuid)

    # Update the MAC
    if not current.has_key(index):
        raise "VF index %u not found for VM %s" % (index, vmuuid)
    c = current[index]
    current[index] = (c[0], c[1], c[2], mac, c[4])

    # Update the config
    _set_vm_assignments(session, vmuuid, current)

    return ""

def change_vf_vlan(session, args):
    """Change the VLAN for a VF already assigned to a VM.
    Use vlan=None to remove VLAN tagging."""

    if not args.has_key("vlan"):
        raise "Need to specify a new VLAN"
    vlan = args["vlan"]
    if vlan.lower() == "none":
        vlan = None

    if not args.has_key("uuid"):
        raise "No VM UUID specified, please use the 'uuid' argument"
    vmuuid = args["uuid"]

    # Specify the VF by index
    if not args.has_key("index"):
        raise "Need to specify the VF index"
    index = int(args["index"])

    # Current assignments
    current = _get_vm_assignments(session, vmuuid)

    # Update the MAC
    if not current.has_key(index):
        raise "VF index %u not found for VM %s" % (index, vmuuid)
    c = current[index]
    current[index] = (c[0], c[1], c[2], c[3], vlan)

    # Update the config
    _set_vm_assignments(session, vmuuid, current)

    return ""

if __name__ == "__main__":
    XenAPIPlugin.dispatch({"enable_iommu":           enable_iommu,
                           "assign_free_vf":         assign_free_vf,
                           "show_summary":           show_summary,
                           "unassign_vf":            unassign_vf,
                           "assign_free_pci_device": assign_free_pci_device,
                           "unassign_pci_device":    unassign_pci_device,
                           "get_vm":                 get_vm,
                           "prep_for_vm":            prep_for_vm,
                           "list_pci_devices":       list_pci_devices,
                           "unassign_all":           unassign_all,
                           "change_vf_vlan":         change_vf_vlan,
                           "change_vf_mac":          change_vf_mac})

