#!/usr/bin/env python
# vsphere-vm-grep --- search for virtual machines by regex on vm properties

# Author: Noah Friedman <friedman@splode.com>
# Created: 2019-05-21
# Public domain

# $Id: vsphere-vm-grep,v 1.4 2019/08/24 03:41:26 friedman Exp $

# Commentary:
# Code:

from   __future__ import print_function
from   pyVmomi    import vim, vmodl

import vspherelib     as vsl
import sys
import re
import functools



def get_args():
    p = vsl.ArgumentParser( loadrc=True )
    p.add( '-v', '--verbose', action='store_true', help='Display matching property and contents' )

    p.add( '-I', '-i', action='store_true', help='Perform case-insensitive matching.')
    p.add( '-L', action='store_true', help='''Make \w, \W, \b, \B, dependent on the current locale.''')
    p.add( '-M', action='store_true', help='''
        "^" matches the beginning of lines (after a newline) as well as the string.
        "$" matches the end of lines (before a newline) as well as the end of the string.''' )
    p.add( '-S', action='store_true', help='"." matches any character at all, including newline.')
    p.add( '-X', action='store_true', help='''Ignore whitespace and comments for nicer looking RE's.''')
    p.add( '-U', action='store_true', help='''Make \w, \W, \b, \B, dependent on the Unicode locale.''')

    p.add( 'regex',                   help='Pattern to search for' )
    p.add( 'properties', nargs='*',   help='vm properties to search (default: config.annotation)' )

    args = p.parse()
    if not args.regex:
        vsl.printerr( 'Specify regex' )
        sys.exit( 1 )
    return args


def matches( regex, obj, name=None, all=False ):
    # Don't descend into other referenced managed objects
    if obj is None or isinstance( obj, (type, bool, vim.ManagedObject) ):
        return

    try:
        if regex.search( obj ):
            return [(name, obj)]
    except TypeError:
        try:
            iter( obj )
        except TypeError:
            obj = vars( obj )

        result = []
        # Sorting the elements makes the results a little more consistent,
        # otherwise what you get depends on hash order.
        for (i, elt) in enumerate( sorted( obj )):
            try:
                val = obj[ elt ]
                subname = str.join( '.', (name, elt))
                m = matches( regex, val, name=subname, all=all )
                if m:
                    result.extend( m )
            except (TypeError, IndexError) as e:
                subname = '{}[{}]'.format( name, i )
                m = matches( regex, elt, name=subname, all=all )
                if m:
                    result.extend( m )
            if result and not all:
                return result
        if result:
            return result

def display( results ):
    for elt in results:
        label, text = elt
        try:
            text = vsl.fold_text( text, maxlen=70, indent=8 )
            if text.find( '\n' ) >= 0:
                text = '\n        ' + text
            print( '    {}: {}'.format( label, text ))
        except AttributeError:
            text = re.sub( '\n+', '\n    ', str( text ), flags=re.M )
            print( '    {} = {}'.format( label, text ))
    print( '' )


def get_obj_props_best_effort( vsi, orig_proplist ):
    # Since you can't specify properties that are subproperties of arrays,
    # we keep trying to include parent properties until everything succeeds.
    # We have to filter the results when this happens.
    proplist = list( orig_proplist )  # make copy, since we modify
    while proplist:
        try:
            vmlist = vsi.get_obj_props( [vim.VirtualMachine], proplist )
        except vmodl.query.InvalidProperty as e:
            culprit = e.name
            proplist.remove( culprit )

            up = culprit.rfind( '.' )
            if up >= 0:
                proplist.append( culprit[ 0 : up ])
        else: # no exception occured
            return vmlist, proplist

def filter_nonrequested( result, requested ):
    want = []
    for elt in result:
        name = re.sub( '\[[0-9]+\]', '', elt[0] )
        if (name in requested
            # the name is a child of any of the requested props
            or any( name.startswith( req ) and name[ len( req ) ] == '.'
                    for req in requested )):
            want.append( elt )
    return want

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

    greplist = args.properties
    if not greplist:
        greplist.append( 'config.annotation' )
    proplist = [ 'name', ]
    proplist.extend( greplist )

    flags = 0
    if args.I: flags |= re.I
    if args.L: flags |= re.L
    if args.M: flags |= re.M
    if args.S: flags |= re.S
    if args.X: flags |= re.X
    if args.U: flags |= re.U
    regex  = re.compile( args.regex, flags=flags )

    vmlist, actual_proplist = get_obj_props_best_effort( vsi, proplist )
    proplists_equiv = proplist == actual_proplist  # eqv, but not eq
    if proplists_equiv:
        iter_proplist = greplist
    else:
        actual_proplist.remove( 'name' )
        iter_proplist = actual_proplist

    keyfn = functools.cmp_to_key( lambda a, b: cmp( a[ 'name' ], b[ 'name' ] ))
    vmlist.sort( key=keyfn )
    for vm in vmlist:
        firstmatch = True
        for prop in iter_proplist:
            try:
                result = matches( regex, vm[ prop ], name=prop, all=args.verbose )
                if result:
                    # Even if we're just testing for any match, we have to
                    # filter out matches from other parts of the proplist
                    # tree that the user didn't request a match on.
                    if iter_proplist is not greplist:
                        result = filter_nonrequested( result, greplist )
                        # And then if there are no results left, it was a
                        # false positive.
                        if not result:
                            continue

                    if firstmatch:
                        print( vm[ 'name' ] )
                        firstmatch = False
                    if args.verbose:
                        display( result )
                    else:
                        break
            except KeyError:
                pass

##########

if __name__ == '__main__':
    main()

# eof