def _translate_single_frame_for_real_length(self, data, opcode,
oldpayloadlength, data_len,
old_real_packet_size):
"""
returns the real payload length and real packet size plus the modified data
"""
return_data = data.copy()
payloadlength = oldpayloadlength
real_packet_size = old_real_packet_size
if opcode == WebsocketCommon.OP_CODE_PING or opcode == WebsocketCommon.OP_CODE_PONG or opcode == WebsocketCommon.OP_CODE_CLOSING:
raise WebsocketInvalidFrameError(
"Some frames cannot have longer payload length", data)
if payloadlength == 126:
real_packet_size += 2
size_bytearray = bytearray()
size_bytearray.append(0)
size_bytearray.append(return_data[0])
size_bytearray.append(return_data[1])
return_data = return_data[2:]
payloadlength = WebsocketHelper.bytes_to_int(size_bytearray)
else:
real_packet_size += 8
size_bytearray = bytearray()
for i in range(8):
size_bytearray.append(return_data[i])
return_data = return_data[8:]
payloadlength = WebsocketHelper.bytes_to_int(size_bytearray)
return payloadlength, real_packet_size, return_data
def translate_handshake(self, data):
"""
return Handshake, could be ServerHandshake or ClientHandshake
:param data:
:return: None if the header is imcomplete,
"""
return_data = data
head_line, return_data = WebsocketHelper.read_string_line(return_data)
if head_line is None:
# do nothing and raise WebsocketIncompletePacketError to indicate that this is a imcomplete header
raise WebsocketIncompletePacketError(
'error translate_handshake: imcomplete header')
assert isinstance(head_line, str)
first_line_tokens = head_line.split(' ', 2)
if len(first_line_tokens) != 3:
print("error parsing first line, invalid handshake")
raise WebsocketDecodeError("invalid header field")
try:
handshake = None
if self.role == WebsocketCommon.ROLE_CLIENT:
handshake = self._translate_handshake_client(
first_line_tokens, head_line)
assert handshake is None or isinstance(handshake,
ServerHandshake)
elif self.role == WebsocketCommon.ROLE_SERVER:
handshake = self._translate_handshake_server(
first_line_tokens, head_line)
assert handshake is None or isinstance(handshake,
ClientHandshake)
else:
raise ValueError("error invalid role")
except WebsocketDecodeError as e:
print("error header line not match")
raise WebsocketDecodeError("header line not match")
key_value_line, return_data = WebsocketHelper.read_string_line(
return_data)
while key_value_line is not None and len(key_value_line) > 0:
# None means incomplete, len == 0 means end of line
pair = key_value_line.split(':', 1)
if len(pair) != 2:
raise WebsocketDecodeError("Invalid key value line")
# TODO currently not support one key appear in multiple lines
print(
"WS_DRAFT: putting key value while translating handshake {}, {}"
.format(pair[0], pair[1]))
handshake.put(pair[0].strip(), pair[1].strip())
key_value_line, return_data = WebsocketHelper.read_string_line(
return_data)
# TODO should we keep it intact ?
if key_value_line is None:
# keep the data intact
return None, data
return handshake, return_data
def do_write(self, key, event_type):
print("WS_SERVER: do_write invoked, ready to write")
ws_impl = key.data
assert isinstance(ws_impl, WebsocketImpl)
# flush message for this ws_impl to network
WebsocketHelper.batch(ws_impl)
# reset the interested event to READ
new_key = self.selector.modify(ws_impl.wrapped_socket,
selectors.EVENT_READ, ws_impl)
# update the ws_impl's key
ws_impl.set_key(new_key)
def create_text_frames(self, message):
"""
Currently only supports string type message
:param message:
:return:
"""
assert isinstance(message, str)
text_frame = TextFrame()
text_frame.payload = WebsocketHelper.str_to_utf8_bytearray(message)
return [text_frame]
def post_process_handshake_request_as_client(self,
handshake: ClientHandshake):
# TODO put header fields
handshake.put(WebsocketCommon.UPGRADE, "websocket")
handshake.put(WebsocketCommon.CONNECTION, WebsocketCommon.UPGRADE)
random_bytes = WebsocketHelper.random_bytearray(16)
# TODO will this be transformed to base64 string ?
random_base64_string = base64.b64encode(random_bytes).decode("ascii")
handshake.put(WebsocketCommon.SEC_WEB_SOCKET_KEY, random_base64_string)
handshake.put(WebsocketCommon.SEC_WEB_SOCKET_VERSION, "13")
# TODO skip extensions for now
# TODO skip protocols for now
return handshake
def do_read(self, key, event_type):
assert event_type & selectors.EVENT_READ
ws_impl = key.data
# retrieve the socket from ws_impl
sock = ws_impl.wrapped_socket
# read data from socket
data = WebsocketHelper.read(ws_impl, sock)
# print("WS_SERVER: received packet from client, string format {}".format(data.decode('ascii')))
assert data is not None
size = len(data)
# assert size > 0
if size == 0:
print("WS_SERVER: received empty packets from socket")
return
ws_impl.put_inqueue(data)
self.assign_executor(ws_impl)
def test_read_string_line(self):
line = 'dfsdfasdf \r\n asdf'
s, line = WebsocketHelper.read_string_line(bytearray(line.encode('ascii')))
self.assertTrue(s == bytearray('dfsdfasdf '.encode('ascii')))
self.assertTrue(line == bytearray(' asdf'.encode('ascii')))
def create_bytearray_for_frame(self, frame):
assert isinstance(frame, Frame)
payload = frame.payload
result = bytearray()
# TODO wait mask is a must support for Client
# TODO give mask support to Sam or Ruikai
# mask = self.role == WebsocketCommon.ROLE_CLIENT
mask = False
size_bytes = self._get_size_bytes(payload)
opt_code = frame.op
one_int8 = -128 if frame.fin else 0
one = opt_code | one_int8
if frame.rsv1:
one = one | self._get_rsv_byte(1)
if frame.rsv2:
one = one | self._get_rsv_byte(2)
if frame.rsv3:
one = one | self._get_rsv_byte(3)
# add the first bytes
one_bytes = WebsocketHelper.int8_to_bytes(one)
result.extend(one_bytes)
# calculate the second bytes
payload_length = len(frame.payload)
payload_length_bytearray = WebsocketHelper.int_to_bytes_for_size(
payload_length, size_bytes)
assert len(payload_length_bytearray) == size_bytes
if size_bytes == 1:
length_plus_mask = payload_length | self._get_mask_int(mask)
length_plus_mask_bytes = WebsocketHelper.int8_to_bytes(
length_plus_mask)
result.extend(length_plus_mask_bytes)
elif size_bytes == 2:
length_plus_mask = 126 | self._get_mask_int(mask)
length_plus_mask_bytes = WebsocketHelper.int8_to_bytes(
length_plus_mask)
result.extend(length_plus_mask_bytes)
result.extend(payload_length_bytearray)
elif size_bytes == 8:
length_plus_mask = 127 | self._get_mask_int(mask)
length_plus_mask_bytes = WebsocketHelper.int8_to_bytes(
length_plus_mask)
result.extend(length_plus_mask_bytes)
result.extend(payload_length_bytearray)
else:
print("Error impossible these many bytes")
if mask:
# TODO implemented by Ruikai and Sam
mask_key = bytearray()
mask_key.extend(
WebsocketHelper.int_to_bytes_for_size(
random.randint(-pow(2, 31),
pow(2, 31) - 1), 4))
result.extend(mask_key)
i = 0
for i in range(len(payload)):
result.extend(
WebsocketHelper.int_to_bytes_for_size(
(payload[i] ^ mask_key[(i % 4)]), 4))
else:
result.extend(payload)
return result
def _translate_single_frame(self, data):
return_data = data.copy()
assert isinstance(return_data, bytearray)
data_len = len(data)
real_packet_size = 2
if data_len < real_packet_size:
# Imcomplete data
return None, data
# TODO we don't put any size constraints on frames
b1 = return_data[0]
print("WS_DRAFT: first byte is {0:b} {1}".format(b1, b1))
fin = b1 >> 8 != 0
rsv1 = (b1 & 0x40) != 0
rsv2 = (b1 & 0x20) != 0
rsv3 = (b1 & 0x10) != 0
b2 = return_data[1]
print("WS_DRAFT: second byte is {0:b} {1}".format(b2, b2))
mask = (b2 & -128) != 0
# TODO will this actually work ?
payloadlength = b2 & ~128
# payloadlength = b2 & 127
opcode = self._to_opcode(b1 & 15)
print("WS_DRAFT: payload length is {0:b} {1}".format(
payloadlength, payloadlength))
print("WS_DRAFT: opcode is {0:b} {1}".format(b1 & 15, b1 & 15))
return_data = return_data[2:]
if payloadlength < 0 or payloadlength > 125:
# TODO the first two bits should not be passed in
print("after truncate the first 2 bytes : {}".format(
WebsocketHelper.bytearray_to_ascii_string(return_data)))
try:
payloadlength, real_packet_size, return_data = self._translate_single_frame_for_real_length(
return_data, opcode, payloadlength, data_len,
real_packet_size)
except WebsocketInvalidFrameError as e:
raise e
if data_len < real_packet_size:
# Imcomplete data
return None, data
real_packet_size += 4 if mask else 0
real_packet_size += payloadlength
payload = bytearray
if data_len < real_packet_size:
# Imcomplete data
return None, data
if mask:
# TODO give this to Sam or Ruikai
mask_key = return_data[:4]
return_data = return_data[4:]
payload = bytearray()
for i in range(payloadlength):
payload.extend(
WebsocketHelper.int_to_bytes_for_size(
return_data[i] ^ mask_key[i % 4], 4))
return_data = return_data[payloadlength:]
else:
payload = return_data[:payloadlength]
return_data = return_data[payloadlength:]
frame = Frame.frame_factory_get(opcode)
frame.fin = fin
frame.rsv1 = rsv1
frame.rsv2 = rsv2
frame.rsv3 = rsv3
# TODO skip extension and protocol validation
frame.payload = payload
# TODO which will it use ?
is_valid = frame.is_valid()
if is_valid:
print("after decode the bytes left are {}".format(
len(return_data)))
return frame, return_data
else:
raise WebsocketInvalidFrameError("Invalid frame", data)