Ansible Module - VMWare Update Guest PCI Device

Ansible Module to enable PCI Passthrough for a VM after the Host Device is configured for Passthrough.

#!/usr/bin/env python

ANSIBLE_METADATA = {
    'metadata_version': '1.1',
    'status': ['preview'],
    'supported_by': 'community'
}

DOCUMENTATION = '''
---
module: vmware_update_guest_pci_device

short_description: Module to update PCI Passthrough device for a guest

version_added: "2.4"

description:
    - "Module to update PCI Passthrough device for a guest, based on device_name or device_id"

options:
    hostname:
        description:
            - vSphere service to connect to
        required: true
    username:
        description:
            - username to connect to vSphere hostname
        required: true
    password:
        description:
            - username password
        required: true
    esxi_hostname:
        desciption:
            - Esxi host to make changes to
    device_name:
        description:
            - Device name(s) to enable passthrough for
        required: false
    device_id:
        description:
            - Device id to enable passthrough for \
                device_id takes precedence over device_name
        required: false
    vm_name:
        description:
            - VM that you want to update the PCI \
                passthrough for
        required: false
    device_uuid:
        description:
            - VM device uuid that you want to update the PCI \
                passthrough for, this takes president over \
                vm_name
        required: false
    validate_certs:
        description:
            - Enable ssl hostname certificate verification
        required: false
        default: true
    port:
        description:
            - vSphere port to connect to
        required: false
        default: 443
    state:
        description:
            - Set Passthrough on -> present, off -> absent \
                only changes state, if state isn't set
        required: true
'''

EXAMPLES = '''
# Enable PCI Passthrough on {{ esxi_hostname }} for 'PCI Device Name'
- name: Enable Passthrough
update_vlan_ids:
    hostname: {{ vsphere_host }}
    username: {{ user_name }}
    password: {{ admin_pass }}
    esxi_hostname: {{ esxi_hostname }}
    device_name: 'PCI Device Name'
    vm_name: {{ vm_name }}
    state: 'present'
    validate_certs: "{{ validate_certs }}"

# Disable PCI Passthrough on {{ esxi_hostname }} for 'PCI Device Name'
- name: Disable Passthrough
update_vlan_ids:
    hostname: {{ vsphere_host }}
    username: {{ user_name }}
    password: {{ admin_pass }}
    esxi_hostname: {{ esxi_hostname }}
    device_name: 'PCI Device Name'
    vm_name: {{ vm_name }}
    state: 'absent'
    validate_certs: "{{ validate_certs }}"


# Enable PCI Passthrough on {{ esxi_hostname }} for '0000:0f:00.0'
- name: Enable Passthrough
update_vlan_ids:
    hostname: {{ vsphere_host }}
    username: {{ user_name }}
    password: {{ admin_pass }}
    esxi_hostname: {{ esxi_hostname }}
    device_id: '0000:0f:00.0'
    state: 'present'
    validate_certs: "{{ validate_certs }}"

# Disable PCI Passthrough on {{ esxi_hostname }} for '0000:0f:00.0'
- name: Enable Passthrough
update_vlan_ids:
    hostname: {{ vsphere_host }}
    username: {{ user_name }}
    password: {{ admin_pass }}
    esxi_hostname: {{ esxi_hostname }}
    device_id: '0000:0f:00.0'
    state: 'absent'
    validate_certs: "{{ validate_certs }}"
'''

RETURN = '''
module_args:
    description: Module arguments that are passed in
    type: array
    returned: always
changed:
    description: Set to true if any port is changed, false if no port is changed
    type: bool
    returned: always
changed_pci_ids:
    description: Ids of changed PCI devices
    type: array
    returned: always
msg:
    description: Msg is passed if error or warning
    type: str
    returned: On Error or Warning
'''

from ansible.module_utils.basic import AnsibleModule
from pyVim import connect
from pyVmomi import vmodl
from pyVmomi import vim


def get_device_by_name(device_name, host):
    """ Find devices by name in a host

    Keyword arguments:
    device_name -- the device name to search for
    host -- the host to search

    If no devices found
        return []
    else
        return array of devices found
    """
    obj = []
    for device in host.hardware.pciDevice:
        if device.deviceName == device_name:
            obj.append(device)
    return obj


