#!/usr/bin/python

from multiprocessing import Process
import os
import sys, getopt

verbose = 1
use_pkill = False

#------------------
# main
#------------------
def Usage():
	print ""
	print "Usage:"
	print "       kill_remote_processes [options] args"
	print ""
	print " This script is used to send a system signal to all processes"
	print "with the specified name. It simply runs \"killall\" (or pkill)"
	print "on each remote node passing it \"args\" as the argument list."
	print "The remote invocation is done via \"ssh\". The list of nodes"
	print "to which the command is sent may be specified explicitly on"
	print "the command line or via an input file containing one node name"
	print "per line."
	print ""
	print "options:"
	print ""
	print "  -h                  Print usage statement"
	print "  -q                  Quiet (print less output)"
	print "  -p                  Use pkill instead of killall (see below)"
	print "  -f file             Name of file from which to get node names"
	print "  -n node1,node2,...  Specify remote nodes to execute command"
	print ""
	print "Note that multiple -f and -n arguments can be given. If a node"
	print "appears more than once then it will have the command executed"
	print "more than once."
	print ""
	print "By default, the 'killall' command is used on the remote node"
	print "to find and kill the process. This will only look at the"
	print "command itself and not the entire command line. This is"
	print "generally considered safer and so it is the default. If you"
	print "need to identify the process(es) by some other string on the"
	print "command line, use the '-p' option and the 'pkill -f' command"
	print "will be used on the remote node instead."
	print ""
	print "Examples:"
	print ""
	print "    kill_remote_processes -n L3_nodes.txt start_L3_proc"
	print ""
	print "This will run \"killall start_L3_proc\" on each node listed"
	print "in the file \"L3nodes.txt\". Since killall sends the SIGINT"
	print "by default, then that is the signal that will get sent. To"
	print "send a SIGKILL, add the argument \"-9\" to the command."
	print ""
	print "Be aware too that if one started a remote process using"
	print "\"start_remote_processes\" and the given command was a script,"
	print "then the script may have started a process with a different"
	print "name and that is what should be used here."
	print ""
	print "   kill_remote_processes -n gluon44 eb_start -9"
	print ""
	print "This will send a SIGKILL signal to all \"eb_start\" processes"
	print "on the node gluon44."
	exit()

#------------------
# RunProg
#------------------
def RunProg(node, cmd):
	os.system("ssh %s \"%s\"" % (node, cmd))

#------------------
# main
#------------------
def main(argv):

	global verbose,use_pkill

	if (len(sys.argv) < 2):
		print "\nTOO FEW ARGUMENTS!"
		Usage()

	nodes = []
	try:
		opts, args = getopt.getopt(argv,"hqpf:u:n:",["help","file=","nodes="])
	except getopt.GetoptError,exc:
		print "\nError parsing options!"
		print exc.msg 
		Usage()
	for opt, arg in opts:
		if opt in ("-h", "--help"):
			Usage()
		elif opt in ("-q", "--quiet"):
			verbose=0
		elif opt in ("-p", "--pkill"):
			use_pkill=True
		elif opt in ("-f", "--file"):
			with open(arg) as f:
				nodes.extend(f.read().splitlines())
		elif opt in ("-n", "--nodes"):
			nodes.extend(arg.split(','))
		elif opt in ("-u"):
			args.insert(0,arg)
			args.insert(0,opt)


	# Command is contained in the non-option arguments list "args"
	if(len(args)<1):
		print "\nYOU MUST SPECIFY A COMMAND TO RUN!"
		Usage()

	command = "killall "
	if use_pkill : command = "pkill -f "
	for a in args:
		command += a+" "
	command = command.strip()

	# Check for empty node names. This can happen if a file is specified and
	# has extra carriage returns.
	tmp_nodes = []
	for n in nodes:
		tmp = n.strip()
		if( len(tmp)>0 ): tmp_nodes.append(tmp)
	nodes = tmp_nodes
	if verbose>=1 : print "   command: '%s'" % (command)
	if verbose>=1 : print "num. nodes: %d" % (len(nodes))

	procs = []
	for n in nodes:
		p = Process(target=RunProg, args=(n,command))
		if verbose>=1 : print "Running command on: %s" % (n)
		p.start()
		procs.append(p)

	for p in procs:
		p.join()


# The following "if" is for Windows
if __name__ == "__main__":
    main(sys.argv[1:])