/
/
/
1"""
2Lightweight protobuf encoder/decoder for Yandex externalCommandBypass.
3
4Adapted from AlexxIT/YandexStation (MIT license).
5"""
6
7from __future__ import annotations
8
9import base64
10
11
12class Protobuf:
13 """Minimal protobuf wire-format parser."""
14
15 def __init__(self, raw: str | bytes) -> None:
16 """Initialize from base64 string or raw bytes."""
17 self.raw = base64.b64decode(raw) if isinstance(raw, str) else raw
18 self.pos = 0
19
20 def read(self, length: int) -> bytes:
21 """Read exactly `length` bytes from buffer."""
22 if self.pos + length > len(self.raw):
23 msg = "Read past end of buffer"
24 raise IndexError(msg)
25 self.pos += length
26 return self.raw[self.pos - length : self.pos]
27
28 def read_byte(self) -> int:
29 """Read a single byte."""
30 res = self.raw[self.pos]
31 self.pos += 1
32 return res
33
34 def read_varint(self) -> int:
35 """Read a variable-length integer."""
36 res = 0
37 shift = 0
38 while True:
39 b = self.read_byte()
40 res += (b & 0x7F) << shift
41 if b & 0x80 == 0:
42 break
43 shift += 7
44 return res
45
46 def read_bytes(self) -> bytes:
47 """Read a length-prefixed byte string."""
48 length = self.read_varint()
49 return self.read(length)
50
51 def read_dict(self) -> dict[int, object]:
52 """Parse the buffer into a tagâvalue dict."""
53 res: dict[int, object] = {}
54 while self.pos < len(self.raw):
55 b = self.read_varint()
56 typ = b & 0b111
57 tag = b >> 3
58
59 if typ == 0: # VARINT
60 v: object = self.read_varint()
61 elif typ == 1: # I64
62 v = self.read(8)
63 elif typ == 2: # LEN
64 raw_bytes = self.read_bytes()
65 try:
66 nested = Protobuf(raw_bytes)
67 parsed = nested.read_dict()
68 # Only accept as nested message if ALL bytes were consumed
69 v = parsed if nested.pos == len(nested.raw) and parsed else raw_bytes
70 except Exception:
71 v = raw_bytes
72 elif typ == 5: # I32
73 v = self.read(4)
74 else:
75 msg = f"Unsupported protobuf wire type: {typ}"
76 raise NotImplementedError(msg)
77
78 if tag in res:
79 existing = res[tag]
80 if isinstance(existing, list):
81 existing.append(v)
82 else:
83 res[tag] = [existing, v]
84 else:
85 res[tag] = v
86
87 return res
88
89
90def _append_varint(b: bytearray, i: int) -> None:
91 while i >= 0x80:
92 b.append(0x80 | (i & 0x7F))
93 i >>= 7
94 b.append(i)
95
96
97def loads(raw: str | bytes) -> dict[int, object]:
98 """Decode protobuf wire format to dict."""
99 return Protobuf(raw).read_dict()
100
101
102def dumps(data: dict[int, str]) -> bytes:
103 """Encode dict to protobuf wire format (string values only)."""
104 b = bytearray()
105 for tag, value in data.items():
106 _append_varint(b, tag << 3 | 2)
107 encoded = value.encode()
108 _append_varint(b, len(encoded))
109 b.extend(encoded)
110 return bytes(b)
111