#!/usr/bin/env python

# Copyright 2016 Abdul Anshad <abdulanshad33@gmail.com>
#
#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.

"""
Written by Abdul Anshad
Github: https://github.com/Abdul-Anshad-A
Email: abdulanshad33@gmail.com

Credits:
Thanks to "reuben.13@gmail.com" for the initial code.

Note: Example code For testing purposes only
vSphere Python SDK program to perform snapshot operations.
"""

import atexit
from time import sleep
import argparse
import sys
import time
import ssl
from pyVmomi import vim, vmodl
from pyVim.task import WaitForTask
from pyVim import connect
from pyVim.connect import Disconnect, SmartConnect, GetSi




def get_args():
	parser =  argparse.ArgumentParser(description='Arguments for creating VM from template')
	#operating system arg
	parser.add_argument('--host',
					 required= True,
					 action='store',
					 help='vCenter server IP or hostname')
	parser.add_argument('--user',
					 required= True,
					 action='store',
					 help='vCenter server username')
	parser.add_argument('--pwd',
					 required= True,
					 action='store',
					 help='vCenter server password')
	parser.add_argument('--name',
					 required= True,
					 action='store',
					 help='VM name')
	parser.add_argument('--action',
					 required= True,
					 action='store',
					 help='Snapshot action create/remove/revert ')
	parser.add_argument('--snapshot_name',
					 required= False,
					 action='store',
					 help='Snapshot name')
	parser.add_argument('--snapshot_id',
					 required= False,
					 action='store',
					 help='Snapshot id')
	args = parser.parse_args()

	return args



def get_obj(content, vimtype, name):
	"""
	 Get the vsphere object associated with a given text 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 get_snapshots_by_name_recursively(snapshots, snapname):
	snap_obj = []
	for snapshot in snapshots:
		if snapshot.name == snapname:
			snap_obj.append(snapshot)
		else:
			snap_obj = snap_obj + get_snapshots_by_name_recursively(
									snapshot.childSnapshotList, snapname)
	return snap_obj

def get_snapshots_by_id_recursively(snapshots, snapid):
	snap_obj = []
	for snapshot in snapshots:
		if snapshot.id == snapid:
			snap_obj.append(snapshot)
		else:
			snap_obj = snap_obj + get_snapshots_by_id_recursively(
									snapshot.childSnapshotList, snapid)
	return snap_obj


def get_current_snap_obj(snapshots, snapob):
	snap_obj = []
	for snapshot in snapshots:
		if snapshot.snapshot == snapob:
			snap_obj.append(snapshot)
		snap_obj = snap_obj + get_current_snap_obj(	snapshot.childSnapshotList, snapob)
	return snap_obj


def main():
	args =get_args()


	inputs = {'vcenter_ip': args.host,
		  'vcenter_password': args.pwd,
		  'vcenter_user': args.user,
		  'vm_name': args.name,
		  # operation in 'create/remove/revert'

		  'operation': args.action,
		  'snapshot_name': args.snapshot_name,
		  'snapshot_id': args.snapshot_id,
		  'ignore_ssl': True
		  }
	si = None

	#print("Trying to connect to VCENTER SERVER . . .")

	context = None
	if inputs['ignore_ssl'] and hasattr(ssl, "_create_unverified_context"):
		context = ssl._create_unverified_context()

	si = connect.Connect(inputs['vcenter_ip'], 443,
						 inputs['vcenter_user'], inputs[
							 'vcenter_password'],
						 sslContext=context)

	atexit.register(Disconnect, si)

	#print("Connected to VCENTER SERVER !")

	content = si.RetrieveContent()

	operation = inputs['operation']
	vm_name = inputs['vm_name']

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

	if not vm:
		print("ERROR\nVirtual Machine %s doesn't exists" % vm_name)
		sys.exit()

	if operation != 'create' and vm.snapshot is None:
		print("ERROR\nVirtual Machine %s doesn't have any snapshots" % vm.name)
		sys.exit()

	if operation == 'create':
		snapshot_name = inputs['snapshot_name']
		description = "Automated Snapshot"
		dumpMemory = False
		quiesce = False

		#print("Creating snapshot %s for virtual machine %s" % (snapshot_name, vm.name))
		result = wait_for_task(vm.CreateSnapshot(snapshot_name, description, dumpMemory, quiesce))
		if (result == "true"):
			print("DONE")
		else:
			print("ERROR")
	elif operation in ['remove', 'revert']:
		snapshot_id = int(inputs['snapshot_id'])
		snap_obj = get_snapshots_by_id_recursively(vm.snapshot.rootSnapshotList, snapshot_id)
		#print(snap_obj);
		# if len(snap_obj) is 0; then no snapshots with specified name
		if len(snap_obj) == 1:
			snap_obj = snap_obj[0].snapshot
			if operation == 'remove':
				#print("Removing snapshot %s" % snapshot_name)
				result = wait_for_task(snap_obj.RemoveSnapshot_Task(True))
				if (result == "true"):
					print("DONE")
				else:
					print("ERROR")
			else:
				#print("Reverting to snapshot %s" % snapshot_name)
				result = wait_for_task(snap_obj.RevertToSnapshot_Task())
				if (result == "true"):
					print("DONE")
				else:
					print("ERROR")
		else:
			print("MESSAGE\nNo snapshots found with the id " + str(snapshot_id) + " provided on VM " + vm_name)


	else:
		print("Specify operation in "
			  "create/remove/revert")


def wait_for_task(task):
	""" wait for a vCenter task to finish """
	task_done = False
	while not task_done:
		if task.info.state == 'success':
			task_done = True
			return "true"
		if task.info.state == 'error':
			task_done = True
			return "false"
# Start program
if __name__ == "__main__":
	main()