
'''
Default included TCP publishers.

@author: Michael Eddington
@version: $Id: tcp.py 2591 2011-11-11 11:27:57Z meddingt $
'''

#
# Copyright (c) Michael Eddington
#
# Permission is hereby granted, free of charge, to any person obtaining a copy 
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights 
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 
# copies of the Software, and to permit persons to whom the Software is 
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in	
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
#

# Authors:
#   Michael Eddington (mike@phed.org)

# $Id: tcp.py 2591 2011-11-11 11:27:57Z meddingt $

import socket, time, sys
from Peach.Engine.engine import Engine
from Peach.Engine.common import SoftException, PeachException
from Peach.publisher import Publisher
from Peach.publisher import Timeout
from Peach.publisher import PublisherSoftException
import Peach

def Debug(msg):
	if Peach.Engine.engine.Engine.debug:
		print msg

#class Timeout(SoftException):
#	def __init__(self, msg):
#		self.msg = msg
#	
#	def __str__(self):
#		return self.msg

class Tcp(Publisher):
	'''
	A simple TCP client publisher.
	'''
	
	
	def __init__(self, host, port, timeout = 0.25, throttle = 0):
		'''
		@type	host: string
		@param	host: Remote host
		@type	port: number
		@param	port: Remote port
		@type	timeout: number
		@param	timeout: How long to wait for reponse
		@type	throttle: number
		@param	throttle: How long to wait between connections
		'''
		Publisher.__init__(self)
		self._host = host
		
		try:
			self._port = int(port)
		except:
			raise PeachException("The Tcp publisher parameter for port was not a valid number.")
		
		try:
			self._timeout = float(timeout)
		except:
			raise PeachException("The Tcp publisher parameter for timeout was not a valid number.")
		
		try:
			self._throttle = float(throttle)
		except:
			raise PeachException("The Tcp publisher parameter for throttle was not a valid number.")
		
		self._socket = None
		
	def start(self):
		pass
	
	def stop(self):
		self.close()
	
	def connect(self):
		'''
		Create connection.
		'''
		self.close()
		
		if self._throttle > 0:
			time.sleep(self._throttle)
		
		# Try connecting many times
		# before we crash.
		for i in range(30):
			try:
				self._socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
				self._socket.connect((self._host,self._port))
				exception = None
				break
			
			except:
				self._socket = None
				exception = sys.exc_info()
			
			# Wait half sec and try again
			time.sleep(1)
		
		if self._socket == None:
			value = ""
			try:
				value = str(exception[1])
			except:
				pass
			
			raise PeachException("TCP onnection attempt failed: %s" % value)
		
		self.buff = ""
		self.pos = 0
	
	def close(self):
		'''
		Close connection if open.
		'''
		try:
			if self._socket != None:
				self._socket.close()
		finally:
			self._socket = None
		
		self.buff = ""
		self.pos = 0
	
	def send(self, data):
		'''
		Send data via sendall.
		
		@type	data: string
		@param	data: Data to send
		'''
		
		if Peach.Engine.engine.Engine.debug:
			print ">>>>>>>>>>>>>>>>"
			print "tcp.Tcp.send():"
			self.hexPrint(data)
		
		try:
			self._socket.sendall(data)
		except:
			if Peach.Engine.engine.Engine.debug:
				print "Tcp: Sendall failed: " + str(sys.exc_info()[1])
			raise PublisherSoftException("sendall failed: " + str(sys.exc_info()[1]))
	
	def _receiveBySize(self, size):
		'''
		This is now a buffered receiver.
		
		@rtype: string
		@return: received data.
		'''
		
		# Do we already a have it?
		if size+self.pos < len(self.buff):
			ret = self.buff[self.pos:self.pos+size]
			self.pos += size
			return ret
		
		# Only ask for the diff of what we don't already have
		diffSize = (self.pos+size) - len(self.buff)
		
		try:
			if Peach.Engine.engine.Engine.debug:
				print "Asking for %d, need %d, have %d" % (size, diffSize, len(self.buff)-self.pos)
			
			self._socket.settimeout(self._timeout)
			ret = self._socket.recv(diffSize)
			
			if not ret:
				# Socket was closed
				if Peach.Engine.engine.Engine.debug:
					print "Socket is closed"
				
				raise PublisherSoftException("Socket is closed")
			
			if Peach.Engine.engine.Engine.debug:
				print "<<<<<<<<<<<<<<<<<"
				print "tcp.Tcp.receive():"
				self.hexPrint(ret)
			
			self.buff += ret
			
		except socket.error, e:
			if str(e).find('The socket operation could not complete without blocking') != -1:
				if Peach.Engine.engine.Engine.debug:
					print "timed out waiting for data"
				
				raise Timeout("Timed out waiting for data [%d:%d:%d:%d]" % (len(self.buff), (size+self.pos), size, diffSize))
			
			elif str(e).find('An existing connection was forcibly') != -1:
				if Peach.Engine.engine.Engine.debug:
					print "Socket was closed!"
				
				raise PublisherSoftException("Socket is closed")
			
			else:
				if Peach.Engine.engine.Engine.debug:
					print "recv failed: " + str(sys.exc_info()[1])
				
				raise PublisherSoftException("recv failed: " + str(sys.exc_info()[1]))
		
		finally:
			self._socket.settimeout(None)
		
		ret = self.buff[self.pos:]
		self.pos = len(self.buff)
		return ret
	
	def _receiveByAvailable(self):
		'''
		Receive as much as possible prior to timeout.
		
		@rtype: string
		@return: received data.
		'''
		self._socket.settimeout(self._timeout)
		
		try:
			ret = self._socket.recv(4096)
			
			if not ret:
				raise PublisherSoftException("Socket is closed")
			
			if Peach.Engine.engine.Engine.debug:
				print "<<<<<<<<<<<<<<<<<"
				print "tcp.Tcp.receive():"
				self.hexPrint(ret)
			
			self.buff += ret
			
		except socket.error, e:
			if str(e).find('The socket operation could not complete without blocking') == -1:
				pass
			
			else:
				raise PublisherSoftException("recv failed: " + str(sys.exc_info()[1]))
		
		self._socket.settimeout(None)
		
		ret = self.buff[self.pos:]
		self.pos = len(self.buff)
		return ret