1import base64
2from engineio import json as _json
3
4(OPEN, CLOSE, PING, PONG, MESSAGE, UPGRADE, NOOP) = (0, 1, 2, 3, 4, 5, 6)
5packet_names = ['OPEN', 'CLOSE', 'PING', 'PONG', 'MESSAGE', 'UPGRADE', 'NOOP']
6
7binary_types = (bytes, bytearray)
8
9
10class Packet(object):
11    """Engine.IO packet."""
12
13    json = _json
14
15    def __init__(self, packet_type=NOOP, data=None, encoded_packet=None):
16        self.packet_type = packet_type
17        self.data = data
18        if isinstance(data, str):
19            self.binary = False
20        elif isinstance(data, binary_types):
21            self.binary = True
22        else:
23            self.binary = False
24        if self.binary and self.packet_type != MESSAGE:
25            raise ValueError('Binary packets can only be of type MESSAGE')
26        if encoded_packet is not None:
27            self.decode(encoded_packet)
28
29    def encode(self, b64=False):
30        """Encode the packet for transmission."""
31        if self.binary:
32            if b64:
33                encoded_packet = 'b' + base64.b64encode(self.data).decode(
34                    'utf-8')
35            else:
36                encoded_packet = self.data
37        else:
38            encoded_packet = str(self.packet_type)
39            if isinstance(self.data, str):
40                encoded_packet += self.data
41            elif isinstance(self.data, dict) or isinstance(self.data, list):
42                encoded_packet += self.json.dumps(self.data,
43                                                  separators=(',', ':'))
44            elif self.data is not None:
45                encoded_packet += str(self.data)
46        return encoded_packet
47
48    def decode(self, encoded_packet):
49        """Decode a transmitted package."""
50        self.binary = isinstance(encoded_packet, binary_types)
51        b64 = not self.binary and encoded_packet[0] == 'b'
52        if b64:
53            self.binary = True
54            self.packet_type = MESSAGE
55            self.data = base64.b64decode(encoded_packet[1:])
56        else:
57            if self.binary and not isinstance(encoded_packet, bytes):
58                encoded_packet = bytes(encoded_packet)
59            if self.binary:
60                self.packet_type = MESSAGE
61                self.data = encoded_packet
62            else:
63                self.packet_type = int(encoded_packet[0])
64                try:
65                    self.data = self.json.loads(encoded_packet[1:])
66                    if isinstance(self.data, int):
67                        # do not allow integer payloads, see
68                        # github.com/miguelgrinberg/python-engineio/issues/75
69                        # for background on this decision
70                        raise ValueError
71                except ValueError:
72                    self.data = encoded_packet[1:]
73