codekingpro/portable-devtools

1import functools2from engineio import json as _json3 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-616    # num_attachments: ASCII encoded, only if num_attachments != 017    # '-': only if num_attachments != 018    # namespace, followed by a ',': only if namespace != '/'19    # id: ASCII encoded, only if id is not None20    # data: JSON dump of data payload21 22    uses_binary_events = True23    json = _json24 25    def __init__(self, packet_type=EVENT, data=None, namespace=None, id=None,26                 binary=None, encoded_packet=None):27        self.packet_type = packet_type28        self.data = data29        self.namespace = namespace30        self.id = id31        if self.uses_binary_events and \32                (binary or (binary is None and self._data_is_binary(33                    self.data))):34            if self.packet_type == EVENT:35                self.packet_type = BINARY_EVENT36            elif self.packet_type == ACK:37                self.packet_type = BINARY_ACK38            else:39                raise ValueError('Packet does not support binary payload.')40        self.attachment_count = 041        self.attachments = []42        if encoded_packet:43            self.attachment_count = self.decode(encoded_packet) or 044 45    def encode(self):46        """Encode the packet for transmission.47 48        If the packet contains binary elements, this function returns a list49        of packets where the first is the original packet with placeholders for50        the binary components and the remaining ones the binary attachments.51        """52        encoded_packet = str(self.packet_type)53        if self.packet_type == BINARY_EVENT or self.packet_type == BINARY_ACK:54            data, attachments = self._deconstruct_binary(self.data)55            encoded_packet += str(len(attachments)) + '-'56        else:57            data = self.data58            attachments = None59        if self.namespace is not None and self.namespace != '/':60            encoded_packet += self.namespace + ','61        if self.id is not None:62            encoded_packet += str(self.id)63        if data is not None:64            encoded_packet += self.json.dumps(data, separators=(',', ':'))65        if attachments is not None:66            encoded_packet = [encoded_packet] + attachments67        return encoded_packet68 69    def decode(self, encoded_packet):70        """Decode a transmitted package.71 72        The return value indicates how many binary attachment packets are73        necessary to fully decode the packet.74        """75        ep = encoded_packet76        try:77            self.packet_type = int(ep[0:1])78        except TypeError:79            self.packet_type = ep80            ep = ''81        self.namespace = None82        self.data = None83        ep = ep[1:]84        dash = ep.find('-')85        attachment_count = 086        if dash > 0 and ep[0:dash].isdigit():87            if dash > 10:88                raise ValueError('too many attachments')89            attachment_count = int(ep[0:dash])90            ep = ep[dash + 1:]91        if ep and ep[0:1] == '/':92            sep = ep.find(',')93            if sep == -1:94                self.namespace = ep95                ep = ''96            else:97                self.namespace = ep[0:sep]98                ep = ep[sep + 1:]99            q = self.namespace.find('?')100            if q != -1:101                self.namespace = self.namespace[0:q]102        if ep and ep[0].isdigit():103            i = 1104            end = len(ep)105            while i < end:106                if not ep[i].isdigit() or i >= 100:107                    break108                i += 1109            self.id = int(ep[:i])110            ep = ep[i:]111            if len(ep) > 0 and ep[0].isdigit():112                raise ValueError('id field is too long')113        if ep:114            self.data = self.json.loads(ep)115        return attachment_count116 117    def add_attachment(self, attachment):118        if self.attachment_count <= len(self.attachments):119            raise ValueError('Unexpected binary attachment')120        self.attachments.append(attachment)121        if self.attachment_count == len(self.attachments):122            self.reconstruct_binary(self.attachments)123            return True124        return False125 126    def reconstruct_binary(self, attachments):127        """Reconstruct a decoded packet using the given list of binary128        attachments.129        """130        self.data = self._reconstruct_binary_internal(self.data,131                                                      self.attachments)132 133    def _reconstruct_binary_internal(self, data, attachments):134        if isinstance(data, list):135            return [self._reconstruct_binary_internal(item, attachments)136                    for item in data]137        elif isinstance(data, dict):138            if data.get('_placeholder') and 'num' in data:139                return attachments[data['num']]140            else:141                return {key: self._reconstruct_binary_internal(value,142                                                               attachments)143                        for key, value in data.items()}144        else:145            return data146 147    def _deconstruct_binary(self, data):148        """Extract binary components in the packet."""149        attachments = []150        data = self._deconstruct_binary_internal(data, attachments)151        return data, attachments152 153    def _deconstruct_binary_internal(self, data, attachments):154        if isinstance(data, bytes):155            attachments.append(data)156            return {'_placeholder': True, 'num': len(attachments) - 1}157        elif isinstance(data, list):158            return [self._deconstruct_binary_internal(item, attachments)159                    for item in data]160        elif isinstance(data, dict):161            return {key: self._deconstruct_binary_internal(value, attachments)162                    for key, value in data.items()}163        else:164            return data165 166    def _data_is_binary(self, data):167        """Check if the data contains binary components."""168        if isinstance(data, bytes):169            return True170        elif isinstance(data, list):171            return functools.reduce(172                lambda a, b: a or b, [self._data_is_binary(item)173                                      for item in data], False)174        elif isinstance(data, dict):175            return functools.reduce(176                lambda a, b: a or b, [self._data_is_binary(item)177                                      for item in data.values()],178                False)179        else:180            return False181 182    def _to_dict(self):183        d = {184            'type': self.packet_type,185            'data': self.data,186            'nsp': self.namespace,187        }188        if self.id:189            d['id'] = self.id190        return d191