Ansible Module - VMWare Update Host PCI Passthrough

Ansible Module to enable PCI Passthrough for a Device on VMWare Host.

#!/usr/bin/env python

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

DOCUMENTATION = '''
---
module: vmware_update_host_pci_passthrough

short_description: Module to update PCI Passthrough device for a host

version_added: \"2.4\"

description:
    - \"Module to update PCI Passthrough device for a host, 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:
        description:
            - 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
    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'
    state: 'present'
    validate_certs: "{{ validate_certs }}"

# Disable 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\'
    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
import time


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 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),
        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']))

        content = service_instance.RetrieveContent()

        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['device_id']):
            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)

        hostPciPassthruSystems = host.configManager.pciPassthruSystem

        changed_pci_ids = []
        for deviceToSet in devicesToSet:
            for pciInfo in hostPciPassthruSystems.pciPassthruInfo:
                if deviceToSet.id == pciInfo.id:
                    if module.params['state'].lower() == 'absent' and pciInfo.passthruEnabled == True:
                        configs = []
                        config = vim.HostPciPassthruConfig()
                        config.passthruEnabled = False
                        config.id = deviceToSet.id
                        configs.append(config)
                        hostPciPassthruSystems.UpdatePassthruConfig(configs)

                        # wait 60 second for change to take effect
                        # TODO should change this to recheck if change has been updated
                        #      rather than sleep for 60 seconds
                        time.sleep( 60 )

                        result['changed'] = True
                        changed_pci_ids.append(config.id)
                    elif module.params['state'].lower() == 'present' \
                        and pciInfo.passthruEnabled == False \
                            and pciInfo.passthruCapable == True:
                        configs = []
                        config = vim.HostPciPassthruConfig()
                        config.passthruEnabled = True
                        config.id = deviceToSet.id
                        configs.append(config)
                        hostPciPassthruSystems.UpdatePassthruConfig(configs)

                        # wait 60 second for change to take effect
                        # TODO should change this to recheck if change has been updated
                        #      rather than sleep for 60 seconds
                        time.sleep( 60 )

                        result['changed'] = True
                        changed_pci_ids.append(config.id)
                    elif pciInfo.passthruCapable == False:
                        module.warn_json(msg='Device doesn\'t all passthru' + pciInfo.id, **result)

        result['changed_pci_ids'] = changed_pci_ids

    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():
    run_module()


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