#!/usr/bin/env python

from rox import g, saving
import rox
import fcntl

try:
	from rox import processes
except ImportError:
	rox.croak(_('Sorry, this version of Archive requires ROX-Lib 1.9.3 or later'))

import sys, os

class ChildError(Exception):
	"Raised when the child process reports an error."

class ChildKilled(saving.AbortSave):
	"Raised when child died due to calling the kill method."
	def __init__(self):
		saving.AbortSave.__init__(self, _("Operation aborted at user's request"))

def escape(text):
	"""Return text with \ and ' escaped"""
	return text.replace("\\", "\\\\").replace("'", "\\'")

def Tmp(mode = 'w+b'):
	"Create a seekable, randomly named temp file (deleted automatically after use)."
	import tempfile
	try:
		return tempfile.NamedTemporaryFile(mode, suffix = '-archive')
	except:
		# python2.2 doesn't have NamedTemporaryFile...
		pass

	import random
	name = tempfile.mktemp(`random.randint(1, 1000000)` + '-archive')

	fd = os.open(name, os.O_RDWR|os.O_CREAT|os.O_EXCL, 0700)
	tmp = tempfile.TemporaryFileWrapper(os.fdopen(fd, mode), name)
	tmp.name = name
	return tmp

def keep_on_exec(fd):
	fcntl.fcntl(fd, fcntl.F_SETFD, 0)

class PipeThroughCommand(processes.Process):
	def __init__(self, command, src, dst):
		"""Execute 'command' with src as stdin and writing to stream
		dst. src must be a fileno() stream, but dst need not be.
		Either stream may be None if input or output is not required.
		Call the wait() method to wait for the command to finish."""

		assert src is None or hasattr(src, 'fileno')

		processes.Process.__init__(self)

		self.command = command
		self.dst = dst
		self.src = src
		self.tmp_stream = None

		self.callback = None
		self.killed = 0
		self.errors = ""

		self.start()

	def pre_fork(self):
		# Output to 'dst' directly if it's a fileno stream. Otherwise,
		# send output to a temporary file.
		assert self.tmp_stream is None

		if self.dst:
			if hasattr(self.dst, 'fileno'):
				self.dst.flush()
				self.tmp_stream = self.dst
			else:
				self.tmp_stream = Tmp()

	def start_error(self):
		"""Clean up effects of pre_fork()."""
		self.tmp_stream = None

	def child_run(self):
		src = self.src

		if src:
			os.dup2(src.fileno(), 0)
			keep_on_exec(0)
			os.lseek(0, 0, 0)	# OpenBSD needs this, dunno why
		if self.dst:
			os.dup2(self.tmp_stream.fileno(), 1)
			keep_on_exec(1)
	
		if os.system(self.command) == 0:
			os._exit(0)	# No error code or signal
		os._exit(1)
	
	def parent_post_fork(self):
		if self.dst and self.tmp_stream is self.dst:
			self.tmp_stream = None
	
	def got_error_output(self, data):
		self.errors += data
	
	def child_died(self, status):
		errors = self.errors.strip()

		err = None

		if self.killed:
			err = ChildKilled
		elif errors:
			err = ChildError(_("Errors from command '%s':\n%s") % (self.command, errors))
		elif status != 0:
			err = ChildError(_("Command '%s' returned an error code!") % self.command)

		# If dst wasn't a fileno stream, copy from the temp file to it
		if not err and self.tmp_stream:
			self.tmp_stream.seek(0)
			self.dst.write(self.tmp_stream.read())
		self.tmp_stream = None

		self.callback(err)
	
	def wait(self):
		"""Run a recursive mainloop until the command terminates.
		Raises an exception on error."""
		done = []
		def set_done(exception):
			done.append(exception)
			g.mainquit()
		self.callback = set_done
		while not done:
			g.mainloop()
		exception, = done
		if exception:
			raise exception
	
	def kill(self):
		self.killed = 1
		processes.Process.kill(self)
	
def test():
	"Check that this module works."

	def show():
		error = sys.exc_info()[1]
		print "(error reported was '%s')" % error
	
	def pipe_through_command(command, src, dst): PipeThroughCommand(command, src, dst).wait()

	print "Test escape()..."

	assert escape('''  a test ''') == '  a test '
	assert escape('''  "a's test" ''') == '''  "a\\'s test" '''
	assert escape('''  "a\\'s test" ''') == '''  "a\\\\\\'s test" '''

	print "Test Tmp()..."
	
	file = Tmp()
	file.write('Hello')
	print >>file, ' ',
	file.flush()
	os.write(file.fileno(), 'World')

	file.seek(0)
	assert file.read() == 'Hello World'

	print "Test pipe_through_command():"

	print "Try an invalid command..."
	try:
		pipe_through_command('bad_command_1234', None, None)
		assert 0
	except ChildError:
		show()
	else:
		assert 0

	print "Try a valid command..."
	pipe_through_command('exit 0', None, None)
	
	print "Writing to a non-fileno stream..."
	from cStringIO import StringIO
	a = StringIO()
	pipe_through_command('echo Hello', None, a)
	assert a.getvalue() == 'Hello\n'

	print "Reading from a stream to a StringIO..."
	file.seek(1)
	pipe_through_command('cat', file, a)
	assert a.getvalue() == 'Hello\nello World'

	print "Writing to a fileno stream..."
	file.seek(0)
	file.truncate(0)
	pipe_through_command('echo Foo', None, file)
	file.seek(0)
	assert file.read() == 'Foo\n'

	print "Read and write fileno streams..."
	src = Tmp()
	src.write('123')
	src.seek(0)
	file.seek(0)
	file.truncate(0)
	pipe_through_command('cat', src, file)
	file.seek(0)
	assert file.read() == '123'

	print "Detect non-zero exit value..."
	try:
		pipe_through_command('exit 1', None, None)
	except ChildError:
		show()
	else:
		assert 0
	
	print "Detect writes to stderr..."
	try:
		pipe_through_command('echo one >&2; sleep 2; echo two >&2', None, None)
	except ChildError:
		show()
	else:
		assert 0

	print "Check tmp file is deleted..."
	name = file.name
	assert os.path.exists(name)
	file = None
	assert not os.path.exists(name)

	print "Check we can kill a runaway proces..."
	ptc = PipeThroughCommand('sleep 100; exit 1', None, None)
	def stop():
		ptc.kill()
	g.timeout_add(2000, stop)
	try:
		ptc.wait()
		assert 0
	except ChildKilled:
		pass
	
	print "All tests passed!"

if __name__ == '__main__':
	test()
