Preserve undecodable archives instead of aborting the read

keynote_parser/codec.py

@@ -1,7 +1,9 @@
+import base64
 import importlib
 import struct
 import sys
 import traceback
+import warnings
 from functools import partial
 
 import snappy
@@ -38,6 +40,88 @@ def import_version(version: str = LATEST_VERSION):
     )
 
 
+class UnknownArchiveWarning(UserWarning):
+    """Raised when an archive can't be decoded and is preserved verbatim instead.
+
+    This is a warning rather than an error so that a single unrecognised message
+    type - often in an incidental file like Index/CalculationEngine.iwa - doesn't
+    prevent the rest of a document from being read. To treat these as fatal:
+
+        warnings.simplefilter("error", keynote_parser.codec.UnknownArchiveWarning)
+    """
+
+
+# Documents can contain thousands of archives of the same unknown type; only
+# warn about each (file, type) pair once so the output stays readable.
+_WARNED_ABOUT = set()
+
+
+def _describe(klass):
+    """A short, human-readable name for a message class, for use in warnings."""
+    descriptor = getattr(klass, "DESCRIPTOR", None)
+    return (
+        getattr(descriptor, "full_name", None)
+        or getattr(klass, "__name__", None)
+        or repr(klass)
+    )
+
+
+def _warn_about_unknown_archive(message, filename, type_id):
+    key = (filename, type_id)
+    if key in _WARNED_ABOUT:
+        return
+    _WARNED_ABOUT.add(key)
+    if filename:
+        message = "%s (in %s)" % (message, filename)
+    warnings.warn(message, UnknownArchiveWarning, stacklevel=2)
+
+
+class UnknownArchive(object):
+    """An archive that could not be decoded, holding its original bytes.
+
+    Keynote documents are read and written whole, so an archive we can't
+    interpret still has to survive a pack/unpack cycle untouched. This holds the
+    raw payload and writes it back byte-for-byte, which keeps round-trips lossless
+    even though the contents aren't introspectable or replaceable.
+    """
+
+    # Deliberately not a real Protobuf type name, so it can't collide with an
+    # entry in NAME_CLASS_MAP.
+    PBTYPE = "keynote_parser.UnknownArchive"
+
+    def __init__(self, type_id, data):
+        self.type_id = type_id
+        self.data = data
+
+    def __eq__(self, other):
+        return (
+            isinstance(other, UnknownArchive)
+            and self.type_id == other.type_id
+            and self.data == other.data
+        )
+
+    def __repr__(self):
+        return "<%s type=%s length=%d>" % (
+            self.__class__.__name__,
+            self.type_id,
+            len(self.data),
+        )
+
+    def to_dict(self):
+        return {
+            "_pbtype": self.PBTYPE,
+            "type": self.type_id,
+            "base64Data": base64.b64encode(self.data).decode("ascii"),
+        }
+
+    @classmethod
+    def from_dict(cls, _dict):
+        return cls(int(_dict["type"]), base64.b64decode(_dict["base64Data"]))
+
+    def SerializeToString(self):
+        return self.data
+
+
 class IWAFile(object):
     def __init__(self, chunks, filename=None):
         self.chunks = chunks
@@ -205,36 +289,53 @@ def from_buffer(cls, buf, filename=None, version: str = LATEST_VERSION):
 
         n = 0
         for message_info in archive_info.message_infos:
+            # message_info.length delimits this archive regardless of whether we
+            # can decode it, so an undecodable archive never desynchronises the
+            # ones that follow it.
+            message_payload = payload[n : n + message_info.length]
+            n += message_info.length
+
             try:
                 if message_info.type == 0 and archive_info.should_merge and payloads:
                     base_message = archive_info.message_infos[
                         message_info.base_message_index
                     ]
-                    klass = partial(
-                        ProtobufPatch.FromString,
-                        message_info,
-                        import_version(version)[0][base_message.type],
-                    )
+                    base_klass = import_version(version)[0][base_message.type]
+                    klass = partial(ProtobufPatch.FromString, message_info, base_klass)
+                    # `klass` is a functools.partial here, whose repr includes the
+                    # entire message_info; not something to put in a warning.
+                    description = "patch to %s" % _describe(base_klass)
                 else:
                     klass = import_version(version)[0][message_info.type]
