strawmelonjuice.com
Forking what is left of ZeroNet and hopefully adding an AT Proto Frontend/Proxy
1import os
2import struct
3import io
4
5import msgpack
6import msgpack.fallback
7
8
9def msgpackHeader(size):
10 if size <= 2 ** 8 - 1:
11 return b"\xc4" + struct.pack("B", size)
12 elif size <= 2 ** 16 - 1:
13 return b"\xc5" + struct.pack(">H", size)
14 elif size <= 2 ** 32 - 1:
15 return b"\xc6" + struct.pack(">I", size)
16 else:
17 raise Exception("huge binary string")
18
19
20def stream(data, writer):
21 packer = msgpack.Packer(use_bin_type=True)
22 writer(packer.pack_map_header(len(data)))
23 for key, val in data.items():
24 writer(packer.pack(key))
25 if isinstance(val, io.IOBase): # File obj
26 max_size = os.fstat(val.fileno()).st_size - val.tell()
27 size = min(max_size, val.read_bytes)
28 bytes_left = size
29 writer(msgpackHeader(size))
30 buff = 1024 * 64
31 while 1:
32 writer(val.read(min(bytes_left, buff)))
33 bytes_left = bytes_left - buff
34 if bytes_left <= 0:
35 break
36 else: # Simple
37 writer(packer.pack(val))
38 return size
39
40
41class FilePart(object):
42 __slots__ = ("file", "read_bytes", "__class__")
43
44 def __init__(self, *args, **kwargs):
45 self.file = open(*args, **kwargs)
46 self.__enter__ == self.file.__enter__
47
48 def __getattr__(self, attr):
49 return getattr(self.file, attr)
50
51 def __enter__(self, *args, **kwargs):
52 return self.file.__enter__(*args, **kwargs)
53
54 def __exit__(self, *args, **kwargs):
55 return self.file.__exit__(*args, **kwargs)
56
57
58# Don't try to decode the value of these fields as utf8
59bin_value_keys = ("hashfield_raw", "peers", "peers_ipv6", "peers_onion", "body", "sites", "bin")
60
61
62def objectDecoderHook(obj):
63 global bin_value_keys
64 back = {}
65 for key, val in obj:
66 if type(key) is bytes:
67 key = key.decode("utf8")
68 if key in bin_value_keys or type(val) is not bytes or len(key) >= 64:
69 back[key] = val
70 else:
71 back[key] = val.decode("utf8")
72 return back
73
74
75def getUnpacker(fallback=False, decode=True):
76 if fallback: # Pure Python
77 unpacker = msgpack.fallback.Unpacker
78 else:
79 unpacker = msgpack.Unpacker
80
81 extra_kwargs = {"max_buffer_size": 5 * 1024 * 1024}
82 if msgpack.version[0] >= 1:
83 extra_kwargs["strict_map_key"] = False
84
85 if decode: # Workaround for backward compatibility: Try to decode bin to str
86 unpacker = unpacker(raw=True, object_pairs_hook=objectDecoderHook, **extra_kwargs)
87 else:
88 unpacker = unpacker(raw=False, **extra_kwargs)
89
90 return unpacker
91
92
93def pack(data, use_bin_type=True):
94 return msgpack.packb(data, use_bin_type=use_bin_type)
95
96
97def unpack(data, decode=True):
98 unpacker = getUnpacker(decode=decode)
99 unpacker.feed(data)
100 return next(unpacker)
101