#!/usr/bin/python3
#
#
#   Uses the oDATA API to query information on a Citrix XenApp/XenDesktop environment.
#
#

#Header Info
__author__ = 'Wagner Garcez'
__email__ = 'wagner.garcez@interop.com.br'
__license__ = 'GPLv3'


#import modules
try: 
	import re
	import sys
	import json
	import signal
	import requests
	from requests_ntlm import HttpNtlmAuth
	from optparse import OptionParser, OptionGroup
except Exception as e:
	print ('UNKNOWN - {}'.format(e)) 
	sys.exit(opme['UNKNOWN'])

#fs = {  1:'None',
#		2:'FailedToStart',
#		3:'StuckOnBoot',
#		4:'Unregistered',
#		5:'MaximumLoad' # return API: MaxCapacity , return Director: MaximumLoad
#}


# get parameters
def parse_args():
	parser = OptionParser()
	parser.add_option('-b', '--broker'	 , dest = 'broker',		default = None , help = 'Broker machine | controller')
	parser.add_option('-m', '--machine'	 , dest = 'machine',	default = None , help = 'Machine name')
	parser.add_option('-u', '--username' , dest = 'user_name',	default = None , help = 'User Name for API connection')
	parser.add_option('-p', '--password' , dest = 'password',	default = None , help = 'Passwors for API connection')
	parser.add_option('-v', '--version'  , dest = 'version',	default = 'v1' , help = 'Use v1 in XenDesktop 7')
	parser.add_option('-o', '--option'   , dest = 'option',		default = None , help = 'The type of data you want to query')
	parser.add_option('-w', '--warning'  , dest = 'warning',	default = None , help = 'Warning threshold')
	parser.add_option('-c', '--critical' , dest = 'critical',	default = None , help = 'Critical threshold')
	parser.add_option('-V', '--verbose'  , dest = 'verbose',	default = False, help = 'Verbose Mode', action = 'store_true')
	(options, args) = parser.parse_args()
	if len(sys.argv) < 2: usage()
	test = validate_args(options)
	if test != '':
		print (test)
		sys.exit(opme['UNKNOWN'])
	else:
		return options
	
# validar argumentos recebidos 
def validate_args(o):
	miss = ''
	if (o.broker is None or o.user_name is None or o.password is None or o.machine is None):
		miss = 'UNKNOWN - Missing argument'
		if o.broker is None:
			miss += ' -b <broker>'
		if o.user_name is None:
			miss += ' -u <username>'
		if o.password is None:
			miss += ' -p <password>'
		if o.machine is None:
			miss += ' -m <machine>'
		return miss
	if (o.warning is None or o.critical is None) and (o.option == 'load_index' or o.option == 'load_index' or o.option == 'session_count'):
		miss += "UNKNOWN - Missing argument -w <warning> or -c <critical"
		return miss
	return miss

# help : executar o script sem parametro			
def usage():
	usage = '''
	Usage:
		{}script.py -b [broker] -u [user] -p [password] -o [option] -m [machine_name] -h [help]
			
		-b  ... Name of broker machine
		-u  ... User name to connect XenAPI
		-p  ... Password to connecet XenAPI
		-w  ... Warning thresholds 
		-c  ... Critical thresholds
		-m  ... Machine Name
		-o  ... Options of check (default: none)
			: load_index
			: session_count
			: maintenance_mode
			: fault_state
	'''.format(sys.argv[0])
	print (usage)
	sys.exit(opme['UNKNOWN'])

# funcoes auxiliares
def sigint_handler(signum, frame):
        print ('Stop pressing the CTRL+C!')
        sys.exit(opme['UNKNOWN'])

# printa na tela e encerra o script
def q(obj):
	print (obj)
	sys.exit(opme['OK'])

# mode verboso
def vprint(obj):
	if options.verbose:
		print (str(obj))
	return

# percorre a lista para descobrir o nome do grupo 
def search_desktop_group_by_id(dg, i):
	for d in dg['value']:
		if d['Id'] == i:
			return d['Name']
		else:
			pass
	print ('UNKNOWN - Not Found Desktop Group Id')
	sys.exit(opme['UNKNOWN'])

def validate_fault_state (fault_state, d):
	perf = " | 'fault_state'="+ str(fault_state)+";;;; 'group_name'="+ str(d) + ";;;;"
	vprint (' Fault State: {} \n'.format(fault_state))
	if fault_state == 'None':
		return 'OK - Server {}[{}] is fine ...{}'.format(options.machine, d, perf), opme['OK']
	else:
		return 'CRITICAL - Server {}[{}].  Failure type: {}{}'.format(options.machine, d, fault_state, perf), opme['CRITICAL']

