# Copyright (C) 2005-2021 Splunk Inc. All Rights Reserved. 

# Core Python imports

import sys
import math
from xml.dom import minidom
import re

# Splunk imports

from splunk import mergeHostPath, util
from splunk.auth import getSessionKey
import splunk.entity as en

# Append SA-Hydra-inframon/bin to the Python path

from splunk.clilib.bundle_paths import make_splunkhome_path
sys.path.append(make_splunkhome_path(['etc', 'apps', 'SA-Hydra-inframon', 'bin']))

# SA-Hydra-inframon imports
from hydra_inframon import six
from hydra_inframon import setupLogger
from hydra_inframon.hydra_common import HydraGatewayAdapter
from hydra_inframon.models import HydraNodeStanza, SplunkStoredCredential

# TA-vmware-inframon imports

import ta_vmware_inframon.simple_vsphere_utils as vsu
from ta_vmware_inframon.models import TAVMwareCollectionStanza

#from test_autoeventgen import generate_moids, distribute_autogenerated_hierarchy_cache
#Setup Logger
logger = setupLogger(log_format='%(asctime)s %(levelname)s [TAVMwareHierarchyAgent] %(message)s', log_name="ta_vmware_inframon_hierarchy_agent.log", logger_name="ta_vmware_inframon_hierarchy_agent")

# It can be put in inframon_ta_vmware_collection.conf file but DCN does not need these values, hence we hard coded here
MAX_HOSTS_IN_VC =  1000
MAX_POWER_ON_VMS_IN_VC = 10000
MAX_REGISTER_VMS_IN_VC = 15000
MAX_HOST_IN_CLUSTER = 32
MAX_VMS_IN_HOST = 512
MAX_VMS_IN_CLUSTER = 4000

def get_node_adapters(local_host_path, local_session_key, pool_name):
	"""
	Given the local session key return an iterable of HydraGatewayAdapters to
	all Hydra Nodes in the TA Vmware app context for which we have valid access.
	Any failed logins will be logged but otherwise will not impact data collection.

	@type local_host_path: str
	@param local_host_path: path to the local splunkd mgmt interface
	@type local_session_key: str
	@param local_session_key: valid splunk session key to the local splunkd instance
	@type pool_name: str
	@param pool_name: pool name

	@rtype: list
	@return: iterable of authenticated HydraGatewayAdapters to all DCN's
	"""
	#Establish node list
	node_stanzas = HydraNodeStanza.all(host_path=local_host_path, sessionKey=local_session_key)
	node_stanzas._owner = "nobody" #self.asset_owner
	node_stanzas = node_stanzas.filter_by_app("Splunk_TA_vmware_inframon")

	#Iterate on all nodes, checking if alive and sorting appropriately
	adapter_list = []
	for node_stanza in node_stanzas:
		if node_stanza.pool_name == pool_name:
			password = SplunkStoredCredential.get_password(node_stanza.name, node_stanza.user, "Splunk_TA_vmware_inframon", session_key=local_session_key, host_path=local_host_path)
			if isinstance(node_stanza.gateway_port, int):
				gateway_port = node_stanza.gateway_port
			else:
				gateway_port = 8008
			gateway_uri = node_stanza.name.rstrip("/0123456789") + str(gateway_port)


			try:
				node_session_key = getSessionKey(node_stanza.user, password, hostPath=node_stanza.name)
				adapter_list.append(HydraGatewayAdapter(node_stanza.name, node_session_key, gateway_uri))
			except Exception as e:
				logger.exception("[get_node_adapters] failed to establish gateway adapter for node=%s of pool=%s due to error=%s", node_stanza.name, pool_name, str(e))

	return adapter_list

def get_virtual_centers(local_host_path, local_session_key, pool_name):
	"""
	Given the local session key return an iterable of tuples of vsphere address,
	service user name, and password to the user.

	@type local_host_path: str
	@param local_host_path: path to the local splunkd mgmt interface
	@type local_session_key: str
	@param local_session_key: valid splunk session key to the local splunkd instance
	@type pool_name: str
	@param pool_name: pool name

	@rtype: list
	@return: iterable of tuples of the form (<vspehere address>, <username>, <password>)
	"""
	#Establish node list
	collection_stanzas = TAVMwareCollectionStanza.all(host_path=local_host_path, sessionKey=local_session_key)
	collection_stanzas._owner = "nobody" #self.asset_owner
	collection_stanzas = collection_stanzas.filter_by_app("Splunk_TA_vmware_inframon")

	#Iterate on all nodes, checking if alive and sorting appropriately
	vc_list = []
	for collection_stanza in collection_stanzas:
		if collection_stanza.pool_name == pool_name:
			## Used autoeventgen property if fake auto event is generated to large scale (maximize vc limit)
			if collection_stanza.autoeventgen is None:
				isautoeventgen = False
			else:
				isautoeventgen = collection_stanza.autoeventgen
			if collection_stanza.realm is None:
				for target in collection_stanza.target:
					password = SplunkStoredCredential.get_password(target, collection_stanza.username, "Splunk_TA_vmware_inframon", session_key=local_session_key, host_path=local_host_path)
					vc_list.append((target, collection_stanza.username, password, isautoeventgen))
			else:
				password = SplunkStoredCredential.get_password(collection_stanza.realm, collection_stanza.username, "Splunk_TA_vmware_inframon", session_key=local_session_key, host_path=local_host_path)
				for target in collection_stanza.target:
					vc_list.append((target, collection_stanza.username, password, isautoeventgen))

	return vc_list



