Tie conmux better to autotest

I am sending this out as it is what we use to keep track of our systems that are in conmux and have other options like power. While using labels isn't the best way it is our stop gap solution for the time being. 


Risk: low
Visibility: Good for admins of conmux server
A console check script that will also apply labels like console or remote-power to an autotest server.
Currently each host is added individually due to our CLI tools, I am planning on updating this once we have a stable CLI.


Signed-off-by: Scott Zawalski <scottz@google.com>




git-svn-id: http://test.kernel.org/svn/autotest/trunk@1603 592f7852-d20e-0410-864c-8624ca9c26a4
diff --git a/conmux/contrib/console_check.py b/conmux/contrib/console_check.py
new file mode 100644
index 0000000..eb0f37d
--- /dev/null
+++ b/conmux/contrib/console_check.py
@@ -0,0 +1,303 @@
+#!/usr/bin/python2.4
+
+_author_ = 'Scott Zawalski (scottz@google.com)'
+
+"""Console check script to be used with conmux.
+
+   Checks if machines are not only connected to conmux but also
+   responding in an expected way
+
+   Supports options to show all, good, bad, unknown and add them 
+   to autotest as well. 
+
+   *In order for the power update option to work you have to have
+   access to the etc directory of the conmux server
+"""
+
+import sys,  pexpect, commands, os
+from optparse import OptionParser
+
+
+def main(argv):
+	consoles = {}
+	consoles['good'] = []
+	consoles['bad'] = []
+	consoles['unknown'] = []
+	# 0, 1, 2 status
+	STATUS = [ 'good', 'bad', 'unknown']
+	parser = OptionParser()
+	parser.add_option('--conmux-server', dest="conmux_server",
+	                 default='localhost',
+	                 help="Conmux server to connect to")
+	parser.add_option('--conmux-dir', dest="conmux_dir",
+	                 default='/usr/local/conmux',
+	                 help="Conmux server to connect to")
+	parser.add_option('--console-binary', dest="console_binary",
+	                 default='/usr/local/conmux/bin/console',
+	                 help="Conmux console binary location")
+	parser.add_option('--autotest-cli-dir', dest="autotest_cli_dir",
+	                 default='/usr/local/autotest/cli',
+	                 help="Autotest CLI dir")
+	parser.add_option('--add-hosts',
+	                  action="store_true", dest="add_hosts",
+	                  default=False,
+	                  help="If host not on autotest server try to add it")
+	parser.add_option('--power-label', dest="power_label",
+	                 default='remote-power',
+	                 help="Label to add to hosts that support hard reset")
+	parser.add_option('--console-label', dest="console_label",
+	                 default='console',
+	                 help="Label to add to hosts that support console")
+	parser.add_option('--update-console-label',
+	                  action="store_true", dest="update_console_label",
+	                  default=False,
+	                  help="Update console label on autotest server")
+	parser.add_option('--update-power-label',
+	                  action="store_true", dest="update_power_label",
+	                  default=False,
+	                  help="Update power label on autotest server" +\
+				"*Note this runs then exists no consoles are checked")
+	parser.add_option('--verbose',
+	                  action="store_true", dest="verbose",
+	                  default=False,
+	                  help="Verbose output")
+	parser.add_option('--show-bad',
+	                  action="store_true", dest="show_bad",
+	                  default=False,
+	                  help="Show consoles that are no longer functioning")
+	parser.add_option('--show-good',
+	                  action="store_true", dest="show_good",
+	                  default=False,
+	                  help="Show consoles that are functioning properly")
+	parser.add_option('--show-unknown',
+	                  action="store_true", dest="show_unknown",
+	                  default=False,
+	                  help="Show consoles that are in an unknown state")
+	parser.add_option('--show-all',
+	                  action="store_true", dest="show_all",
+	                  default=False,
+	                  help="Show status of all consoles")
+	options, args = parser.parse_args()
+	if len(argv) == 2 and options.verbose:
+		parser.print_help()
+		return 1
+	elif len(argv) < 2:
+		parser.print_help()
+		return 1
+
+	if options.update_power_label:
+		remove_create_label(options.power_label,
+		                    options.autotest_cli_dir)
+		update_power_label(options.power_label, options.conmux_dir,
+		                   options.autotest_cli_dir, options.add_hosts)
+		return
+	print options.console_binary
+	if not os.path.exists(options.console_binary):
+		print "Error %s does not exist, please specify another path" %\
+		      options.console_binary
+		return 1
+	hosts = get_console_hosts(options.console_binary, options.conmux_server)
+	for host in hosts:
+		rc = check_host(host, options.console_binary)
+		if options.verbose is True:
+			print "%s status: %s" % (host, STATUS[rc])
+		consoles[STATUS[rc]].append(host)
+
+	if options.show_all:
+		for status in consoles:
+			print "--- %s ---" % status
+			for host in consoles[status]:
+				print host
+	if options.show_good:
+		print "--- good ---"
+		for host in consoles['good']:
+			print host
+	if options.show_bad:
+		print "--- bad ---"
+		for host in consoles['bad']:
+			print host
+	if options.show_unknown:
+		print "--- unknown ---"
+		for host in consoles['unknown']:
+			print host
+
+	if options.update_console_label:
+		remove_create_label(options.console_label,
+		                    options.autotest_cli_dir)
+		update_console_label(options.console_label, consoles['good'],
+		                     options.autotest_cli_dir, options.add_hosts)
+		
+
+def update_console_label(console_label, consoles, cli_dir, add_hosts=False):
+	"""Update CONSOLE_LABEL on your autotest server.
+	   This removes the label and recreates it, then populating the label 
+	   with all the machines your conmux server knows about. 
+	   
+	   *Note If the hosts do not exist they are created.
+	   Args:
+		console_label:
+		string, describes the autotest label to add to machines.
+		consoles: 
+		list, all the consoles that have confirmed console support.
+	"""
+	# TODO: Update to new CLI and change logic until then
+	# this is the best way to ensure a machine is added i.e. one at a time
+
+	for host in consoles:
+		if not host_label_add(host, console_label, cli_dir):
+			# Try to create host
+			if add_hosts:
+				if host_create(host, cli_dir):
+					host_label_add(host, power_label,
+					               cli_dir)
+				else:
+					print "Unable to add host " + host
+
+
+def update_power_label(power_label, conmux_dir, cli_dir, add_hosts=False):
+	"""Look in CONSOLE_DIR/etc and grab known power commands
+	   Then remove POWER_LABEL and add machines to that label
+	"""
+	# remove label and add it
+	for host in hard_reset_hosts(conmux_dir):
+		rc = label_add_host(host, power_label, cli_dir)
+		if not rc:
+			# Try to create the host
+			if add_hosts:
+				if host_create(host, cli_dir):
+					rc = label_add_host(host, power_label,
+							    cli_dir)
+				else:
+					print "Unable to add host " + host
+
+
+def hard_reset_hosts(conmux_dir): 
+	"""Go through conmux dir and find hosts that have reset commands"""
+	config_dir = os.path.join(conmux_dir, "etc")
+	hosts = []
+	for file in os.listdir(config_dir):
+		if not file.endswith(".cf"):
+			continue
+		file_path = os.path.join(config_dir, file)
+		try:
+			try:
+				f = open(file_path)
+				for line in f:
+					if "reset" in line:
+						hosts.append(file.rstrip(".cf"))
+			except IOError:
+				pass
+		finally:
+			f.close()
+	return hosts
+
+
+def host_create(host, cli_dir):
+	"""Create a host
+	   Return:
+		True, if successfuly false if failed
+	"""
+	cmd = "%s/host-create %s" % (cli_dir, host)
+	status, output = commands.getstatusoutput(cmd)
+	return status == 0
+
+
+def label_add_host(host, label, cli_dir):
+	"""Add a host to a label"""
+	host_cmd = "%s/label-add-hosts %s %s" % (cli_dir, label, host)
+	(status, output) = commands.getstatusoutput(host_cmd)
+	if status != 0:
+		return False
+
+	return True
+
+
+def remove_create_label(label, cli_dir):
+	"""Remove and recreate a given label"""
+	cmd = "%s/label-rm %s" % (cli_dir, label)
+	status, output = commands.getstatusoutput(cmd)
+	if status != 0:
+		raise Exception("Error deleting label: " + label)
+
+	cmd = "%s/label-create %s" % (cli_dir, label)
+	status, output = commands.getstatusoutput(cmd)
+	if status != 0:
+		raise Exception("Error creating label: " + label + output)
+
+	return True
+
+
+def get_console_hosts(console_binary, conmux_server):
+	"""Use console to collect console hosts and return a list.
+ 
+	   Args:
+		console_binary:
+		string, location of the conmux console binary
+		conmux_server:
+		string, hostname of the conmux server
+	
+           Returns:
+		A List of console conmux is currently running on.
+	"""
+
+	hosts_list = []
+	cmd = "%s --list %s" % (console_binary, conmux_server)
+	for line in commands.getoutput(cmd).split('\n'):
+		host = (line.split(' '))[0]
+		hosts_list.append(host)
+
+	return hosts_list
+
+
+def check_host(host, console_binary):
+	"""Check hosts for common errors and return the status.
+
+	   Args:
+		host:
+		string, the console host identifier
+
+		console_binary:
+		string, location of the conmux console binary
+	   Returns:
+		int, 0: Machine state is good
+		int, 1: Machine state is bad
+		int, 2: Machine state is unknown
+	"""
+	RESPONSES = [ host + ' login:',
+		      'ENOENT entry not found',
+		      'login:',
+		      'Connection refused',
+		      '<<<NOT CONNECTED>>>',
+		      'Authentication failure',
+		      'Give root password for maintenance', ]
+
+	cmd = '%s %s' % (console_binary, host)
+	shell = pexpect.spawn(cmd)
+
+	shell.send('\r\n')
+	shell.send('\r\n')
+	shell.send('\r\n')
+	try:
+		# May need to increase the timeout but good so far
+		response = shell.expect(RESPONSES, 1)
+	except pexpect.TIMEOUT:
+		shell.sendline('~$')
+		shell.expect('>')
+		shell.sendline('quit')
+		return 1
+	except pexpect.EOF:
+		# unknown error
+		shell.sendline('~$')
+		shell.expect('>')
+		shell.sendline('quit')
+		return 2
+	# TODO: Change actions based on what server returned
+	if response == 0:
+		# OK response
+		return 0
+	else:
+		return 1
+
+
+if __name__ == '__main__':
+	main(sys.argv)