bazarr/libs/socketio/packet.py

178 lines
6.6 KiB
Python
Raw Normal View History

2020-01-30 01:07:26 +00:00
import functools
import json as _json
(CONNECT, DISCONNECT, EVENT, ACK, CONNECT_ERROR, BINARY_EVENT, BINARY_ACK) = \
2020-01-30 01:07:26 +00:00
(0, 1, 2, 3, 4, 5, 6)
packet_names = ['CONNECT', 'DISCONNECT', 'EVENT', 'ACK', 'CONNECT_ERROR',
2020-01-30 01:07:26 +00:00
'BINARY_EVENT', 'BINARY_ACK']
class Packet(object):
"""Socket.IO packet."""
# the format of the Socket.IO packet is as follows:
#
# packet type: 1 byte, values 0-6
# num_attachments: ASCII encoded, only if num_attachments != 0
# '-': only if num_attachments != 0
# namespace: only if namespace != '/'
# ',': only if namespace and one of id and data are defined in this packet
# id: ASCII encoded, only if id is not None
# data: JSON dump of data payload
json = _json
def __init__(self, packet_type=EVENT, data=None, namespace=None, id=None,
binary=None, encoded_packet=None):
self.packet_type = packet_type
self.data = data
self.namespace = namespace
self.id = id
if binary or (binary is None and self._data_is_binary(self.data)):
if self.packet_type == EVENT:
self.packet_type = BINARY_EVENT
elif self.packet_type == ACK:
self.packet_type = BINARY_ACK
else:
raise ValueError('Packet does not support binary payload.')
self.attachment_count = 0
self.attachments = []
if encoded_packet:
self.attachment_count = self.decode(encoded_packet)
def encode(self):
"""Encode the packet for transmission.
If the packet contains binary elements, this function returns a list
of packets where the first is the original packet with placeholders for
the binary components and the remaining ones the binary attachments.
"""
encoded_packet = str(self.packet_type)
2020-01-30 01:07:26 +00:00
if self.packet_type == BINARY_EVENT or self.packet_type == BINARY_ACK:
data, attachments = self._deconstruct_binary(self.data)
encoded_packet += str(len(attachments)) + '-'
2020-01-30 01:07:26 +00:00
else:
data = self.data
attachments = None
needs_comma = False
if self.namespace is not None and self.namespace != '/':
encoded_packet += self.namespace
needs_comma = True
if self.id is not None:
if needs_comma:
encoded_packet += ','
needs_comma = False
encoded_packet += str(self.id)
2020-01-30 01:07:26 +00:00
if data is not None:
if needs_comma:
encoded_packet += ','
encoded_packet += self.json.dumps(data, separators=(',', ':'))
if attachments is not None:
encoded_packet = [encoded_packet] + attachments
return encoded_packet
def decode(self, encoded_packet):
"""Decode a transmitted package.
The return value indicates how many binary attachment packets are
necessary to fully decode the packet.
"""
ep = encoded_packet
try:
self.packet_type = int(ep[0:1])
except TypeError:
self.packet_type = ep
ep = ''
self.namespace = None
self.data = None
ep = ep[1:]
dash = ep.find('-')
attachment_count = 0
if dash > 0 and ep[0:dash].isdigit():
attachment_count = int(ep[0:dash])
ep = ep[dash + 1:]
if ep and ep[0:1] == '/':
sep = ep.find(',')
if sep == -1:
self.namespace = ep
ep = ''
else:
self.namespace = ep[0:sep]
ep = ep[sep + 1:]
q = self.namespace.find('?')
if q != -1:
self.namespace = self.namespace[0:q]
if ep and ep[0].isdigit():
self.id = 0
while ep and ep[0].isdigit():
self.id = self.id * 10 + int(ep[0])
ep = ep[1:]
if ep:
self.data = self.json.loads(ep)
return attachment_count
def add_attachment(self, attachment):
if self.attachment_count <= len(self.attachments):
raise ValueError('Unexpected binary attachment')
self.attachments.append(attachment)
if self.attachment_count == len(self.attachments):
self.reconstruct_binary(self.attachments)
return True
return False
def reconstruct_binary(self, attachments):
"""Reconstruct a decoded packet using the given list of binary
attachments.
"""
self.data = self._reconstruct_binary_internal(self.data,
self.attachments)
def _reconstruct_binary_internal(self, data, attachments):
if isinstance(data, list):
return [self._reconstruct_binary_internal(item, attachments)
for item in data]
elif isinstance(data, dict):
if data.get('_placeholder') and 'num' in data:
return attachments[data['num']]
else:
return {key: self._reconstruct_binary_internal(value,
attachments)
for key, value in data.items()}
2020-01-30 01:07:26 +00:00
else:
return data
def _deconstruct_binary(self, data):
"""Extract binary components in the packet."""
attachments = []
data = self._deconstruct_binary_internal(data, attachments)
return data, attachments
def _deconstruct_binary_internal(self, data, attachments):
if isinstance(data, bytes):
2020-01-30 01:07:26 +00:00
attachments.append(data)
return {'_placeholder': True, 'num': len(attachments) - 1}
elif isinstance(data, list):
return [self._deconstruct_binary_internal(item, attachments)
for item in data]
elif isinstance(data, dict):
return {key: self._deconstruct_binary_internal(value, attachments)
for key, value in data.items()}
2020-01-30 01:07:26 +00:00
else:
return data
def _data_is_binary(self, data):
"""Check if the data contains binary components."""
if isinstance(data, bytes):
2020-01-30 01:07:26 +00:00
return True
elif isinstance(data, list):
return functools.reduce(
lambda a, b: a or b, [self._data_is_binary(item)
for item in data], False)
elif isinstance(data, dict):
return functools.reduce(
lambda a, b: a or b, [self._data_is_binary(item)
for item in data.values()],
2020-01-30 01:07:26 +00:00
False)
else:
return False