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