#!/usr/bin/python2

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

# Experimental HTML5 WebSockets client for my proposed scheme. Not compatible with current draft specification.
# Includes basic proxy support. For more complex proxying use http://proxychains.sourceforge.net/
# May have issues with older versions of Python. v2.4 or above recommended.

import asyncore, asynchat, socket
import base64, time

# Server name and port
HOST = '127.0.0.1'
PORT = 81

# Set to proxy server name/ip or empty string for no proxy
PROXY = ''
PROXY_PORT = 8080

# Set proxy credentials or leave as empty strings
PROXY_USER = ''
PROXY_PASS = ''

# Size (in chars) and number of async data blocks to send 
CHUNK_SIZE = 100
MAX_CHUNKS = 100

BUFFER_SIZE = CHUNK_SIZE

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

	def __init__(self, host, port):
		asynchat.async_chat.__init__(self)
		# Client state
		self.host = host
		self.port = port
		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
		# open connection to server
		self.create_socket(socket.AF_INET, socket.SOCK_STREAM)
		if PROXY:
			self.connect( (PROXY, PROXY_PORT) )
		else:
			self.connect( (host, port) )
	
	def push(self, data):
		print "=> " + data.rstrip()
		asynchat.async_chat.push(self, data)
		
	def handle_connect(self):
		status("Connected")
		# Initial client headers
		self.headers = {
			'Upgrade': 'WebSocket/1.0',
			'Connection': 'keep-alive',
			'Host': '%s:%s' % (self.host, self.port)
		}
		# Switch to websocket
		#self.push('OPTIONS * HTTP/1.1')
		self.push('GET http://%s:%s/ HTTP/1.1\r\n'% (self.host, self.port))

		# Proxy headers
		if PROXY:
			self.push('Proxy-Connection: keep-alive\r\n')
			if PROXY_USER:
				userpass = base64.b64encode('%s:%s'%(PROXY_USER,PROXY_PASS))
				self.push('Proxy-Authorization: Basic %s\r\n' % userpass)
		
		# Server request headers
		for key, value in self.headers.items():
			self.push(key + ': ' + value + '\r\n')
		self.push('\r\n')
		#self.flush()

	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 server HTTP headers
			self.reading_headers = False
			self.parse_headers(data)
			status("End of headers reached")
			if self.is_websocket:
				status("Valid WebSocket response received")
				# 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('c' * CHUNK_SIZE)
				
				# Close connection when all data has been sent
				#self.close_when_done()
			else:
				status("Failed WebSocket handshake")
				self.handle_close()
			
	def parse_headers(self,headers):
		header_lines = headers.split("\r\n")
		# check for valid response
		response_line = header_lines[0]
		if not response_line[:12] in ('HTTP/1.0 101','HTTP/1.1 101'):
			status("Unsupported response %s" % response_line)
		else:
			for line in header_lines[1:]:
				if line[:8].lower() == 'upgrade:':
					upgrade_to = line[8:].lower().strip()
					if upgrade_to == 'websocket/1.0':
						status("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_expt(self):
		self.handle_error()
		
	def handle_error(self):
		status("Network error")
		self.handle_close()

	def handle_close(self):
		status("Connection closed")
		self.close()
# Test

if __name__=="__main__":
	try:
		# connect to server
		if PROXY:
			details = (HOST,PORT,PROXY,PROXY_PORT)
			status("WebSocket client connecting to %s:%s via %s:%s ..." % details)
		else:
			status("WebSocket client connecting to %s:%s ..." % (HOST,PORT))
		c = WebSocketClient(HOST,PORT)
		# poll
		asyncore.loop(timeout=2)
	except KeyboardInterrupt:
		status("Ctrl+C pressed. Shutting down.")