# -*- Mode: Python; tab-width: 4 -*-
#	$Id: consumer.py,v 1.5 1996/03/22 04:34:16 rushing Exp $
#	Author: Sam Rushing <rushing@nightmare.com>

'Consumer fifo support for async chat sockets'

import os
import socket
import asyncore
import asynchat

import string

import sys
import tb

def top_of_traceback (tb):
	while tb.tb_next:
		tb = tb.tb_next
	return tb

# A 'function chain' is used to represent a future computation, as a
# series of functions to be called.  Every function in the chain
# (except the last) represents an asynchronous sub-task that needs to
# be performed, before we can finally return a result to the user.
# Except for the last function, each one is called with an extra
# 'continuation' as its first argument.  This represents the rest of
# the function chain, and will be called explicitly when ready for the
# next stage of the computation.

class function_chain:

	def __init__ (self, *funs):
		# Should we try to be smart here, and if the first
		# fun in funs is another function chain, simply copy
		# its fun list?  Or does it matter?
		self.funs = funs

	def __call__ (self, *args):
		fun = self.funs[0]
		self.funs = self.funs[1:]
		if self.funs:
			apply (fun, (self,)+args)
		else:
			apply (fun, args)

	# cut to the chase - call the last function in the
	# chain - usually used to abort because of a failed
	# sub-task.

	def abort (self, *args):
		apply (self.funs[-1], args)

# A delayed function object will evaluate itself when
# called.  It is used to package up a function and its
# arguments for later execution.

class delayed_function:
	def __init__ (self, fun, args):
		self.fun = fun
		self.args = args
	def __call__ (self, *args):
		# the given args are ignored, if any.
		# should we do something else?
		return apply (self.fun, self.args)

# a better way of returning a boolean result value
# for example - "return status_report (0, 'lost connection')"

class status_report:

	def __init__ (self, yesno, reason):
		self.reason = reason
		self.yesno = yesno

	def __nonzero__ (self):
		if self.yesno:
			return 1
		else:
			return 0

	def __repr__ (self):
		if self.yesno:
			s = 'success'
		else:
			s = 'failure'
		return '<%s "%s">' % (s, self.reason)

# This is a consumer object for single-line commands.  It's used to
# package up the user's reaction (as a function) to the response to a
# command being sent to a server.

class simple_consumer:
	def __init__ (self,
				  # done_fun (result)
				  done_fun,
				  terminator='\r\n'):
		self.terminator = terminator
		self.done_fun = done_fun
		self.result = ''

	def collect (self, data):
		self.result = self.result + data

	def found_terminator (self):
		self.done_fun (self.result)
		# to indicate we're done
		return 1

	def abort (self, why):
		self.done_fun (why)

import regex

# A consumer object used for multi-line commands.  If the first line
# (or whatever is terminated by <response_terminator>) that comes back
# matches good_result_reg, then switch to <command_terminator>, and
# collect the multi-line output, calling the user's callback for each
# line, until we see <command_terminator>, at which point we call the
# user's <done_fun>

class multi_line_consumer:
	def __init__ (self,
				  # done_fun (success)
				  done_fun,
				  good_result_reg,
				  line_collect_fun,
				  response_terminator='\r\n',
				  line_terminator='\r\n',
				  command_terminator='\r\n.\r\n'):
		self.good_result_reg = good_result_reg
		self.line_collect_fun = line_collect_fun
		self.done_fun = done_fun
		self.terminator = response_terminator
		self.line_terminator = line_terminator
		self.command_terminator = command_terminator
		self.data = ''
		# indicates we're waiting for the result line still
		self.good = 0

	def collect (self, data):
		self.data = self.data + data
		# if we're collecting output
		if self.good:
			lines = string.splitfields (self.data, self.line_terminator)
			# check for an incomplete line at the end...
			if lines[-1] == '':
				# a final empty string from splitfields indicates that
				# line_terminator was found at the end of the input string
				# i.e., no incomplete line.
				self.data = ''
			else:
				self.data = lines[-1]
			# either way, ignore it this pass
			lines = lines[:-1]
			# call the user's per-line callback function
			for line in lines:
				self.line_collect_fun (line)

	def found_terminator (self):
		if not self.good:
			if self.good_result_reg.match (self.data) == len(self.data):
				self.data = ''
				self.good = 1
				self.terminator = self.command_terminator
				return 0
			else:
				# report failure, and return 1 so that this consumer
				# is removed from the fifo.
				self.done_fun (status_report (0, self.data))
				return 1
		else:
			# if there's still a line waiting to be processed, don't
			# forget about it!
			if self.data:
				self.line_collect_fun (self.data)
			# ??? should self.data be returned here ???
			self.done_fun (status_report (1, self.data))
			return 1

	def abort (self, why):
		self.done_fun (why)

# ======================================================================
# Encapsulates the behavior of a typical internet client, with both
# single and multi-line response to commands. The response to a user
# command is passed back to the user via callback functions.
#
# Upon connection, a <simple_consumer> is optionally pushed unto the
# fifo that will handle the initial greeting from the server.
#
# FIXME: need a user interface for error and disconnect events.
# Some servers will send output unexpectedly (INN sends a warning just
# before killing an idle connection).
# ======================================================================

