#!/usr/bin/env python3
from pyVmomi import vim
import pyVim.connect
import atexit
import argparse
import re
import siteconf
 
site_conf = siteconf.read()


def get_args():
    parser = argparse.ArgumentParser(
        description='Set VM network')

    parser.add_argument('-s', '--site',
                        required=True,
                        action='store',
                        help='name of the site to connect to')
    parser.add_argument('-d', '--datastore',
                        required=True,
                        action='store',
                        help='datastore name')
    parser.add_argument('-W', '--warning',
                        required=False,
                        action='store',
                        default='20%',
                        help='warning if less free space (nnn%, or nnn[GT]B, default: TB)')
    parser.add_argument('-C', '--critical',
                        required=False,
                        action='store',
                        default='10%',
                        help='critical if less free space')

    args = parser.parse_args()
    return args

def get_datastore(content, name):
    obj_view = content.viewManager.CreateContainerView(
        content.rootFolder, [vim.Datastore], True)
    ds_list = obj_view.view
    obj_view.Destroy()
    for ds in ds_list:
        if not ds.summary.accessible:
            continue
        if ds.summary.name == name:
            return ds
    return None

def less_free(ds, minimum):
    cap = ds.summary.capacity
    free = ds.summary.freeSpace
    ratio = free / cap
    if m := re.match("([0-9.]+)%$", minimum):
        min_ratio = float(m[1]) / 100
        #print(f"min_ratio: {min_ratio}")
        if ratio < min_ratio:
            return True
        else:
            return False
    elif m := re.match("([0-9.]+)([GT]B?)$", minimum):
        min_n, unit = float(m[1]), m[2]
        if unit[0] == 'G':
            min_n *= 1024**3
        if unit[0] == 'T':
            min_n *= 1024**4
    else:
        min_n = float(min_n) * 1024**4 # TB
    #print(f"min_n: {min_n}")
    if ds.summary.freeSpace < min_n:
        return True
    else:
        return False


def main():
    args = get_args()

    if not args.site in site_conf:
        print(f"site not found: {args.site}")
        exit(1)

    host = site_conf[args.site]['vcenter']
    user = site_conf[args.site]['vcenter_user']
    passwd = site_conf[args.site]['vcenter_passwd']

    si = pyVim.connect.SmartConnect(
        host = host,
        user = user,
        pwd = passwd,
        disableSslCertValidation = True
        )
    atexit.register(pyVim.connect.Disconnect, si)
    content = si.RetrieveContent()
    ds = get_datastore(content, args.datastore)
    if ds == None:
        print(f"UNKNOWN: datastore {args.site} | {args.datastore} not found.")
        exit(3)
    cap = ds.summary.capacity / 1024**4
    free = ds.summary.freeSpace / 1024**4
    ratio = free / cap
    text = f"{ds.summary.name} {ratio:.0%} free: {free:.3f} TB / {cap:.3f} TB"
    if less_free(ds, args.critical):
        print(f"CRITICAL: {text}")
        exit(2)
    elif less_free(ds, args.warning):
        print(f"WARNING: {text}")
        exit(1)
    else:
        print(f"OK: {text}")


if __name__ == "__main__":
    main()

# vim: set tabstop=4 shiftwidth=4 expandtab smarttab: