init

pwr/formats.py

@@ -0,0 +1,286 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Callable, Iterator
+
+from .compression import open_decompressed_reader
+from .proto import (
+    BlockHash,
+    BsdiffHeader,
+    Control,
+    ManifestBlockHash,
+    ManifestHeader,
+    PatchHeader,
+    SignatureHeader,
+    SyncHeader,
+    SyncHeaderType,
+    SyncOp,
+    SyncOpType,
+    Wound,
+    WoundsHeader,
+)
+from .wire import (
+    MANIFEST_MAGIC,
+    PATCH_MAGIC,
+    SIGNATURE_MAGIC,
+    WOUNDS_MAGIC,
+    WireError,
+    WireReader,
+)
+
+ContainerDecoder = Callable[[bytes], Any]
+
+
+def _split_decoder(decoder: ContainerDecoder | tuple[ContainerDecoder | None, ContainerDecoder | None] | None):
+    if decoder is None:
+        return None, None
+    if isinstance(decoder, tuple):
+        if len(decoder) != 2:
+            raise ValueError("container decoder tuple must have two elements")
+        return decoder
+    return decoder, decoder
+
+
+@dataclass
+class FilePatch:
+    sync_header: SyncHeader
+    sync_ops: list[SyncOp] | None = None
+    bsdiff_header: BsdiffHeader | None = None
+    bsdiff_controls: list[Control] | None = None
+
+    def is_rsync(self) -> bool:
+        return self.sync_ops is not None
+
+    def is_bsdiff(self) -> bool:
+        return self.bsdiff_header is not None
+
+
+class PatchReader:
+    def __init__(self, stream, *, container_decoder: ContainerDecoder | tuple[ContainerDecoder | None, ContainerDecoder | None] | None = None):
+        self._stream = stream
+        self._raw_wire = WireReader(stream)
+        self._wire: WireReader | None = None
+
+        self.header: PatchHeader | None = None
+        self.target_container_raw: bytes | None = None
+        self.source_container_raw: bytes | None = None
+        self.target_container: Any | None = None
+        self.source_container: Any | None = None
+
+        self._init_stream(container_decoder)
+
+    @classmethod
+    def open(cls, path: str, *, container_decoder: ContainerDecoder | tuple[ContainerDecoder | None, ContainerDecoder | None] | None = None):
+        return cls(open(path, "rb"), container_decoder=container_decoder)
+
+    def _init_stream(self, container_decoder):
+        self._raw_wire.expect_magic(PATCH_MAGIC)
+        self.header = self._raw_wire.read_message(PatchHeader)
+
+        decompressed = open_decompressed_reader(self._stream, self.header.compression)
+        self._wire = WireReader(decompressed)
+
+        target_decoder, source_decoder = _split_decoder(container_decoder)
+        self.target_container_raw, self.target_container = self._read_container(target_decoder)
+        self.source_container_raw, self.source_container = self._read_container(source_decoder)
+
+    def _read_container(self, decoder: ContainerDecoder | None):
+        assert self._wire is not None
+        raw = self._wire.read_message_bytes()
+        parsed = decoder(raw) if decoder else None
+        return raw, parsed
+
+    def iter_file_entries(self) -> Iterator[FilePatch]:
+        assert self._wire is not None
+        file_index = 0
+        while True:
+            try:
+                sync_header = self._wire.read_message(SyncHeader)
+            except EOFError:
+                return
+
+            if sync_header.file_index is None:
+                sync_header.file_index = file_index
+            else:
+                file_index = int(sync_header.file_index)
+
+            header_type = sync_header.type
+            if header_type is None:
+                header_type = SyncHeaderType.RSYNC
+            else:
+                try:
+                    header_type = SyncHeaderType(int(header_type))
+                except ValueError:
+                    raise WireError(f"unknown sync header type: {header_type}")
+
+            if header_type == SyncHeaderType.BSDIFF:
+                bsdiff_header = self._wire.read_message(BsdiffHeader)
+                controls: list[Control] = []
+                while True:
+                    ctrl = self._wire.read_message(Control)
+                    controls.append(ctrl)
+                    if ctrl.eof:
+                        break
+
+                sentinel = self._wire.read_message(SyncOp)
+                if sentinel.type != SyncOpType.HEY_YOU_DID_IT:
+                    raise WireError("expected BSDIFF sentinel SyncOp")
+
+                yield FilePatch(
+                    sync_header=sync_header,
+                    bsdiff_header=bsdiff_header,
+                    bsdiff_controls=controls,
+                )
+                file_index += 1
+                continue
+
+            if header_type != SyncHeaderType.RSYNC:
+                raise WireError(f"unsupported sync header type: {header_type}")
+
+            ops: list[SyncOp] = []
+            while True:
+                op = self._wire.read_message(SyncOp)
+                if op.type == SyncOpType.HEY_YOU_DID_IT:
+                    break
+                ops.append(op)
+
+            yield FilePatch(sync_header=sync_header, sync_ops=ops)
+            file_index += 1
+
+    def close(self) -> None:
+        if self._wire:
+            self._wire.close()
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc, tb):
+        self.close()
+
+
+class SignatureReader:
+    def __init__(self, stream, *, container_decoder: ContainerDecoder | None = None):
+        self._stream = stream
+        self._raw_wire = WireReader(stream)
+        self._wire: WireReader | None = None
+
+        self.header: SignatureHeader | None = None
+        self.container_raw: bytes | None = None
+        self.container: Any | None = None
+
+        self._init_stream(container_decoder)
+
+    @classmethod
+    def open(cls, path: str, *, container_decoder: ContainerDecoder | None = None):
+        return cls(open(path, "rb"), container_decoder=container_decoder)
+
+    def _init_stream(self, container_decoder):
+        self._raw_wire.expect_magic(SIGNATURE_MAGIC)
+        self.header = self._raw_wire.read_message(SignatureHeader)
+
+        decompressed = open_decompressed_reader(self._stream, self.header.compression)
+        self._wire = WireReader(decompressed)
+
+        self.container_raw = self._wire.read_message_bytes()
+        self.container = container_decoder(self.container_raw) if container_decoder else None
+
+    def iter_block_hashes(self) -> Iterator[BlockHash]:
+        assert self._wire is not None
+        while True:
+            try:
+                yield self._wire.read_message(BlockHash)
+            except EOFError:
+                return
+
+    def close(self) -> None:
+        if self._wire:
+            self._wire.close()
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc, tb):
+        self.close()
+
+
+class ManifestReader:
+    def __init__(self, stream, *, container_decoder: ContainerDecoder | None = None):
+        self._stream = stream
+        self._raw_wire = WireReader(stream)
+        self._wire: WireReader | None = None
+
+        self.header: ManifestHeader | None = None
+        self.container_raw: bytes | None = None
+        self.container: Any | None = None
+
+        self._init_stream(container_decoder)
+
+    @classmethod
+    def open(cls, path: str, *, container_decoder: ContainerDecoder | None = None):
+        return cls(open(path, "rb"), container_decoder=container_decoder)
+
+    def _init_stream(self, container_decoder):
+        self._raw_wire.expect_magic(MANIFEST_MAGIC)
+        self.header = self._raw_wire.read_message(ManifestHeader)
+
+        decompressed = open_decompressed_reader(self._stream, self.header.compression)
+        self._wire = WireReader(decompressed)
+
+        self.container_raw = self._wire.read_message_bytes()
+        self.container = container_decoder(self.container_raw) if container_decoder else None
+
+    def iter_block_hashes(self) -> Iterator[ManifestBlockHash]:
+        assert self._wire is not None
+        while True:
+            try:
+                yield self._wire.read_message(ManifestBlockHash)
+            except EOFError:
+                return
+
+    def close(self) -> None:
+        if self._wire:
+            self._wire.close()
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc, tb):
+        self.close()
+
+
+class WoundsReader:
+    def __init__(self, stream, *, container_decoder: ContainerDecoder | None = None):
+        self._stream = stream
+        self._wire = WireReader(stream)
+
+        self.header: WoundsHeader | None = None
+        self.container_raw: bytes | None = None
+        self.container: Any | None = None
+
+        self._init_stream(container_decoder)
+
+    @classmethod
+    def open(cls, path: str, *, container_decoder: ContainerDecoder | None = None):
+        return cls(open(path, "rb"), container_decoder=container_decoder)
+
+    def _init_stream(self, container_decoder):
+        self._wire.expect_magic(WOUNDS_MAGIC)
+        self.header = self._wire.read_message(WoundsHeader)
+        self.container_raw = self._wire.read_message_bytes()
+        self.container = container_decoder(self.container_raw) if container_decoder else None
+
+    def iter_wounds(self) -> Iterator[Wound]:
+        while True:
+            try:
+                yield self._wire.read_message(Wound)
+            except EOFError:
+                return
+
+    def close(self) -> None:
+        self._wire.close()
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc, tb):
+        self.close()