+                    description = _describe(klass)
             except KeyError:
-                raise NotImplementedError(
-                    "Don't know how to parse Protobuf message type "
-                    + str(message_info.type)
+                _warn_about_unknown_archive(
+                    "Don't know how to parse Protobuf message type %s; "
+                    "preserving it verbatim. Slide content is unaffected unless "
+                    "it lives in this archive." % message_info.type,
+                    filename,
+                    message_info.type,
                 )
+                payloads.append(UnknownArchive(message_info.type, message_payload))
+                continue
+
             try:
-                message_payload = payload[n : n + message_info.length]
                 if hasattr(klass, "FromString"):
                     output = klass.FromString(message_payload)
                 else:
                     output = klass(message_payload)
             except Exception as e:
-                raise ValueError(
-                    "Failed to deserialize %s payload of length %d: %s"
-                    % (klass, message_info.length, e)
+                _warn_about_unknown_archive(
+                    "Failed to deserialize %s of length %d (%s: %s); "
+                    "preserving it verbatim."
+                    % (description, message_info.length, type(e).__name__, e),
+                    filename,
+                    message_info.type,
                 )
+                payloads.append(UnknownArchive(message_info.type, message_payload))
+                continue
+
             payloads.append(output)
-            n += message_info.length
 
         return cls(archive_info, payloads), payload[n:]
 
@@ -311,6 +412,8 @@ def _work_around_protobuf_max_float_handling(_dict):
 
 
 def dict_to_message(_dict, version: str = LATEST_VERSION):
+    if _dict.get("_pbtype") == UnknownArchive.PBTYPE:
+        return UnknownArchive.from_dict(_dict)
     _type = _dict["_pbtype"]
     del _dict["_pbtype"]
     _dict = _work_around_protobuf_max_float_handling(_dict)

tests/test_unknown_archives.py

@@ -0,0 +1,143 @@
+"""Tests for graceful handling of archives we can't decode.
+
+Apple adds Protobuf message types between Keynote releases, so a document
+written by a newer Keynote than the one we have protos for will contain
+archives we don't recognise. Those must not abort the whole read - see #60,
+#64 and #70, where a single unknown type in Index/CalculationEngine.iwa made
+`ls`, `cat` and `unpack` unusable on the entire document.
+"""
+
+import copy
+import warnings
+
+import pytest
+
+from keynote_parser import codec
+
+SIMPLE_FILENAME = "./tests/data/simple-oneslide.iwa"
+
+
+@pytest.fixture(autouse=True)
+def _reset_warning_dedupe():
+    # codec only warns once per (file, type); clear that between tests.
+    codec._WARNED_ABOUT.clear()
+
+
+def _first_archive_type(filename):
+    with open(filename, "rb") as f:
+        file = codec.IWAFile.from_buffer(f.read(), filename)
+    return file.chunks[0].archives[0].header.message_infos[0].type
+
+
+def _first_archive_payload(filename):
+    """The exact on-disk bytes of the first archive in `filename`."""
+    with open(filename, "rb") as f:
+        decompressed = b"".join(codec.IWACompressedChunk._decompress_all(f.read()))
+    archive_info, payload = codec.get_archive_info_and_remainder(decompressed)
+    return payload[: archive_info.message_infos[0].length]
+
+
+def _patch_mapping(monkeypatch, id_name_map):
+    _, name_class_map, archive_info = codec.import_version()
+    monkeypatch.setattr(
+        codec,
+        "import_version",
+        lambda *a, **k: (id_name_map, name_class_map, archive_info),
+    )
+
+
+def _read_with_type_unmapped(filename, type_id, monkeypatch):
+    """Read `filename` as though `type_id` were absent from the mapping."""
+    id_name_map = codec.import_version()[0]
+    _patch_mapping(monkeypatch, {k: v for k, v in id_name_map.items() if k != type_id})
+    with open(filename, "rb") as f:
+        data = f.read()
+    return codec.IWAFile.from_buffer(data, filename), data
+
+
+def test_unknown_message_type_warns_instead_of_raising(monkeypatch):
+    type_id = _first_archive_type(SIMPLE_FILENAME)
+    with pytest.warns(codec.UnknownArchiveWarning, match=str(type_id)):
+        file, _ = _read_with_type_unmapped(SIMPLE_FILENAME, type_id, monkeypatch)
+    assert file is not None
+
+
+def test_unknown_message_type_is_preserved_verbatim(monkeypatch):
+    type_id = _first_archive_type(SIMPLE_FILENAME)
+    expected = _first_archive_payload(SIMPLE_FILENAME)
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore", codec.UnknownArchiveWarning)
+        file, _ = _read_with_type_unmapped(SIMPLE_FILENAME, type_id, monkeypatch)
+
+        unknown = file.chunks[0].archives[0].objects[0]
+        assert isinstance(unknown, codec.UnknownArchive)
+        assert unknown.type_id == type_id
+
+        # The undecoded archive must be written back byte-for-byte.
+        assert unknown.data == expected
+        assert unknown.SerializeToString() == expected
+
+        # ...and the file as a whole must still round-trip.
+        assert codec.IWAFile.from_buffer(file.to_buffer()).to_dict() == file.to_dict()
+
+
+def test_unknown_archive_survives_a_yaml_roundtrip(monkeypatch):
+    type_id = _first_archive_type(SIMPLE_FILENAME)
+    expected = _first_archive_payload(SIMPLE_FILENAME)
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore", codec.UnknownArchiveWarning)
+        file, _ = _read_with_type_unmapped(SIMPLE_FILENAME, type_id, monkeypatch)
+
+        as_dict = file.to_dict()
+        assert as_dict["chunks"][0]["archives"][0]["objects"][0]["_pbtype"] == (
+            codec.UnknownArchive.PBTYPE
+        )
+
+        # Unpack -> YAML -> pack must not lose the bytes we couldn't decode.
+        # (from_dict consumes the dict it's handed, so hand it a copy.)
+        reparsed = codec.IWAFile.from_dict(copy.deepcopy(as_dict))
+        assert reparsed.chunks[0].archives[0].objects[0].data == expected
+        assert reparsed.to_dict() == as_dict
+
+
+def test_undecodable_payload_is_preserved_verbatim(monkeypatch):
+    """A mapped type whose payload won't parse should degrade the same way."""
+
+    class Undecodable:
+        @staticmethod
+        def FromString(data):
+            raise ValueError("nope")
+
+    type_id = _first_archive_type(SIMPLE_FILENAME)
+    expected = _first_archive_payload(SIMPLE_FILENAME)
+    broken = dict(codec.import_version()[0])
+    broken[type_id] = Undecodable
+    _patch_mapping(monkeypatch, broken)
+
+    with open(SIMPLE_FILENAME, "rb") as f:
+        original = f.read()
+
+    with pytest.warns(codec.UnknownArchiveWarning, match="Failed to deserialize"):
+        file = codec.IWAFile.from_buffer(original, SIMPLE_FILENAME)
+
+    unknown = file.chunks[0].archives[0].objects[0]
+    assert isinstance(unknown, codec.UnknownArchive)
+    assert unknown.data == expected
+
+
+def test_unknown_archives_can_be_made_fatal(monkeypatch):
+    """Callers who'd rather fail loudly can promote the warning to an error."""
+    type_id = _first_archive_type(SIMPLE_FILENAME)
+    with warnings.catch_warnings():
+        warnings.simplefilter("error", codec.UnknownArchiveWarning)
+        with pytest.raises(ValueError):
+            _read_with_type_unmapped(SIMPLE_FILENAME, type_id, monkeypatch)
+
+
+def test_only_warns_once_per_file_and_type(monkeypatch):
+    type_id = _first_archive_type(SIMPLE_FILENAME)
+    with warnings.catch_warnings(record=True) as caught:
+        warnings.simplefilter("always", codec.UnknownArchiveWarning)
+        _read_with_type_unmapped(SIMPLE_FILENAME, type_id, monkeypatch)
+    relevant = [w for w in caught if w.category is codec.UnknownArchiveWarning]
+    assert len(relevant) == 1