Will check if message is valid UTF8

Author: Pithikos

websocket.py

@@ -6,6 +6,7 @@
 from base64 import b64encode
 from hashlib import sha1
 
+
 if sys.version_info[0] < 3 :
 	from SocketServer import ThreadingMixIn, TCPServer, StreamRequestHandler
 else:
@@ -102,6 +103,7 @@ def handler_to_client(self, handler):
 				return client
 
 
+
 class WebSocketHandler(StreamRequestHandler):
 
 	def __init__(self, socket, addr, server):
@@ -122,13 +124,14 @@ def handle(self):
 				self.read_next_message()
 
 	def read_next_message(self):
+		
 		b1 = self.rfile.read(1)
 		b2 = self.rfile.read(1)
 		FIN    = ord(b1) & 0b10000000
 		OPCODE = ord(b1) & 0b00001111
 		MASKED = ord(b2) & 0b10000000
 		LENGTH = ord(b2) & 0b01111111
-		
+
 		if not b1:
 			print("Client closed connection.")
 			self.keep_alive = 0
@@ -191,35 +194,47 @@ def send_text(self, message):
 		Fragmented(=continuation) messages are not being used since their usage
 		is needed in very limited cases - when we don't know the payload length.
 		'''
-		
+
 		#   0 = '0x00' = 0b00000000
 		# 125 = '0x7d' = 0b01111101
 		# 126 = '0x7e' = 0b01111110
 		# 127 = '0x7f' = 0b01111111
 		# 128 = '0x80' = 0b10000000
 		# 129 = '0x81' = 0b10000001
-		length = len(message)
+
+		# assure message is of acceptable format and encode to UTF-8 if needed
+		encoded_message = None
+		if isinstance(message, bytes):
+			message = try_decode_UTF8(message) # this is slower but assures we have UTF-8
+			if not message:
+				print("Can\'t send message, message is not valid UTF-8")
+				return False
+		else:
+			print('Can\'t send message, message has to be a string or bytes. Given type is %s' % type(message))
+			return False
+		encoded_message = encode_to_UTF8(message)
+		length = len(encoded_message)
 
 		# normal payload
 		if length <= 125:
 			#print("sending single frame of size %s", length)
 			self.request.send(b'\x81')
 			self.request.send(chr(length).encode())
-			self.request.send(unicode(message, 'UTF-8'))
+			self.request.send(encoded_message)
 
 		# extended payload
 		elif length >= 126 and length <= 65535:
 			#print("sending extended frame of size %s", length)
 			self.request.send(b'\x81\x7e')
 			self.request.send(struct.pack(">H", length)) # MUST be 16bits
-			self.request.send(unicode(message, 'UTF-8'))
+			self.request.send(encoded_message)
 
 		# huge extended payload
 		elif length < 18446744073709551616:
 			#print("sending extended frame of size %s", length)
 			self.request.send(b'\x81\x7f')
 			self.request.send(struct.pack(">Q", length)) # MUST be 64bits
-			self.request.send(unicode(message, 'UTF-8'))
+			self.request.send(encoded_message)
 
 	def handshake(self):
 		message = self.request.recv(1024).decode().strip()
@@ -238,15 +253,15 @@ def handshake(self):
 		self.handshake_done = self.request.send(response.encode())
 		self.valid_client = True
 		self.server._new_client_(self)
-		
+
 	def make_handshake_response(self, key):
 		return \
 		  'HTTP/1.1 101 Switching Protocols\r\n'\
 		  'Upgrade: websocket\r\n'              \
 		  'Connection: Upgrade\r\n'             \
 		  'Sec-WebSocket-Accept: %s\r\n'        \
 		  '\r\n' % self.calculate_response_key(key)
-		
+
 	def calculate_response_key(self, key):
 		GUID = '258EAFA5-E914-47DA-95CA-C5AB0DC85B11'
 		hash = sha1(key.encode() + GUID.encode())
@@ -257,6 +272,29 @@ def finish(self):
 		self.server._client_left_(self)
 
 
+
+def encode_to_UTF8(data):
+	try:
+		return data.encode('UTF-8')
+	except UnicodeEncodeError as e:
+		print("Could not encode data to UTF-8 -- %s" % e)
+		return False
+	except Exception as e:
+		raise(e)
+		return False
+
+
+
+def try_decode_UTF8(data):
+	try:
+		return data.decode('utf-8')
+	except UnicodeDecodeError:
+		return False
+	except Exception as e:
+		raise(e)
+		
+
+
 # This is only for testing purposes
 class DummyWebsocketHandler(WebSocketHandler):
     def __init__(self, *_):