vmware_automate.py

from pyVim.connect import SmartConnectNoSSL, Disconnect
from pyVim.task import WaitForTask
from pyVmomi import vim, vmodl
import argparse
import requests
import time
from datetime import datetime, timedelta
import os
import re
from sys import stderr

import config

def setup():
    return {'esxi': config.ESXI, 'guest': config.GUEST, 'naming_scheme': config.VM_NAMING_SCHEME}

def parse_args():
    parser = argparse.ArgumentParser()

    parser.add_argument('file', help='A file to upload to a guest on the configured ESXI host')

    args = parser.parse_args()

    return args

class ESXI_Host():
    def __init__(self, host, user, password, naming_scheme=None):
        self.host = host
        self._user = user
        self._password = password
        self.naming_scheme = naming_scheme
        if self._setup() is False:
            raise Exception('Failed to setup host connection!')

    def __str__(self):
        return 'ESXI_Host<{}@{}>'.format(self._user, self.host)

    def _setup(self):
        service_instance = self._connect()
        if service_instance is None:
            raise Exception('Could not connect to the ESXI host ({})!'.format(self.host))

        content = service_instance.content
        return content

    def _connect(self):
        try:
            c = SmartConnectNoSSL(host=self.host, user=self._user, pwd=self._password)
            print('Connected to \'{}\' as \'{}\'!'.format(self.host, self._user)) 
            return c
        except:
            print('Unable to connect to \'{}\' as \'{}\'!'.format(self.host, self._user), file=stderr)
            return None

    def getDatacenter(self):
        content = self._setup()
        return content.rootFolder.childEntity[0]

    def getVMs(self):
        datacenter = self.getDatacenter()
        vms = datacenter.vmFolder.childEntity
        if self.naming_scheme:
            vms = list(filter(lambda x: len(x.name.split('.')) == len(self.naming_scheme.split('.')), vms))
        return vms

    def lsVMs(self):
        return [vm.name for vm in self.getVMs()]

    def getHost(self):
        datacenter = self.getDatacenter()
        host = datacenter.hostFolder.childEntity[0].host[0]
        return host

    def getPCIDevices(self):
        host = self.getHost()
        return host.hardware.pciDevice

    def getPCIPassthru(self):
        host = self.getHost()
        return host.configManager.pciPassthruSystem

    def togglePassthru(self, passthru, id: str):
        ids = list(c.id for c in list(passthru.pciPassthruInfo))
        idx = ids.index(id)
        id = passthru.pciPassthruInfo[idx].id
        passthru_state = not passthru.pciPassthruInfo[idx].passthruActive
        config = vim.host.PciPassthruConfig(id=id, passthruEnabled=passthru_state)
        passthru.UpdatePassthruConfig([config])
        print("Passthru toggled for '{0}' to '{1}'!".format(id, passthru_state))
        return True

    def findVM(self, name):
        for vm in self.getVMs():
            if vm.name == name:
                return vm
        return None

