#!/usr/bin/env python3
from pyVmomi import vim
import pyVim.connect
import atexit
import argparse
from pprint import pprint
import siteconf

site_conf = siteconf.read()

def get_args():
    parser = argparse.ArgumentParser(
        description='Clone VM from template')

    parser.add_argument('-s', '--site',
                        required=True,
                        action='store',
                        help='Name of the site to create VM on')

    parser.add_argument('-t', '--template',
                        required=False,
                        action='store',
                        help='Name of the template you wish to clone')

    parser.add_argument('-v', '--vm-name',
                        required=True,
                        action='store',
                        help='Name of the VM you wish to make')

    parser.add_argument('-c', '--cpu',
                        required=False,
                        type=int,
                        action='store',
                        default=1,
                        help='Set number of CPUs')

    parser.add_argument('-m', '--memory',
                        required=False,
                        type=int,
                        action='store',
                        default=2048,
                        help='[MB] Set memory size')

    parser.add_argument('-d', '--disk-size',
                        required=False,
                        type=int,
                        action='store',
                        default=0,
                        help='[GB] Resize LVM disk to this size')

    parser.add_argument('--datastore',
                        required=False,
                        action='store',
                        default=None,
                        help='target datastore')

    args = parser.parse_args()
    args.disk_size = args.disk_size * 1024 * 1024
    return args


def wait_for_task(task):
    """ wait for a vCenter task to finish """
    task_done = False
    while not task_done:
        if task.info.state == 'success':
            return task.info.result

        if task.info.state == 'error':
            print("there was an error")
            task_done = True


def get_obj(content, vimtype, name):
    """
    Return an object by name, if name is None the
    first found object is returned
    """
    obj = None
    container = content.viewManager.CreateContainerView(
        content.rootFolder, vimtype, True)
    for c in container.view:
        if name:
            if c.name == name:
                obj = c
                break
        else:
            obj = c
            break

    return obj


def clone_vm(content, template, vm_name, datacenter_name, datastore_name,
             cluster_or_host):
    datacenter = get_obj(content, [vim.Datacenter], datacenter_name)
    destfolder = datacenter.vmFolder
    datastore = get_obj(content, [vim.Datastore], datastore_name)
    cluster = get_obj(content, [vim.ClusterComputeResource], cluster_or_host)
    if cluster:
        resource_pool = cluster.resourcePool
    else:
        host = get_obj(content, [vim.ComputeResource], cluster_or_host)
        resource_pool = host.resourcePool

    relospec = vim.vm.RelocateSpec()
    relospec.datastore = datastore
    relospec.pool = resource_pool
    clonespec = vim.vm.CloneSpec()
    clonespec.location = relospec
    clonespec.powerOn = False

    print("cloning VM...")
    task = template.Clone(folder=destfolder, name=vm_name, spec=clonespec)
    wait_for_task(task)

def extend_disk(content, vm_name, extend_to_kb):
    vm = get_obj(content, [vim.VirtualMachine], vm_name)
    for dev in vm.config.hardware.device:
        #print("dev "+ dev.deviceInfo.label)
        if hasattr(dev.backing, 'fileName'):
            size_kb = dev.capacityInKB
            if size_kb > 5 * 1024 * 1024: # lvm disk
                if extend_to_kb > size_kb:
                    disk_spec = vim.vm.device.VirtualDeviceSpec()
                    disk_spec.operation = vim.vm.device.VirtualDeviceSpec.Operation.edit
                    disk_spec.device = vim.vm.device.VirtualDisk()
                    disk_spec.device.key = dev.key
                    disk_spec.device.backing = vim.vm.device.VirtualDisk.FlatVer2BackingInfo()
                    disk_spec.device.backing.fileName = dev.backing.fileName
                    disk_spec.device.backing.diskMode = dev.backing.diskMode
                    disk_spec.device.controllerKey = dev.controllerKey
                    disk_spec.device.unitNumber = dev.unitNumber
                    disk_spec.device.capacityInKB = extend_to_kb
                    dev_changes = []
                    dev_changes.append(disk_spec)
                    spec = vim.vm.ConfigSpec()
                    spec.deviceChange = dev_changes
                    print("extending disk '"+ dev.deviceInfo.label +"' to "+ str(extend_to_kb) +" kB")
                    task = vm.ReconfigVM_Task(spec=spec)
                    wait_for_task(task)

def set_cpu_mem(content, vm_name, cpu, mem_mb):
    vm = get_obj(content, [vim.VirtualMachine], vm_name)
    spec = vim.vm.ConfigSpec()
    spec.numCPUs = cpu
    spec.memoryMB = mem_mb
    print("Set cpu="+ str(cpu) +", memory="+ str(mem_mb) +" MB")
    task = vm.ReconfigVM_Task(spec=spec)
    wait_for_task(task)

def power_on(content, vm_name):
    vm = get_obj(content, [vim.VirtualMachine], vm_name)
    print("Powering on '"+ vm_name +"'")
    task = vm.PowerOn()
    wait_for_task(task)

def main():
    args = get_args()

    if not args.site in site_conf:
        print(f"site not found: {args.site}")
        exit(1)
    site = site_conf[args.site];
    #pprint(site)

    si = pyVim.connect.SmartConnect(
        host = site["vcenter"],
        user = site["vcenter_user"],
        pwd = site["vcenter_passwd"],
        disableSslCertValidation = True
        )

    atexit.register(pyVim.connect.Disconnect, si)

    content = si.RetrieveContent()
    template_name = site["template"]
    if args.template is not None:
        template_name = args.template
    template = get_obj(content, [vim.VirtualMachine], template_name)
    if not template:
        print("template not found")
        exit(1)

    vm = get_obj(content, [vim.VirtualMachine], args.vm_name)
    if vm:
        print("VM '"+ args.vm_name +"' exists.")
        exit(1)

    datastore = site["datastore"]
    if args.datastore is not None:
        datastore = args.datastore
    clone_vm(content, template, args.vm_name, site["datacenter"],
        datastore, site["cluster_or_host"])

    set_cpu_mem(content, args.vm_name, args.cpu, args.memory)
    extend_disk(content, args.vm_name, args.disk_size)
    power_on(content, args.vm_name)


if __name__ == "__main__":
    main()

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