python-pages 0.1.0__py3-none-any.whl

pages/protobuf.py

@@ -0,0 +1,162 @@
+"""Small lossless protobuf wire parser for schema discovery and mutation."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Iterable
+
+
+class ProtobufError(ValueError):
+    pass
+
+
+def read_varint(data: bytes, pos: int = 0) -> tuple[int, int]:
+    value = 0
+    shift = 0
+    for _ in range(10):
+        if pos >= len(data):
+            raise ProtobufError("truncated varint")
+        byte = data[pos]
+        pos += 1
+        value |= (byte & 0x7F) << shift
+        if byte < 0x80:
+            return value, pos
+        shift += 7
+    raise ProtobufError("varint is longer than 10 bytes")
+
+
+def encode_varint(value: int) -> bytes:
+    if value < 0:
+        value &= (1 << 64) - 1
+    out = bytearray()
+    while value >= 0x80:
+        out.append((value & 0x7F) | 0x80)
+        value >>= 7
+    out.append(value)
+    return bytes(out)
+
+
+@dataclass
+class Field:
+    number: int
+    wire_type: int
+    value: int | bytes
+    original: bytes = b""
+    dirty: bool = False
+    children: "Message | None" = None
+
+    def encode(self) -> bytes:
+        if not self.dirty and self.original:
+            return self.original
+        key = encode_varint((self.number << 3) | self.wire_type)
+        if self.wire_type == 0:
+            return key + encode_varint(int(self.value))
+        if self.wire_type == 1:
+            return key + int(self.value).to_bytes(8, "little")
+        if self.wire_type == 2:
+            payload = self.children.encode() if self.children is not None else bytes(self.value)
+            return key + encode_varint(len(payload)) + payload
+        if self.wire_type == 5:
+            return key + int(self.value).to_bytes(4, "little")
+        raise ProtobufError(f"unsupported wire type {self.wire_type}")
+
+    def set_bytes(self, value: bytes) -> None:
+        if self.wire_type != 2:
+            raise TypeError("not a length-delimited field")
+        self.value = value
+        self.children = None
+        self.dirty = True
+
+    def set_varint(self, value: int) -> None:
+        if self.wire_type != 0:
+            raise TypeError("not a varint field")
+        self.value = value
+        self.dirty = True
+
+    def set_fixed32(self, value: int) -> None:
+        if self.wire_type != 5:
+            raise TypeError("not a fixed32 field")
+        self.value = value
+        self.dirty = True
+
+
+@dataclass
+class Message:
+    fields: list[Field] = field(default_factory=list)
+
+    @classmethod
+    def parse(cls, data: bytes) -> "Message":
+        fields: list[Field] = []
+        pos = 0
+        while pos < len(data):
+            start = pos
+            key, pos = read_varint(data, pos)
+            number, wire = key >> 3, key & 7
+            if number == 0:
+                raise ProtobufError(f"field zero at offset {start}")
+            if wire == 0:
+                value, pos = read_varint(data, pos)
+            elif wire == 1:
+                if pos + 8 > len(data):
+                    raise ProtobufError("truncated fixed64")
+                value = int.from_bytes(data[pos : pos + 8], "little")
+                pos += 8
+            elif wire == 2:
+                length, pos = read_varint(data, pos)
+                if pos + length > len(data):
+                    raise ProtobufError("truncated length-delimited field")
+                value = data[pos : pos + length]
+                pos += length
+            elif wire == 5:
+                if pos + 4 > len(data):
+                    raise ProtobufError("truncated fixed32")
+                value = int.from_bytes(data[pos : pos + 4], "little")
+                pos += 4
+            else:
+                raise ProtobufError(f"unsupported wire type {wire} at {start}")
+            fields.append(Field(number, wire, value, data[start:pos]))
+        return cls(fields)
+
+    def encode(self) -> bytes:
+        return b"".join(f.encode() for f in self.fields)
+
+    def get(self, number: int) -> list[Field]:
+        return [f for f in self.fields if f.number == number]
+
+    def one(self, number: int, default=None):
+        items = self.get(number)
+        return items[0] if items else default
+
+    def add_varint(self, number: int, value: int) -> Field:
+        result = Field(number, 0, value, dirty=True)
+        self.fields.append(result)
+        return result
+
+    def add_bytes(self, number: int, value: bytes) -> Field:
+        result = Field(number, 2, value, dirty=True)
+        self.fields.append(result)
+        return result
+
+    def add_fixed32(self, number: int, value: int) -> Field:
+        result = Field(number, 5, value, dirty=True)
+        self.fields.append(result)
+        return result
+
+    def iter_strings(self, min_length: int = 2) -> Iterable[tuple[Field, str]]:
+        for item in self.fields:
+            if item.wire_type != 2 or len(item.value) < min_length:
+                continue
+            try:
+                text = bytes(item.value).decode("utf-8")
+            except UnicodeDecodeError:
+                continue
+            if text.isprintable() or any(c in text for c in "\n\t"):
+                yield item, text
+
+
+def try_message(data: bytes) -> Message | None:
+    try:
+        message = Message.parse(data)
+    except ProtobufError:
+        return None
+    return message if message.fields else None

pages/snappy.py

@@ -0,0 +1,140 @@
+"""Raw Snappy and iWork's small custom framing.
+
+iWork uses blocks consisting of byte 0 followed by a 24-bit little-endian
+compressed length and a raw (not standard framed-stream) Snappy block.
+"""
+
+from __future__ import annotations
+
+
+class SnappyError(ValueError):
+    pass
+
+
+def _read_varint(data: bytes, pos: int = 0) -> tuple[int, int]:
+    value = 0
+    shift = 0
+    while pos < len(data) and shift < 35:
+        byte = data[pos]
+        pos += 1
+        value |= (byte & 0x7F) << shift
+        if byte < 0x80:
+            return value, pos
+        shift += 7
+    raise SnappyError("invalid Snappy length varint")
+
+
+def _varint(value: int) -> bytes:
+    out = bytearray()
+    while value >= 0x80:
+        out.append((value & 0x7F) | 0x80)
+        value >>= 7
+    out.append(value)
+    return bytes(out)
+
+
+def decompress_raw(data: bytes) -> bytes:
+    expected, pos = _read_varint(data)
+    out = bytearray()
+    while pos < len(data):
+        tag = data[pos]
+        pos += 1
+        kind = tag & 3
+        if kind == 0:
+            n = tag >> 2
+            if n < 60:
+                length = n + 1
+            else:
+                width = n - 59
+                if width > 4 or pos + width > len(data):
+                    raise SnappyError("invalid literal length")
+                length = int.from_bytes(data[pos : pos + width], "little") + 1
+                pos += width
+            if pos + length > len(data):
+                raise SnappyError("truncated literal")
+            out += data[pos : pos + length]
+            pos += length
+            continue
+
+        if kind == 1:
+            if pos >= len(data):
+                raise SnappyError("truncated COPY_1")
+            length = ((tag >> 2) & 7) + 4
+            offset = ((tag & 0xE0) << 3) | data[pos]
+            pos += 1
+        elif kind == 2:
+            if pos + 2 > len(data):
+                raise SnappyError("truncated COPY_2")
+            length = (tag >> 2) + 1
+            offset = int.from_bytes(data[pos : pos + 2], "little")
+            pos += 2
+        else:
+            if pos + 4 > len(data):
+                raise SnappyError("truncated COPY_4")
+            length = (tag >> 2) + 1
+            offset = int.from_bytes(data[pos : pos + 4], "little")
+            pos += 4
+        if offset <= 0 or offset > len(out):
+            raise SnappyError(f"invalid copy offset {offset}")
+        for _ in range(length):
+            out.append(out[-offset])
+
+    if len(out) != expected:
+        raise SnappyError(f"decoded {len(out)} bytes, expected {expected}")
+    return bytes(out)
+
+
+def compress_raw_literal(data: bytes) -> bytes:
+    """Create a valid raw Snappy block using literals only.
+
+    It is deliberately simple: ZIP will deflate the outer .pages archive, and
+    correctness/preservation is more important here than IWA-level ratio.
+    """
+    out = bytearray(_varint(len(data)))
+    pos = 0
+    while pos < len(data):
+        length = min(len(data) - pos, 65536)
+        n = length - 1
+        if n < 60:
+            out.append(n << 2)
+        else:
+            width = max(1, (n.bit_length() + 7) // 8)
+            out.append((59 + width) << 2)
+            out += n.to_bytes(width, "little")
+        out += data[pos : pos + length]
+        pos += length
+    return bytes(out)
+
+
+def decompress_iwork(data: bytes) -> tuple[bytes, list[int]]:
+    out = bytearray()
+    decoded_sizes: list[int] = []
+    pos = 0
+    while pos < len(data):
+        if pos + 4 > len(data) or data[pos] != 0:
+            raise SnappyError(f"bad iWork frame at offset {pos}")
+        length = int.from_bytes(data[pos + 1 : pos + 4], "little")
+        pos += 4
+        if pos + length > len(data):
+            raise SnappyError("truncated iWork frame")
+        block = decompress_raw(data[pos : pos + length])
+        out += block
+        decoded_sizes.append(len(block))
+        pos += length
+    return bytes(out), decoded_sizes
+
+
+def compress_iwork(data: bytes, decoded_sizes: list[int] | None = None) -> bytes:
+    if not decoded_sizes or sum(decoded_sizes) != len(data):
+        decoded_sizes = [min(65536, len(data) - pos) for pos in range(0, len(data), 65536)]
+    out = bytearray()
+    pos = 0
+    for size in decoded_sizes:
+        compressed = compress_raw_literal(data[pos : pos + size])
+        if len(compressed) >= 1 << 24:
+            raise SnappyError("compressed iWork block exceeds 24-bit length")
+        out.append(0)
+        out += len(compressed).to_bytes(3, "little")
+        out += compressed
+        pos += size
+    return bytes(out)