def collect_vms_by_host(target, user, password):
	"""
	Given a virtual center and the user and password to access it return a
	data structure of the form { <host moid> : <vm moid>[]}

	@type target: str
	@param target: valid vsphere address
	@type user: str
	@param user: valid user to that vsphere address
	@type password: str
	@param password: password for the passed user

	@rtype: dict
	@return: dict of host moid to iterable of vm moid
	"""

	vss = vsu.vSphereService(target, username=user, password=password)
	response = vss.get_obj_list([{'type':'VirtualMachine','all':'false', 'pathSet':'summary.runtime.powerState'},
								{'type':'VirtualMachine','all':'false', 'pathSet':'summary.runtime.host'}],
							{'type':'Folder', 'moid':"group-d1"})
	vss.logout()
	resp_xml = minidom.parseString(response)
	vms_by_host = {}
	for returnval in resp_xml.getElementsByTagName("returnval"):
		#Get the MOR object, there will only ever be one
		obj =  returnval.getElementsByTagName('obj')[0]

		#Convert property sets into a dictionary
		propsets = returnval.getElementsByTagName('propSet')
		prop_dict = {}
		for propset in propsets:
			name = None
			val = None
			for node in propset.childNodes:
				if node.tagName == "name":
					name = node.firstChild.data
				elif node.tagName == "val":
					val = node.firstChild.data
			prop_dict[name] = val

		#Check on props and add to list if appropriate
		if prop_dict.get("summary.runtime.powerState", "") == "poweredOn":
			host = prop_dict.get("summary.runtime.host", "orphan")
			if vms_by_host.get(host, None) is None:
				vms_by_host[host] = []
			vms_list = vms_by_host[host]

			#Add vm moid to the list
			vms_list.append(str(obj.firstChild.data))

	return vms_by_host

def get_datastore_info(target, user, password):
	"""
	Given a virtual center and the user and password to access it return a
	data structure of the form { <datastore uuid> : (<datastore moid>, <datastore name>)}

	@type target: str
	@param target: valid vsphere address
	@type user: str
	@param user: valid user to that vsphere address
	@type password: str
	@param password: password for the passed user

	@rtype: dict
	@return: dict of datastore uuid to datastore moid and datastore name
	"""
	try:
		vss = vsu.vSphereService(target, username=user, password=password)
		response = vss.get_obj_list([{'type':'Datastore','all':'false', 'pathSet':'summary.url'},
									 {'type':'Datastore','all':'false', 'pathSet':'name'}],
									{'type':'Folder', 'moid':"group-d1"})
		vss.logout()
		ds_info_by_uuid = {}
		resp_xml = minidom.parseString(response)

		for returnval in resp_xml.getElementsByTagName("returnval"):
				#Get the MOR object, there will only ever be one
			obj =  returnval.getElementsByTagName('obj')[0]

			#Convert property sets into a dictionary
			propsets = returnval.getElementsByTagName('propSet')
			prop_dict = {}
			for propset in propsets:
				name = None
				val = None
				for node in propset.childNodes:
					if node.tagName == "name":
						name = node.firstChild.data
					elif node.tagName == "val":
						val = node.firstChild.data
				prop_dict[name] = val

			#Check on props and add to list if appropriate
			ds_url = prop_dict.get("summary.url", "")
			ds_name = prop_dict.get("name", "")
			match = re.match('ds:\/\/\/vmfs\/volumes\/(.*?)\/[\"\n\s\t\b]?.*$', ds_url)
			if match:
				ds_uuid = str(match.group(1))
				ds_info_by_uuid[ds_uuid] = (str(obj.firstChild.data), ds_name)
		return ds_info_by_uuid
	except Exception as e:
		logger.error("Failed to get datastore uuid-moid mapping for vc:%s", target)
	

