
# $Id: remote.py,v 1.16 2004/01/17 18:50:46 cran Exp $
#
# Copyright (c) 2002 Sebastian Stark
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
# 1. Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
# 2. Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE AUTHOR SEBASTIAN STARK
# ``AS IS'' AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
# A PARTICULAR PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE AUTHOR
# OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.


"""Various remote methods for tentakel

Currently, only the "ssh" method is implemented."""

from error import *
import sys
import commands
import threading, Queue
import tpg

class FormatString(tpg.Parser):
	r"""

	token escape	: '\\[\\nt]'	str ;
	token fmtchar	: '%[%dos]'	str ;
	token char	: '.'		str ;

	START/e -> FORMAT/e ;

	FORMAT/f ->
					$ f = ""
			( escape/e	$ f = f + self.getEscape(e)
			| fmtchar/fc	$ f = f + self.getSpecialChar(fc)
			| char/c	$ f = f + c
			)*
	;
	"""

	def __init__(self, formatMap=None):
		tpg.Parser.__init__(self)
		self.formatMap = formatMap
	
	def getEscape(self, e):
		return {
			r"\\": r"\\",
			r"\n": "\n",
			r"\t": "\t"
		}[e]

	def getSpecialChar(self,s):
		return self.formatMap[s]


class RemoteCommand(threading.Thread):
	"""SSH remote execution object"""

	def __init__(self, sshpath="ssh", destination=None, user=None):
		threading.Thread.__init__(self)
		self.setDaemon(1)
		self.sshpath = sshpath
		self.user = user
		self.destination = destination
		self.commandQueue = Queue.Queue()
		self.resultQueue = Queue.Queue()
		self.start()

	def execute(self, command):
		"""Execute a command in this thread"""
		
		self.commandQueue.put(command)

	def run(self):
		while 1:
			s = '%s %s@%s "%s"' % (self.sshpath, self.user, self.destination, self.commandQueue.get())
			status, output = commands.getstatusoutput(s)
			# shift 8 bits right to strip signal number from status
			self.resultQueue.put((status >> 8, output))

class RemoteCollator:
	"""This class is meant to hold RemoteCommand instances"""

	def __init__(self, conf, groupName):
		self.clear()
		self.useConf(conf, groupName)
		self.formatter = FormatString()

	def clear(self):
		self.remoteobjects = []

	def useConf(self, conf, groupName):
		save = self
		self.clear()
	        try:
			user = conf.getParam("user", group=groupName)
			sshpath = conf.getParam("ssh_path", group=groupName)
	                for host in conf.getExpandedGroup(groupName):
	                        self.add(RemoteCommand(destination=host, user=user, sshpath=sshpath))
			self.format = conf.getParam("format", group=groupName)
	        except KeyError:
			self = save
	                warn("unknown group: '%s'" % groupName)

	def getDestinations(self):
		return [x.destination for x in self.remoteobjects]

	def add(self,obj):
		if isinstance(obj,RemoteCommand):
			self.remoteobjects.append(obj)
		else:
			pass
	
	def remove(self,obj):
		self.remoteobjects.remove(obj)

	def formatRecord(self, formatMap=None):
		"""Apply a format mapping to the format string

		Outputs the format with formatting expressions replaced
		by values taken from formatMap. The formatMap must contain
		translations for the formatting expressions. For example:

		  formatMap = { r"%%": "%", r"%d": "something" }

		The real translation happens in the formatter object
		which is an instance of the FormatString parser class.
		
		"""

		self.formatter.formatMap = formatMap
		return self.formatter(self.format)

	def execAll(self, command):
		"Execute command on all remote objects"

		# TODO: think about this a bit
		for obj in self.remoteobjects:
			obj.execute(command)

	def displayAll(self):
		"Display the results of all remote objects"

		# TODO: output results when they're ready, not later
		for obj in self.remoteobjects:
			status, output = obj.resultQueue.get()
			formatMap = {
					r"%%": "%",
					r"%d": obj.destination,
					r"%o": output,
					r"%s": str(status)
				}
			sys.stdout.write(self.formatRecord(formatMap))


if __name__ == "__main__":

	failures = 0
	print "self testing..."

	print "### instantiate RemoteCommand:"
	import pwd, os
	user = pwd.getpwuid(os.geteuid())[0]
	r = RemoteCommand(user=user, destination="localhost")
	if isinstance(r, RemoteCommand):
		print "OK"
	else:
		print "-> failed <-"
		failures += 1

	print "### execute command:"
	v = (0, "test123")
	r.execute("echo test123")
	res = r.resultQueue.get()
	if res == v:
		print "OK"
	else:
		print "-> failed <-"
		print "read", res, "but should be", v
		failures += 1

	if failures:
		print "self test: encountered", failure, "failures."
		sys.exit(1)
	else:
		print "self test: no failures."
		sys.exit(0)


