Repository navigation
Expand file tree
/
Copy pathpbuf.py
More file actions
206 lines (183 loc) · 6.29 KB
/
Copy pathpbuf.py
File metadata and controls
206 lines (183 loc) · 6.29 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
from __future__ import annotations
import codecs
import enum
import re
from refinery.lib.structures import StructReader
from refinery.lib.types import Param
from refinery.units import Arg
from refinery.units.formats import JSONEncoderUnit
class WireType(enum.IntEnum):
VARINT = 0
I64 = 1
I32 = 5
LEN = 2
SGROUP = 3
EGROUP = 4
class ProtoBufReader(StructReader[memoryview]):
try_repeated = False
def varint(self):
return self.read_7bit_encoded_int(64)
def _same_type(self, a, b) -> bool:
if type(a) is not type(b):
return False
if isinstance(a, dict):
if set(a) != set(b):
return False
for key in a:
if not self._same_type(a[key], b[key]):
return False
return True
def _decode_as_bytes(self, msg):
if msg is None:
return True
if isinstance(msg, (bytes, bytearray, memoryview, float)):
return True
if isinstance(msg, list):
return len(msg) <= 1
if isinstance(msg, dict):
if len(msg) > 1:
return False
if not msg:
return True
return self._decode_as_bytes(next(iter(msg.values())))
else:
return False
def _as_map(self, value: list[dict]):
if not isinstance(value, list):
return value
switch = False
as_map = {}
key_nr = None
val_nr = None
key_01 = None
val_01 = None
for entry in value:
if not isinstance(entry, dict):
return value
if not len(entry) == 2:
return value
if not as_map:
key_nr, val_nr = entry.keys()
try:
key = entry[key_nr]
val = entry[val_nr]
except KeyError:
return value
if not as_map:
key_01 = key
val_01 = val
elif not self._same_type(key, key_01):
return value
elif not self._same_type(val, val_01):
return value
if key not in as_map:
as_map[key] = val
continue
elif switch:
break
else:
switch = True
key_nr, val_nr = val_nr, key_nr
key_01, val_01 = val_01, key_01
key, val = val, key
temp_map = {}
for k, v in as_map.items():
if v in temp_map:
return value
temp_map[v] = k
if key in temp_map:
return value
temp_map[key] = val
as_map = temp_map
else:
return as_map
def read_key_value_pair(self):
nr, wt = divmod(self.varint(), 8)
return nr, WireType(wt)
def read_message(self, gid: int | None = None):
def insert(key, val):
if key in msg:
box = msg[key]
if isinstance(box, list):
box.append(val)
else:
msg[key] = [box, val]
else:
msg[key] = val
msg = {}
while not self.eof:
nr, wt = self.read_key_value_pair()
if nr not in range(1, 536_870_911):
raise ValueError
if wt == WireType.EGROUP:
if nr == gid:
break
raise ValueError
if wt == WireType.SGROUP:
insert(nr, self.read_message(nr))
elif wt == WireType.VARINT:
insert(nr, self.varint())
elif wt == WireType.I64:
insert(nr, self.f64())
elif wt == WireType.I32:
insert(nr, self.f32())
elif wt == WireType.LEN:
size = self.varint()
blob = self.read_exactly(size)
data = blob
wire = ProtoBufReader(blob)
try:
if re.fullmatch(R'[\s!-~]+', data := codecs.decode(blob, 'utf8')):
insert(nr, data)
continue
except UnicodeDecodeError:
pass
try:
data = wire.read_message()
except Exception:
if self.try_repeated and any(b & 0x80 for b in blob):
wire.seekset(0)
data = []
try:
while not wire.eof:
data.append(wire.varint())
except (EOFError, OverflowError):
data = blob
if self._decode_as_bytes(data):
data = blob
insert(nr, data)
else:
raise TypeError
for nr, value in msg.items():
msg[nr] = self._as_map(value)
return msg
class pbuf(JSONEncoderUnit):
"""
Convert ProtoBuf messages to JSON.
Deserialization is ambiguous without the definition file, so the output is partly based on
heuristics. Some fields like fixed integers are never recovered, fixed 32-bit and 64-bit
data types are always recovered as floating point numbers. For variable length data, the
unit first attempts to decode the data as a printable UTF-8 string. If this fails, it will
attempt to deserialize it as ProtoBuf. If this also fails and the corresponding option is
set, it will try to reconstruct a sequence of repeated variable length integers. The final
fallback is to return the body as a byte string.
"""
def __init__(
self,
try_repeated: Param[bool, Arg.Switch('-r',
help='Try to detect and decode repeated integer fields.')] = False,
encode=None,
digest=None,
arrays=False,
):
super().__init__(
encode=encode,
digest=digest,
arrays=arrays,
try_repeated=try_repeated
)
def process(self, data):
reader = ProtoBufReader(memoryview(data))
reader.try_repeated = self.args.try_repeated
message = reader.read_message()
return self.to_json(message)