11130Szelenkov@nginx.comimport base64 21130Szelenkov@nginx.comimport hashlib 31130Szelenkov@nginx.comimport itertools 41477Szelenkov@nginx.comimport random 51477Szelenkov@nginx.comimport select 61477Szelenkov@nginx.comimport struct 71477Szelenkov@nginx.com 81635Szelenkov@nginx.comimport pytest 9*2616Szelenkov@nginx.com 102491Szelenkov@nginx.comfrom unit.applications.proto import ApplicationProto 111130Szelenkov@nginx.com 121130Szelenkov@nginx.comGUID = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11" 131130Szelenkov@nginx.com 141130Szelenkov@nginx.com 152491Szelenkov@nginx.comclass ApplicationWebsocket(ApplicationProto): 161130Szelenkov@nginx.com 171130Szelenkov@nginx.com OP_CONT = 0x00 181130Szelenkov@nginx.com OP_TEXT = 0x01 191130Szelenkov@nginx.com OP_BINARY = 0x02 201130Szelenkov@nginx.com OP_CLOSE = 0x08 211130Szelenkov@nginx.com OP_PING = 0x09 221130Szelenkov@nginx.com OP_PONG = 0x0A 231130Szelenkov@nginx.com CLOSE_CODES = [1000, 1001, 1002, 1003, 1007, 1008, 1009, 1010, 1011] 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 331262Szelenkov@nginx.com def upgrade(self, headers=None): 341262Szelenkov@nginx.com key = None 351262Szelenkov@nginx.com 361262Szelenkov@nginx.com if headers is None: 371262Szelenkov@nginx.com key = self.key() 381262Szelenkov@nginx.com headers = { 391130Szelenkov@nginx.com 'Host': 'localhost', 401130Szelenkov@nginx.com 'Upgrade': 'websocket', 411130Szelenkov@nginx.com 'Connection': 'Upgrade', 421130Szelenkov@nginx.com 'Sec-WebSocket-Key': key, 431625Smax.romanov@nginx.com 'Sec-WebSocket-Protocol': 'chat, phone, video', 441130Szelenkov@nginx.com 'Sec-WebSocket-Version': 13, 451262Szelenkov@nginx.com } 461262Szelenkov@nginx.com 472212Szelenkov@nginx.com sock = self.get( 482073Szelenkov@nginx.com headers=headers, 492073Szelenkov@nginx.com no_recv=True, 502073Szelenkov@nginx.com ) 511130Szelenkov@nginx.com 521246Szelenkov@nginx.com resp = '' 531444Szelenkov@nginx.com while True: 541444Szelenkov@nginx.com rlist = select.select([sock], [], [], 60)[0] 551444Szelenkov@nginx.com if not rlist: 562330Szelenkov@nginx.com pytest.fail("Can't read response from server.") 571444Szelenkov@nginx.com 581246Szelenkov@nginx.com resp += sock.recv(4096).decode() 591246Szelenkov@nginx.com 601848Szelenkov@nginx.com if resp.startswith('HTTP/') and '\r\n\r\n' in resp: 611246Szelenkov@nginx.com resp = self._resp_to_dict(resp) 621246Szelenkov@nginx.com break 631246Szelenkov@nginx.com 641130Szelenkov@nginx.com return (resp, sock, key) 651130Szelenkov@nginx.com 661130Szelenkov@nginx.com def apply_mask(self, data, mask): 671130Szelenkov@nginx.com return bytes(b ^ m for b, m in zip(data, itertools.cycle(mask))) 681130Szelenkov@nginx.com 691150Szelenkov@nginx.com def serialize_close(self, code=1000, reason=''): 701130Szelenkov@nginx.com return struct.pack('!H', code) + reason.encode('utf-8') 711130Szelenkov@nginx.com 721424Szelenkov@nginx.com def frame_read(self, sock, read_timeout=60): 73*2616Szelenkov@nginx.com def recv_bytes(sock, bytes_len): 741130Szelenkov@nginx.com data = b'' 751444Szelenkov@nginx.com while True: 761444Szelenkov@nginx.com rlist = select.select([sock], [], [], read_timeout)[0] 771444Szelenkov@nginx.com if not rlist: 781444Szelenkov@nginx.com # For all current cases if the "read_timeout" was changed 791444Szelenkov@nginx.com # than test do not expect to get a response from server. 801444Szelenkov@nginx.com if read_timeout == 60: 812330Szelenkov@nginx.com pytest.fail("Can't read response from server.") 821444Szelenkov@nginx.com break 831444Szelenkov@nginx.com 84*2616Szelenkov@nginx.com data += sock.recv(bytes_len - len(data)) 851154Szelenkov@nginx.com 86*2616Szelenkov@nginx.com if len(data) == bytes_len: 871130Szelenkov@nginx.com break 881130Szelenkov@nginx.com 891130Szelenkov@nginx.com return data 901130Szelenkov@nginx.com 911130Szelenkov@nginx.com frame = {} 921130Szelenkov@nginx.com 931848Szelenkov@nginx.com (head1,) = struct.unpack('!B', recv_bytes(sock, 1)) 941848Szelenkov@nginx.com (head2,) = struct.unpack('!B', recv_bytes(sock, 1)) 951130Szelenkov@nginx.com 961130Szelenkov@nginx.com frame['fin'] = bool(head1 & 0b10000000) 971130Szelenkov@nginx.com frame['rsv1'] = bool(head1 & 0b01000000) 981130Szelenkov@nginx.com frame['rsv2'] = bool(head1 & 0b00100000) 991130Szelenkov@nginx.com frame['rsv3'] = bool(head1 & 0b00010000) 1001130Szelenkov@nginx.com frame['opcode'] = head1 & 0b00001111 1011130Szelenkov@nginx.com frame['mask'] = head2 & 0b10000000 1021130Szelenkov@nginx.com 1031130Szelenkov@nginx.com length = head2 & 0b01111111 1041130Szelenkov@nginx.com if length == 126: 1051130Szelenkov@nginx.com data = recv_bytes(sock, 2) 1061848Szelenkov@nginx.com (length,) = struct.unpack('!H', data) 1071130Szelenkov@nginx.com elif length == 127: 1081130Szelenkov@nginx.com data = recv_bytes(sock, 8) 1091848Szelenkov@nginx.com (length,) = struct.unpack('!Q', data) 1101130Szelenkov@nginx.com 1111130Szelenkov@nginx.com if frame['mask']: 1121130Szelenkov@nginx.com mask_bits = recv_bytes(sock, 4) 1131130Szelenkov@nginx.com 1141156Szelenkov@nginx.com data = b'' 1151156Szelenkov@nginx.com 1161156Szelenkov@nginx.com if length != 0: 1171156Szelenkov@nginx.com data = recv_bytes(sock, length) 1181156Szelenkov@nginx.com 1191130Szelenkov@nginx.com if frame['mask']: 1201130Szelenkov@nginx.com data = self.apply_mask(data, mask_bits) 1211130Szelenkov@nginx.com 1221130Szelenkov@nginx.com if frame['opcode'] == self.OP_CLOSE: 1231130Szelenkov@nginx.com if length >= 2: 1241848Szelenkov@nginx.com (code,) = struct.unpack('!H', data[:2]) 1251130Szelenkov@nginx.com reason = data[2:].decode('utf-8') 1261130Szelenkov@nginx.com if not (code in self.CLOSE_CODES or 3000 <= code < 5000): 1271596Szelenkov@nginx.com pytest.fail('Invalid status code') 1281130Szelenkov@nginx.com frame['code'] = code 1291130Szelenkov@nginx.com frame['reason'] = reason 1301130Szelenkov@nginx.com elif length == 0: 1311130Szelenkov@nginx.com frame['code'] = 1005 1321130Szelenkov@nginx.com frame['reason'] = '' 1331130Szelenkov@nginx.com else: 1341596Szelenkov@nginx.com pytest.fail('Close frame too short') 1351130Szelenkov@nginx.com 1361130Szelenkov@nginx.com frame['data'] = data 1371130Szelenkov@nginx.com 1381130Szelenkov@nginx.com if frame['mask']: 1391596Szelenkov@nginx.com pytest.fail('Received frame with mask') 1401130Szelenkov@nginx.com 1411130Szelenkov@nginx.com return frame 1421130Szelenkov@nginx.com 1431130Szelenkov@nginx.com def frame_to_send( 1441130Szelenkov@nginx.com self, 1451130Szelenkov@nginx.com opcode, 1461130Szelenkov@nginx.com data, 1471130Szelenkov@nginx.com fin=True, 1481130Szelenkov@nginx.com length=None, 1491130Szelenkov@nginx.com rsv1=False, 1501130Szelenkov@nginx.com rsv2=False, 1511130Szelenkov@nginx.com rsv3=False, 1521130Szelenkov@nginx.com mask=True, 1531130Szelenkov@nginx.com ): 1541130Szelenkov@nginx.com frame = b'' 1551130Szelenkov@nginx.com 1561130Szelenkov@nginx.com if isinstance(data, str): 1571130Szelenkov@nginx.com data = data.encode('utf-8') 1581130Szelenkov@nginx.com 1591130Szelenkov@nginx.com head1 = ( 1601130Szelenkov@nginx.com (0b10000000 if fin else 0) 1611130Szelenkov@nginx.com | (0b01000000 if rsv1 else 0) 1621130Szelenkov@nginx.com | (0b00100000 if rsv2 else 0) 1631130Szelenkov@nginx.com | (0b00010000 if rsv3 else 0) 1641130Szelenkov@nginx.com | opcode 1651130Szelenkov@nginx.com ) 1661130Szelenkov@nginx.com 1671130Szelenkov@nginx.com head2 = 0b10000000 if mask else 0 1681130Szelenkov@nginx.com 1691130Szelenkov@nginx.com data_length = len(data) if length is None else length 1701130Szelenkov@nginx.com if data_length < 126: 1711130Szelenkov@nginx.com frame += struct.pack('!BB', head1, head2 | data_length) 1721130Szelenkov@nginx.com elif data_length < 65536: 1731130Szelenkov@nginx.com frame += struct.pack('!BBH', head1, head2 | 126, data_length) 1741130Szelenkov@nginx.com else: 1751130Szelenkov@nginx.com frame += struct.pack('!BBQ', head1, head2 | 127, data_length) 1761130Szelenkov@nginx.com 1771130Szelenkov@nginx.com if mask: 1781130Szelenkov@nginx.com mask_bits = struct.pack('!I', random.getrandbits(32)) 1791130Szelenkov@nginx.com frame += mask_bits 1801130Szelenkov@nginx.com 1811130Szelenkov@nginx.com if mask: 1821130Szelenkov@nginx.com frame += self.apply_mask(data, mask_bits) 1831130Szelenkov@nginx.com else: 1841130Szelenkov@nginx.com frame += data 1851130Szelenkov@nginx.com 1861130Szelenkov@nginx.com return frame 1871130Szelenkov@nginx.com 1881130Szelenkov@nginx.com def frame_write(self, sock, *args, **kwargs): 1891130Szelenkov@nginx.com chopsize = kwargs.pop('chopsize') if 'chopsize' in kwargs else None 1901130Szelenkov@nginx.com 1911130Szelenkov@nginx.com frame = self.frame_to_send(*args, **kwargs) 1921130Szelenkov@nginx.com 1931130Szelenkov@nginx.com if chopsize is None: 1941152Szelenkov@nginx.com try: 1951152Szelenkov@nginx.com sock.sendall(frame) 1961152Szelenkov@nginx.com except BrokenPipeError: 1971152Szelenkov@nginx.com pass 1981130Szelenkov@nginx.com 1991130Szelenkov@nginx.com else: 2001130Szelenkov@nginx.com pos = 0 2011130Szelenkov@nginx.com frame_len = len(frame) 2021150Szelenkov@nginx.com while pos < frame_len: 2031130Szelenkov@nginx.com end = min(pos + chopsize, frame_len) 2041151Szelenkov@nginx.com try: 2051151Szelenkov@nginx.com sock.sendall(frame[pos:end]) 2061151Szelenkov@nginx.com except BrokenPipeError: 2071151Szelenkov@nginx.com end = frame_len 2081130Szelenkov@nginx.com pos = end 2091130Szelenkov@nginx.com 210*2616Szelenkov@nginx.com def message(self, sock, mes_type, message, fragmention_size=None, **kwargs): 2111130Szelenkov@nginx.com message_len = len(message) 2121130Szelenkov@nginx.com 2131130Szelenkov@nginx.com if fragmention_size is None: 2141130Szelenkov@nginx.com fragmention_size = message_len 2151130Szelenkov@nginx.com 2161130Szelenkov@nginx.com if message_len <= fragmention_size: 217*2616Szelenkov@nginx.com self.frame_write(sock, mes_type, message, **kwargs) 2181130Szelenkov@nginx.com return 2191130Szelenkov@nginx.com 2201130Szelenkov@nginx.com pos = 0 221*2616Szelenkov@nginx.com op_code = mes_type 2221150Szelenkov@nginx.com while pos < message_len: 2231130Szelenkov@nginx.com end = min(pos + fragmention_size, message_len) 2241150Szelenkov@nginx.com fin = end == message_len 2252073Szelenkov@nginx.com self.frame_write(sock, op_code, message[pos:end], fin=fin, **kwargs) 2261130Szelenkov@nginx.com op_code = self.OP_CONT 2271130Szelenkov@nginx.com pos = end 2281130Szelenkov@nginx.com 2291433Szelenkov@nginx.com def message_read(self, sock, read_timeout=60): 2301130Szelenkov@nginx.com frame = self.frame_read(sock, read_timeout=read_timeout) 2311130Szelenkov@nginx.com 2321150Szelenkov@nginx.com while not frame['fin']: 2331130Szelenkov@nginx.com temp = self.frame_read(sock, read_timeout=read_timeout) 2341130Szelenkov@nginx.com frame['data'] += temp['data'] 2351130Szelenkov@nginx.com frame['fin'] = temp['fin'] 2361130Szelenkov@nginx.com 2371130Szelenkov@nginx.com return frame 238