def get_device_by_id(device_id, host):
    """ Find devices by id in a host

    Keyword arguments:
    device_id -- the device id to search for
    host -- the host to search

    If no devices found
        return []
    else
        return array of devices found
    """
    obj = []
    for device in host.hardware.pciDevice:
        if device.id == device_id:
            obj.append(device)
    return obj


def get_pci_device(devices, device_id):
    """ Find PCI device by device_id

    Keyword arguments:
    device_id -- the device id to search for
    devices -- the devices to search

    If no devices found
        return []
    else
        return array of PCI devices found
    """
    obj = []
    for device in devices:
        if type(device.backing) == (vim.vm.device.VirtualPCIPassthrough.DeviceBackingInfo) \
            and device.backing.id and device.backing.id == device_id:
            obj.append(device)

    return obj


def wait_for_task(task, module):
    """ Wait for a vCenter task to finish

    Keyword arguments:
    task -- the current scheduled task
    module -- reference to ansible module, so we
        can send info back to ansible
    """
    task_done = False
    while not task_done:
        if task.info.state == 'success':

            return task.info.result

        if task.info.state == 'error':
            module.warn('There was an error updating...' + str(task.info.error))
            task_done = True


def get_obj(content, vimtype, name):
    """ Get view for a particular type and name

    Keyword arguments:
    content -- reference to vmware connection obj
    vimtype -- type of obj to search for
    name -- name of obj to search for

    If view not found
        return None
    else
        return view
    """
    obj = None
    container = content.viewManager.CreateContainerView(
        content.rootFolder, vimtype, True)
    for c in container.view:
        if c.name == name:
            obj = c
            break
    return obj


