diff --git a/machine/corpora/dbl_bundle_text_corpus.py b/machine/corpora/dbl_bundle_text_corpus.py index 33aa62cd..87cc3ae3 100644 --- a/machine/corpora/dbl_bundle_text_corpus.py +++ b/machine/corpora/dbl_bundle_text_corpus.py @@ -7,6 +7,7 @@ from ..scripture import ENGLISH_VERSIFICATION from ..scripture.verse_ref import Versification from ..utils.typeshed import StrPath +from ..utils.zip_entry_utils import open_bounded_stream from .scripture_text_corpus import ScriptureTextCorpus from .usx_zip_text import UsxZipText @@ -16,7 +17,7 @@ class DblBundleTextCorpus(ScriptureTextCorpus): def __init__(self, filename: StrPath) -> None: with ZipFile(filename, "r") as archive: - with archive.open("metadata.xml", "r") as stream: + with open_bounded_stream(archive, "metadata.xml") as stream: doc = ElementTree.parse(stream) version = doc.getroot().get("version", "2.0") parts = version.split(".", maxsplit=3) @@ -27,7 +28,7 @@ def __init__(self, filename: StrPath) -> None: (zi for zi in archive.filelist if os.path.basename(zi.filename) == "versification.vrs"), None ) if versification_entry is not None: - with archive.open(versification_entry, "r") as stream: + with open_bounded_stream(archive, versification_entry) as stream: abbr = doc.getroot().findtext("./identification/abbreviation", "") versification = Versification.parse( TextIOWrapper(stream, encoding="utf-8-sig"), "versification.vrs", fallback_name=abbr diff --git a/machine/corpora/zip_entry_stream_container.py b/machine/corpora/zip_entry_stream_container.py index a7f40a70..f719a95d 100644 --- a/machine/corpora/zip_entry_stream_container.py +++ b/machine/corpora/zip_entry_stream_container.py @@ -5,6 +5,7 @@ from zipfile import ZipFile from ..utils.typeshed import StrPath +from ..utils.zip_entry_utils import open_bounded_stream from .stream_container import StreamContainer @@ -20,7 +21,8 @@ def __exit__(self, type: Any, value: Any, traceback: Any) -> None: self.close() def open_stream(self) -> BinaryIO: - return BytesIO(self._archive.read(self._entry)) + with open_bounded_stream(self._archive, self._entry) as stream: + return BytesIO(stream.read()) def close(self) -> None: self._archive.close() diff --git a/machine/corpora/zip_paratext_project_file_handler.py b/machine/corpora/zip_paratext_project_file_handler.py index 3f430048..ca075c97 100644 --- a/machine/corpora/zip_paratext_project_file_handler.py +++ b/machine/corpora/zip_paratext_project_file_handler.py @@ -4,6 +4,7 @@ from typing import BinaryIO, Optional, cast from zipfile import ZipFile +from ..utils.zip_entry_utils import open_bounded_stream from .paratext_project_file_handler import ParatextProjectFileHandler from .usfm_stylesheet import UsfmStylesheet @@ -27,7 +28,8 @@ def find(self, extension: str) -> Optional[str]: def open(self, file_name: str) -> Optional[BinaryIO]: for actual_entry_name in self._archive.namelist(): if actual_entry_name.lower() == file_name.lower(): - return BytesIO(self._archive.read(actual_entry_name)) + with open_bounded_stream(self._archive, actual_entry_name) as stream: + return BytesIO(stream.read()) return None def create_stylesheet(self, file_name: str) -> UsfmStylesheet: diff --git a/machine/utils/zip_entry_utils.py b/machine/utils/zip_entry_utils.py new file mode 100644 index 00000000..9d1c69cc --- /dev/null +++ b/machine/utils/zip_entry_utils.py @@ -0,0 +1,28 @@ +from typing import IO, Union +from zipfile import BadZipFile, ZipFile, ZipInfo + +DEFAULT_MAX_UNCOMPRESSED_SIZE = 100 * 1024 * 1024 # 100 MB +DEFAULT_MAX_COMPRESSION_RATIO = 100.0 # 100:1 ratio limit + + +def open_bounded_stream( + archive: ZipFile, + entry: Union[str, ZipInfo], + max_uncompressed_size: int = DEFAULT_MAX_UNCOMPRESSED_SIZE, + max_compression_ratio: float = DEFAULT_MAX_COMPRESSION_RATIO, +) -> IO[bytes]: + """Opens a zip entry for reading, refusing entries that are too large or too highly compressed. + + Raises `BadZipFile` if the entry's declared size or compression ratio exceeds the limits. `ZipFile` never yields + more than the declared size, so an entry whose header understates its size fails the CRC check with `BadZipFile`. + """ + if isinstance(entry, str): + entry = archive.getinfo(entry) + + if entry.file_size > max_uncompressed_size: + raise BadZipFile("Entry uncompressed size exceeds maximum allowed limit.") + + if entry.compress_size > 0 and entry.file_size / entry.compress_size > max_compression_ratio: + raise BadZipFile("Compression ratio exceeds safe threshold.") + + return archive.open(entry, "r") diff --git a/tests/utils/test_zip_entry_utils.py b/tests/utils/test_zip_entry_utils.py new file mode 100644 index 00000000..e22e3e1d --- /dev/null +++ b/tests/utils/test_zip_entry_utils.py @@ -0,0 +1,37 @@ +from io import BytesIO +from zipfile import ZIP_DEFLATED, BadZipFile, ZipFile + +from pytest import raises + +from machine.utils.zip_entry_utils import open_bounded_stream + + +def create_zip(name: str, content: bytes) -> ZipFile: + buffer = BytesIO() + with ZipFile(buffer, "w", ZIP_DEFLATED) as archive: + archive.writestr(name, content) + return ZipFile(buffer, "r") + + +def test_open_bounded_stream_returns_readable_stream() -> None: + with create_zip("test.txt", b"Hello World") as archive: + with open_bounded_stream(archive, "test.txt", max_uncompressed_size=100) as stream: + assert stream.read() == b"Hello World" + + +def test_open_bounded_stream_empty_entry() -> None: + with create_zip("empty.txt", b"") as archive: + with open_bounded_stream(archive, "empty.txt") as stream: + assert stream.read() == b"" + + +def test_open_bounded_stream_header_size_exceeded() -> None: + with create_zip("large.txt", bytes(200)) as archive: + with raises(BadZipFile): + open_bounded_stream(archive, "large.txt", max_uncompressed_size=100) + + +def test_open_bounded_stream_compression_ratio_exceeded() -> None: + with create_zip("bomb.txt", bytes(10_000)) as archive: + with raises(BadZipFile): + open_bounded_stream(archive, "bomb.txt", max_uncompressed_size=20_000, max_compression_ratio=2.0)