#!/usr/bin/python

import getopt, os, sys, signal, socket, SocketServer

class Functions:
	"Contains a set of methods for gathering data from the server." 
	def __init__(self):
		self.nagios_statd_version = 3.12

		# As of right now, the commands are for df, who, proc, uptime, and swap.
		commandlist = {}
		commandlist['AIX'] = ("df -Ik","who | wc -l","ps ax","uptime","lsps -sl | grep -v Paging | awk '{print $2}' | cut -f1 -d%")
		commandlist['BSD/OS'] = ("df","who | wc -l","ps -ax","uptime",None)
		commandlist['CYGWIN_NT-5.0'] = ("df -P",None,"ps -s -W | awk '{printf(\"%6s%6s%3s%6s%s\\n\",$1,$2,\"  S\",\"  0:00\",substr($0,22))}'",None,None)
		commandlist['CYGWIN_NT-5.1'] = commandlist['CYGWIN_NT-5.0']
		commandlist['Darwin'] = ("df -k -t nodevfs,volfs,synthfs,fdesc | grep -v \"automount -\"","who | wc -l","ps ax","uptime",None)
		commandlist['FreeBSD'] = ("df -k","who | wc -l","ps ax","uptime","swapinfo | awk '$1!~/^Device/{print $5}'")
		commandlist['HP-UX'] = ("bdf -l","who -q | grep \"#\"","ps -el","uptime",None)
		commandlist['IRIX'] = ("df -kP","who -q | grep \"#\"","ps -e -o \"pid tty state time comm\"","/usr/bsd/uptime",None)
		commandlist['IRIX64'] = commandlist['IRIX']
		commandlist['Linux'] = ("df -P","who -q | grep \"#\"","ps ax","uptime","free | awk '$1~/^Swap:/{print ($3/$2)*100}'")
		commandlist['NetBSD'] = ("df -k","who | wc -l","ps ax","uptime","swapctl -l | awk '$1!~/^Device/{print $5}'")
		commandlist['NEXTSTEP'] = ("df","who | /usr/ucb/wc -l","ps -ax","uptime",None)
		commandlist['OpenBSD'] = ("df -k","who | wc -l","ps -ax","uptime","swapctl -l | awk '$1!~/^Device/{print $5}'")
		commandlist['OSF1'] = ("df -P","who -q | grep \"#\"","ps ax","uptime",None)
		commandlist['SCO-SV'] = ("df -Bk","who -q | grep \"#\"","ps -el -o \"pid tty s time args\"","uptime",None)
		commandlist['SunOS'] = ("df -k","who -q | grep \"#\"","ps -e -o \"pid tty s time comm\"","uptime","swap -s | tr -d -s -c [:digit:][:space:] | nawk '{print ($3/($3+$4))*100}'")
		commandlist['UNIXWARE2'] = ("/usr/ucb/df","who -q | grep \"#\"","ps -el | awk '{printf(\"%6d%9s%2s%5s %s\\n\",$5,substr($0, 61, 8),$2,substr($0,69,5),substr($0,75))}","echo `uptime`, load average: 0.00, `sar | awk '{oldidle=idle;idle=$5} END {print 100-oldidle}'`,0.00",None)
	
		# Now to make commandlist with the correct one for your OS.
		try:
			self.commandlist = commandlist[os.uname()[0]]
		except KeyError:
			print "Your platform isn't supported by nagios-statd - exiting."
			sys.exit(3)

	# Below are the functions that the client can call.
	def disk(self):
		return self.__run(0)

	def proc(self):
		return self.__run(2)

	def swap(self):
		return self.__run(4)

	def uptime(self):
		return self.__run(3)

	def user(self):
		return self.__run(1)

	def version(self):
		i = "nagios-statd " + str(self.nagios_statd_version)
		return i

	def __run(self,cmdnum):
		# Stupid workaround because of Redhat's incorrect warnings
		if os.uname()[0] == 'Linux':
			outputfh = os.popen(self.commandlist[cmdnum])
			output = outputfh.read()
			try:
				returnvalue = outputfh.close()
			except:
				returnvalue = None

		else:
			# Unmask SIGCHLD so popen can detect the return status (temporarily)
			signal.signal(signal.SIGCHLD, signal.SIG_DFL)
			outputfh = os.popen(self.commandlist[cmdnum])
			output = outputfh.read()
			try:
				returnvalue = outputfh.close()
			except IOError:   # Darwin (Mac OS X) always generates IOErrors here - some thread handling problem.
				returnvalue = None
			signal.signal(signal.SIGCHLD, signal.SIG_IGN)

		if (returnvalue):
			return "ERROR"
		else:
			return output

class NagiosStatd(SocketServer.StreamRequestHandler):
	"Handles connection initialization and data transfer (as daemon)"
	def handle(self):
		# Check to see if user is allowed
		if self.__notallowedhost():
			self.wfile.write(self.error)
			return 1
		if not hasattr(self,"generichandler"):
			self.generichandler = GenericHandler(self.rfile,self.wfile)
		self.generichandler.run()

	def __notallowedhost(self):
		"Compares list of allowed users to client's IP address."
		if hasattr(self.server,"allowedhosts") == 0:
			return 0 
		for i in self.server.allowedhosts:
			if i == self.client_address[0]:  # Address is in list
				return 0
			try: # Do an IP lookup of host in blocked list
				i_ip = socket.gethostbyname(i)
			except:
				self.error = "ERROR DNS lookup of blocked host \"%s\" failed. Denying by default." % i
				return 1
			if i_ip != i:        # If address in list isn't an IP
				if socket.getfqdn(i) == socket.getfqdn(self.client_address[0]):
					return 0
		self.error = "ERROR Client is not among hosts allowed to connect."
		return 1