class typical_internet_client (asynchat.async_chat):

	ac_in_buffer_size	= 16384
	ac_out_buffer_size	= 16384

	def __init__ (self):
		self.create_socket (socket.AF_INET, socket.SOCK_STREAM)
		asynchat.async_chat.__init__ (self)
		self.consumer_fifo = asynchat.fifo()
		self.command_fifo = asynchat.fifo()
		self.saved_terminator = '\r\n'
		self.set_terminator('\r\n')
	
##	def send (self, data):
##		import win32ui
##		win32ui.SetStatusText (data[:-2], 1)
##		asynchat.async_chat.push (self, data)

	# ==================================================
	# consumer fifo
	# ==================================================

	def push_consumer (self, consumer):
		#? empty = len(self.consumer_fifo)
		self.consumer_fifo.push (consumer)

	def idle (self):
		return (len(self.consumer_fifo) == 0)

	def collect_incoming_data (self, data):
		if len(self.consumer_fifo):
			try:
				self.consumer_fifo.first().collect(data)
			except:
				tb = top_of_traceback (sys.exc_traceback)
				fname = tb.tb_frame.f_code.co_filename
				line_no = tb.tb_frame.f_lineno
				self.abort (
					status_report (
						0,
						('exception raised: line %d in %s' % (
							line_no,
							fname
							),
						 sys.exc_type,
						 sys.exc_value,
						 sys.exc_traceback
						 )
						)
					)
				return
		else:
			self.handle_unexpected (data)

	# overridable
	def handle_unexpected (self, data):
		print '%s: unexpected incoming data: %s' % (self, data)

	def handle_error (self, *why):
		self.abort (status_report (0, why))

	# we override this in order to make sure that we
	# use the current terminator of the first consumer
	# in the fifo.

	def get_terminator (self):
		if len(self.consumer_fifo):
			return self.consumer_fifo.first().terminator
		else:
			# reasonable default.
			print 'terminator defaulting!'
			return '\r\n'

	def found_terminator (self):
		if len(self.consumer_fifo):
			consumer = self.consumer_fifo.first()
			try:
				done = consumer.found_terminator()
			except:
				tb = top_of_traceback (sys.exc_traceback)
				fname = tb.tb_frame.f_code.co_filename
				line_no = tb.tb_frame.f_lineno
				self.abort (
					('exception raised: line %d in %s' % (
						line_no,
						fname
						),
					 sys.exc_type,
					 sys.exc_value,
					 sys.exc_traceback
					 )
					)
				return
			if done:
				self.consumer_fifo.pop()
		else:
			# ignore it for now.  usually the disconnect message
			pass
		# check the command queue
		if len(self.consumer_fifo) < self.get_max_outstanding():
			self.pop_command_queue()

	# ==================================================

	def go (self, host, port, greeting_fun=None):
		if greeting_fun:
			self.push_consumer (simple_consumer (greeting_fun))
		self.connect ((host, port))

	def abort (self, why):
		# close the connection first!
		self.close()
		# tell each queued consumer to abort.
		while 1:
			more, c = self.consumer_fifo.pop()
			if not more:
				break
			else:
				c.abort (why)
		# handle the consumers that have been queued,
		# but which have not yet had their commands sent.
		while 1:
			more, stuff = self.command_fifo.pop()
			if not more:
				break
			else:
				command, consumer = stuff	
				consumer.abort (why)

	# ==================================================
	# command queue management
	# ==================================================
	
	# We manage a queue of pending commands, along with
	# their consumers.  The maximum number of outstanding
	# commands can be overridden for finer control over
	# the server-abusive nature of the client.

	# places a limit on the number of 'outstanding'
	# commands - those for which a command has been sent
	# (using self.push()), but which have not yet received
	# a reply.
	max_outstanding = 4

	# overridable accessor
	def get_max_outstanding (self):
		return self.max_outstanding

	# send or queue a command and its consumer
	def send_command (self, command, consumer):
		self.command_fifo.push ((command, consumer))
		# if we're not too backed up, go ahead and send
		# the command at the front of the queue (which
		# may very well be the one we just pushed).
		if len(self.consumer_fifo) < self.get_max_outstanding():
			# this will always succeed due to the previous line
			self.pop_command_queue()

	def pop_command_queue (self):
		bool, stuff = self.command_fifo.pop()
		if bool:
			(command, consumer) = stuff
			self.push (command+'\r\n')
			self.push_consumer (consumer)

	# command == None implies that we expect unsolicited data
	# from the connection, (for example, a greeting string)
	def send_single_line_command (self, command, done_fun):
		con = simple_consumer (done_fun)
		if command:
			self.send_command (command, con)
		else:
			self.push_consumer (con)

	def send_multi_line_command (self,
								 command,
								 good_result_regex,
								 line_collect_fun,
								 done_fun):
		con = multi_line_consumer (
			done_fun,
			good_result_regex,
			line_collect_fun
			)
		if command:
			self.send_command (command, con)
		else:
			self.push_consumer (con)


