feat(sources): MOBI and AZW3 books as sources

surfsense_local/backend/modules/documents/kindle_books/__init__.py

@@ -0,0 +1,2 @@
+"""A MOBI or AZW3 book as both the API and the ingest worker see it: its
+container, and whether an upload is one the reader can open."""

surfsense_local/backend/modules/documents/kindle_books/container.py

@@ -0,0 +1,123 @@
+"""The Palm database a MOBI or AZW3 book is stored in, and its record 0 headers."""
+
+import struct
+from dataclasses import dataclass, field
+from itertools import pairwise
+
+NO_RECORD = 0xFFFFFFFF
+BOOK_TYPES = (b"BOOKMOBI", b"TEXtREAd")
+PALM_HEADER_SIZE = 78
+COMPRESSION_NONE, COMPRESSION_PALMDOC, COMPRESSION_HUFFCDIC = 1, 2, 17480
+EXTH_KF8_BOUNDARY = 121
+# A text record unpacks to 4 KB. Codes that keep copying, or dictionary entries
+# that double at each level, could reach gigabytes; past this the file is damaged.
+MAX_UNPACKED_RECORD = 64 * 1024
+
+
+class KindleBookError(ValueError):
+    """The bytes are not a book this reader can open."""
+
+
+class DrmProtectedError(KindleBookError):
+    pass
+
+
+def u16(data: bytes, at: int) -> int:
+    return struct.unpack_from(">H", data, at)[0]
+
+
+def u32(data: bytes, at: int) -> int:
+    return struct.unpack_from(">I", data, at)[0]
+
+
+def record_offsets(head: bytes) -> list[int]:
+    """Each record's start, from the database header and record table."""
+    if len(head) < PALM_HEADER_SIZE or head[60:68] not in BOOK_TYPES:
+        raise KindleBookError("not a Kindle book")
+    count = u16(head, 76)
+    if len(head) < PALM_HEADER_SIZE + 8 * count:
+        raise KindleBookError("the book is cut short")
+    offsets = [u32(head, PALM_HEADER_SIZE + 8 * i) for i in range(count)]
+    if not offsets or any(a > b for a, b in pairwise(offsets)):
+        raise KindleBookError("the book's record table is out of order")
+    return offsets
+
+
+def records(data: bytes) -> list[bytes]:
+    offsets = record_offsets(data)
+    if offsets[-1] > len(data):
+        raise KindleBookError("the book is cut short")
+    ends = [*offsets[1:], len(data)]
+    return [data[start:end] for start, end in zip(offsets, ends, strict=True)]
+
+
+@dataclass(frozen=True)
+class BookHeader:
+    compression: int
+    text_records: int
+    encryption: int
+    encoding: str = "cp1252"
+    version: int = 0
+    trailing_flags: int = 0
+    huff_record: int = NO_RECORD
+    huff_count: int = 0
+    frag_index: int = NO_RECORD
+    skel_index: int = NO_RECORD
+    exth: dict[int, bytes] = field(default_factory=dict)
+
+    @property
+    def kf8_boundary(self) -> int | None:
+        """Where a combined file's KF8 part starts, if it has one."""
+        value = self.exth.get(EXTH_KF8_BOUNDARY)
+        if value is None or len(value) < 4 or u32(value, 0) == NO_RECORD:
+            return None
+        return u32(value, 0)
+
+
+def has_mobi_header(record0: bytes) -> bool:
+    return record0[16:20] == b"MOBI"
+
+
+def book_header(record0: bytes) -> BookHeader:
+    """Record 0: a PalmDOC header, then for a MOBI book its MOBI and EXTH headers."""
+    if len(record0) < 16:
+        raise KindleBookError("the book is cut short")
+    compression, text_records = u16(record0, 0), u16(record0, 8)
+    if not has_mobi_header(record0):
+        # A plain PalmDOC text: bytes 12-13 are its encryption, as for MOBI.
+        return BookHeader(compression, text_records, u16(record0, 12))
+    length = u32(record0, 20)
+    encoding = {1252: "cp1252", 65001: "utf-8"}.get(u32(record0, 28))
+    if encoding is None:
+        raise KindleBookError("the book's text encoding is unknown")
+    version = u32(record0, 36)
+    kf8 = version >= 8
+    return BookHeader(
+        compression=compression,
+        text_records=text_records,
+        encryption=u16(record0, 12),
+        encoding=encoding,
+        version=version,
+        # The field exists only in a header at least this long.
+        trailing_flags=u32(record0, 240) if length >= 0xE4 else 0,
+        huff_record=u32(record0, 112),
+        huff_count=u32(record0, 116),
+        frag_index=u32(record0, 248) if kf8 else NO_RECORD,
+        skel_index=u32(record0, 252) if kf8 else NO_RECORD,
+        exth=_exth(record0, 16 + length) if u32(record0, 128) & 0x40 else {},
+    )
+
+
+def _exth(record0: bytes, at: int) -> dict[int, bytes]:
+    if record0[at : at + 4] != b"EXTH":
+        return {}
+    entries: dict[int, bytes] = {}
+    pos = at + 12
+    for _ in range(u32(record0, at + 8)):
+        kind, size = u32(record0, pos), u32(record0, pos + 4)
+        if size < 8:
+            break
+        # The first wins: authors repeat, the boundary does not.
+        entries.setdefault(kind, record0[pos + 8 : pos + size])
+        pos += size
+    return entries

surfsense_local/backend/modules/documents/kindle_books/upload_check.py

