#!/usr/bin/python2

# Copyright Shannon Baker 2008 (shannon@arc.net.au).
# Public Domain. Free for all uses.

import asyncore, asynchat, socket, time

HOST = '127.0.0.1'
PORT = 81

CHUNK_SIZE = 100
MAX_CHUNKS = 100

BUFFER_SIZE = CHUNK_SIZE

def status(msg):
	print "==", msg
	
class WebSocketChannel(asynchat.async_chat):

	def __init__(self, server, conn, addr):
		asynchat.async_chat.__init__(self, conn)
		# Server state
		self.conn = conn
		self.addr = addr
		self.ibuffer = []
		self.obuffer = ""
		self.ac_in_buffer_size = BUFFER_SIZE
		self.ac_out_buffer_size = BUFFER_SIZE
		self.set_terminator("\r\n\r\n")
		self.reading_headers = True
		self.handling = False
		self.is_websocket = False
		self.sent_packets = 0

	def push(self, data):
		print "=> " + data.rstrip()
		asynchat.async_chat.push(self, data)

	def collect_incoming_data(self, data):
		if data:
			print "<= " + data.replace("\n","\n<= ")
			if self.is_websocket:
				# Running in async mode
				self.ibuffer.append(data)
				self.handle_incoming_websocket_data(data)
			else:
				# Buffer headers
				self.ibuffer.append(data)
	
	def handle_incoming_websocket_data(self, data):
		# echo data sent by server
		#print "server>> " + data
		pass

	def send_websocket_data(self, data):
		status("Sending WebSocket data...")
		while self.sent_packets < MAX_CHUNKS:
			self.push(data)
			self.sent_packets += 1	
			#self.flush()
					
	def found_terminator(self):
		# gather the data pieces collected so far
		data = "".join(self.ibuffer)
		self.ibuffer = []
		
		if self.reading_headers:
			# reached end of client HTTP headers
			self.reading_headers = False
			self.parse_headers(data)
			print "End of headers reached"
			self.ibuffer = []
			if self.is_websocket:
				print "Valid WebSocket request received"
				self.push('HTTP/1.1 101 Switching Protocols\r\n')
				self.push('Upgrade: WebSocket/1.0\r\n')
				self.push('Connection: keep-alive\r\n\r\n')
				self.flush()
				
				# switch to async block mode (read/write raw data CHUNK_SIZE bytes at a time)
				self.set_terminator(CHUNK_SIZE)
				
				# Send data burst to other end
				self.send_websocket_data('s' * CHUNK_SIZE)
				
				# Close connection when all data has been sent
				#self.close_when_done()
			else:
				print "Failed WebSocket handshake"
			
	def parse_headers(self,headers):
		print "Parsing headers ..."
		header_lines = headers.split("\r\n")
		# check for valid head request
		request_line = header_lines[0]
		print request_line
		if request_line != "GET http://%s:%s/ HTTP/1.1"%(HOST,PORT):
			print "Unsupported request %s" % request_line
		#else:
		if True:
			for line in header_lines[1:]:
				print line
				if line[:5].lower() == 'host:':
					print "Host header found"
				if line[:8].lower() == 'upgrade:':
					upgrade_to = line[8:].lower().strip()
					if upgrade_to == 'websocket/1.0':
						print "Valid WebSocket/1.0 upgrade header found"
						self.is_websocket = True

	def flush(self):
		while self.producer_fifo or self.ac_out_buffer:
			self.initiate_send()

	def handle_connect(self):
		print "Client connected"
		
	def handle_expt(self):
		self.handle_error()

	def handle_error(self):
		print "Transfer error"
		self.close()
		
	def handle_close(self):
		print "Client disconnected"
		self.close()
		
class WebSocketServer(asyncore.dispatcher):

	def __init__(self, host, port):
		asyncore.dispatcher.__init__(self)
		s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
		s.setblocking(False)
		self.set_socket(s, asyncore.socket_map)
		self.set_reuse_addr()
		self.bind((host, port))
		self.listen(5)

	def handle_accept(self):
		conn, addr = self.accept()
		WebSocketChannel(self, conn, addr)

# Test

if __name__=="__main__":
	try:
		# launch the server on the specified port
		s = WebSocketServer(HOST,PORT)
		print "WebSocket server listening on port", PORT, "of host", HOST, "..."
		asyncore.loop(timeout=2,map=asyncore.socket_map)
	except KeyboardInterrupt:
		print "Ctrl+C pressed. Shutting down."