102 lines
2.8 KiB
Python
102 lines
2.8 KiB
Python
import os
|
|
import struct
|
|
import io
|
|
|
|
import msgpack
|
|
import msgpack.fallback
|
|
|
|
|
|
def msgpackHeader(size):
|
|
if size <= 2 ** 8 - 1:
|
|
return b"\xc4" + struct.pack("B", size)
|
|
elif size <= 2 ** 16 - 1:
|
|
return b"\xc5" + struct.pack(">H", size)
|
|
elif size <= 2 ** 32 - 1:
|
|
return b"\xc6" + struct.pack(">I", size)
|
|
else:
|
|
raise Exception("huge binary string")
|
|
|
|
|
|
def stream(data, writer):
|
|
packer = msgpack.Packer(use_bin_type=True)
|
|
writer(packer.pack_map_header(len(data)))
|
|
for key, val in data.items():
|
|
writer(packer.pack(key))
|
|
if isinstance(val, io.IOBase): # File obj
|
|
max_size = os.fstat(val.fileno()).st_size - val.tell()
|
|
size = min(max_size, val.read_bytes)
|
|
bytes_left = size
|
|
writer(msgpackHeader(size))
|
|
buff = 1024 * 64
|
|
while 1:
|
|
writer(val.read(min(bytes_left, buff)))
|
|
bytes_left = bytes_left - buff
|
|
if bytes_left <= 0:
|
|
break
|
|
else: # Simple
|
|
writer(packer.pack(val))
|
|
return size
|
|
|
|
|
|
class FilePart(object):
|
|
__slots__ = ("file", "read_bytes", "__class__")
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
self.file = open(*args, **kwargs)
|
|
self.__enter__ == self.file.__enter__
|
|
|
|
def __getattr__(self, attr):
|
|
return getattr(self.file, attr)
|
|
|
|
def __enter__(self, *args, **kwargs):
|
|
return self.file.__enter__(*args, **kwargs)
|
|
|
|
def __exit__(self, *args, **kwargs):
|
|
return self.file.__exit__(*args, **kwargs)
|
|
|
|
|
|
# Don't try to decode the value of these fields as utf8
|
|
bin_value_keys = ("hashfield_raw", "peers", "peers_ipv6", "peers_onion", "body", "sites", "bin")
|
|
|
|
|
|
def objectDecoderHook(obj):
|
|
global bin_value_keys
|
|
back = {}
|
|
for key, val in obj:
|
|
if type(key) is bytes:
|
|
key = key.decode("utf8")
|
|
if key in bin_value_keys or type(val) is not bytes or len(key) >= 64:
|
|
back[key] = val
|
|
else:
|
|
back[key] = val.decode("utf8")
|
|
return back
|
|
|
|
|
|
def getUnpacker(fallback=False, decode=True):
|
|
if fallback: # Pure Python
|
|
unpacker = msgpack.fallback.Unpacker
|
|
else:
|
|
unpacker = msgpack.Unpacker
|
|
|
|
extra_kwargs = {"max_buffer_size": 5 * 1024 * 1024}
|
|
if msgpack.version[0] >= 1:
|
|
extra_kwargs["strict_map_key"] = False
|
|
|
|
if decode: # Workaround for backward compatibility: Try to decode bin to str
|
|
unpacker = unpacker(raw=True, object_pairs_hook=objectDecoderHook, **extra_kwargs)
|
|
else:
|
|
unpacker = unpacker(raw=False, **extra_kwargs)
|
|
|
|
return unpacker
|
|
|
|
|
|
def pack(data, use_bin_type=True):
|
|
return msgpack.packb(data, use_bin_type=use_bin_type)
|
|
|
|
|
|
def unpack(data, decode=True):
|
|
unpacker = getUnpacker(decode=decode)
|
|
unpacker.feed(data)
|
|
return next(unpacker)
|
|
|