def validate_maintenance_mode (s):
	perf = " | 'maintenance_status'="+ str(s)+";;;;"
	vprint('\n Maintenance Mode : {}\n'.format(s))
	if not s: 
		return 'OK - Server {} not in maintenanceis mode{}'.format(options.machine, perf), opme['OK']
	else:
		return 'WARNING - Server {} in maintenance mode{}'.format(options.machine, perf), opme['WARNING']

def validate_session_count (s, o):
	# verbose
	vprint (' Session Count: {}'.format(s))
	vprint (' Warning Value : {}'.format(options.warning))
	vprint (' Critical Value: {}\n'.format(options.critical))
	perf = " | 'sessions'="+ str(s)+";"+options.warning+";"+options.critical+";;"
	if s < int(options.warning):
		return 'OK - ' + str(o) + perf, opme['OK']
	elif s < int(options.critical):
		return 'WARNING - ' + str(o) + perf, opme['WARNING']
	else:
		return 'CRITICAL - ' + str(o) + perf, opme['CRITICAL']

def validate_load (l, o):
	# verbose
	vprint (' Load Citrix   : {}'.format(l))
	vprint (' Warning Value : {}'.format(options.warning))
	vprint (' Critical Value: {}\n '.format(options.critical))
	perf = " | 'load'="+ str(l)+";"+options.warning+";"+options.critical+";;"
	if l < int(options.warning):
		return 'OK - ' + str(o) + perf, opme['OK']
	elif l < int(options.critical):
		return 'WARNING - ' + str(o) + perf, opme['WARNING']
	else:
		return 'CRITICAL - ' + str(o) + perf, opme['CRITICAL']



def get_info (data):
	hostname = data['DnsName'].split('.',1)[0] 
	domain   = data['DnsName'].split('.',1)[1] 
	vprint ('\nHostname : {}'.format(data['DnsName'].split('.',1)[0]))
	vprint ('Domain	   : {}'.format(data['DnsName'].split('.',1)[1]))
	vprint ('Group Name : {}'.format(data['DesktopGroupName']))
	vprint ('OS         : {}'.format(data['OSType']))
	vprint ('Sessions   : {}'.format(data['CurrentSessionCount']))
	vprint ('Maintenance: {}'.format(data['IsInMaintenanceMode']))
	vprint ('Fault State: {}'.format(fs[data['FaultState']]))
	out = " Server {} Info ... | 'hostname={} domain={} groupname={} os={} session_count={} maintenance_mode={} fault_state={} ".format(data['DnsName'], data['DnsName'].split('.',1)[0],data['DnsName'].split('.',1)[1],data['DesktopGroupName'], data['OSType'], data['CurrentSessionCount'], data['IsInMaintenanceMode'], fs[data['FaultState']] )
	return out, opme['UNKNOWN']	

def get_fault_state (u, s, d, g):
	fs = {  1:'None', 
			2:'FailedToStart',
			3:'StuckOnBoot',
			4:'Unregistered',
			5:'MaximumLoad' # return API: MaxCapacity , return Director: MaximumLoad
	}
	if d in fs: # valida se o codigo retornado faz parte das opcoes de status conhecidas		
		return validate_fault_state(fs[d], g)	# status exist
	else:
		print ('CRITICAL - Server {} with state Unknown: {}'.format(fs[d]), g)
		sys.exit(opme['CRITICAL'])

def get_maintenance_mode (s):
	# valida se o dado recebido é um valor booleano
	if isinstance(s, bool):
		return validate_maintenance_mode(s)
	else:
		output = 'UNKNOWN - Value {} is not boolean'.format(s)
		return output, opme['UNKNOWN']

def get_session_count (session_count):
	# validar se numero de sessoes eh um inteiro
	if isinstance(session_count, int):
		output = 'Session count of server {} is {}'.format(options.machine, session_count)
		return validate_session_count(session_count, output)	
	else:
		output = 'UNKNOWN - Value {} is not number'.format(session_count)
		return output, opme['UNKNOWN']

def get_load_by_id (u, s, i):
	url = u + 'LoadIndexes?$filter=Id eq {}&$format=json'.format(i)
	r = (s.get(url))
	json_data = json.loads(r.text)
	try:
		load = json_data['value'][0]['EffectiveLoadIndex']
	except Exception as e:
		return 'UNKNOWN - %s' + str(e), 3
	# testar se eh um numero inteiro
	if isinstance(load, int):
		if load >= 0:
			load = int(load / 100)
		output = 'Citrix Load of server {} is {}'.format(options.machine, load) 
		# aplicar metricas
		return validate_load(load, output)
	else:
		output = 'UNKNOWN - Value {} is not number'.format(load)
		return output, opme['UNKNOWN']

