python-pages 0.1.0__py3-none-any.whl

pages/pictures.py

@@ -0,0 +1,549 @@
+"""Middle-layer access to inline TSD image attachments and package media."""
+
+from __future__ import annotations
+
+from hashlib import sha1
+import binascii
+from pathlib import Path
+import struct
+from typing import TYPE_CHECKING
+import zlib
+
+from .components import GraphClone, clone_object_graph, remap_payload_data_references
+from .metadata import component_name
+from .protobuf import Message
+from .text import LocatedObject, parse_attribute_table, parse_reference
+
+if TYPE_CHECKING:
+    from .pages import PagesDocument
+
+
+TYPE_ATTACHMENT = 2003
+TYPE_IMAGE = 3005
+
+STORAGE_TABLE_ATTACHMENTS = 9
+IMAGE_DATA = 11
+IMAGE_THUMBNAIL_DATA = 12
+IMAGE_NATURAL_SIZE = 9
+
+PACKAGE_DATAS = 4
+DATA_IDENTIFIER = 1
+DATA_DIGEST = 2
+DATA_PREFERRED_FILENAME = 3
+DATA_FILENAME = 4
+DATA_MATERIALIZED_LENGTH = 18
+DATA_ATTRIBUTES = 10
+DATA_ATTRIBUTES_IMAGE = 100
+IMAGE_ATTRIBUTES_PIXEL_SIZE = 1
+
+
+def _reference(message: Message, number: int) -> int | None:
+    field = message.one(number)
+    return None if field is None else parse_reference(bytes(field.value))
+
+
+class ImageStore:
+    """Resolved ImageArchive plus its full-size and thumbnail DataInfo entries."""
+
+    def __init__(self, document: "PagesDocument", image: LocatedObject):
+        self.document = document
+        self.image = image
+        self.data_identifier = _reference(image.message, IMAGE_DATA)
+        self.thumbnail_identifier = _reference(image.message, IMAGE_THUMBNAIL_DATA)
+        if self.data_identifier is None:
+            raise ValueError(f"TSD.ImageArchive {image.segment.identifier} has no image data")
+
+    def _data_info(self, identifier: int):
+        package = self.document.package_metadata()
+        for raw in package.message.get(PACKAGE_DATAS):
+            info = Message.parse(bytes(raw.value))
+            field = info.one(DATA_IDENTIFIER)
+            if field is not None and int(field.value) == identifier:
+                return raw, info
+        raise KeyError(f"PackageMetadata has no DataInfo {identifier}")
+
+    def _member_name(self, identifier: int) -> str:
+        _, info = self._data_info(identifier)
+        filename = info.one(DATA_FILENAME) or info.one(DATA_PREFERRED_FILENAME)
+        if filename is None:
+            raise ValueError(f"DataInfo {identifier} has no package filename")
+        return "Data/" + bytes(filename.value).decode("utf-8", "replace")
+
+    @property
+    def filename(self) -> str:
+        _, info = self._data_info(self.data_identifier)
+        filename = info.one(DATA_PREFERRED_FILENAME) or info.one(DATA_FILENAME)
+        return "" if filename is None else bytes(filename.value).decode("utf-8", "replace")
+
+    @property
+    def blob(self) -> bytes:
+        name = self._member_name(self.data_identifier)
+        data = self.document.package_data(name)
+        if data is None:
+            raise KeyError(f"image package member {name!r} is missing")
+        return data
+
+    @property
+    def thumbnail_blob(self) -> bytes | None:
+        if self.thumbnail_identifier is None:
+            return None
+        name = self._member_name(self.thumbnail_identifier)
+        return self.document.package_data(name)
+
+    @property
+    def natural_size(self) -> tuple[float, float] | None:
+        field = self.image.message.one(IMAGE_NATURAL_SIZE)
+        if field is None:
+            return None
+        size = Message.parse(bytes(field.value))
+        width, height = size.one(1), size.one(2)
+        if width is None or height is None or width.wire_type != 5 or height.wire_type != 5:
+            return None
+        return tuple(
+            struct.unpack("<f", int(item.value).to_bytes(4, "little"))[0]
+            for item in (width, height)
+        )
+
+    @staticmethod
+    def _float_bits(value: float) -> int:
+        return int.from_bytes(struct.pack("<f", float(value)), "little")
+
+    def _set_natural_size(self, width: int, height: int) -> None:
+        field = self.image.message.one(IMAGE_NATURAL_SIZE)
+        size = Message.parse(bytes(field.value)) if field is not None else Message()
+        for number, value in ((1, width), (2, height)):
+            item = size.one(number)
+            if item is None:
+                size.add_fixed32(number, self._float_bits(value))
+            else:
+                item.set_fixed32(self._float_bits(value))
+        if field is None:
+            self.image.message.add_bytes(IMAGE_NATURAL_SIZE, size.encode())
+        else:
+            field.set_bytes(size.encode())
+
+    @classmethod
+    def _set_pixel_size(cls, info: Message, width: int, height: int) -> None:
+        attributes_field = info.one(DATA_ATTRIBUTES)
+        if attributes_field is None:
+            return
+        attributes = Message.parse(bytes(attributes_field.value))
+        image_field = attributes.one(DATA_ATTRIBUTES_IMAGE)
+        if image_field is None:
+            return
+        image_attributes = Message.parse(bytes(image_field.value))
+        size_field = image_attributes.one(IMAGE_ATTRIBUTES_PIXEL_SIZE)
+        if size_field is None:
+            return
+        size = Message.parse(bytes(size_field.value))
+        for number, value in ((1, width), (2, height)):
+            item = size.one(number)
+            if item is None:
+                size.add_fixed32(number, cls._float_bits(value))
+            else:
+                item.set_fixed32(cls._float_bits(value))
+        size_field.set_bytes(size.encode())
+        image_field.set_bytes(image_attributes.encode())
+        attributes_field.set_bytes(attributes.encode())
+
+    def _replace_data(
+        self, identifier: int, data: bytes, pixel_size: tuple[int, int] | None = None
+    ) -> None:
+        raw_info, info = self._data_info(identifier)
+        if pixel_size is not None:
+            self._set_pixel_size(info, *pixel_size)
+        digest = info.one(DATA_DIGEST)
+        if digest is None:
+            info.add_bytes(DATA_DIGEST, sha1(data).digest())
+        else:
+            digest.set_bytes(sha1(data).digest())
+        length = info.one(DATA_MATERIALIZED_LENGTH)
+        if length is None:
+            info.add_varint(DATA_MATERIALIZED_LENGTH, len(data))
+        else:
+            length.set_varint(len(data))
+        raw_info.set_bytes(info.encode())
+        self.document.set_package_data(self._member_name(identifier), data)
+
+    @staticmethod
+    def _png_size(data: bytes) -> tuple[int, int]:
+        if not data.startswith(b"\x89PNG\r\n\x1a\n"):
+            raise ValueError("fixture-safe image replacement currently requires PNG data")
+        if len(data) < 24 or data[12:16] != b"IHDR":
+            raise ValueError("PNG replacement has no complete IHDR")
+        pixel_size = struct.unpack(">II", data[16:24])
+        if not all(pixel_size):
+            raise ValueError("PNG replacement dimensions must be non-zero")
+        return pixel_size
+
+    def replace(self, data: bytes, thumbnail_data: bytes | None = None) -> None:
+        pixel_size = self._png_size(data)
+        selected_thumbnail = thumbnail_data
+        if self.thumbnail_identifier is not None and selected_thumbnail is None:
+            # A small source is already a valid editing thumbnail. For larger
+            # sources, preserve Pages' existing 256px thumbnail rather than
+            # violating the distinct full/thumbnail dimension contract.
+            if max(pixel_size) <= 256:
+                selected_thumbnail = data
+        if selected_thumbnail is not None:
+            thumbnail_size = self._png_size(selected_thumbnail)
+            scale = min(1.0, 256.0 / max(pixel_size))
+            expected = tuple(max(1, round(value * scale)) for value in pixel_size)
+            if thumbnail_size != expected:
+                raise ValueError(
+                    f"thumbnail dimensions {thumbnail_size} do not match expected {expected}"
+                )
+        self._set_natural_size(*pixel_size)
+        self._replace_data(self.data_identifier, data, pixel_size)
+        if self.thumbnail_identifier is not None and selected_thumbnail is not None:
+            self._replace_data(self.thumbnail_identifier, selected_thumbnail)
+
+
+def _set_varint(message: Message, number: int, value: int) -> None:
+    field = message.one(number)
+    if field is None:
+        message.add_varint(number, value)
+    else:
+        field.set_varint(value)
+
+
+def _set_bytes(message: Message, number: int, value: bytes) -> None:
+    field = message.one(number)
+    if field is None:
+        message.add_bytes(number, value)
+    else:
+        field.set_bytes(value)
+
+
+def _png_chunks(data: bytes) -> list[tuple[bytes, bytes]]:
+    ImageStore._png_size(data)
+    chunks: list[tuple[bytes, bytes]] = []
+    position = 8
+    while position + 12 <= len(data):
+        length = struct.unpack(">I", data[position : position + 4])[0]
+        kind = data[position + 4 : position + 8]
+        end = position + 12 + length
+        if end > len(data):
+            raise ValueError("PNG has a truncated chunk")
+        payload = data[position + 8 : position + 8 + length]
+        chunks.append((kind, payload))
+        position = end
+        if kind == b"IEND":
+            break
+    return chunks
+
+
+def _paeth(left: int, above: int, upper_left: int) -> int:
+    estimate = left + above - upper_left
+    distances = (
+        abs(estimate - left),
+        abs(estimate - above),
+        abs(estimate - upper_left),
+    )
+    return (left, above, upper_left)[distances.index(min(distances))]
+
+
+def _thumbnail_png(data: bytes) -> bytes:
+    """Create a dependency-free nearest-neighbor PNG with a 256px maximum."""
+    width, height = ImageStore._png_size(data)
+    if max(width, height) <= 256:
+        return data
+    chunks = _png_chunks(data)
+    ihdr = next(payload for kind, payload in chunks if kind == b"IHDR")
+    _, _, depth, color_type, compression, filtering, interlace = struct.unpack(
+        ">IIBBBBB", ihdr
+    )
+    if depth != 8 or color_type not in (2, 6) or compression or filtering or interlace:
+        raise ValueError(
+            "automatic thumbnail creation supports non-interlaced 8-bit RGB/RGBA PNG"
+        )
+    channels = 3 if color_type == 2 else 4
+    encoded = zlib.decompress(
+        b"".join(payload for kind, payload in chunks if kind == b"IDAT")
+    )
+    stride = width * channels
+    rows: list[bytearray] = []
+    position = 0
+    previous = bytearray(stride)
+    for _ in range(height):
+        filter_type = encoded[position]
+        position += 1
+        raw = encoded[position : position + stride]
+        position += stride
+        row = bytearray(stride)
+        for index, value in enumerate(raw):
+            left = row[index - channels] if index >= channels else 0
+            above = previous[index]
+            upper_left = previous[index - channels] if index >= channels else 0
+            if filter_type == 0:
+                predictor = 0
+            elif filter_type == 1:
+                predictor = left
+            elif filter_type == 2:
+                predictor = above
+            elif filter_type == 3:
+                predictor = (left + above) // 2
+            elif filter_type == 4:
+                predictor = _paeth(left, above, upper_left)
+            else:
+                raise ValueError(f"PNG uses unknown scanline filter {filter_type}")
+            row[index] = (value + predictor) & 0xFF
+        rows.append(row)
+        previous = row
+
+    scale = 256.0 / max(width, height)
+    target_width = max(1, round(width * scale))
+    target_height = max(1, round(height * scale))
+    resized = bytearray()
+    for y in range(target_height):
+        source_y = min(height - 1, int(y * height / target_height))
+        resized.append(0)
+        row = rows[source_y]
+        for x in range(target_width):
+            source_x = min(width - 1, int(x * width / target_width))
+            start = source_x * channels
+            resized.extend(row[start : start + channels])
+
+    def chunk(kind: bytes, payload: bytes) -> bytes:
+        return (
+            struct.pack(">I", len(payload))
+            + kind
+            + payload
+            + struct.pack(">I", binascii.crc32(kind + payload) & 0xFFFFFFFF)
+        )
+
+    result = bytearray(b"\x89PNG\r\n\x1a\n")
+    result.extend(
+        chunk(
+            b"IHDR",
+            struct.pack(
+                ">IIBBBBB", target_width, target_height, 8, color_type, 0, 0, 0
+            ),
+        )
+    )
+    for kind, payload in chunks:
+        if kind in (b"pHYs", b"sRGB", b"gAMA", b"cHRM", b"iCCP"):
+            result.extend(chunk(kind, payload))
+    result.extend(chunk(b"IDAT", zlib.compress(bytes(resized))))
+    result.extend(chunk(b"IEND", b""))
+    return bytes(result)
+
+
+def _data_info(document: "PagesDocument", identifier: int) -> Message:
+    for raw in document.package_metadata().message.get(PACKAGE_DATAS):
+        info = Message.parse(bytes(raw.value))
+        field = info.one(DATA_IDENTIFIER)
+        if field is not None and int(field.value) == identifier:
+            return info
+    raise KeyError(f"PackageMetadata has no DataInfo {identifier}")
+
+
+def _append_component_data_references(
+    document: "PagesDocument",
+    old_image_identifier: int,
+    new_image_identifier: int,
+    data_identifiers: dict[int, int],
+) -> None:
+    package = document.package_metadata().message
+    for raw_component in package.get(3):
+        component = Message.parse(bytes(raw_component.value))
+        if component_name(component) != "Document":
+            continue
+        original = list(component.get(7))
+        for raw in original:
+            reference = Message.parse(bytes(raw.value))
+            data = reference.one(1)
+            if data is None or int(data.value) not in data_identifiers:
+                continue
+            data.set_varint(data_identifiers[int(data.value)])
+            for raw_object in reference.get(2):
+                object_reference = Message.parse(bytes(raw_object.value))
+                identifier = object_reference.one(1)
+                if identifier is not None and int(identifier.value) == old_image_identifier:
+                    identifier.set_varint(new_image_identifier)
+                    raw_object.set_bytes(object_reference.encode())
+            encoded = reference.encode()
+            if all(bytes(item.value) != encoded for item in component.get(7)):
+                component.add_bytes(7, encoded)
+        raw_component.set_bytes(component.encode())
+        return
+    raise ValueError("PackageMetadata has no Document component")
+
+
+def _set_logical_size(image: Message, width: float, height: float) -> None:
+    def set_size(size: Message) -> None:
+        for number, value in ((1, width), (2, height)):
+            bits = ImageStore._float_bits(value)
+            field = size.one(number)
+            if field is None:
+                size.add_fixed32(number, bits)
+            else:
+                field.set_fixed32(bits)
+
+    original = image.one(4)
+    original_size = Message.parse(bytes(original.value)) if original is not None else Message()
+    set_size(original_size)
+    if original is None:
+        image.add_bytes(4, original_size.encode())
+    else:
+        original.set_bytes(original_size.encode())
+
+    drawable_field = image.one(1)
+    if drawable_field is None:
+        return
+    drawable = Message.parse(bytes(drawable_field.value))
+    geometry_field = drawable.one(1)
+    if geometry_field is None:
+        return
+    geometry = Message.parse(bytes(geometry_field.value))
+    size_field = geometry.one(2)
+    size = Message.parse(bytes(size_field.value)) if size_field is not None else Message()
+    set_size(size)
+    if size_field is None:
+        geometry.add_bytes(2, size.encode())
+    else:
+        size_field.set_bytes(size.encode())
+    geometry_field.set_bytes(geometry.encode())
+    drawable_field.set_bytes(drawable.encode())
+
+
+def _logical_size(image: Message) -> tuple[float, float] | None:
+    field = image.one(4)
+    if field is None:
+        return None
+    size = Message.parse(bytes(field.value))
+    values = []
+    for number in (1, 2):
+        item = size.one(number)
+        if item is None or item.wire_type != 5:
+            return None
+        values.append(struct.unpack("<f", int(item.value).to_bytes(4, "little"))[0])
+    return values[0], values[1]
+
+
+def clone_body_image(
+    document: "PagesDocument",
+    storage: LocatedObject,
+    data: bytes,
+    filename: str,
+    *,
+    width: float | None = None,
+    height: float | None = None,
+) -> tuple[ImageStore, GraphClone]:
+    """Clone a body image graph, then install independent full/thumbnail data."""
+    stores = body_image_stores(document, storage)
+    if not stores:
+        raise ValueError("document has no Pages-authored image graph to use as a template")
+    template = stores[0]
+    objects = document.objects()
+    template_attachment = None
+    field = storage.message.one(STORAGE_TABLE_ATTACHMENTS)
+    if field is not None:
+        for run in parse_attribute_table(bytes(field.value)):
+            attachment = objects.get(run.object_id) if run.object_id is not None else None
+            if attachment is None or attachment.info.type_id != TYPE_ATTACHMENT:
+                continue
+            if _reference(attachment.message, 1) == template.image.segment.identifier:
+                template_attachment = attachment
+                break
+    if template_attachment is None:
+        raise ValueError("image template attachment was not found")
+
+    clone = clone_object_graph(
+        document,
+        template_attachment.segment.identifier,
+        include=lambda obj: obj.filename == template.image.filename
+        and obj.info.type_id in (TYPE_IMAGE, 3097),
+    )
+    cloned_image_id = _reference(clone.root.message, 1)
+    cloned_image = document.objects().get(cloned_image_id) if cloned_image_id is not None else None
+    if cloned_image is None or cloned_image.info.type_id != TYPE_IMAGE:
+        raise ValueError("cloned attachment has no ImageArchive")
+
+    package = document.package_metadata().message
+    existing_data_ids = [
+        int(Message.parse(bytes(raw.value)).one(DATA_IDENTIFIER).value)
+        for raw in package.get(PACKAGE_DATAS)
+    ]
+    full_id = max(existing_data_ids, default=0) + 1
+    thumbnail_id = full_id + 1
+    data_identifiers = {
+        template.data_identifier: full_id,
+        template.thumbnail_identifier: thumbnail_id,
+    }
+    data_identifiers = {old: new for old, new in data_identifiers.items() if old is not None}
+    thumbnail = _thumbnail_png(data)
+    stem = Path(filename).stem or "image"
+    preferred_full = f"{stem}.png"
+    preferred_thumbnail = f"{stem}-small.png"
+    stored_full = f"{stem}-{full_id}.png"
+    stored_thumbnail = f"{stem}-small-{thumbnail_id}.png"
+    for old, new, preferred, stored in (
+        (template.data_identifier, full_id, preferred_full, stored_full),
+        (template.thumbnail_identifier, thumbnail_id, preferred_thumbnail, stored_thumbnail),
+    ):
+        if old is None:
+            continue
+        info = _data_info(document, old)
+        _set_varint(info, DATA_IDENTIFIER, new)
+        _set_bytes(info, DATA_PREFERRED_FILENAME, preferred.encode("utf-8"))
+        _set_bytes(info, DATA_FILENAME, stored.encode("utf-8"))
+        package.add_bytes(PACKAGE_DATAS, info.encode())
+
+    remap_payload_data_references(cloned_image, data_identifiers)
+    _append_component_data_references(
+        document,
+        template.image.segment.identifier,
+        cloned_image.segment.identifier,
+        data_identifiers,
+    )
+    document.set_package_data("Data/" + stored_full, data)
+    document.set_package_data("Data/" + stored_thumbnail, thumbnail)
+    store = ImageStore(document, cloned_image)
+    store.replace(data, thumbnail)
+
+    pixel_width, pixel_height = ImageStore._png_size(data)
+    if width is None and height is None:
+        template_size = _logical_size(cloned_image.message)
+        if template_size is not None:
+            maximum = max(template_size)
+            if pixel_width >= pixel_height:
+                width = maximum
+                height = maximum * pixel_height / pixel_width
+            else:
+                height = maximum
+                width = maximum * pixel_width / pixel_height
+    if width is not None or height is not None:
+        logical_width = (
+            float(width)
+            if width is not None
+            else float(height) * pixel_width / pixel_height
+        )
+        logical_height = (
+            float(height)
+            if height is not None
+            else float(width) * pixel_height / pixel_width
+        )
+        _set_logical_size(cloned_image.message, logical_width, logical_height)
+    return store, clone
+
+
+def body_image_stores(document: "PagesDocument", storage: LocatedObject) -> list[ImageStore]:
+    """Resolve body attachment runs to ImageArchive objects in text order."""
+    field = storage.message.one(STORAGE_TABLE_ATTACHMENTS)
+    if field is None:
+        return []
+    objects = document.objects()
+    result: list[ImageStore] = []
+    seen: set[int] = set()
+    for run in parse_attribute_table(bytes(field.value)):
+        attachment = objects.get(run.object_id) if run.object_id is not None else None
+        if attachment is None or attachment.info.type_id != TYPE_ATTACHMENT:
+            continue
+        image_id = _reference(attachment.message, 1)
+        image = objects.get(image_id) if image_id is not None else None
+        if image is None or image.info.type_id != TYPE_IMAGE or image_id in seen:
+            continue
+        seen.add(image_id)
+        result.append(ImageStore(document, image))
+    return result