1# -*- coding: utf-8 -*- 2import os 3import struct 4 5from ws4py.framing import Frame, OPCODE_CONTINUATION, OPCODE_TEXT, \ 6 OPCODE_BINARY, OPCODE_CLOSE, OPCODE_PING, OPCODE_PONG 7from ws4py.compat import unicode, py3k 8 9__all__ = ['Message', 'TextMessage', 'BinaryMessage', 'CloseControlMessage', 10 'PingControlMessage', 'PongControlMessage'] 11 12class Message(object): 13 def __init__(self, opcode, data=b'', encoding='utf-8'): 14 """ 15 A message is a application level entity. It's usually built 16 from one or many frames. The protocol defines several kind 17 of messages which are grouped into two sets: 18 19 * data messages which can be text or binary typed 20 * control messages which provide a mechanism to perform 21 in-band control communication between peers 22 23 The ``opcode`` indicates the message type and ``data`` is 24 the possible message payload. 25 26 The payload is held internally as a a :func:`bytearray` as they are 27 faster than pure strings for append operations. 28 29 Unicode data will be encoded using the provided ``encoding``. 30 """ 31 self.opcode = opcode 32 self._completed = False 33 self.encoding = encoding 34 35 if isinstance(data, unicode): 36 if not encoding: 37 raise TypeError("unicode data without an encoding") 38 data = data.encode(encoding) 39 elif isinstance(data, bytearray): 40 data = bytes(data) 41 elif not isinstance(data, bytes): 42 raise TypeError("%s is not a supported data type" % type(data)) 43 44 self.data = data 45 46 def single(self, mask=False): 47 """ 48 Returns a frame bytes with the fin bit set and a random mask. 49 50 If ``mask`` is set, automatically mask the frame 51 using a generated 4-byte token. 52 """ 53 mask = os.urandom(4) if mask else None 54 return Frame(body=self.data, opcode=self.opcode, 55 masking_key=mask, fin=1).build() 56 57 def fragment(self, first=False, last=False, mask=False): 58 """ 59 Returns a :class:`ws4py.framing.Frame` bytes. 60 61 The behavior depends on the given flags: 62 63 * ``first``: the frame uses ``self.opcode`` else a continuation opcode 64 * ``last``: the frame has its ``fin`` bit set 65 * ``mask``: the frame is masked using a automatically generated 4-byte token 66 """ 67 fin = 1 if last is True else 0 68 opcode = self.opcode if first is True else OPCODE_CONTINUATION 69 mask = os.urandom(4) if mask else None 70 return Frame(body=self.data, 71 opcode=opcode, masking_key=mask, 72 fin=fin).build() 73 74 @property 75 def completed(self): 76 """ 77 Indicates the the message is complete, meaning 78 the frame's ``fin`` bit was set. 79 """ 80 return self._completed 81 82 @completed.setter 83 def completed(self, state): 84 """ 85 Sets the state for this message. Usually 86 set by the stream's parser. 87 """ 88 self._completed = state 89 90 def extend(self, data): 91 """ 92 Add more ``data`` to the message. 93 """ 94 if isinstance(data, bytes): 95 self.data += data 96 elif isinstance(data, bytearray): 97 self.data += bytes(data) 98 elif isinstance(data, unicode): 99 self.data += data.encode(self.encoding) 100 else: 101 raise TypeError("%s is not a supported data type" % type(data)) 102 103 def __len__(self): 104 return len(self.__unicode__()) 105 106 def __str__(self): 107 if py3k: 108 return self.data.decode(self.encoding) 109 return self.data 110 111 def __unicode__(self): 112 return self.data.decode(self.encoding) 113 114class TextMessage(Message): 115 def __init__(self, text=None): 116 Message.__init__(self, OPCODE_TEXT, text) 117 118 @property 119 def is_binary(self): 120 return False 121 122 @property 123 def is_text(self): 124 return True 125 126class BinaryMessage(Message): 127 def __init__(self, bytes=None): 128 Message.__init__(self, OPCODE_BINARY, bytes, encoding=None) 129 130 @property 131 def is_binary(self): 132 return True 133 134 @property 135 def is_text(self): 136 return False 137 138 def __len__(self): 139 return len(self.data) 140 141class CloseControlMessage(Message): 142 def __init__(self, code=1000, reason=''): 143 data = b"" 144 if code: 145 data += struct.pack("!H", code) 146 if reason is not None: 147 if isinstance(reason, unicode): 148 reason = reason.encode('utf-8') 149 data += reason 150 151 Message.__init__(self, OPCODE_CLOSE, data, 'utf-8') 152 self.code = code 153 self.reason = reason 154 155 def __str__(self): 156 if py3k: 157 return self.reason.decode('utf-8') 158 return self.reason 159 160 def __unicode__(self): 161 return self.reason.decode(self.encoding) 162 163class PingControlMessage(Message): 164 def __init__(self, data=None): 165 Message.__init__(self, OPCODE_PING, data) 166 167class PongControlMessage(Message): 168 def __init__(self, data): 169 Message.__init__(self, OPCODE_PONG, data) 170