def get_desktop_groups(u, s):
    url = u + 'DesktopGroups?$select=Id,Name&$format=json'
    r = (s.get(url))
    json_data = json.loads(r.text)
    if ((json_data['value'])):
        return json_data
    else:
        print ("UNKNOWN - Unable to collect desktop group Ids")
        sys.exit(opme['UNKNOWN'])

def get_machine_info(u, s, m):
    selects = 'DnsName,IPAddress,IsInMaintenanceMode,FaultState,OSType,CurrentLoadIndexId,DesktopGroupId,CurrentSessionCount'
    url = u + 'Machines?$filter=startswith(DnsName, \'{}\')&$select={}&$format=json'.format(m, selects)
    try:
        r = (s.get(url))
        r.raise_for_status()
    except requests.exceptions.HTTPError as errh:
        print ("UNKNOWN - Http Error:",errh)
        sys.exit(opme['UNKNOWN'])
    except requests.exceptions.ConnectionError as errc:
        print ("UNKNOWN - Error Connecting:",errc)
        sys.exit(opme['UNKNOWN'])
    except requests.exceptions.Timeout as errt:
        print ("UNKNOWN - Timeout Error:",errt)
        sys.exit(opme['UNKNOWN'])
    except requests.exceptions.RequestException as err:
        print ("UNKNOWN - Something is wrong [get_machine_info]",err)
        sys.exit(opme['UNKNOWN'])
    json_data = json.loads(r.text)
    if ((json_data['value'])):
        return json_data
    else:
        print ("UNKNOWN - Unable to collect information on machine '{}'".format(options.machine))
        sys.exit(opme['UNKNOWN'])

def get_index_group_id (dg, dgi):
	index = 0
	for i in dgi:
		if i['DesktopGroupId'] is not None: 
			return index
		else:
			index += 1
			pass # group ID nulo = dois registros referente ao mesmo server
	else:
		print ('UNKNOWN - Not Found Desktop Group Id')
		sys.exit(opme['UNKNOWN'])


# main function
def main():
	indice = 0
	url_api = 'http://{}/Citrix/Monitor/OData/{}/Data/'.format(options.broker, options.version)
	# sessao, autenticacao 
	session = requests.Session()
	session.auth  = HttpNtlmAuth(options.user_name, options.password)
	# json com informacoes sobre a maquina alvo
	data = get_machine_info (url_api, session, options.machine)
	# json com informacoes com o nome e o ID de todos os grupos criados 
	desktop_groups = get_desktop_groups (url_api, session)
	vprint ('Desktop Groups:\n-------------')
	vprint (desktop_groups)
	# indice do grupo pertencente a maquina alvo
	index = get_index_group_id (desktop_groups, data['value'])
	# descobrir o nome do grupo atraves do ID 
	group = search_desktop_group_by_id (desktop_groups, data['value'][index]['DesktopGroupId'])
	# adciona nome do grupo ao json com informacoes da maquina
	data['value'][index].update({'DesktopGroupName': group})
	# verbose
	vprint ('\nAll Registers: \n ----------')
	vprint (data)
	vprint ('\nData Target: \n ----------')
	vprint (data['value'][index])
	vprint ('\n DNS Name: {}'.format(data['value'][index]['DnsName']))
	vprint (' Desktop Group Name: {}'.format(data['value'][index]['DesktopGroupName']))
	vprint (' Operating System: {} \n'.format(data['value'][index]['OSType']))
	opt = options.option
	# opcoes disponiveis no script
	if opt == 'load_index':
		if data['value'][index]['CurrentLoadIndexId'] is not None:
			out, code = get_load_by_id(url_api, session, data['value'][index]['CurrentLoadIndexId'])
		else:
			(out , code ) = "UNKNOWN - Invalid option for machine '" + str(options.machine) + "' " , '3'
	elif opt == 'session_count':
		out, code = get_session_count(data['value'][index]['CurrentSessionCount'])
	elif opt == 'maintenance_mode':
		out, code = get_maintenance_mode(data['value'][index]['IsInMaintenanceMode'])
	elif opt == 'fault_state':
		out, code = get_fault_state(url_api, session, data['value'][index]['FaultState'], data['value'][index]['DesktopGroupName'])
	elif opt == 'info':
		out, code = get_info(data['value'][index])
	else:
		(out , code ) = "UNKNOWN - Invalid option: '-o " + str(opt) + "' " , '3'
	# output
	print (out)
	sys.exit(int(code))

# call main function
if __name__ == '__main__':
	opme = {'OK':0,'WARNING':1,'CRITICAL':2,'UNKNOWN':3}
	options = parse_args()
	signal.signal(signal.SIGINT, sigint_handler)
	main()