cpush

#!/usr/bin/env python2
try:
	import sys, os
	######### set defaultt c3 path ########################
        try:
                def_path = os.environ[ 'C3_PATH' ]
        except KeyError:
                def_path = '/opt/c3-4'
        #######################################################

	sys.path.append( def_path )
	import c3_com_obj, c3_file_obj, c3_sock, socket, re

	######## constants ####################################
	help_info = """c3 version 4.0
	Usage: cpush [OPTIONS] [MACHINE DEFINTIONS] source [target]

	the transfer will fail if the target contains any wild cards

	if target is not specified then the source is placed in the directory from 
	which "source" came from, the source can not have any wildcards if no target
	was specified


		-h, --help	= this screen
		-f, --file	= file containing list of ip's if one is not supplied 
				  then /etc/c3.conf will be used
		-l, --list	= list of files to push (see man page for format
		-i		= interactive mode, ask once before executing
		--head		= execute command on head node does not execute on 
				  compute nodes if specified
		--nolocal	= the source file or directory lies on the head node 
				  of the remote cluster
		-b, --blind	= pushes the entier file (cpush uses rsync to push 
		--all           = execute command on all the nodes
				  on all the clusters in the configuration
				  file, ignores the [MACHINE_DEFINITONS]


		machine definitions are in the form of 
			clusterName: start-end, node number"""

	backslash = re.compile( r"\\" )
	#######################################################


	######## check for arguments  #########################
	if len(sys.argv) == 1:
		print help_info
		sys.exit()
	#######################################################


	######### et default options  #########################
	to_print = 0			#not used in parallel version
	interactive = 0			#prompt before execution
	head_node = 0			#execute only on head node
	filename = "/etc/c3.conf"	#default config file
	options = ""			#options to be passed to a remote cluster
	options_to_pass = ""		#if a remote cluster, options to command
	filelist = ""			#file with list of files to be transfered
	list_of_files = []		#list object of files to transfer
	exclude_pattern = ""		#regex to exclude from push
	pidlist = []                    #list if pid's to wait on   
	blind_set = 0			#if set do a blind push
	nolocal = 0			#default is to push the file from the current machine
	defusername = ""		#default user name
	defc3path = ""			#default path to the C3 files
	isdir = ""
	#######################################################


	######### internal variables ##########################
	cluster_from_file = {}
	all = 0
	file_set = 0
	file_list_set = 0
	exclude_set = 0
	blind_set = 0
	#######################################################



	######### parse command line ##########################

	# first create one large string of the command 
	command_line_list = sys.argv[1:]
	command_line_string = ''
	for item in command_line_list:
		command_line_string = '%s %s' % (command_line_string, item)

	# object used to parse command line
	c3_command = c3_com_obj.c3_command_line( command_line_string )
	 
	# get first option
	try:
		option = c3_command.get_opt()
		while option: # while more options
			if option == '-l' or option == "--list": # list of file to transfer
				if not file_list_set:
					filelist = c3_command.get_opt_string()
					file_list_set = 1
				else:
					print "only one file list may be specified"
					sys.exit()
			elif (option == '-f') or (option == '--file'): # alternate config file
				if not file_set:
					filename = c3_command.get_opt_string()
					file_set = 1
				else:
					print "only one file name can be specified."
					sys.exit()
			elif option == '-h' or option == "--help": # print help info
				print help_info
				sys.exit()
			elif option == '-i': # ask once before executing command
				interactive = 1
			elif option == '--head': # execute only the head node
				head_node = 1
			elif option == '-x' or option == '--exclude': # exlude a pattern from the push
				if blind_set:
					print "can not do a blind push with an exclude"
					sys.exit()
				exclude_set = 1	
				exclude_pattern = c3_command.get_opt_string()
				options_to_pass = options_to_pass + " -x \"" + exclude_pattern + "\""
				options = options + " --exclude \'" + exclude_pattern + "\'"
			elif option == '-d' or option == '--delete':
				options_to_pass = options_to_pass + " -d"
				options = options + " --delete"
			elif option == '--nolocal':
				nolocal = 1
			elif option == '-b' or option == '--blind':
				if exclude_set:
					print "can not do a blind push with an exclude"
					sys.exit()
				blind_set = 1
			elif option == '--isdir':
				isdir = c3_command.get_opt_string()
			elif option == '--all':
				all = 1
			else: # a catch all, option supplied not valid
				print "option ", option, " is not valid"
				sys.exit()
			option = c3_command.get_opt()
	except c3_com_obj.end_of_option: #quit parsing options
		pass

	if isdir != "":
		sock = c3_sock.client_sock()
		if os.path.isdir( isdir ):
			sock.send( 'yes' )
		else:
			sock.send( 'no' )
		sys.exit()

	# create cluster object from command line 
	clusters = c3_com_obj.c3_cluster_list()
	 
	if all == 0:
		clusters = c3_command.get_clusters()
	else:
		c3_command.get_clusters()

	if filelist:
		templistoffiles = open( filelist, "r" ).readlines()
		for line in templistoffiles:
			temp_list = line.strip().split()
			if len( temp_list ) > 2:
				print "error in file list ", filelist, " more than source and destination files on one line."
				sys.exit()
			if len( temp_list ) == 0:
				continue
			list_of_files.append( temp_list )
			try:
				list_of_files[-1][1]
			except IndexError:
				list_of_files[-1].append( list_of_files[-1][0] )
			if list_of_files[-1][0][0] != '/':
				list_of_files[-1][0] = os.getcwd() + "/" + list_of_files[-1][0]
			if list_of_files[-1][-1][0] != '/':
				list_of_files[-1][1] = os.getcwd() + "/" + list_of_files[-1][1]
	else:	
		# get source 
		try:
			source = c3_command.get_opt_string()
		except c3_com_obj.bad_string:
			print "Must supply a source file."
			sys.exit()
		# get target
		try:
			target = c3_command.get_opt_string()
		except c3_com_obj.bad_string:
			if not source[0] == '/':
				target = os.getcwd() +'/' + source
			else:
				target = source
		list_of_files.append( [] )
		list_of_files[0].append( source )
		list_of_files[0].append( target )

	#######################################################


	######### test if ssh or rsh ##########################
	try:
		transport = os.environ[ 'C3_RSH' ]
	except KeyError:
		transport = 'ssh'
	#######################################################

	######### et default username #########################
	try:
		defusername = os.environ[ 'C3_USER' ]
	except KeyError:
		defusername = os.environ[ 'USER' ]
	#######################################################
        ######### set filename  ##########################
        if not file_set:
                try:
                        filename = os.environ[ 'C3_CONF' ]
                except KeyError:
                        filename = '/etc/c3.conf'
        #######################################################


	######### make cluster list object from file ##########
	try: # open file & initialize file parser
		file = c3_file_obj.cluster_def( filename )
	except IOError:
		print "error opening file: ", filename
		sys.exit()

	# note, someday need to change to only get clusters needed, need changes to the
	# file parser object for that, will be done in later versions

	try:
		file.get_next_cluster() # set the default cluster (first cluster in list)
		try:
			if clusters.clusters[0] ==  "/default": # if default cluster set correct cluster name
				clusters.clusters[0] = file.get_cluster_name()
				clusters.node[file.get_cluster_name()] = clusters.node["/default"]
				del clusters.node["/default"]
				clusters.username[file.get_cluster_name()] = clusters.username["/default"]
		except IndexError: #if all is specified will throw
			pass
		while(1): #throws exception when no more clusters
			c_name = file.get_cluster_name() #name of cluster being parsed
			if all == 1:
				clusters.clusters.append( c_name )
				clusters.node[c_name] = []
				clusters.node[c_name].append( "" )
			if c_name in clusters.clusters:
				cluster_from_file[c_name] = {}
				cluster_from_file[c_name]['external'] = file.get_external_head_node()
				cluster_from_file[c_name]['internal'] = file.get_internal_head_node()
				cluster_from_file[c_name]['nodes'] = [] #list of node names from file
				if cluster_from_file[c_name]['external']: #if a direct cluster
					index = 0
					try:
						while(1): # build list of nodes
							node_obj = file.get_next_node()
							cluster_from_file[c_name]['nodes'].append( c3_file_obj.node_obj() )
							cluster_from_file[c_name]['nodes'][index] = node_obj
							index = index + 1
					except c3_file_obj.end_of_cluster:
						pass
			file.get_next_cluster() #repeat untill no more clusters
	except c3_file_obj.no_more_clusters:
		pass
	except c3_file_obj.parse_error, error:
		print error.description
		print "somewhere around ", error.last
		sys.exit()

	#######################################################


	######### execute command on each node in cluster



	# there are two main groups, local and remote clusters
	# in each of those groups there are direct and indirect
	# modes, that is every node specified or a "link". A link
	# on a local cluster is of course invalid.

	# right now the only way I know how to check if a cluster
	# is local is to use a "gethostbyname" and compare it
	# to the head node names (both internel and externel).
	# right now that is acceptable as many tools require 
	# the function to work correctly (ssh being one of them)
	# my want to think about a better way.

	#host_ip = socket.gethostbyname( socket.gethostname() )



        while interactive: # if interactive execution, prompt once before executing
                answer = raw_input( "push " + source + " to " + target + " :" )
                if re.compile( r".*n(o)?.*", re.IGNORECASE ).match( answer ):
                        sys.exit()
                if re.compile( r".*y(es)?.*", re.IGNORECASE).match( answer ):
                        interactive = 0
		

	pid_list_outer = []

	for cluster in clusters.clusters:

		pid = os.fork()
		if pid == 0:

			############ get machine names #############################
			try: 
				local_ip = socket.gethostbyname( socket.gethostname() )
			except socket.error:
				print "Can not resolve local hostname"
				sys.exit()
			try:
				int_ip = socket.gethostbyname( cluster_from_file[cluster]['internal'] )
			except socket.error:
				int_ip = ""
                        except KeyError:
                                print "Cluster ", cluster, " is not in you configuration file."
                                continue
	
			try:
				ext_ip = socket.gethostbyname( cluster_from_file[cluster]['external'] )
			except socket.error:
				ext_ip = ""
			except TypeError:
				ext_ip = int_ip
			############################################################
			
			try:
				if clusters.username[cluster]=="/default":
					username = defusername
				else:
					username = clusters.username[cluster]
			except KeyError: #if --all is specified
				username = defusername
					
			if head_node:   #if only execute on head node , do so
					#will execute on local cluster with ssh also
				for file in list_of_files:
					if blind_set:
						string_to_execute = transport + " " + "-l " + username + " " + node_name + " \'/bin/rm -f " + file[1] + "\'"
						os.system( string_to_execute )
					string_to_execute = "rsync -avz --rsh=" + transport + " " + options + " \'" + file[0] + "\' " + username + "@" + ext_ip + ":\'" + file[1] + "\'"
					os.system( string_to_execute )

			elif cluster_from_file[cluster]['external']: # if a direct cluster
				if ext_ip == local_ip or int_ip == local_ip: # if a local cluster

					if ext_ip == "": # error conditions (just don't execute current cluster)
						print "Can not resolve ", cluster_from_file[cluster]['external']
					elif int_ip == "":
						print "can not resolve ", cluster_from_file[cluster]['internal']

					elif clusters.node[cluster][0] != "" : #range specified on command line
						for node in clusters.node[cluster]: #for each cluster specified on the command line
							try:
								if not cluster_from_file[cluster]['nodes'][node].dead: #if machine is not dead
									node_name = cluster_from_file[cluster]['nodes'][node].name
									pid = os.fork() # execute command on each node in it's own process
									if pid == 0:
										node_name = cluster_from_file[cluster]['nodes'][node].name
										for file in list_of_files:
											if not os.path.isfile( file[0] ):
												print file[0] + " must be a file."
												sys.stdout.flush()
												continue
											if blind_set:
												string_to_execute = transport + " " + "-l " + username + " " + node_name + " \'/bin/rm -f " + file[1] + "\'"
												os.system( string_to_execute )
											string_to_execute = "rsync -avz --rsh=" + transport + " " + options + " \'" + file[0] + "\' " + username + "@" + node_name + ":\'" + file[1] + "\'"
											os.system( string_to_execute )
										os._exit(1)
									pidlist.append(pid)
							except IndexError:
								pass
					else:  # no range specified on command line, do all nodes
						for node in cluster_from_file[cluster]['nodes']:  # for each node in cluster
							if not node.dead: #if node not dead
								pid = os.fork() # execute command in own process
								if pid == 0:
									node_name = node.name
									for file in list_of_files:
										if not os.path.isfile( file[0] ):
											print file[0] + " must be a file."
											sys.stdout.flush()
											continue
										if blind_set:
											string_to_execute = transport + " " + "-l " + username + " " + node_name + " \'/bin/rm -f " + file[1] + "\'"
											os.system( string_to_execute )
										string_to_execute = "rsync -avz --rsh=" + transport + " " + options + " \'" + file[0] + "\' " + username + "@" + node_name + ":\'" + file[1] + "\'"
										os.system( string_to_execute )
									os._exit(1)
								pidlist.append(pid)
				else: # remote cluster
					if ext_ip == "": # error condition
						print "Can not resolve ", cluster_from_file[cluster]['external']
						sys.stdout.flush()
					else:
						# generate temprorary config file
						cluster_def_string = "cluster " + cluster + " {\n"
						cluster_def_string = cluster_def_string + cluster_from_file[cluster]['external'] + ":" + cluster_from_file[cluster]['internal'] + "\n"

			
						if clusters.node[cluster][0] != "" : #range specified on command line
							try:
								for node in clusters.node[cluster]: #for each node specified on command line
									if not cluster_from_file[cluster]['nodes'][node].dead: # if cluster node not dead
										node_name = cluster_from_file[cluster]['nodes'][node].name #add cluster to list
										cluster_def_string = cluster_def_string + node_name + "\n"
							except IndexError:
								pass
						else:  # no range specified on command line, do all nodes
							for node in cluster_from_file[cluster]['nodes']: # for each node in cluster
								if not node.dead: #if node not dead, add to list
										cluster_def_string = cluster_def_string +  node.name + "\n"
						cluster_def_string = cluster_def_string + "}" # close list
										
						#this is an attempt to generate a unique file name. As there is no easy way
						#for me to see the remote machine I mangle a group of relativly unique
						#identifiers. since i prepend the machine ip address if the file does not reside localy
						#it should not remotely, at the very least it should be safe to rewrite the file
						tempfilenamebase = local_ip + "%d" % os.getuid() + "%d" % os.getpid() 
						tempfilename = "/tmp/" + tempfilenamebase
							
						string_to_execute = transport + " -l " + username + " " + ext_ip + " " + def_path + "/cget --head --non_interactive " + tempfilename
						sock = c3_sock.server_sock( string_to_execute )
						answer = sock.recieve()
						sock.close()
						while answer == 'good': #make sure file name is unique on remote machine
							tempfilename = tempfilename + "1"
							string_to_execute = transport + " -l " + username + " " + ext_ip + " " + def_path + "/cget --head --non_interactive " + tempfilename
							sock.__init__( string_to_execute )
							answer = sock.recieve()
							sock.close()
						for line in list_of_files:
							go_nogo = "go"
							if not nolocal:
								remotefilename = tempfilename + os.path.basename(line[0]) 
							else:
								remotefilename = line[0]
								string_to_execute = transport + " -l " + username + " " + ext_ip + " " + def_path + "/cget --head --non_interactive " + remotefilename
								sock = c3_sock.server_sock( string_to_execute )
								answer = sock.recieve()
								sock.close()
								if answer == 'bad': #make sure file name is unique on remote machine
									print "file must exist on remote machine"
									go_nogo = "nogo"

							
							if go_nogo == 'go':	
								FILE = open( tempfilename, 'w' )
								FILE.write( cluster_def_string )
								FILE.close()
								# push file to remote machine
								if not nolocal:
									if not os.path.isfile( line[0] ):
										print line[0] + " must be a file."
										continue
									string_to_execute = "rsync -avz --rsh=" + transport + " " + line[0] + " " + username + "@" + ext_ip + ":" + remotefilename
									os.system( string_to_execute )
							
								string_to_execute = "rsync -v --rsh=" + transport + " " + tempfilename + " " + username + "@" + ext_ip + ":" +  tempfilename
								os.system( string_to_execute )

								# execute commend remotely
								string_to_execute = transport + " -l " + username + " " + ext_ip + " " + def_path + "/cpush --isdir " + line[1]
								sock = c3_sock.server_sock( string_to_execute )
								answer = sock.recieve()
								if answer == 'yes':
									line[1] = line[1] + '/' + os.path.basename(line[0])
								sock.close()
								
								string_to_execute = transport + " -l " + username + " " + ext_ip + " \' " + def_path + "/cpush " + options_to_pass + " -f " + tempfilename + " " + remotefilename + " " + line[1] + "\'"
								os.system( string_to_execute )

								# remove temporary file
								string_to_execute = transport + " -l " + username + " " + ext_ip + " /bin/rm -f " + tempfilename
								os.system( string_to_execute )
								if not nolocal:
									string_to_execute = transport + " -l " + username + " " + ext_ip + " /bin/rm -rf " + remotefilename
									os.system( string_to_execute )
						if os.path.exists( tempfilename):
							os.unlink( tempfilename ) # remove local temp file
					
			else: # indirect clusters
				if int_ip == local_ip:  # can not have a indirect local cluster since if your default cluster
							# is local you would generate a circular reference
					print "error local cluster"
					sys.stdout.flush()
				else: # remote indirect cluster
					if int_ip == "": # error condition
						print "Can not resolve hostname ", cluster_from_file[cluster]['internal']
						sys.stdout.flush()
					else:
						for line in list_of_files:
							go_nogo = "go"
							if not nolocal:
								tempfilename = "/tmp/" + local_ip + "%d" % os.getuid() + "%d" % os.getpid() + os.path.basename(line[0])

								string_to_execute = transport + " -l " + username + " " + ext_ip + " " + def_path + "/cget --head --non_interactive " + tempfilename
								sock = c3_sock.server_sock( string_to_execute )
								answer = sock.recieve()
								sock.close()
								while answer == 'good': #make sure file name is unique on local machine
									filename = filename + "1"
									string_to_execute = transport + " -l " + username + " " + ext_ip + " " + def_path + "/cget --head --non_interactive " + tempfilename
									sock.__init__( string_to_execute )
									answer = sock.recieve()
									sock.close()

							
							else:
								tempfilename = line[0]
								string_to_execute = transport + " -l " + username + " " + ext_ip + " " + def_path + "/cget --head --non_interactive " + tempfilename
								sock = c3_sock.server_sock( string_to_execute )
								answer = sock.recieve()
								sock.close()
								if answer == 'bad': #make sure file name is unique on local machine
									go_nogo = "nogo"
									print "file must exist on remote machine"
							
							if go_nogo == 'go':
								# push file to remote machine
								if not nolocal:
									if not os.path.isfile ( line[0] ):
										print line[0] + " must be a file."	
										continue
									string_to_execute = "rsync -avz --rsh=" + transport + " " + line[0] + " " + username + "@" + int_ip + ":" + tempfilename
									os.system( string_to_execute )
								
								
								node_list = "" # generate new node list from the command line
								if clusters.node[cluster][0] != "" : #range specified on command line
									node_list = ":%d" % clusters.node[cluster].pop(0)

									for node in clusters.node[cluster]:
										node_list = node_list + ", %d" % node	
								# execute command on remote machine
								string_to_execute = transport + " -l " + username + " " + ext_ip + " " + def_path + "/cpush --isdir " + line[1]
								sock = c3_sock.server_sock( string_to_execute )
								answer = sock.recieve()
								if answer == 'yes':
									line[1] = line[1] + '/' + os.path.basename(line[0])
								sock.close()

								string_to_execute = transport + " -l " + username + " " + int_ip + " \' " + def_path + "/cpush " + options_to_pass + " " + node_list + " " + tempfilename + " " + line[1] + "\'"
								os.system( string_to_execute )

								if not nolocal:
									# remove temporary file
									string_to_execute = transport + " -l " + username + " " + int_ip + " /bin/rm -rf " + tempfilename
									os.system( string_to_execute )


			for pid in pidlist: # wait for all processes spawned to finish
				os.waitpid(pid,0)
			os._exit(1)
		pid_list_outer.append(pid)
	for pid in pid_list_outer:
		os.waitpid(pid, 0)
except KeyboardInterrupt:
        print "Keyboard interrupt\n"