class ESXI_Guest(ESXI_Host):
    def __init__(self, host, guest, user, password):
        ESXI_Host.__init__(self, host.host, host._user, host._password)
        self.guest = guest
        self._vm_user = user
        self._vm_password = password

        if self._setupVM() is None:
            raise Exception('Failed to connect to the specified VM ({})'.format(self.guest))

        self.snapshots = self.getSnapshots()

    def __str__(self):
        return "ESXI_Guest<{}@'{}'".format(self._vm_user, self.guest)

    def _setupVM(self):
        vm = self.findVM(self.guest)
        if vm:
            self._auth = vim.vm.guest.NamePasswordAuthentication(username=self._vm_user, password=self._vm_password, interactiveSession=True)
            return vm

        return None

    def _recurseSnapshots(self, snapshot):
        print(snapshot)
        snapshot_data = {'snapshot': snapshot.snapshot, 'name': snapshot.name, 'children': []}

        if len(snapshot.childSnapshotList) == 0:
            return snapshot_data
        for child in snapshot.childSnapshotList:
            child_data = self._recurseSnapshots(child)
            if child_data:
                snapshot_data['children'].append(child_data)
        return snapshot_data

    def _getVM(self):
        return self._setupVM()

    def getSnapshots(self):
        snapshots = [] 
        root = self._getVM().snapshot.rootSnapshotList
        snapshot_ptr = snapshots
        for snapshot in root:
            snapshots.append(self._recurseSnapshots(snapshot))
        return snapshots

    def check_tools_status(self):
        timeout = datetime.now() + timedelta(seconds=600)

        while self._getVM().guest.toolsRunningStatus != 'guestToolsRunning': 
            if datetime.now() > timeout:
                return False
            time.sleep(5)

        while not self._getVM().guest.guestOperationsReady:
            if datetime.now() > timeout:
                return False
            time.sleep(5)

        return True

    def upload(self, src, dest=None, filename=None):
        if dest is None:
            dest = 'C:\\Users\\{}\\Desktop\\{}'.format(self._vm_user, os.path.basename(filename))

        print(dest)

        if self.power_status() != 'running':
            self.power('on')
        while (self.power_status() != 'running'):
            print(self.power_status())
        if self.check_tools_status() is False:
            raise Exception('Unable to reach ok status before timeout!')

        if type(src) is bytes:
            upload_file = src
        else:
            if os.path.isfile(src):
                with open(src, 'rb') as upload_file_fd:
                    upload_file = upload_file_fd.read()

        upload_file_attributes = vim.vm.guest.FileManager.FileAttributes()
        upload_url = self._setup().guestOperationsManager.fileManager.InitiateFileTransferToGuest(self._getVM(), self._auth, dest, upload_file_attributes, len(upload_file), True)
        upload_url = re.sub(r"^https://\*:", "https://" + str(self.host) + ":", str(upload_url))

        status = requests.put(upload_url, data=upload_file, verify=False).status_code
        return True if status == 200 else False

    def power_status(self):
        return self._getVM().guest.guestState

    def getPowerModes(self):
        return ['status', 'on', 'off', 'restart', 'pause', 'terminate']

    def power(self, mode):
        if mode not in self.getPowerModes():
            return {'error': 'invalid power operation'}
        
        status = None
        task = None

        if mode == 'status':
            status = self.power_status()
        elif mode == 'on':
            # permission: VirtualMachine.Interact.PowerOn 
            task = self._getVM().PowerOnVM_Task()
        elif mode == 'off':
            # permission: VirtualMachine.Interact.PowerOff 
            task = self._getVM().PowerOffVM_Task()
        elif mode == 'restart':
            # permission: VirtualMachine.Interact.Reset
            task = self._getVM().ResetVM_Task()
        elif mode == 'pause':
            # permission: VirtualMachine.Interact.Suspend 
            task = self._getVM().SuspendVM_Task()
        elif mode =='terminate':
            # permission: VirtualMachine.Interact.PowerOff
            task = self._getVM().TerminateVM()

        res = {}
        if status is None:
            status = task.info.state
        res['status'] = status
        if task:
            error = task.info.error
            if error is not None:
                res[status] = error.msg

        return res

    def createSnapshot(self, name, memory=False, pause=True):
        print("Creating snapshot for: " + str(self._getVM().name))
        print(self._getVM())
        return self.getVM().CreateSnapshot_Task(name, memory=memory, quiesce=pause)

    def revertToSnapshot(self, name):
        snapshots = self.getSnapshots()
        snapshot = list(filter(lambda x: x.get('name') == name, snapshots)).pop()
        if not snapshot:
            return None

        print(snapshot)
        snapshot = snapshot.get('snapshot')
        task = snapshot.RevertToSnapshot_Task(suppressPowerOn=True)
        print(task)
        print(task.info)
        state = task.info.state
        result = {}
        result['vm'] = task.info.entityName
        result['state'] = state
        if state == 'error':
            result['error'] = task.info.error.msg
        else:
            result['result'] = str(task.info.result)

        print(result)
        return result

    def as_dict(self):
        data = {}
        data['status'] = self.power_status()
        data['user'] = self._vm_user
        data['guest'] = self.guest
        data['snapshots'] = len(self.makeJSON(self.snapshots))
        return data

    def makeJSON(self, obj):
        print(obj)
        if type(obj) is list:
            res = []
            for item in obj:
                res.append(self.makeJSON(item))
        elif type(obj) is dict:
            res = {}
            for key, val in obj.items():
                res[key] = self.makeJSON(val)
        else:
            res = str(obj)

        return res

if __name__ == '__main__':
    args = parse_args()
    if args.file:
        upload_file = args.file

    setup_data = setup()
    default_vms = setup_data['default_vms']
    research_vms = setup_data['research_vms']

    print(default_vms)
    print(research_vms)

    host = ESXI_Host(setup_data['esxi']['host'], setup_data['esxi']['user'], setup_data['esxi']['password'])

    guest = default_vms[0]
    guest = ESXI_Guest(host, guest, setup_data['guest']['user'], setup_data['guest']['password'])

    status = guest.upload(upload_file, dest='C:\\Users\\{}\\Desktop\\{}'.format(setup_data['guest']['user'], os.path.basename(upload_file)))
    print(status)