def distribute_hierarchy_cache(nodes, target, vms_by_host, ds_info_by_uuid):
	"""
	Distribute the given hierarchy dictionary to all nodes in the given set of
	hydra gateway adapters

	@type nodes: list of HydraGatewayAdapter
	@param nodes: list of authenticated gateway adapters to DCN's
	@type target: str
	@param target: target virtual center value
	@type vms_by_host: dict of <host moid> -> list of <vm moid>
	@param vms_by_host: dictionary of parent moid to list of child moids to cache
	@type ds_info_by_uuid: dict of <datastore uuid> -> <datastore moid>, <datastore name>
	@param ds_info_by_uuid: dictionary of datastore uuid to moid and name mapping to cache

	@rtype: None
	@return: None
	"""
	cache_items = []
	for host_moid, vms in six.iteritems(vms_by_host):
		cache_name = "perfhierarchy:" + target + ":" + host_moid
		cache_value = {host_moid: vms}
		cache_items.append((cache_name, cache_value))
	cache_name = "dsinfo:" + target
	cache_value = ds_info_by_uuid
	cache_items.append((cache_name, cache_value))
	for hga in nodes:
		try:
			hga.set_cache_batch(cache_items, expiration=14400)
		except Exception as e:
			logger.exception("[distribute_hierarchy_cache] failed to set cache for target=%s for node=%s due to error=%s", target, hga.splunkd_uri, str(e))

def main(local_host_path, local_session_key, pool_list):
	"""
	The main method to prepare host-VM list for the vCenter Servers and distribute it among nodes of pool from the pool_list

	@type local_host_path: str
	@param local_host_path: path to the local splunkd mgmt interface
	@type local_session_key: str
	@param local_session_key: valid splunk session key to the local splunkd instance
	@type pool_name: str
	@param pool_name: pool name
	"""
	#Get Data Collection Nodes
	for pool_name in pool_list:
		logger.info("Start vmware hierarchy agent for pool:%s", pool_name)
		nodes = get_node_adapters(local_host_path, local_session_key, pool_name)
		if len(nodes) == 0:
			logger.info("could not authenticate with any data collection nodes for pool_name : %s, skipping run of vmware hierarchy agent", pool_name)
			continue

		#Get target Virtual Centers
		vcs = get_virtual_centers(local_host_path, local_session_key, pool_name)
		if len(vcs) == 0:
			logger.info("could not find any configured virtual centers for pool_name : %s, skipping run of vmware hierarchy agent", pool_name)
			continue
		# Capture vc and list existing hosts and vms count so it would be useful if we make
		# Uncomment the following line if you want to generate dummy events using autoeventgen
		#vcs_autogenerated_id_info = []
		#Process and distribute virtual center hierarchy cache
		for target, username, password, isautogenerateid in vcs:
			try:
				vms_by_host = collect_vms_by_host(target, username, password)
				logger.debug("Successfully get hierarchy list for vc:%s of pool:%s", target, pool_name)

				ds_info_by_uuid = get_datastore_info(target, username, password)
				logger.debug("Successfully get datastore uuid-info mapping for vc:%s of pool:%s", target, pool_name)

				# Uncomment the following lines if you want to generate dummy events using autoeventgen
				'''
				if isautogenerateid:
					ids = generate_moids(vms_by_host, logger, target)
					vcs_autogenerated_id_info.append(ids)
					distribute_autogenerated_hierarchy_cache(nodes, target, vcs_autogenerated_id_info, logger)
				'''
				# Distribute to the node
				distribute_hierarchy_cache(nodes, target, vms_by_host, ds_info_by_uuid)
			except vsu.LoginFailure as e:
				logger.error(str(e))
			except vsu.ConnectionFailure as e:
				logger.error(str(e))
			except Exception as e:
				logger.error("Could not get hierarchy list or datastore uuid-moid mapping for vc:%s of pool:%s Error:%s", target, pool_name, str(e))

if __name__ == "__main__":

	local_session_key = sys.stdin.readline().strip("\r\n")
	local_host_path = mergeHostPath()
	# Get enabled Scheduler stanzas
	scheduler_stanzas = en.getEntities("/data/inputs/ta_vmware_collection_scheduler_inframon", "Splunk_TA_vmware_inframon", "nobody", sessionKey=local_session_key, hostPath=local_host_path)
	pool_list = [name for name, config in six.iteritems(scheduler_stanzas) if not util.normalizeBoolean(config.get("disabled"))]

	main(local_host_path, local_session_key, pool_list)
