11130Szelenkov@nginx.comimport random 21130Szelenkov@nginx.comimport base64 31130Szelenkov@nginx.comimport struct 41130Szelenkov@nginx.comimport select 51130Szelenkov@nginx.comimport hashlib 61130Szelenkov@nginx.comimport itertools 71130Szelenkov@nginx.comfrom unit.applications.proto import TestApplicationProto 81130Szelenkov@nginx.com 91130Szelenkov@nginx.comGUID = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11" 101130Szelenkov@nginx.com 111130Szelenkov@nginx.com 121130Szelenkov@nginx.comclass TestApplicationWebsocket(TestApplicationProto): 131130Szelenkov@nginx.com 141130Szelenkov@nginx.com OP_CONT = 0x00 151130Szelenkov@nginx.com OP_TEXT = 0x01 161130Szelenkov@nginx.com OP_BINARY = 0x02 171130Szelenkov@nginx.com OP_CLOSE = 0x08 181130Szelenkov@nginx.com OP_PING = 0x09 191130Szelenkov@nginx.com OP_PONG = 0x0A 201130Szelenkov@nginx.com CLOSE_CODES = [1000, 1001, 1002, 1003, 1007, 1008, 1009, 1010, 1011] 211130Szelenkov@nginx.com 221130Szelenkov@nginx.com def __init__(self, preinit=False): 231130Szelenkov@nginx.com self.preinit = preinit 241130Szelenkov@nginx.com 251130Szelenkov@nginx.com def key(self): 261130Szelenkov@nginx.com raw_key = bytes(random.getrandbits(8) for _ in range(16)) 271130Szelenkov@nginx.com return base64.b64encode(raw_key).decode() 281130Szelenkov@nginx.com 291130Szelenkov@nginx.com def accept(self, key): 301130Szelenkov@nginx.com sha1 = hashlib.sha1((key + GUID).encode()).digest() 311130Szelenkov@nginx.com return base64.b64encode(sha1).decode() 321130Szelenkov@nginx.com 331130Szelenkov@nginx.com def upgrade(self): 341130Szelenkov@nginx.com key = self.key() 351130Szelenkov@nginx.com 361130Szelenkov@nginx.com if self.preinit: 371130Szelenkov@nginx.com self.get() 381130Szelenkov@nginx.com 391130Szelenkov@nginx.com resp, sock = self.get( 401130Szelenkov@nginx.com headers={ 411130Szelenkov@nginx.com 'Host': 'localhost', 421130Szelenkov@nginx.com 'Upgrade': 'websocket', 431130Szelenkov@nginx.com 'Connection': 'Upgrade', 441130Szelenkov@nginx.com 'Sec-WebSocket-Key': key, 451130Szelenkov@nginx.com 'Sec-WebSocket-Protocol': 'chat', 461130Szelenkov@nginx.com 'Sec-WebSocket-Version': 13, 471130Szelenkov@nginx.com }, 481130Szelenkov@nginx.com read_timeout=1, 491130Szelenkov@nginx.com start=True, 501130Szelenkov@nginx.com ) 511130Szelenkov@nginx.com 521130Szelenkov@nginx.com return (resp, sock, key) 531130Szelenkov@nginx.com 541130Szelenkov@nginx.com def apply_mask(self, data, mask): 551130Szelenkov@nginx.com return bytes(b ^ m for b, m in zip(data, itertools.cycle(mask))) 561130Szelenkov@nginx.com 571150Szelenkov@nginx.com def serialize_close(self, code=1000, reason=''): 581130Szelenkov@nginx.com return struct.pack('!H', code) + reason.encode('utf-8') 591130Szelenkov@nginx.com 601138Szelenkov@nginx.com def frame_read(self, sock, read_timeout=10): 611130Szelenkov@nginx.com def recv_bytes(sock, bytes): 621130Szelenkov@nginx.com data = b'' 631130Szelenkov@nginx.com while select.select([sock], [], [], read_timeout)[0]: 641130Szelenkov@nginx.com try: 651130Szelenkov@nginx.com if bytes < 65536: 661130Szelenkov@nginx.com data = sock.recv(bytes) 671130Szelenkov@nginx.com else: 681130Szelenkov@nginx.com data = self.recvall( 691150Szelenkov@nginx.com sock, read_timeout=read_timeout, buff_size=bytes 701130Szelenkov@nginx.com ) 711130Szelenkov@nginx.com break 721130Szelenkov@nginx.com except: 731130Szelenkov@nginx.com break 741130Szelenkov@nginx.com 751130Szelenkov@nginx.com return data 761130Szelenkov@nginx.com 771130Szelenkov@nginx.com frame = {} 781130Szelenkov@nginx.com 791130Szelenkov@nginx.com head1, = struct.unpack('!B', recv_bytes(sock, 1)) 801130Szelenkov@nginx.com head2, = struct.unpack('!B', recv_bytes(sock, 1)) 811130Szelenkov@nginx.com 821130Szelenkov@nginx.com frame['fin'] = bool(head1 & 0b10000000) 831130Szelenkov@nginx.com frame['rsv1'] = bool(head1 & 0b01000000) 841130Szelenkov@nginx.com frame['rsv2'] = bool(head1 & 0b00100000) 851130Szelenkov@nginx.com frame['rsv3'] = bool(head1 & 0b00010000) 861130Szelenkov@nginx.com frame['opcode'] = head1 & 0b00001111 871130Szelenkov@nginx.com frame['mask'] = head2 & 0b10000000 881130Szelenkov@nginx.com 891130Szelenkov@nginx.com length = head2 & 0b01111111 901130Szelenkov@nginx.com if length == 126: 911130Szelenkov@nginx.com data = recv_bytes(sock, 2) 921130Szelenkov@nginx.com length, = struct.unpack('!H', data) 931130Szelenkov@nginx.com elif length == 127: 941130Szelenkov@nginx.com data = recv_bytes(sock, 8) 951130Szelenkov@nginx.com length, = struct.unpack('!Q', data) 961130Szelenkov@nginx.com 971130Szelenkov@nginx.com if frame['mask']: 981130Szelenkov@nginx.com mask_bits = recv_bytes(sock, 4) 991130Szelenkov@nginx.com 1001130Szelenkov@nginx.com data = recv_bytes(sock, length) 1011130Szelenkov@nginx.com if frame['mask']: 1021130Szelenkov@nginx.com data = self.apply_mask(data, mask_bits) 1031130Szelenkov@nginx.com 1041130Szelenkov@nginx.com if frame['opcode'] == self.OP_CLOSE: 1051130Szelenkov@nginx.com if length >= 2: 1061130Szelenkov@nginx.com code, = struct.unpack('!H', data[:2]) 1071130Szelenkov@nginx.com reason = data[2:].decode('utf-8') 1081130Szelenkov@nginx.com if not (code in self.CLOSE_CODES or 3000 <= code < 5000): 1091130Szelenkov@nginx.com self.fail('Invalid status code') 1101130Szelenkov@nginx.com frame['code'] = code 1111130Szelenkov@nginx.com frame['reason'] = reason 1121130Szelenkov@nginx.com elif length == 0: 1131130Szelenkov@nginx.com frame['code'] = 1005 1141130Szelenkov@nginx.com frame['reason'] = '' 1151130Szelenkov@nginx.com else: 1161130Szelenkov@nginx.com self.fail('Close frame too short') 1171130Szelenkov@nginx.com 1181130Szelenkov@nginx.com frame['data'] = data 1191130Szelenkov@nginx.com 1201130Szelenkov@nginx.com if frame['mask']: 1211130Szelenkov@nginx.com self.fail('Received frame with mask') 1221130Szelenkov@nginx.com 1231130Szelenkov@nginx.com return frame 1241130Szelenkov@nginx.com 1251130Szelenkov@nginx.com def frame_to_send( 1261130Szelenkov@nginx.com self, 1271130Szelenkov@nginx.com opcode, 1281130Szelenkov@nginx.com data, 1291130Szelenkov@nginx.com fin=True, 1301130Szelenkov@nginx.com length=None, 1311130Szelenkov@nginx.com rsv1=False, 1321130Szelenkov@nginx.com rsv2=False, 1331130Szelenkov@nginx.com rsv3=False, 1341130Szelenkov@nginx.com mask=True, 1351130Szelenkov@nginx.com ): 1361130Szelenkov@nginx.com frame = b'' 1371130Szelenkov@nginx.com 1381130Szelenkov@nginx.com if isinstance(data, str): 1391130Szelenkov@nginx.com data = data.encode('utf-8') 1401130Szelenkov@nginx.com 1411130Szelenkov@nginx.com head1 = ( 1421130Szelenkov@nginx.com (0b10000000 if fin else 0) 1431130Szelenkov@nginx.com | (0b01000000 if rsv1 else 0) 1441130Szelenkov@nginx.com | (0b00100000 if rsv2 else 0) 1451130Szelenkov@nginx.com | (0b00010000 if rsv3 else 0) 1461130Szelenkov@nginx.com | opcode 1471130Szelenkov@nginx.com ) 1481130Szelenkov@nginx.com 1491130Szelenkov@nginx.com head2 = 0b10000000 if mask else 0 1501130Szelenkov@nginx.com 1511130Szelenkov@nginx.com data_length = len(data) if length is None else length 1521130Szelenkov@nginx.com if data_length < 126: 1531130Szelenkov@nginx.com frame += struct.pack('!BB', head1, head2 | data_length) 1541130Szelenkov@nginx.com elif data_length < 65536: 1551130Szelenkov@nginx.com frame += struct.pack('!BBH', head1, head2 | 126, data_length) 1561130Szelenkov@nginx.com else: 1571130Szelenkov@nginx.com frame += struct.pack('!BBQ', head1, head2 | 127, data_length) 1581130Szelenkov@nginx.com 1591130Szelenkov@nginx.com if mask: 1601130Szelenkov@nginx.com mask_bits = struct.pack('!I', random.getrandbits(32)) 1611130Szelenkov@nginx.com frame += mask_bits 1621130Szelenkov@nginx.com 1631130Szelenkov@nginx.com if mask: 1641130Szelenkov@nginx.com frame += self.apply_mask(data, mask_bits) 1651130Szelenkov@nginx.com else: 1661130Szelenkov@nginx.com frame += data 1671130Szelenkov@nginx.com 1681130Szelenkov@nginx.com return frame 1691130Szelenkov@nginx.com 1701130Szelenkov@nginx.com def frame_write(self, sock, *args, **kwargs): 1711130Szelenkov@nginx.com chopsize = kwargs.pop('chopsize') if 'chopsize' in kwargs else None 1721130Szelenkov@nginx.com 1731130Szelenkov@nginx.com frame = self.frame_to_send(*args, **kwargs) 1741130Szelenkov@nginx.com 1751130Szelenkov@nginx.com if chopsize is None: 1761130Szelenkov@nginx.com sock.sendall(frame) 1771130Szelenkov@nginx.com 1781130Szelenkov@nginx.com else: 1791130Szelenkov@nginx.com pos = 0 1801130Szelenkov@nginx.com frame_len = len(frame) 1811150Szelenkov@nginx.com while pos < frame_len: 1821130Szelenkov@nginx.com end = min(pos + chopsize, frame_len) 183*1151Szelenkov@nginx.com try: 184*1151Szelenkov@nginx.com sock.sendall(frame[pos:end]) 185*1151Szelenkov@nginx.com except BrokenPipeError: 186*1151Szelenkov@nginx.com end = frame_len 1871130Szelenkov@nginx.com pos = end 1881130Szelenkov@nginx.com 1891130Szelenkov@nginx.com def message(self, sock, type, message, fragmention_size=None, **kwargs): 1901130Szelenkov@nginx.com message_len = len(message) 1911130Szelenkov@nginx.com 1921130Szelenkov@nginx.com if fragmention_size is None: 1931130Szelenkov@nginx.com fragmention_size = message_len 1941130Szelenkov@nginx.com 1951130Szelenkov@nginx.com if message_len <= fragmention_size: 1961130Szelenkov@nginx.com self.frame_write(sock, type, message, **kwargs) 1971130Szelenkov@nginx.com return 1981130Szelenkov@nginx.com 1991130Szelenkov@nginx.com pos = 0 2001130Szelenkov@nginx.com op_code = type 2011150Szelenkov@nginx.com while pos < message_len: 2021130Szelenkov@nginx.com end = min(pos + fragmention_size, message_len) 2031150Szelenkov@nginx.com fin = end == message_len 2041150Szelenkov@nginx.com self.frame_write( 2051150Szelenkov@nginx.com sock, op_code, message[pos:end], fin=fin, **kwargs 2061150Szelenkov@nginx.com ) 2071130Szelenkov@nginx.com op_code = self.OP_CONT 2081130Szelenkov@nginx.com pos = end 2091130Szelenkov@nginx.com 2101138Szelenkov@nginx.com def message_read(self, sock, read_timeout=10): 2111130Szelenkov@nginx.com frame = self.frame_read(sock, read_timeout=read_timeout) 2121130Szelenkov@nginx.com 2131150Szelenkov@nginx.com while not frame['fin']: 2141130Szelenkov@nginx.com temp = self.frame_read(sock, read_timeout=read_timeout) 2151130Szelenkov@nginx.com frame['data'] += temp['data'] 2161130Szelenkov@nginx.com frame['fin'] = temp['fin'] 2171130Szelenkov@nginx.com 2181130Szelenkov@nginx.com return frame 219