openwebrx-clone/rxws.py

"""
rxws: WebSocket methods implemented for OpenWebRX

	This file is part of OpenWebRX, 
	an open-source SDR receiver software with a web UI.
	Copyright (c) 2013-2015 by Andras Retzler <randras@sdr.hu>

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

    This program 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 Affero General Public License for more details.

    You should have received a copy of the GNU Affero General Public License
    along with this program.  If not, see <http://www.gnu.org/licenses/>.

"""

import base64
import sha
import select
import code

class WebSocketException(Exception):
	pass

def handshake(myself):
	my_client_id=myself.path[4:]
	my_headers=myself.headers.items()
	my_header_keys=map(lambda x:x[0],my_headers)
	h_key_exists=lambda x:my_header_keys.count(x)
	h_value=lambda x:my_headers[my_header_keys.index(x)][1]
	#print "The Lambdas(tm)"
	#print h_key_exists("upgrade")
	#print h_value("upgrade")
	#print h_key_exists("sec-websocket-key")
	if (not h_key_exists("upgrade")) or not (h_value("upgrade")=="websocket") or (not h_key_exists("sec-websocket-key")):
		raise WebSocketException
	ws_key=h_value("sec-websocket-key")
	ws_key_toreturn=base64.b64encode(sha.new(ws_key+"258EAFA5-E914-47DA-95CA-C5AB0DC85B11").digest())
	#A sample list of keys we get: [('origin', 'http://localhost:8073'), ('upgrade', 'websocket'), ('sec-websocket-extensions', 'x-webkit-deflate-frame'), ('sec-websocket-version', '13'), ('host', 'localhost:8073'), ('sec-websocket-key', 't9J1rgy4fc9fg2Hshhnkmg=='), ('connection', 'Upgrade'), ('pragma', 'no-cache'), ('cache-control', 'no-cache')]
	myself.wfile.write("HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Accept: "+ws_key_toreturn+"\r\nCQ-CQ-de: HA5KFU\r\n\r\n")

def get_header(size):
	#this does something similar: https://github.com/lemmingzshadow/php-websocket/blob/master/server/lib/WebSocket/Connection.php
	ws_first_byte=0b10000010 # FIN=1, OP=2
	if(size>125):
		ws_second_byte=126 # The following two bytes will indicate frame size
		extended_size=chr((size>>8)&0xff)+chr(size&0xff) #Okay, it uses reverse byte order (little-endian) compared to anything else sent on TCP
	else:
		ws_second_byte=size
		#256 bytes binary message in a single unmasked frame | 0x82 0x7E 0x0100 [256 bytes of binary data]
		extended_size=""
	return chr(ws_first_byte)+chr(ws_second_byte)+extended_size

def code_payload(data, masking_key=""):
	# both encode or decode
	if masking_key=="":
		key = (61, 84, 35, 6)
	else:
		key = [ord(i) for i in masking_key]
	encoded=""
	for i in range(0,len(data)):
		encoded+=chr(ord(data[i])^key[i%4])
	return encoded

def xxdg(data):
	output=""
	for i in range(0,len(data)/8):
		output+=xxd(data[i:i+8])
		if i%2: output+="\n"
		else: output+="  "
	return output
		

def xxd(data):
	#diagnostic purposes only
	output=""
	for d in data:
		output+=hex(ord(d))[2:].zfill(2)+" " 
	return output

#for R/W the WebSocket, use recv/send
#for reading the TCP socket, use readsock 
#for writing the TCP socket, use myself.wfile.write and flush

def readsock(myself,size,blocking):
	#http://thenestofheliopolis.blogspot.hu/2011/01/how-to-implement-non-blocking-two-way.html
	if blocking:
		return myself.rfile.read(size)
	else:
		poll = select.poll()
		poll.register(myself.rfile.fileno(), select.POLLIN or select.POLLPRI)
		fd = poll.poll(0) #timeout is 0
		if len(fd):
			f = fd[0]
			if f[1] > 0:
				return myself.rfile.read(size)
	return ""


def recv(myself, blocking=False, debug=False):
	bufsize=70000
	#myself.connection.setblocking(blocking) #umm... we cannot do that with rfile
	if debug: print "ws_recv begin"
	try:
		data=readsock(myself,6,blocking)
		#print "rxws.recv bytes:",xxd(data)	
	except:
		if debug: print "ws_recv error"	
		return ""
	if debug: print "ws_recv recved"
	if(len(data)==0): return ""
	fin=ord(data[0])&128!=0
	is_text_frame=ord(data[0])&15==1
	length=ord(data[1])&0x7f
	data+=readsock(myself,length,blocking)
	#print "rxws.recv length is ",length," (multiple packets together?) len(data) =",len(data)
	has_one_byte_length=length<125
	masked=ord(data[1])&0x80!=0
	#print "len=", length, len(data)-2
	#print "fin, is_text_frame, has_one_byte_length, masked = ", (fin, is_text_frame, has_one_byte_length, masked)
	#print xxd(data)
	if fin and is_text_frame and has_one_byte_length:
		if masked:
			return code_payload(data[6:], data[2:6])
		else:
			return data[2:]

#Useful links for ideas on WebSockets:
#  http://stackoverflow.com/questions/8125507/how-can-i-send-and-receive-websocket-messages-on-the-server-side
#  https://developer.mozilla.org/en-US/docs/WebSockets/Writing_WebSocket_server
#  http://tools.ietf.org/html/rfc6455#section-5.2	


def flush(myself): 
	myself.wfile.flush()
	#or the socket, not the rfile:
	#lR,lW,lX = select.select([],[myself.connection,],[],60)
	

def send(myself, data, begin_id="", debug=0):
	base_frame_size=35000 #could guess by MTU?
	debug=0
	#try:
	while True:
		counter=0
		from_end=len(data)-counter
		if from_end+len(begin_id)>base_frame_size:
			data_to_send=begin_id+data[counter:counter+base_frame_size-len(begin_id)]
			header=get_header(len(data_to_send))
			flush(myself)
			myself.wfile.write(header+data_to_send)
			flush(myself)
			if debug: print "rxws.send ==================== #1 if branch :: from={0} to={1} dlen={2} hlen={3}".format(counter,counter+base_frame_size-len(begin_id),len(data_to_send),len(header))
		else:
			data_to_send=begin_id+data[counter:]
			header=get_header(len(data_to_send))
			flush(myself)
			myself.wfile.write(header+data_to_send)
			flush(myself)
			if debug: print "rxws.send :: #2 else branch :: dlen={0} hlen={1}".format(len(data_to_send),len(header))
			#if debug: print "header:\n"+xxdg(header)+"\n\nws data:\n"+xxdg(data_to_send)
			break
		counter+=base_frame_size-len(begin_id)
	#except:
	#	pass