#!/usr/bin/env python


# Quick and dirty script to configure cpu and memory on VMs
# 11-2-2017
# escapenguin@gmail.com

import atexit
import argparse
import getpass
import ssl
import sys

from pyVmomi import vim, vmodl
from pyVim.connect import SmartConnect, Disconnect
from pyVim.task import WaitForTask

# Disable SSL certificate verification... disable for production
context = ssl.SSLContext(ssl.PROTOCOL_TLSv1)
context.verify_mode = ssl.CERT_NONE

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

	parser.add_argument('-s', '--host',
						required=True,
						action='store',
						help='vSphere server 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='Username 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('--uuid',
						required=False,
						action='store',
						help='vmuuid of vm')

	parser.add_argument('--vcpu_nu',
						required=False,
						action='store',
						help='Number of vcpu to configure on vm')

	parser.add_argument('--mem_size',
						required=False,
						action='store',
						help='Amount of memory (MB) to configure on vm')

	args = parser.parse_args()

	if not args.password:
		args.password = getpass.getpass(
			prompt='Enter user 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 change_vcpu(vm, si, vcpu_nu):
	vcpu_nu = int(vcpu_nu)
	cspec = vim.vm.ConfigSpec()
	cspec.numCPUs = vcpu_nu
	cspec.numCoresPerSocket = 1
	WaitForTask(vm.Reconfigure(cspec))

def change_memory(vm, si, mem_size):
	mem_size = int(mem_size)
	cspec = vim.vm.ConfigSpec()
	cspec.memoryMB = mem_size
	WaitForTask(vm.Reconfigure(cspec))

def main():
	args = get_args()
	try:
		si = SmartConnect(
			host=args.host,
			user=args.user,
			pwd=args.password,
			port=(args.port),
			sslContext=context)
		atexit.register(Disconnect, si)
		content = si.RetrieveContent()
		vm = None
		if args.uuid:
			search_index = si.content.searchIndex
			vm = search_index.FindByUuid(None, args.uuid, True)
		elif args.vm_name:
			#content = si.RetrieveContent()
			vm = get_obj(content, [vim.VirtualMachine], args.vm_name)

		if vm:
			if args.vcpu_nu and args.mem_size:
				print(args.vcpu_nu)
				print(args.mem_size)
				change_vcpu(vm, si, args.vcpu_nu)
				change_memory(vm, si, args.mem_size)
			elif args.mem_size and not args.vcpu_nu:
				print(args.mem_size)
				change_memory(vm, si, args.mem_size)
			elif args.vcpu_nu and not args.mem_size:
				print(args.vcpu_nu)
				change_vcpu(vm, si, args.vcpu_nu)
			else:
				print("VM not found.")
	except IOError as e:
		print(e)
		pass

if __name__ == "__main__":
	main()