@@ -0,0 +1,44 @@
+"""Whether an upload is a Kindle book the reader can open, from its first bytes."""
+
+import struct
+from pathlib import Path
+
+from modules.documents.kindle_books.container import (
+    PALM_HEADER_SIZE,
+    KindleBookError,
+    book_header,
+    record_offsets,
+    u16,
+)
+
+KINDLE_SUFFIXES = frozenset({".mobi", ".azw3"})
+DRM_REFUSAL = "This book is DRM-protected, so it can't be read"
+KFX_REFUSAL = "KFX books from the Kindle store can't be read yet"
+# Amazon's KFX: a DRM envelope, or the bare container.
+_KFX_MAGIC = (b"\xeaDRMION\xee", b"CONT")
+# Record 0 holds the headers; anything past this is not needed to judge the book.
+_RECORD0_READ = 64 * 1024
+
+
+def is_kindle_book(path: Path) -> bool:
+    """True for a readable book; raises with the user's reason for DRM and KFX."""
+    try:
+        with path.open("rb") as book:
+            head = book.read(PALM_HEADER_SIZE)
+            if head.startswith(_KFX_MAGIC):
+                raise ValueError(KFX_REFUSAL)
+            if len(head) < PALM_HEADER_SIZE:
+                return False
+            head += book.read(8 * u16(head, 76))
+            offsets = record_offsets(head)
+            book.seek(offsets[0])
+            record0 = book.read(_RECORD0_READ)
+    except (OSError, KindleBookError, struct.error):
+        return False
+    try:
+        header = book_header(record0)
+    except (KindleBookError, struct.error):
+        return False
+    if header.encryption:
+        raise ValueError(DRM_REFUSAL)
+    return True

surfsense_local/backend/modules/documents/storage.py

@@ -8,6 +8,11 @@
 
 from fastapi import HTTPException, UploadFile, status
 
+from modules.documents.kindle_books.upload_check import (
+    KINDLE_SUFFIXES,
+    is_kindle_book,
+)
+
 MAX_UPLOAD_BYTES = 500 * 1024 * 1024
 MAX_ARCHIVE_ENTRIES = 10_000
 MAX_ARCHIVE_UNPACKED_BYTES = 2 * 1024 * 1024 * 1024
@@ -25,6 +30,8 @@
     ".pptx": "application/vnd.openxmlformats-officedocument.presentationml.presentation",
     ".xlsx": "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
     ".epub": "application/epub+zip",
+    ".mobi": "application/x-mobipocket-ebook",
+    ".azw3": "application/vnd.amazon.mobi8-ebook",
     ".html": "text/html",
     ".htm": "text/html",
     ".csv": "text/csv",
@@ -69,6 +76,8 @@ def validate_upload(path: Path, suffix: str) -> str:
         valid = _is_ooxml(path, suffix)
     elif suffix == ".epub":
         valid = _is_epub(path)
+    elif suffix in KINDLE_SUFFIXES:
+        valid = is_kindle_book(path)
     elif suffix in {".html", ".htm", ".csv", ".md", ".markdown", ".txt", ".text"}:
         valid = _is_utf8_text(path)
     elif suffix in {".png", ".jpg", ".jpeg", ".tif", ".tiff", ".bmp", ".webp"}:

surfsense_local/backend/tests/integration/worker/test_parsing.py

@@ -19,6 +19,7 @@
 from reportlab.lib.styles import getSampleStyleSheet
 from reportlab.platypus import Paragraph, SimpleDocTemplate, Spacer
 
+from tests import kindle_samples
 from worker.ingestion.parser_pack import missing_parser_folders
 
 pytestmark = [pytest.mark.integration, pytest.mark.image]
@@ -189,6 +190,18 @@ def parsed(tmp_path_factory: pytest.TempPathFactory) -> dict[str, str]:
 
     files["novel.epub"] = book(CHAPTERS)
     files["broken-heading.epub"] = book([(BROKEN_HEADING, RECEIPT)])
+    files["novel.azw3"] = kindle_samples.azw3(
+        [[f"<h2>{heading}</h2>", f"<p>{text}</p>"] for heading, text in CHAPTERS]
+    )
+    files["novel.mobi"] = kindle_samples.mobi(
+        "".join(
+            f'<p><font size="7">CHAPTER {n}.</font><br/><font size="7">{heading}</font></p>'
+            f"<p>{text}</p>"
+            for n, (heading, text) in enumerate(CHAPTERS, start=1)
+        ),
+        compression=2,
+    )
+    files["damaged.azw3"] = kindle_samples.azw3([["<p>Lost.</p>"]])[:400]
     files["broken-heading.html"] = (
         f"<html><body><h2>{BROKEN_HEADING}</h2><p>{RECEIPT}</p></body></html>"
     ).encode()
@@ -264,3 +277,26 @@ def test_a_line_break_in_a_heading_reads_as_a_space(
 
     assert "CHAPTER I. Down the Rabbit-Hole" in markdown
     assert "\ue000" not in markdown
+
+
+def test_a_kf8_book_reads_its_chapters_in_order(parsed: dict[str, str]) -> None:
+    """Each chapter's fragments go back into its file, and the files stay in order."""
+    expected = " ".join(f"## {heading}\n\n{text}" for heading, text in CHAPTERS)
+
+    assert words(expected) in words(parsed["novel.azw3"])
+
+
+def test_a_mobi6_book_reads_its_large_type_titles_as_headings(
+    parsed: dict[str, str],
+) -> None:
+    """Its titles are big paragraphs with a line break, as calibre writes them."""
+    markdown = parsed["novel.mobi"]
+
+    for n, (heading, text) in enumerate(CHAPTERS, start=1):
+        assert f"# CHAPTER {n}. {heading}\n\n{text}" in markdown
+    assert "\ue000" not in markdown
+
+
+def test_a_damaged_book_fails_as_unreadable(parsed: dict[str, str]) -> None:
+    """Unreadable fails the source with its reason and is not retried."""
+    assert parsed["damaged.azw3"].startswith("UnreadableFileError: ")