#!/usr/bin/env python

import argparse
import subprocess
from multiprocessing.dummy import Pool as ThreadPool
from os.path import isfile
from shutil import copy

## ConfigParse is different from python2.7 and python3.5
from sys import version as python_version
if -1 != python_version.split(" ")[0].find("2.7"):
	import ConfigParser
	ConfigParser = ConfigParser.ConfigParser
else:
	import configparser
	ConfigParser = configparser.ConfigParser
	## python3.5 replace input with raw_input
	raw_input = input

host_note = """
###############################################################################
###############################################################################
#
# HOST DEFINITIONS
#
###############################################################################
###############################################################################

# Define hosts for the Lenovo XCC machines we'll be monitoring
# Change the host_name, alias, and address to fit your situation
"""

group_note = """
###############################################################################
###############################################################################
#
# HOST GROUP DEFINITIONS
#
###############################################################################
###############################################################################

# Define hostgroups for Lenovo XCC machines
# All hosts that use the linux-server template will automatically be a member of this group
"""

service_note = """
###############################################################################
###############################################################################
#
# SERVICE DEFINITIONS
#
###############################################################################
###############################################################################

# Create services for monitoring
"""

host_group_tmp = """
define hostgroup{
	hostgroup_name    %s  ; The name of the hostgroup
	alias             %s  ; Long name of the group
	members           %s  ; Members in this group
}
"""

service_tmp = """
define service{
	use                 generic-service
	hostgroup_name      %s
	service_description Check %s for Lenovo XCC
	check_command       %s
}
"""

host_tmp = """
define host{
    use             linux-server    ; Inherit default values from a template
    host_name       %s ; The name we're giving to this host
    alias           %s ; A longer name associated with the host
    address         %s  ; IP address of the host
}
"""

command_tmp = """
################################################################################
#
# LENOVO SERVICE CHECK COMMANDS
#
# These are some example service check commands.  They may or may not work on
# your system, as they must be modified for your plugins.  See the HTML 
# documentation on the plugins for examples of how to configure command definitions.
#
################################################################################

define command{
        command_name    check_system_authPriv
        command_line    $USER1$/check_lenovo_xcc.py -H "$HOSTADDRESS$" -v3 --snmp_username $ARG1$ --snmp_security_level authPriv --snmp_aprotocol $ARG2$ --snmp_apassword $ARG3$  --snmp_pprotocol $ARG4$ --snmp_ppassword $ARG5$ --mode $ARG6$ 
        }

define command{
        command_name    check_system_authNoPriv
        command_line    $USER1$/check_lenovo_xcc.py -H "$HOSTADDRESS$" -v3 --snmp_username $ARG1$ --snmp_security_level authNoPriv --snmp_aprotocol $ARG2$ --snmp_apassword $ARG3$ --mode $ARG4$ 
        }

define command{
        command_name    check_system_noAuthNoPriv
        command_line    $USER1$/check_lenovo_xcc.py -H "$HOSTADDRESS$" -v3 --snmp_username $ARG1$ --snmp_security_level noAuthNoPriv --mode $ARG2$ 
        }

"""

def parse_argument():
	"""
	This function just deal with the argument
	"""
	parser = argparse.ArgumentParser()
	parser.add_argument(
		"-f", "--configfile", help="IP list and configuration file")
	arg = parser.parse_args()
	return arg

## Get config file name
def get_config_file():
	"""
	this function just get the dict of arguments
	"""
	args = parse_argument()
	if args.configfile:
		config_file = args.configfile
	else:
		print( "Configuration file is missing for discovery")
		config_file = None
	
	return config_file

## Get config information from config file, include nagios path, created file name, force rewrite file
def get_config_info(discovery_cfg):
	output_path, config_file, force = "./", "Lenovo_XCC.cfg", False
	if "output_path" in discovery_cfg.keys():
		output_path = discovery_cfg["output_path"]
	if "config_file" in discovery_cfg.keys():
		config_file = discovery_cfg["cfg_file_name"]
	if "force" in discovery_cfg.keys():
		if 'true' == discovery_cfg["force"].lower():
			force = True
	return output_path, config_file, force
	

## Ping the IP to test connect state
def isconnect(ip):
	run_watch = subprocess.Popen("ping %s -c 2 -W 1" % ip,
		shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE,universal_newlines=True)
	ret, err_msg = run_watch.communicate()
	if -1 == ret.find("100%") and -1 == ret.find("50%"):
		return ip
	else:
		return "notconnect"

## Use multithread to get connected ip list
def get_connected_iplist(ip_list):
	pool = ThreadPool()
	results = pool.map(isconnect, ip_list)
	ok_list = [str for str in results if str != "notconnect"]
	pool.close()
	pool.join()
	return ok_list

## IP to number
def ip2num(ip):
	ip_segment = [int(x) for x in ip.split('.')]
	return ip_segment[0] << 24 | ip_segment[1] << 16 | ip_segment[2] << 8 | ip_segment[3]

