#!/usr/bin/env python

# a few notes:
# 1. this was built to work with the check_ncpa.py plugin, however - it will likely work with any other check, provided they have valid perfdata
#    containing a portion labeled exactly 'total', and the size that 'total' reports is in GiB
# 2. this will only work with check_ncpa with the default units or GiB specified. any other unit will cause unspecified behavior
# 3. this ONLY will fire off if the state is CRITICAL and the statetype is HARD


# STEPS:
# install the wizard via admin > manage configwizards
# create the following files: /usr/local/nagiosxi/var/las_vegas /usr/local/nagiosxi/var/studio_city with the appropriate password for each - and NO TRAILING NEWLINE
# set the ownership on both files to nagios:nagios


import pyVmomi
from pyVmomi import vim
from pyVmomi import vmodl
from pyVim.connect import SmartConnect, Disconnect, SmartConnectNoSSL
import atexit
import argparse
import getpass
import sys
import time


def get_args():
    parser = argparse.ArgumentParser(description='Arguments for talking to vCenter')

    # event handler arguments
    parser.add_argument('', '--state',
                        required=True,
                        action='store',
                        help='current state of the service. OK, UNKNOWN, WARNING, or CRITICAL')

    parser.add_argument('', '--statetype',
                        required=True,
                        action='store',
                        help='current state type of the service. SOFT or HARD')

    parser.add_argument('', '--perfdata',
                        required=True,
                        action='store',
                        help='the perfdata from the check. expecting at least a "total" unit to indicate existing disk size')

    # extending disk arguments
    parser.add_argument('-H', '--host',
                        required=True,
                        action='store',
                        help='vsphere host to connect to')

    parser.add_argument('-o', '--port',
                        type=int,
                        default=443,
                        action='store',
                        help='port to connect on')

    parser.add_argument('-u', '--user',
                        required=True,
                        action='store',
                        help='user name to use')

    parser.add_argument('-p', '--password',
                        required=False,
                        action='store',
                        help='password to use')

    parser.add_argument('-v', '--vm-name',
                        required=True,
                        action='store',
                        help='name of the vm')

    parser.add_argument('-P', '--path',
                        required=True,
                        help='the path to the disk (in the datastore)')

    parser.add_argument('-s', '--size',
                        required=True,
                        action='store',
                        help='percentage of current disk size to increase by')

    parser.add_argument('-d', '--datastore',
                        required=True,
                        action='store',
                        help='the datastore where the disk lives')

    parser.add_argument('-S', '--skip-ssl',
                        type=int,
                        default=1)

    args = parser.parse_args()

    if not args.password:
        args.password = getpass.getpass(prompt='Enter password: ')

    return args


def get_obj(content, vimtype, name):
    obj = None
    container = content.viewManager.CreateContainerView(content.rootFolder, vimtype, True)

    for c in container.view:
        if c.name == name:
            obj = c
            break

    return obj


def wait_for_task(task):
    while task.info.state not in [vim.TaskInfo.State.success, vim.TaskInfo.State.error]:
        time.sleep(1)


def get_disk_size(perfdata, percentage_to_increase):
    for data in perfdata.split(' '):
        if "'total'" not in data:
            continue
        value = data.split('=')[1].split(';')[0].replace('GiB', '')
        print float(value) * ((float(100) + percentage_to_increase) / 100)


def main():
    args = get_args()

    # this script only executes when the service is critical and all checks have failed
    if args.state != "CRITICAL" or args.statetype != "HARD":
        sys.exit(0)

    try:
        disk_size = get_disk_size(args.perfdata, args.size)
    except:
        print("invalid perfdata supplied!")
        sys.exit(1)

    if args.skip_ssl == 1:
        si = SmartConnectNoSSL(
            host=args.host,
            user=args.user,
            pwd=args.password,
            port=args.port)
    else:
        si = SmartConnect(
            host=args.host,
            user=args.user,
            pwd=args.password,
            port=args.port)

    atexit.register(Disconnect, si)
    content = si.RetrieveContent()

    vm = get_obj(content, [vim.VirtualMachine], args.vm_name)
    virtualDiskManager = si.content.virtualDiskManager

    if not vm:
        print("VM not found!")
        sys.exit(2)

    disk = "[%s] %s" % (args.datastore, args.path)
    
    wait_for_task(virtualDiskManager.ExtendVirtualDisk(disk, None, disk_size, False))


main()