1import functools
2import json as _json
3
4(CONNECT, DISCONNECT, EVENT, ACK, CONNECT_ERROR, BINARY_EVENT, BINARY_ACK) = \
5    (0, 1, 2, 3, 4, 5, 6)
6packet_names = ['CONNECT', 'DISCONNECT', 'EVENT', 'ACK', 'CONNECT_ERROR',
7                'BINARY_EVENT', 'BINARY_ACK']
8
9
10class Packet(object):
11    """Socket.IO packet."""
12
13    # the format of the Socket.IO packet is as follows:
14    #
15    # packet type: 1 byte, values 0-6
16    # num_attachments: ASCII encoded, only if num_attachments != 0
17    # '-': only if num_attachments != 0
18    # namespace: only if namespace != '/'
19    # ',': only if namespace and one of id and data are defined in this packet
20    # id: ASCII encoded, only if id is not None
21    # data: JSON dump of data payload
22
23    json = _json
24
25    def __init__(self, packet_type=EVENT, data=None, namespace=None, id=None,
26                 binary=None, encoded_packet=None):
27        self.packet_type = packet_type
28        self.data = data
29        self.namespace = namespace
30        self.id = id
31        if binary or (binary is None and self._data_is_binary(self.data)):
32            if self.packet_type == EVENT:
33                self.packet_type = BINARY_EVENT
34            elif self.packet_type == ACK:
35                self.packet_type = BINARY_ACK
36            else:
37                raise ValueError('Packet does not support binary payload.')
38        self.attachment_count = 0
39        self.attachments = []
40        if encoded_packet:
41            self.attachment_count = self.decode(encoded_packet)
42
43    def encode(self):
44        """Encode the packet for transmission.
45
46        If the packet contains binary elements, this function returns a list
47        of packets where the first is the original packet with placeholders for
48        the binary components and the remaining ones the binary attachments.
49        """
50        encoded_packet = str(self.packet_type)
51        if self.packet_type == BINARY_EVENT or self.packet_type == BINARY_ACK:
52            data, attachments = self._deconstruct_binary(self.data)
53            encoded_packet += str(len(attachments)) + '-'
54        else:
55            data = self.data
56            attachments = None
57        needs_comma = False
58        if self.namespace is not None and self.namespace != '/':
59            encoded_packet += self.namespace
60            needs_comma = True
61        if self.id is not None:
62            if needs_comma:
63                encoded_packet += ','
64                needs_comma = False
65            encoded_packet += str(self.id)
66        if data is not None:
67            if needs_comma:
68                encoded_packet += ','
69            encoded_packet += self.json.dumps(data, separators=(',', ':'))
70        if attachments is not None:
71            encoded_packet = [encoded_packet] + attachments
72        return encoded_packet
73
74    def decode(self, encoded_packet):
75        """Decode a transmitted package.
76
77        The return value indicates how many binary attachment packets are
78        necessary to fully decode the packet.
79        """
80        ep = encoded_packet
81        try:
82            self.packet_type = int(ep[0:1])
83        except TypeError:
84            self.packet_type = ep
85            ep = ''
86        self.namespace = None
87        self.data = None
88        ep = ep[1:]
89        dash = ep.find('-')
90        attachment_count = 0
91        if dash > 0 and ep[0:dash].isdigit():
92            attachment_count = int(ep[0:dash])
93            ep = ep[dash + 1:]
94        if ep and ep[0:1] == '/':
95            sep = ep.find(',')
96            if sep == -1:
97                self.namespace = ep
98                ep = ''
99            else:
100                self.namespace = ep[0:sep]
101                ep = ep[sep + 1:]
102            q = self.namespace.find('?')
103            if q != -1:
104                self.namespace = self.namespace[0:q]
105        if ep and ep[0].isdigit():
106            self.id = 0
107            while ep and ep[0].isdigit():
108                self.id = self.id * 10 + int(ep[0])
109                ep = ep[1:]
110        if ep:
111            self.data = self.json.loads(ep)
112        return attachment_count
113
114    def add_attachment(self, attachment):
115        if self.attachment_count <= len(self.attachments):
116            raise ValueError('Unexpected binary attachment')
117        self.attachments.append(attachment)
118        if self.attachment_count == len(self.attachments):
119            self.reconstruct_binary(self.attachments)
120            return True
121        return False
122
123    def reconstruct_binary(self, attachments):
124        """Reconstruct a decoded packet using the given list of binary
125        attachments.
126        """
127        self.data = self._reconstruct_binary_internal(self.data,
128                                                      self.attachments)
129
130    def _reconstruct_binary_internal(self, data, attachments):
131        if isinstance(data, list):
132            return [self._reconstruct_binary_internal(item, attachments)
133                    for item in data]
134        elif isinstance(data, dict):
135            if data.get('_placeholder') and 'num' in data:
136                return attachments[data['num']]
137            else:
138                return {key: self._reconstruct_binary_internal(value,
139                                                               attachments)
140                        for key, value in data.items()}
141        else:
142            return data
143
144    def _deconstruct_binary(self, data):
145        """Extract binary components in the packet."""
146        attachments = []
147        data = self._deconstruct_binary_internal(data, attachments)
148        return data, attachments
149
150    def _deconstruct_binary_internal(self, data, attachments):
151        if isinstance(data, bytes):
152            attachments.append(data)
153            return {'_placeholder': True, 'num': len(attachments) - 1}
154        elif isinstance(data, list):
155            return [self._deconstruct_binary_internal(item, attachments)
156                    for item in data]
157        elif isinstance(data, dict):
158            return {key: self._deconstruct_binary_internal(value, attachments)
159                    for key, value in data.items()}
160        else:
161            return data
162
163    def _data_is_binary(self, data):
164        """Check if the data contains binary components."""
165        if isinstance(data, bytes):
166            return True
167        elif isinstance(data, list):
168            return functools.reduce(
169                lambda a, b: a or b, [self._data_is_binary(item)
170                                      for item in data], False)
171        elif isinstance(data, dict):
172            return functools.reduce(
173                lambda a, b: a or b, [self._data_is_binary(item)
174                                      for item in data.values()],
175                False)
176        else:
177            return False
178