#!/usr/bin/python
#
""" wtConnection.py
 
 author michael branton

 This script is free software; you can redistribute it and/or
 modify it under the terms of the GNU Lesser General Public
 License as published by the Free Software Foundation; either
 version 2.1 of the License, or (at your option) any later version.

 This library is distributed in the hope that it will be useful,
 but WITHOUT ANY WARRANTY; without even the implied warranty of
 MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 Lesser General Public License for more details.

 You should have received a copy of the GNU Lesser General Public
 License along with this script; if not, write to the Free Software
 Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
"""

import socket

class wtConnection:

	"""encapulates basic socket methods needed to talk to a
	wireTAP server (http://wiretap.stetson.edu)"""

        def __init__(self,*port):

		""" Class constructor. Initializes a socket connection
		using a given port. If no port is specified, one is 
		chosen at random. Defaults to RAW protocol"""

		self.sockexists = 'False'
		self.setHost('')
		self.setHostPort('')
		self.setProtocol('')
		self.block()
		if len(port) != 0:
			self.setClientPort(port[0])
		else:
			self.setClientPort(0)

	def create(self):

		"""private method """
		#create a local socket to communicate on
		if self.sockexists == 'True':
			self.s.close()

		if self.prot=='UDP':
        		self.s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
		elif self.prot=='TCP':
			self.s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
		else:
			self.s = socket.socket(socket.AF_INET, socket.SOCK_RAW)
			
		if self.cport!=0:
        		self.s.bind(("",self.cport))
		self.s.setblocking(self.blocking)
		self.sockexists='True'

	def setProtocol(self,prot):

		""" supported protocols are 'UDP', 'TCP' and 'RAW'. 'RAW' is 
		the default."""

		self.prot=prot
		#if we already have a connection, we
		#need to close it and create a new one
		if self.sockexists=='True':
			self.close()
			self.create()

	def setHost(self,host):

		""" set host to communicate with"""

		self.host=host

	def setHostPort(self,port):

		""" set port on host to use for connection """

		self.hport=port

	def setClientPort(self,port):

		""" set port on client to use for connection """
		self.cport=port
		#if we already have a connection, we
		#need to close it and create a new one
		if self.sockexists=='True':
			self.close()
			self.create()

	def send(self,message):

		"""send data out to  host"""

		if self.host!='':
        		self.s.sendto(message, (self.host,self.hport))

	def receive(self):

		"""get data back"""

        	result = self.s.recvfrom(1024)
        	return result

	def close(self):

		"""end it all by closing the socket connection. Returns
		status of attempted socket close"""

		result = self.s.close()
		return result

	def block(self):
		"""set socket to blocking"""

		self.blocking=1

	def nonblock(self):
		"""set socket to non-blocking"""

		self.blocking=0