def testUtf8Validator(self): state = validate_utf8(six.b('\xf0\x90\x80\x80')) self.assertEqual(state, True) state = validate_utf8(six.b('\xce\xba\xe1\xbd\xb9\xcf\x83\xce\xbc\xce\xb5\xed\xa0\x80edited')) self.assertEqual(state, False) state = validate_utf8(six.b('')) self.assertEqual(state, True)
def format(self): """ format this object to string(byte array) to send data to server. """ if any(x not in (0, 1) for x in [self.fin, self.rsv1, self.rsv2, self.rsv3]): raise ValueError("not 0 or 1") if self.opcode not in ABNF.OPCODES: raise ValueError("Invalid OPCODE") length = len(self.data) if length >= ABNF.LENGTH_63: raise ValueError("data is too long") frame_header = chr(self.fin << 7 | self.rsv1 << 6 | self.rsv2 << 5 | self.rsv3 << 4 | self.opcode) if length < ABNF.LENGTH_7: frame_header += chr(self.mask << 7 | length) frame_header = six.b(frame_header) elif length < ABNF.LENGTH_16: frame_header += chr(self.mask << 7 | 0x7e) frame_header = six.b(frame_header) frame_header += struct.pack("!H", length) else: frame_header += chr(self.mask << 7 | 0x7f) frame_header = six.b(frame_header) frame_header += struct.pack("!Q", length) if not self.mask: return frame_header + self.data else: mask_key = self.get_mask_key(4) return frame_header + self._get_masked(mask_key)
def recv_line(sock): line = [] while True: c = recv(sock, 1) line.append(c) if c == six.b("\n"): break return six.b("").join(line)
def testRecvWithSimpleFragmentation(self): sock = ws.WebSocket() s = sock.sock = SockMock() # OPCODE=TEXT, FIN=0, MSG="Brevity is " s.add_packet(six.b("\x01\x8babcd#\x10\x06\x12\x08\x16\x1aD\x08\x11C")) # OPCODE=CONT, FIN=1, MSG="the soul of wit" s.add_packet(six.b("\x80\x8fabcd\x15\n\x06D\x12\r\x16\x08A\r\x05D\x16\x0b\x17")) data = sock.recv() self.assertEqual(data, "Brevity is the soul of wit") with self.assertRaises(ws.WebSocketConnectionClosedException): sock.recv()
def testRecv(self): # TODO: add longer frame data sock = ws.WebSocket() s = sock.sock = SockMock() something = six.b("\x81\x8fabcd\x82\xe3\xf0\x87\xe3\xf1\x80\xe5\xca\x81\xe2\xc5\x82\xe3\xcc") s.add_packet(something) data = sock.recv() self.assertEqual(data, "こんにちは") s.add_packet(six.b("\x81\x85abcd)\x07\x0f\x08\x0e")) data = sock.recv() self.assertEqual(data, "Hello")
def testSend(self): # TODO: add longer frame data sock = ws.WebSocket() sock.set_mask_key(create_mask_key) s = sock.sock = HeaderSockMock("data/header01.txt") sock.send("Hello") self.assertEqual(s.sent[0], six.b("\x81\x85abcd)\x07\x0f\x08\x0e")) sock.send("こんにちは") self.assertEqual(s.sent[1], six.b("\x81\x8fabcd\x82\xe3\xf0\x87\xe3\xf1\x80\xe5\xca\x81\xe2\xc5\x82\xe3\xcc")) sock.send(u"こんにちは") self.assertEqual(s.sent[1], six.b("\x81\x8fabcd\x82\xe3\xf0\x87\xe3\xf1\x80\xe5\xca\x81\xe2\xc5\x82\xe3\xcc")) sock.send("x" * 127)
def testRecvTimeout(self): sock = ws.WebSocket() s = sock.sock = SockMock() s.add_packet(six.b("\x81")) s.add_packet(socket.timeout()) s.add_packet(six.b("\x8dabcd\x29\x07\x0f\x08\x0e")) s.add_packet(socket.timeout()) s.add_packet(six.b("\x4e\x43\x33\x0e\x10\x0f\x00\x40")) with self.assertRaises(ws.WebSocketTimeoutException): sock.recv() with self.assertRaises(ws.WebSocketTimeoutException): sock.recv() data = sock.recv() self.assertEqual(data, "Hello, World!") with self.assertRaises(ws.WebSocketConnectionClosedException): sock.recv()
def close(self, status=STATUS_NORMAL, reason=six.b(""), timeout=3): """ Close Websocket object status: status code to send. see STATUS_XXX. reason: the reason to close. This must be string. timeout: timeout until receive a close frame. If None, it will wait forever until receive a close frame. """ if self.connected: if status < 0 or status >= ABNF.LENGTH_16: raise ValueError("code is invalid range") try: self.connected = False self.send( struct.pack('!H', status) + reason, ABNF.OPCODE_CLOSE) sock_timeout = self.sock.gettimeout() self.sock.settimeout(timeout) try: frame = self.recv_frame() if isEnabledForError(): recv_status = struct.unpack("!H", frame.data[0:2])[0] if recv_status != STATUS_NORMAL: error("close status: " + repr(recv_status)) except: pass self.sock.settimeout(sock_timeout) self.sock.shutdown(socket.SHUT_RDWR) except: pass self.shutdown()
def testRecvWithProlongedFragmentation(self): sock = ws.WebSocket() s = sock.sock = SockMock() # OPCODE=TEXT, FIN=0, MSG="Once more unto the breach, " s.add_packet(six.b("\x01\x9babcd.\x0c\x00\x01A\x0f\x0c\x16\x04B\x16\n\x15" "\rC\x10\t\x07C\x06\x13\x07\x02\x07\tNC")) # OPCODE=CONT, FIN=0, MSG="dear friends, " s.add_packet(six.b("\x00\x8eabcd\x05\x07\x02\x16A\x04\x11\r\x04\x0c\x07" "\x17MB")) # OPCODE=CONT, FIN=1, MSG="once more" s.add_packet(six.b("\x80\x89abcd\x0e\x0c\x00\x01A\x0f\x0c\x16\x04")) data = sock.recv() self.assertEqual( data, "Once more unto the breach, dear friends, once more") with self.assertRaises(ws.WebSocketConnectionClosedException): sock.recv()
def testRecvWithFireEventOfFragmentation(self): sock = ws.WebSocket(fire_cont_frame=True) s = sock.sock = SockMock() # OPCODE=TEXT, FIN=0, MSG="Brevity is " s.add_packet(six.b("\x01\x8babcd#\x10\x06\x12\x08\x16\x1aD\x08\x11C")) # OPCODE=CONT, FIN=0, MSG="Brevity is " s.add_packet(six.b("\x00\x8babcd#\x10\x06\x12\x08\x16\x1aD\x08\x11C")) # OPCODE=CONT, FIN=1, MSG="the soul of wit" s.add_packet(six.b("\x80\x8fabcd\x15\n\x06D\x12\r\x16\x08A\r\x05D\x16\x0b\x17")) _, data = sock.recv_data() self.assertEqual(data, six.b("Brevity is ")) _, data = sock.recv_data() self.assertEqual(data, six.b("Brevity is ")) _, data = sock.recv_data() self.assertEqual(data, six.b("the soul of wit")) # OPCODE=CONT, FIN=0, MSG="Brevity is " s.add_packet(six.b("\x80\x8babcd#\x10\x06\x12\x08\x16\x1aD\x08\x11C")) with self.assertRaises(ws.WebSocketException): sock.recv_data() with self.assertRaises(ws.WebSocketConnectionClosedException): sock.recv()
def testRecvWithFragmentationAndControlFrame(self): sock = ws.WebSocket() sock.set_mask_key(create_mask_key) s = sock.sock = SockMock() # OPCODE=TEXT, FIN=0, MSG="Too much " s.add_packet(six.b("\x01\x89abcd5\r\x0cD\x0c\x17\x00\x0cA")) # OPCODE=PING, FIN=1, MSG="Please PONG this" s.add_packet(six.b("\x89\x90abcd1\x0e\x06\x05\x12\x07C4.,$D\x15\n\n\x17")) # OPCODE=CONT, FIN=1, MSG="of a good thing" s.add_packet(six.b("\x80\x8fabcd\x0e\x04C\x05A\x05\x0c\x0b\x05B\x17\x0c" "\x08\x0c\x04")) data = sock.recv() self.assertEqual(data, "Too much of a good thing") with self.assertRaises(ws.WebSocketConnectionClosedException): sock.recv() self.assertEqual( s.sent[0], six.b("\x8a\x90abcd1\x0e\x06\x05\x12\x07C4.,$D\x15\n\n\x17"))
def testInternalRecvStrict(self): sock = ws.WebSocket() s = sock.sock = SockMock() s.add_packet(six.b("foo")) s.add_packet(socket.timeout()) s.add_packet(six.b("bar")) # s.add_packet(SSLError("The read operation timed out")) s.add_packet(six.b("baz")) with self.assertRaises(ws.WebSocketTimeoutException): sock.frame_buffer.recv_strict(9) # if six.PY2: # with self.assertRaises(ws.WebSocketTimeoutException): # data = sock._recv_strict(9) # else: # with self.assertRaises(SSLError): # data = sock._recv_strict(9) data = sock.frame_buffer.recv_strict(9) self.assertEqual(data, six.b("foobarbaz")) with self.assertRaises(ws.WebSocketConnectionClosedException): sock.frame_buffer.recv_strict(1)
def mask(mask_key, data): """ mask or unmask data. Just do xor for each byte mask_key: 4 byte string(byte). data: data to mask/unmask. """ if data is None: data = "" if isinstance(mask_key, six.text_type): mask_key = six.b(mask_key) if isinstance(data, six.text_type): data = six.b(data) _m = array.array("B", mask_key) _d = array.array("B", data) return _mask(_m, _d)
def send_close(self, status=STATUS_NORMAL, reason=six.b("")): """ send close data to the server. status: status code to send. see STATUS_XXX. reason: the reason to close. This must be string or bytes. """ if status < 0 or status >= ABNF.LENGTH_16: raise ValueError("code is invalid range") self.connected = False self.send(struct.pack('!H', status) + reason, ABNF.OPCODE_CLOSE)
def testClose(self): sock = ws.WebSocket() sock.sock = SockMock() sock.connected = True sock.close() self.assertEqual(sock.connected, False) sock = ws.WebSocket() s = sock.sock = SockMock() sock.connected = True s.add_packet(six.b('\x88\x80\x17\x98p\x84')) sock.recv() self.assertEqual(sock.connected, False)
def recv_strict(self, bufsize): shortage = bufsize - sum(len(x) for x in self.recv_buffer) while shortage > 0: # Limit buffer size that we pass to socket.recv() to avoid # fragmenting the heap -- the number of bytes recv() actually # reads is limited by socket buffer and is relatively small, # yet passing large numbers repeatedly causes lots of large # buffers allocated and then shrunk, which results in # fragmentation. bytes_ = self.recv(min(16384, shortage)) self.recv_buffer.append(bytes_) shortage -= len(bytes_) unified = six.b("").join(self.recv_buffer) if shortage == 0: self.recv_buffer = [] return unified else: self.recv_buffer = [unified[bufsize:]] return unified[:bufsize]
def testRecvContFragmentation(self): sock = ws.WebSocket() s = sock.sock = SockMock() # OPCODE=CONT, FIN=1, MSG="the soul of wit" s.add_packet(six.b("\x80\x8fabcd\x15\n\x06D\x12\r\x16\x08A\r\x05D\x16\x0b\x17")) self.assertRaises(ws.WebSocketException, sock.recv)