class GenericHandler:
	def __init__(self,rfile=sys.stdin,wfile=sys.stdout):
		# Create functions object
		self.functions = Functions()
		self.rfile = rfile
		self.wfile = wfile

	def run(self):
		# Get the request from the client
		line = self.rfile.readline()
		line = line.strip()

		# Check for appropriate requests from client
		if len(line) == 0:
			self.wfile.write("ERROR No function requested from client.")
			return 1

		# Call the appropriate function
		try:
			output = getattr(self.functions,line)()
		except AttributeError:
			error = "ERROR Function \"" + line + "\" does not exist."
			self.wfile.write(error)
			return 1
		except TypeError:
			error = "ERROR Function \"" + line + "\" not supported on this platform."
			self.wfile.write(error)
			return 1

		# Send output
		if output.isspace():
			error = "ERROR Function \"" + line + "\" returned no information."
			self.wfile.write(error)
			return 1
		elif output == "ERROR":
			error = "ERROR Function \"" + line + "\" exited abnormally."
			self.wfile.write(error)
		else:
			for line in output:
				self.wfile.write(line)


class Initialization:
	"Methods for interacting with user - initial code entry point."
	def __init__(self):
		self.port = 1040
		self.ip = ''
		# Run this through Functions initially, to make sure the platform is supported.
		i = Functions()
		del(i)

	def getoptions(self):
		"Parses command line"
		try:
			opts, args = getopt.getopt(sys.argv[1:], "a:b:ip:P:Vh", ["allowedhosts=","bindto=","inetd","port=","pid=","version","help"])
		except getopt.GetoptError, (msg, opt):
			print sys.argv[0] + ": " + msg
			print "Try '" + sys.argv[0] + " --help' for more information."
			sys.exit(3)
		for option,value in opts:
			if option in ("-a","--allowedhosts"):
				value = value.replace(" ","")
				self.allowedhosts = value.split(",")
			elif option in ("-b","--bindto"):
				self.ip = value
			elif option in ("-i","--inetd"):
				self.runfrominetd = 1
			elif option in ("-p","--port"):
				self.port = int(value)
			elif option in ("-P","--pid"):
				self.pidfile = value
			elif option	in ("-V","--version"):
				self.version()
				sys.exit(3)
			elif option in ("-h","--help"):
				self.usage()

	def main(self):
		# Retrieve command line options
		self.getoptions()

		# Just splat to stdout if we're running under inetd
		if hasattr(self,"runfrominetd"):
			server = GenericHandler()
			server.run()
			sys.exit(0)

		# Check to see if the port is available
		try:
			s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
			s.bind((self.ip, self.port))
			s.close()
			del(s)
		except socket.error:
			print "Port %s is already in use.  Unable to bind - exiting." % self.port
			sys.exit(2)

		# Detach from terminal
		if os.fork() == 0:

			# Make this the controlling process
			os.setsid()

			# Be polite and chdir to /
			os.chdir('/')

			# Try to close all open filehandles
			for i in range(0,256):
				try: os.close(i)
				except: pass

			# Redirect the offending filehandles
			sys.stdin = open('/dev/null','r')
			sys.stdout = open('/dev/null','w')
			sys.stderr = open('/dev/null','w')

			# Set the path
			os.environ["PATH"] = "/bin:/usr/bin:/usr/local/bin:/usr/sbin"

			# Reap children automatically
			signal.signal(signal.SIGCHLD, signal.SIG_IGN)

			# Save pid if user requested it
			if hasattr(self,"pidfile"):
				self.savepid(self.pidfile)

			# Create a forking TCP/IP server and start processing	
			server = SocketServer.ForkingTCPServer((self.ip,self.port),NagiosStatd)
			if hasattr(self,"allowedhosts"):
				server.allowedhosts = self.allowedhosts
			server.serve_forever()

		# Get rid of the parent 
		else:
			sys.exit(0)

	def savepid(self,file):
		try:
			fh = open(file,"w")
			fh.write(str(os.getpid()))
			fh.close()
		except:
			print "Unable to save PID file - exiting."
			sys.exit(2)

	def usage(self):
		print "Usage: " + sys.argv[0] + " [OPTION]"
		print "nagios-statd daemon - remote UNIX system monitoring tool for Nagios.\n"
		print "-a, --allowedhosts=HOSTS   Comma delimited list of IPs/hosts allowed to connect."
		print "-b, --bindto=IP            IP address for the daemon to bind to."
		print "-i, --inetd                Run from inetd."
		print "-p, --port=PORT            Port to listen on."
		print "-P, --pid=FILE             Save pid to FILE."
		print "-V, --version              Output version information and exit."
		print "  -h, --help               Print this help and exit."
		sys.exit(3)

	def version(self):
		i = Functions()
		print "nagios-statd %.2f" % i.nagios_statd_version
		print "Written by April King (april@twoevils.org).\n"
		print "This is free software.  There is NO warranty; not even for MERCHANTABILITY or"
		print "FITNESS FOR A PARTICULAR PURPOSE."
		print "\nNagios is a trademark of Ethan Galstad."

if __name__ == "__main__":
	# Check to see if running Python 2.x+ / needed because getfqdn() is Python 2.0+ only
	if (int(sys.version[0]) < 2):
		print "nagios-statd requires Python version 2.0 or greater."
		sys.exit(3)

	i = Initialization()
	i.main()