## Number to ip
def num2ip(num):
	ip = ['', '', '', '']
	ip[3] = (num & 0xff)
	ip[2] = (num & 0xff00) >> 8
	ip[1] = (num & 0xff0000) >> 16
	ip[0] = (num & 0xff000000) >> 24
	return '%s.%s.%s.%s' % (ip[0], ip[1], ip[2], ip[3])

# Get ip list from 3 format: 
# format1: 198.126.0.0/24
# format2: 198.126.0.1-198.126.0.100
# format3: 198.126.0.1(just one ip address)
def getiplist_impl(ip_range):
	if -1 != ip_range.find("-"):
		index = ip_range.find("-")
		try:
			l_endpoint = ip2num(ip_range[:index])
			r_endpoint = ip2num(ip_range[index + 1:])
			if l_endpoint >= r_endpoint:
				print("The right endpoint must be greater than the left endpoint(%s)" % ip_range)
				return None
			else:
				return [num2ip(ip) for ip in range(l_endpoint, r_endpoint + 1)]
		except:
			print("IP range %s is invalid, please set a valid one in configuration file" % ip_range)
			return None
	elif -1 != ip_range.find("/"):
		index = ip_range.find("/")
		ip_addr = ip_range[:index]
		mask = ip_range[index + 1:]
		if 0 != ip2num(ip_addr) & (0xffffffff >> int(mask)):
			print("IP range %s is invalid, please set a valid one in configuration file" % ip_range)
			return None
		else:
			l_endpoint =  ip2num(ip_addr)
			r_endpoint = l_endpoint + (0xffffffff >> int(mask))
			return [num2ip(ip) for ip in range(l_endpoint, r_endpoint + 1)]
	else:
		return [ip_range]

## Get ip list from ip range
def getiplist(ip_str):
	iplist = []
	iprange_list = ip_str.split(",")
	iprange_list = [iprange.strip() for iprange in iprange_list if '' != iprange.strip()]
	for iprange in iprange_list:
		ret_iplist = getiplist_impl(iprange)
		if ret_iplist:
			iplist += ret_iplist
	return iplist

## Get snmpwalk options list, this function write for thread and map function
def get_snmpwalk_options_list(connected_list, authlevel, user, aprotocol, apassword, pprotocol, ppassword):
	options_list = []
	for ip in connected_list:
		if 'authpriv' == authlevel.lower():
			snmpwalk_options = " -l authPriv -u %s -a %s -A %s -x %s -X %s %s" % (user, aprotocol, apassword, pprotocol, ppassword, ip)
		elif 'authnopriv' == authlevel.lower():
			snmpwalk_options = " -l authNoPriv -u %s -a %s -A %s %s" % (user, aprotocol, apassword, ip)
		elif 'noauthnopriv' == authlevel.lower():
			snmpwalk_options = " -l noAuthNoPriv -u %s %s" % (user, ip)
		else:
			break
		options_list.append(snmpwalk_options)
	return options_list

## Run snmpwalk command to find Lenovo XCC machines
def ispurley(snmpwalk_options):
	"""
	this function will judgment the machine is purley
	"""
	try:
		ip = snmpwalk_options.split(" ")[-1]
	except:
		return "notxcc"
	##Command: snmpwalk -v3 -l authPriv -u SNMPUSER -a SHA -A SYS2009health -x DES -X SYS2009health 196.198.0.1 1.3.6.1.4.1.19046.11.1.3.7.0 -t 0.5 
	cmd = "snmpwalk -v3 %s 1.3.6.1.4.1.19046.11.1.3.7.0 -t 0.5 | grep 'INTEGER:' | wc -l " % snmpwalk_options
	run_watch = subprocess.Popen(
		cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True)
	ret, err_msg = run_watch.communicate()
	if "0" != ret.strip():
		return ip
	else:
		return "notxcc"

## Use multithread to get Lenovo XCC ip list
def get_xcc_list(snmpwalk_options_list):
	pool = ThreadPool()
	if 0 == len(snmpwalk_options_list):
		return []
	results = pool.map(ispurley, snmpwalk_options_list)
	xcc_list = [str for str in results if str != "notxcc"]
	pool.close()
	pool.join()
	return xcc_list

