#!/usr/bin/env python
# vsphere-folder-tree --- show datacenter vm folder hierarchy, optionally with VMs

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

# $Id: vsphere-folder-tree,v 1.8 2018/10/04 15:51:05 friedman Exp $

# Commentary:

# TODO: display other objects besides virtual machines in verbose+all mode.

# Code:

from __future__ import print_function
from pyVmomi    import vim
import vspherelib as vsl
import re
import sys

poweredOff = vim.VirtualMachine.PowerState.poweredOff

slash  = '/'
indent = '    '

def get_args():
    p = vsl.ArgumentParser( loadrc=True )
    p.add_bool( '-a', '--all',     help='Display all folder hierarchies; default is vm-only' )
    p.add_bool( '-v', '--verbose', help='Display extended info' )
    return p.parse()

def print_hier( folder_to_path, path_to_folder, vm_table ):
    def print_vms( folder, depth ):
        try:
            vmlist = map( lambda vm: vm_table[ vm.id ],
                          vsl.get_seq_type( folder.childEntity, vim.VirtualMachine ))
        except AttributeError:
            return
        # Display virtual machines before templates,
        # but within each group display in alphabetical order.
        vmlist.sort( lambda a, b: (
            cmp( a[ 'config.template' ], b[ 'config.template' ] )
            or cmp( a[ 'name' ].lower(), b[ 'name' ].lower() )))
        ws = indent * depth
        for vm in vmlist:
            if vm[ 'config.template' ]:
                flag = '+ '
            elif vm[ 'runtime.powerState' ] == poweredOff:
                flag = 'o '
            else:
                flag = '  '
            print( ws, flag, vm[ 'name' ], sep='' )

    def print_folder( path, depth ):
        last = path[ path.rfind( slash ) + 1 : ]
        print( depth * indent, last, slash, sep='' )

        folder = path_to_folder[ path ]
        del path_to_folder[ path ]
        if isinstance( folder, vim.Datacenter ):
            # Datastore objects have no childEntity attr
            subpaths = [ folder_to_path[ obj ] for obj in
                         ( getattr( folder, attr ) for attr in
                           ( 'datastoreFolder',
                             'hostFolder',
                             'networkFolder',
                             'vmFolder') ) ]
        else:
            subpaths = [ folder_to_path[ obj ] for obj in
                         vsl.get_seq_type( folder.childEntity, vim.Folder ) ]
        subpaths.sort()
        for subpath in subpaths:
            print_folder( subpath, depth + 1 )

        if vm_table:
            print_vms( folder, depth + 1 )

    path_list = path_to_folder.keys()
    path_list.sort( cmp=lambda a, b: cmp( a.lower(), b.lower() ) )
    for path in path_list:
        if path in path_to_folder:
            depth = path.count( slash ) - 1
            print_folder( path, depth )


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

    if args.all:
        path_to_folder = vsi.path_to_folder_map().copy()
        folder_to_path = vsi.folder_to_path_map()
    else:
        path_to_folder = vsi.path_to_subfolder_map( 'vm' )
        folder_to_path = vsl.inverted_dict( path_to_folder )

    if args.verbose:
        props = [ 'name', 'config.template', 'runtime.powerState' ]
        vm_table = { vm[ 'obj' ].id : vm for vm in
                     vsi.get_obj_props( [vim.VirtualMachine], props ) }
        print_hier( folder_to_path, path_to_folder, vm_table )
    else:
        path_list = sorted( path_to_folder.keys() )
        for path in path_list:
            depth = path.count( slash ) - 1
            last = path[ path.rfind( slash ) + 1 : ]
            print( depth * indent, last, slash, sep='' )


if __name__ == '__main__':
    main()

# vsphere-folder-tree ends here