def run_module():
    """ Define available arguments/parameters a user can pass to the module

    See doc at top of file for details of input/outputs
    """
    module_args = dict(
        hostname=dict(type='str', required=True),
        username=dict(type='str', required=True),
        password=dict(type='str', required=True, no_log=True),
        validate_certs=dict(type='bool', required=False, default=True),
        port=dict(type='int', required=False, default=443),
        esxi_hostname=dict(type='str', required=True),
        device_name=dict(type='str', required=False),
        device_id=dict(type='str', required=False),
        device_uuid=dict(type='str', required=False),
        vm_name=dict(type='str', required=True),
        vm_shutdown=dict(type='bool', required=False, default=True),
        state=dict(choices=['present', 'absent'], required=True)
    )

    module = AnsibleModule(
        argument_spec=module_args,
        supports_check_mode=True
    )

    result = dict(
        changed=False,
    )

    if module.check_mode:
        module.exit_json(**result)

    try:
        if module.params['validate_certs']:
            service_instance = connect.SmartConnect(host=module.params['hostname'],
                                                    user=module.params['username'],
                                                    pwd=module.params['password'],
                                                    port=int(module.params['port']))
        else:
            service_instance = connect.SmartConnectNoSSL(host=module.params['hostname'],
                                                        user=module.params['username'],
                                                        pwd=module.params['password'],
                                                        port=int(module.params['port']))

        search_index = service_instance.content.searchIndex

        content = service_instance.RetrieveContent()


        if module.params['device_uuid']:
            vm = search_index.FindByUuid(None, module.params['device_uuid'], True, True)
        else:
            vm = get_obj(content, [vim.VirtualMachine], module.params['vm_name'])

        if not vm:
            module.fail_json(msg='VM not found device_uuid=' + str(module.params['device_uuid'])\
                + ' vm_name=' + str(module.params['vm_name']), **result)

        objview = content.viewManager.CreateContainerView(content.rootFolder,
                                                        [vim.HostSystem],
                                                        True)
        esxi_hosts = objview.view
        objview.Destroy()

        host = None
        for esxi_host in esxi_hosts:
            if module.params['esxi_hostname'] == esxi_host.name:
                host = esxi_host

        if not host:
            module.fail_json(msg='Host not found: ' + module.params['esxi_hostname'], **result)

        if not (module.params['device_name'] or module.params['vm_name']):
            module.fail_json(msg='Please specify device_id or device_name...', **result)

        devicesToSet = []
        if module.params['device_id']:
            devicesToSet = get_device_by_id(module.params['device_id'], host)

            if not devicesToSet:
                module.fail_json(msg='Device not found... ' + module.params['device_id'], **result)
        elif module.params['device_name']:
            devicesToSet = get_device_by_name(module.params['device_name'], host)

        if not devicesToSet:
            module.fail_json(msg='Device not found...' + module.params['device_name'], **result)


        systemSpecs = vim.EnvironmentBrowserConfigOptionQuerySpec()
        guestId = []
        guestId.append(vm.summary.config.guestId)
        systemSpecs.guestId = guestId
        systemSpec = vm.environmentBrowser.QueryConfigOptionEx(systemSpecs)
        queryConfigTarget = vm.environmentBrowser.QueryConfigTarget()

        for deviceToSet in devicesToSet:
            systemId = None
            for pciPassthroughDevices in queryConfigTarget.pciPassthrough:
                if pciPassthroughDevices.pciDevice.id == deviceToSet.id:
                    systemId = pciPassthroughDevices.systemId

            if not systemId:
                module.fail_json(msg='systemId not found for device ' + deviceToSet.id, **result)


            if module._verbosity >= 1:
                module.warn("SystemId: " + systemId)
                module.warn("VendorId: " + str(deviceToSet.vendorId))
                module.warn("Id: " + deviceToSet.id)
                module.warn("DeviceId: " + str(hex(deviceToSet.deviceId).replace('Ox', '')))
                module.warn("DeviceName: " + deviceToSet.deviceName)

            spec = vim.VirtualMachineConfigSpec()
            deviceChanges = []
            deviceChange = vim.VirtualDeviceConfigSpec()
            device = vim.VirtualPCIPassthrough()
            backing = vim.VirtualPCIPassthroughDeviceBackingInfo()
            backing.systemId = systemId
            backing.vendorId = deviceToSet.vendorId
            backing.id = deviceToSet.id
            backing.deviceId = str(hex(deviceToSet.deviceId).replace('Ox', ''))
            backing.deviceName = str(deviceToSet.deviceName)

            device.backing = backing

            deviceInfo = vim.Description()

            device.deviceInfo = deviceInfo

            deviceChange.device = device

            setdevices = get_pci_device(vm.config.hardware.device, deviceToSet.id)

            if not setdevices and module.params['state'] == 'present':
                deviceChange.operation = 'add'

                deviceChanges.append(deviceChange)
                spec.deviceChange = deviceChanges

                cpuFeatureMasks = []
                cpuFeatureMask = vim.VirtualMachineCpuIdInfoSpec()
                cpuFeatureMask.operation = 'add'
                cpuFeatureMasks.append(cpuFeatureMask)

                if module.params['vm_shutdown']:
                    if format(vm.runtime.powerState) == "poweredOn":
                        task = vm.PowerOffVM_Task()
                        wait_for_task(task, module)

                task = vm.ReconfigVM_Task(spec)
                wait_for_task(task, module)
                result['changed'] = True

            for setdevice in setdevices:
                print("here")

                if module.params['state'] == 'absent':
                    deviceChange.operation = 'remove'
                    deviceChange.device.key = setdevice.key

                    deviceChanges.append(deviceChange)
                    spec.deviceChange = deviceChanges

                    cpuFeatureMasks = []
                    cpuFeatureMask = vim.VirtualMachineCpuIdInfoSpec()
                    cpuFeatureMask.operation = 'add'
                    cpuFeatureMasks.append(cpuFeatureMask)

                    if module.params['vm_shutdown']:
                        if format(vm.runtime.powerState) == "poweredOn":
                            task = vm.PowerOffVM_Task()
                            wait_for_task(task, module)

                    task = vm.ReconfigVM_Task(spec)
                    wait_for_task(task, module)
                    result['changed'] = True

    except vmodl.MethodFault as error:
        module.fail_json(msg='Caught vmodl fault: ' + error.message, **result)
    except EnvironmentError as error:
        module.fail_json(msg='Connection error: ' + error.strerror, **result)

    module.exit_json(**result)


def main():
    """ Call the real method to do work """
    run_module()


# Start program
if __name__ == "__main__":
    main()