11246Szelenkov@nginx.comimport re 21130Szelenkov@nginx.comimport random 31130Szelenkov@nginx.comimport base64 41130Szelenkov@nginx.comimport struct 51130Szelenkov@nginx.comimport select 61130Szelenkov@nginx.comimport hashlib 71130Szelenkov@nginx.comimport itertools 81130Szelenkov@nginx.comfrom unit.applications.proto import TestApplicationProto 91130Szelenkov@nginx.com 101130Szelenkov@nginx.comGUID = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11" 111130Szelenkov@nginx.com 121130Szelenkov@nginx.com 131130Szelenkov@nginx.comclass TestApplicationWebsocket(TestApplicationProto): 141130Szelenkov@nginx.com 151130Szelenkov@nginx.com OP_CONT = 0x00 161130Szelenkov@nginx.com OP_TEXT = 0x01 171130Szelenkov@nginx.com OP_BINARY = 0x02 181130Szelenkov@nginx.com OP_CLOSE = 0x08 191130Szelenkov@nginx.com OP_PING = 0x09 201130Szelenkov@nginx.com OP_PONG = 0x0A 211130Szelenkov@nginx.com CLOSE_CODES = [1000, 1001, 1002, 1003, 1007, 1008, 1009, 1010, 1011] 221130Szelenkov@nginx.com 231130Szelenkov@nginx.com def __init__(self, preinit=False): 241130Szelenkov@nginx.com self.preinit = preinit 251130Szelenkov@nginx.com 261130Szelenkov@nginx.com def key(self): 271130Szelenkov@nginx.com raw_key = bytes(random.getrandbits(8) for _ in range(16)) 281130Szelenkov@nginx.com return base64.b64encode(raw_key).decode() 291130Szelenkov@nginx.com 301130Szelenkov@nginx.com def accept(self, key): 311130Szelenkov@nginx.com sha1 = hashlib.sha1((key + GUID).encode()).digest() 321130Szelenkov@nginx.com return base64.b64encode(sha1).decode() 331130Szelenkov@nginx.com 341262Szelenkov@nginx.com def upgrade(self, headers=None): 351262Szelenkov@nginx.com key = None 361262Szelenkov@nginx.com 371262Szelenkov@nginx.com if headers is None: 381262Szelenkov@nginx.com key = self.key() 391262Szelenkov@nginx.com headers = { 401130Szelenkov@nginx.com 'Host': 'localhost', 411130Szelenkov@nginx.com 'Upgrade': 'websocket', 421130Szelenkov@nginx.com 'Connection': 'Upgrade', 431130Szelenkov@nginx.com 'Sec-WebSocket-Key': key, 441130Szelenkov@nginx.com 'Sec-WebSocket-Protocol': 'chat', 451130Szelenkov@nginx.com 'Sec-WebSocket-Version': 13, 461262Szelenkov@nginx.com } 471262Szelenkov@nginx.com 481262Szelenkov@nginx.com _, sock = self.get( 491262Szelenkov@nginx.com headers=headers, 501246Szelenkov@nginx.com no_recv=True, 511130Szelenkov@nginx.com start=True, 521130Szelenkov@nginx.com ) 531130Szelenkov@nginx.com 541246Szelenkov@nginx.com resp = '' 551424Szelenkov@nginx.com while select.select([sock], [], [], 60)[0]: 561246Szelenkov@nginx.com resp += sock.recv(4096).decode() 571246Szelenkov@nginx.com 581246Szelenkov@nginx.com if ( 591246Szelenkov@nginx.com re.search('101 Switching Protocols', resp) 601246Szelenkov@nginx.com and resp[-4:] == '\r\n\r\n' 611246Szelenkov@nginx.com ): 621246Szelenkov@nginx.com resp = self._resp_to_dict(resp) 631246Szelenkov@nginx.com break 641246Szelenkov@nginx.com 651130Szelenkov@nginx.com return (resp, sock, key) 661130Szelenkov@nginx.com 671130Szelenkov@nginx.com def apply_mask(self, data, mask): 681130Szelenkov@nginx.com return bytes(b ^ m for b, m in zip(data, itertools.cycle(mask))) 691130Szelenkov@nginx.com 701150Szelenkov@nginx.com def serialize_close(self, code=1000, reason=''): 711130Szelenkov@nginx.com return struct.pack('!H', code) + reason.encode('utf-8') 721130Szelenkov@nginx.com 731424Szelenkov@nginx.com def frame_read(self, sock, read_timeout=60): 741130Szelenkov@nginx.com def recv_bytes(sock, bytes): 751130Szelenkov@nginx.com data = b'' 761130Szelenkov@nginx.com while select.select([sock], [], [], read_timeout)[0]: 771154Szelenkov@nginx.com data += sock.recv(bytes - len(data)) 781154Szelenkov@nginx.com 791154Szelenkov@nginx.com if len(data) == bytes: 801130Szelenkov@nginx.com break 811130Szelenkov@nginx.com 821130Szelenkov@nginx.com return data 831130Szelenkov@nginx.com 841130Szelenkov@nginx.com frame = {} 851130Szelenkov@nginx.com 861130Szelenkov@nginx.com head1, = struct.unpack('!B', recv_bytes(sock, 1)) 871130Szelenkov@nginx.com head2, = struct.unpack('!B', recv_bytes(sock, 1)) 881130Szelenkov@nginx.com 891130Szelenkov@nginx.com frame['fin'] = bool(head1 & 0b10000000) 901130Szelenkov@nginx.com frame['rsv1'] = bool(head1 & 0b01000000) 911130Szelenkov@nginx.com frame['rsv2'] = bool(head1 & 0b00100000) 921130Szelenkov@nginx.com frame['rsv3'] = bool(head1 & 0b00010000) 931130Szelenkov@nginx.com frame['opcode'] = head1 & 0b00001111 941130Szelenkov@nginx.com frame['mask'] = head2 & 0b10000000 951130Szelenkov@nginx.com 961130Szelenkov@nginx.com length = head2 & 0b01111111 971130Szelenkov@nginx.com if length == 126: 981130Szelenkov@nginx.com data = recv_bytes(sock, 2) 991130Szelenkov@nginx.com length, = struct.unpack('!H', data) 1001130Szelenkov@nginx.com elif length == 127: 1011130Szelenkov@nginx.com data = recv_bytes(sock, 8) 1021130Szelenkov@nginx.com length, = struct.unpack('!Q', data) 1031130Szelenkov@nginx.com 1041130Szelenkov@nginx.com if frame['mask']: 1051130Szelenkov@nginx.com mask_bits = recv_bytes(sock, 4) 1061130Szelenkov@nginx.com 1071156Szelenkov@nginx.com data = b'' 1081156Szelenkov@nginx.com 1091156Szelenkov@nginx.com if length != 0: 1101156Szelenkov@nginx.com data = recv_bytes(sock, length) 1111156Szelenkov@nginx.com 1121130Szelenkov@nginx.com if frame['mask']: 1131130Szelenkov@nginx.com data = self.apply_mask(data, mask_bits) 1141130Szelenkov@nginx.com 1151130Szelenkov@nginx.com if frame['opcode'] == self.OP_CLOSE: 1161130Szelenkov@nginx.com if length >= 2: 1171130Szelenkov@nginx.com code, = struct.unpack('!H', data[:2]) 1181130Szelenkov@nginx.com reason = data[2:].decode('utf-8') 1191130Szelenkov@nginx.com if not (code in self.CLOSE_CODES or 3000 <= code < 5000): 1201130Szelenkov@nginx.com self.fail('Invalid status code') 1211130Szelenkov@nginx.com frame['code'] = code 1221130Szelenkov@nginx.com frame['reason'] = reason 1231130Szelenkov@nginx.com elif length == 0: 1241130Szelenkov@nginx.com frame['code'] = 1005 1251130Szelenkov@nginx.com frame['reason'] = '' 1261130Szelenkov@nginx.com else: 1271130Szelenkov@nginx.com self.fail('Close frame too short') 1281130Szelenkov@nginx.com 1291130Szelenkov@nginx.com frame['data'] = data 1301130Szelenkov@nginx.com 1311130Szelenkov@nginx.com if frame['mask']: 1321130Szelenkov@nginx.com self.fail('Received frame with mask') 1331130Szelenkov@nginx.com 1341130Szelenkov@nginx.com return frame 1351130Szelenkov@nginx.com 1361130Szelenkov@nginx.com def frame_to_send( 1371130Szelenkov@nginx.com self, 1381130Szelenkov@nginx.com opcode, 1391130Szelenkov@nginx.com data, 1401130Szelenkov@nginx.com fin=True, 1411130Szelenkov@nginx.com length=None, 1421130Szelenkov@nginx.com rsv1=False, 1431130Szelenkov@nginx.com rsv2=False, 1441130Szelenkov@nginx.com rsv3=False, 1451130Szelenkov@nginx.com mask=True, 1461130Szelenkov@nginx.com ): 1471130Szelenkov@nginx.com frame = b'' 1481130Szelenkov@nginx.com 1491130Szelenkov@nginx.com if isinstance(data, str): 1501130Szelenkov@nginx.com data = data.encode('utf-8') 1511130Szelenkov@nginx.com 1521130Szelenkov@nginx.com head1 = ( 1531130Szelenkov@nginx.com (0b10000000 if fin else 0) 1541130Szelenkov@nginx.com | (0b01000000 if rsv1 else 0) 1551130Szelenkov@nginx.com | (0b00100000 if rsv2 else 0) 1561130Szelenkov@nginx.com | (0b00010000 if rsv3 else 0) 1571130Szelenkov@nginx.com | opcode 1581130Szelenkov@nginx.com ) 1591130Szelenkov@nginx.com 1601130Szelenkov@nginx.com head2 = 0b10000000 if mask else 0 1611130Szelenkov@nginx.com 1621130Szelenkov@nginx.com data_length = len(data) if length is None else length 1631130Szelenkov@nginx.com if data_length < 126: 1641130Szelenkov@nginx.com frame += struct.pack('!BB', head1, head2 | data_length) 1651130Szelenkov@nginx.com elif data_length < 65536: 1661130Szelenkov@nginx.com frame += struct.pack('!BBH', head1, head2 | 126, data_length) 1671130Szelenkov@nginx.com else: 1681130Szelenkov@nginx.com frame += struct.pack('!BBQ', head1, head2 | 127, data_length) 1691130Szelenkov@nginx.com 1701130Szelenkov@nginx.com if mask: 1711130Szelenkov@nginx.com mask_bits = struct.pack('!I', random.getrandbits(32)) 1721130Szelenkov@nginx.com frame += mask_bits 1731130Szelenkov@nginx.com 1741130Szelenkov@nginx.com if mask: 1751130Szelenkov@nginx.com frame += self.apply_mask(data, mask_bits) 1761130Szelenkov@nginx.com else: 1771130Szelenkov@nginx.com frame += data 1781130Szelenkov@nginx.com 1791130Szelenkov@nginx.com return frame 1801130Szelenkov@nginx.com 1811130Szelenkov@nginx.com def frame_write(self, sock, *args, **kwargs): 1821130Szelenkov@nginx.com chopsize = kwargs.pop('chopsize') if 'chopsize' in kwargs else None 1831130Szelenkov@nginx.com 1841130Szelenkov@nginx.com frame = self.frame_to_send(*args, **kwargs) 1851130Szelenkov@nginx.com 1861130Szelenkov@nginx.com if chopsize is None: 1871152Szelenkov@nginx.com try: 1881152Szelenkov@nginx.com sock.sendall(frame) 1891152Szelenkov@nginx.com except BrokenPipeError: 1901152Szelenkov@nginx.com pass 1911130Szelenkov@nginx.com 1921130Szelenkov@nginx.com else: 1931130Szelenkov@nginx.com pos = 0 1941130Szelenkov@nginx.com frame_len = len(frame) 1951150Szelenkov@nginx.com while pos < frame_len: 1961130Szelenkov@nginx.com end = min(pos + chopsize, frame_len) 1971151Szelenkov@nginx.com try: 1981151Szelenkov@nginx.com sock.sendall(frame[pos:end]) 1991151Szelenkov@nginx.com except BrokenPipeError: 2001151Szelenkov@nginx.com end = frame_len 2011130Szelenkov@nginx.com pos = end 2021130Szelenkov@nginx.com 2031130Szelenkov@nginx.com def message(self, sock, type, message, fragmention_size=None, **kwargs): 2041130Szelenkov@nginx.com message_len = len(message) 2051130Szelenkov@nginx.com 2061130Szelenkov@nginx.com if fragmention_size is None: 2071130Szelenkov@nginx.com fragmention_size = message_len 2081130Szelenkov@nginx.com 2091130Szelenkov@nginx.com if message_len <= fragmention_size: 2101130Szelenkov@nginx.com self.frame_write(sock, type, message, **kwargs) 2111130Szelenkov@nginx.com return 2121130Szelenkov@nginx.com 2131130Szelenkov@nginx.com pos = 0 2141130Szelenkov@nginx.com op_code = type 2151150Szelenkov@nginx.com while pos < message_len: 2161130Szelenkov@nginx.com end = min(pos + fragmention_size, message_len) 2171150Szelenkov@nginx.com fin = end == message_len 2181150Szelenkov@nginx.com self.frame_write( 2191150Szelenkov@nginx.com sock, op_code, message[pos:end], fin=fin, **kwargs 2201150Szelenkov@nginx.com ) 2211130Szelenkov@nginx.com op_code = self.OP_CONT 2221130Szelenkov@nginx.com pos = end 2231130Szelenkov@nginx.com 224*1433Szelenkov@nginx.com def message_read(self, sock, read_timeout=60): 2251130Szelenkov@nginx.com frame = self.frame_read(sock, read_timeout=read_timeout) 2261130Szelenkov@nginx.com 2271150Szelenkov@nginx.com while not frame['fin']: 2281130Szelenkov@nginx.com temp = self.frame_read(sock, read_timeout=read_timeout) 2291130Szelenkov@nginx.com frame['data'] += temp['data'] 2301130Szelenkov@nginx.com frame['fin'] = temp['fin'] 2311130Szelenkov@nginx.com 2321130Szelenkov@nginx.com return frame 233