#!/usr/bin/env python
# vsphere-vm-power --- virtual machine power on/off/suspend operations

# Author: Noah Friedman <friedman@splode.com>
# Created: 2018-03-22
# Public domain

# $Id: vsphere-vm-power,v 1.17 2019/08/23 22:59:53 friedman Exp $

# Commentary:
# Code:

from   __future__ import print_function
from   pyVmomi    import vim, vmodl
import vspherelib     as vsl
import sys

cmds = { 'on'          : 'PowerOn',
         'off'         : 'PowerOff',
         'terminate'   : 'TerminateVM',   # kill -9 vmx process

         'vm-reboot'   : 'RebootGuest',
         'vm-suspend'  : 'StandbyGuest',  # vmsvc/hibernate ?
         'vm-shutdown' : 'ShutdownGuest',

         'reset'       : 'Reset',
         'suspend'     : 'Suspend',
         'resume'      : 'PowerOn', }


def get_args():
    p = vsl.ArgumentParser( loadrc=True )
    p.add( 'cmd', choices=cmds,   help='Command to perform on virtual machines' )
    p.add( 'vm',  nargs='+',      help='VM name(s) or pattern(s)' )
    return p.parse()


# TODO: handle questions and arrange for answering if interactive.
class powerOnCallback( object ):
    def __init__( self ):
        self.error = 0
        self.att   = []  # attempted
        self.natt  = []  # not attempted
        self.rec   = []  # recommendations
        # subtasks spawned for attempted power on of individual VMs
        self.tasklist = []

    def entry( self, change, *rest ):
        state = None
        res   = None
        if   change.name == 'info.state':
            state = change.val
        elif change.name == 'info.result':
            res   = change.val
        elif change.name == 'info':
            state = change.val.state
            try:
                res = change.val.result
            except AttributeError:
                pass

        try:
            for elt in res.attempted:
                # For vms managed by DRS, no subtask is returned.
                try:
                    self.tasklist.append( elt.task )
                except AttributeError:
                    self.att.append( elt )
            self.natt.extend( res.notAttempted )
            self.rec.extend(  res.recommendations)
        except AttributeError:
            pass

        if state in [vim.TaskInfo.State.success, vim.TaskInfo.State.error]:
            if state == vim.TaskInfo.State.error:
                self.error = 1

            for elt in self.att:
                print( elt.vm.name, 'Success', sep=': ' )

            for elt in self.natt:
                vsl.printerr( elt.vm.name, elt.fault.msg )

            # This needs work.
            # Recommendations occur when DRS is enabled but set to manual
            # approval.  In the case of recommendations with no warnings,
            # just approve them.  The cases with warning we ought to
            # require some interactive approval unless some --auto-approve
            # option is provided.
            for elt in self.rec:
                if elt.warningText:
                    vsl.printerr( 'recommendation', elt )
                else:
                    vsl.printerr( elt.target.name, 'applying manual recommendations' )
                    for action in elt.action:
                        desc = action._wsdlName
                        plc = '{} => {}'.format(
                            action.target.name,
                            action.targetHost.name )
                        vsl.printerr( desc, plc )
                    elt.target.ApplyRecommendation( elt.key )


class monitorChangeCallback( object ):
    def __init__( self, vmlist ):
        self.succ     = 1
        self.vm_table = { vm : {} for vm in vmlist }

    def entry( self, change, objSet, *rest ):
        name    = change.name
        val     = change.val
        vm      = objSet.obj
        vm_prop = self.vm_table[ vm ]

        if objSet.kind == 'enter':
            vm_prop[ name ] = val
        elif vm_prop.get( name, None ) != val:
            if vsl.debug:
                print( '{}: {}: {} => {}'.format(
                    vm.name, name, vm_prop[ name ], val ))
            else:
                print( vm.name, 'Success', sep=': ' )
            del self.vm_table[ vm ]

        if not self.vm_table:
            return self.succ


# Walk up the chain of parent objects from the vm until we find a datacenter.
def obj_datacenter( obj ):
    while obj:
        if isinstance( obj, vim.Datacenter ):
            return obj
        obj = obj.parent

# For data centers, VMware recommends using this interface as of api 5.1
# rather than powering on VMs directly one at a time, because the latter
# method doesn't provide manual DRS recommendations.  This method also
# involves fewer round trips with the server to launch each task.  (The
# multi-vm method is available with standalone hosts as well.)
def power_on( vsi, vmlist ):
    dc = obj_datacenter( vmlist[0] )
    tasklist = dc.PowerOnMultiVM_Task( vmlist )
    if tasklist:
        callback = powerOnCallback()
        succ = vsi.taskwait( tasklist, callback=callback.entry, printsucc=False )
        if callback.tasklist:
            succ2 = vsi.taskwait( callback.tasklist )
            if not succ2:
                sys.exit( 1 )
        if not succ or callback.error:
            sys.exit( 1 )

def main():
    args = get_args()
    vsi  = vsl.vmomiConnect( args )

    vmlist = vsi.search_by_name( args.vm )
    if not vmlist:
        return

    # esxi4.x doesn't handle PowerOnMultiVM_Task
    if ( args.cmd in ['on', 'resume']
         and vsi.si.content.about.apiVersion >= '5.5' ):
        return power_on( vsi, vmlist )

    op = getattr( type( vmlist[ 0 ] ), cmds[ args.cmd ] )
    tasklist = []
    monilist = []
    for vm in vmlist:
        try:
            task = op( vm )
            if task:
                tasklist.append( task )
            else:
                monilist.append( vm )
        except vmodl.MethodFault as e:
            vsl.printerr( vm.name, e.msg )
        except Exception as e:
            vsl.printerr( 'Caught Exception', str( e ))

    tsucc, msucc = 1, 1
    if tasklist:
        tsucc = vsi.taskwait( tasklist )
    if monilist:
        if args.cmd in ['vm-reboot']:
            proplist = [ 'runtime.bootTime' ]
        else:
            proplist = [ 'runtime.powerState' ]

        msucc = vsi.monitor_property_changes(
            monilist, proplist,
            monitorChangeCallback( monilist ).entry )

    if not (tsucc and msucc):
        sys.exit( 1 )


if __name__ == '__main__':
    try:
        main()
    except KeyboardInterrupt:
        pass

# vsphere-vm-power ends here