### This function will create config file for nagios
### include: host definition, group definition and service definition
def create_nagios_cfg_file(input_file):
	ip_range, user, alias, authlevel, snmp_aprotocol, apassword, snmp_pprotocol, ppassword = \
	'0.0.0.0', 'user_tmp', '', 'noAuthNoPriv', 'SHA', '', 'DES', ''
	group_cfg_str, services_cfg_str, host_cfg_str = '', '', ''
	output_file, output_path, cmd_file, force = "Lenovo_XCC.cfg", "./", "lenovo_snmp_command.cfg", False
	
	# system-health, power,cpu,storage,temperature,fans,voltage,memory
	services_list = ['system-health', 'power', 'cpu', 'storage', 'temperature', 'fans', 'voltage', 'memory']
	conf = ConfigParser()
	conf.read(cfg_file)

	print("Discovering systems, please wait...")

	# Get all the group name
	group_list = conf.sections()
	for group in group_list:
		if "discovery_cfg" == group:
			dis_cfg = dict(conf.items("discovery_cfg"))
			output_path, output_file, force = get_config_info(dis_cfg)
			continue
		command_str = ""
		print("Discovering group '%s'." % group)
		items = dict(conf.items(group))

		if "ip" in items.keys():
			ip_range = items['ip']
		else:
			print("[ERROR]IP range is missing for %s" % group)
			continue

		if 'alias' in items.keys():
			alias = items['alias']
		else:
			alias = group

		if "user" in items.keys():
			user = items["user"]
		else:
			print("[ERROR]SNMPv3 User is missing for %s" % group)
			continue

		if "authlevel" in items.keys():
			authlevel = items["authlevel"]
		
		if "snmp_aprotocol" in items.keys():
			snmp_aprotocol = items["snmp_aprotocol"]

		if "apassword" in items.keys():
			apassword = items["apassword"]

		if "snmp_pprotocol" in items.keys():
			snmp_pprotocol = items['snmp_pprotocol']

		if "ppassword" in items.keys():
			ppassword = items["ppassword"]

		if "services" in items.keys():
			services_list = items['services'].split(",")
			services_list = [service.strip() for service in services_list if '' != service.strip()]
		
		## Get all ip list in the ip range
		all_ip_list = getiplist(ip_range)

		## Get list of ip which is connected via ping
		connected_ip_list = get_connected_iplist(all_ip_list)
		if 0 == len(connected_ip_list):
			print("0 out of %s are successfully discovered in group %s" % (len(all_ip_list), group))
			continue

		## Get list of snmpwalk options which is used for judge XCC machine
		## This code is used for thread map
		# options_list = get_snmpwalk_options_list(connected_ip_list)
		snmpwalk_options_list = get_snmpwalk_options_list(connected_ip_list, authlevel, user, snmp_aprotocol, apassword, snmp_pprotocol, ppassword)

		## Get list of ip which is XCC's address
		xcc_ip_list = get_xcc_list(snmpwalk_options_list)
		if 0 != len(xcc_ip_list):
			print(xcc_ip_list)
		print("%s out of %s are successfully discovered in group %s" % (len(xcc_ip_list), len(all_ip_list), group))

		## Get the host config
		for xcc_ip in xcc_ip_list:
			host_cfg_str += host_tmp % (group+'_'+xcc_ip, group+'_'+xcc_ip, xcc_ip)

		## Get the group config
		## Get the members of this group
		members_list = [group +"_" + xcc_ip for xcc_ip in xcc_ip_list]
		group_members = ','.join(members_list)
		if group_members:
			group_cfg_str += host_group_tmp % (group, alias, group_members)

		## Get the services config
		## get command for services
		if 'authpriv' == authlevel.lower():
			if '' == apassword or '' == ppassword:
				print("Both of authentication and privacy protocol password are necessary for authPriv level")
				continue
			else:
				command_str = "%s!%s!%s!%s!%s!%s" % ("check_system_authPriv", user, snmp_aprotocol, apassword, snmp_pprotocol, ppassword)
		elif 'authnopriv' == authlevel.lower():
			if '' == apassword:
				print("Authentication protocol password is necessary for authPriv level")
				continue
			else:
				command_str = "%s!%s!%s!%s" % ("check_system_authNoPriv", user, snmp_aprotocol, apassword)
		elif 'noauthnopriv' == authlevel.lower():
			command_str = "%s!%s" % ("check_system_noAuthNoPriv", user)

		## add services 
		if services_list:
			services_cfg_str += ("\n#### Services for %s ####\n" % group)
		for service in services_list:
			services_cfg_str += service_tmp % (group, service, command_str + "!" + service)

	if host_cfg_str and group_cfg_str and services_cfg_str:
		can_write_cfg = False
		if '/' == output_path[-1]:
			ng_cfg_file = output_path + output_file
			ng_cmd_file = output_path + cmd_file
		else:
			ng_cfg_file = output_path + "/" + output_file
			ng_cmd_file = output_path + "/" + cmd_file

		if not force and isfile(ng_cfg_file):
			ret = raw_input('File %s exists. Do you want to overwrite it? (Y/N) ' % ng_cfg_file)
			if "y" == ret.lower():
				can_write_cfg = True
		else:
			can_write_cfg = True

		if can_write_cfg:
			try:
				fh = open(ng_cfg_file, 'w+')
				fh.write(host_note + host_cfg_str + group_note + group_cfg_str + service_note + services_cfg_str)
				fh.close()
			except:
				print("Failed in writing configuration file")
			try:
				fh_cmd = open(ng_cmd_file, 'w+')
				fh_cmd.write(command_tmp)
				fh_cmd.close()
				print("Configuration files generated successfully:\n- %s\n- %s" % ((output_path+"lenovo_snmp_command.cfg"),(output_path + output_file)))
			except:
				print("Failed in writing lenovo_snmp_command.cfg ")
	else:
		print("No configuration files generated.")

if __name__ == "__main__":
	cfg_file = get_config_file()
	if cfg_file:
		create_nagios_cfg_file(cfg_file)
		pass
	else:
		exit(1)