diff --git a/Tests/test_file_blp.py b/Tests/test_file_blp.py index 1e2f20c407b..9f2de8f982e 100644 --- a/Tests/test_file_blp.py +++ b/Tests/test_file_blp.py @@ -4,7 +4,7 @@ import pytest -from PIL import Image +from PIL import BlpImagePlugin, Image from .helper import ( assert_image_equal, @@ -19,6 +19,7 @@ def test_load_blp1() -> None: assert_image_equal_tofile(im, "Tests/images/blp/blp1_jpeg.png") with Image.open("Tests/images/blp/blp1_jpeg2.blp") as im: + assert im.mode == "RGBA" im.load() @@ -37,6 +38,13 @@ def test_load_blp2_dxt1a() -> None: assert_image_equal_tofile(im, "Tests/images/blp/blp2_dxt1a.png") +def test_invalid_file() -> None: + invalid_file = "Tests/images/flower.jpg" + + with pytest.raises(BlpImagePlugin.BLPFormatError): + BlpImagePlugin.BlpImageFile(invalid_file) + + def test_save(tmp_path: Path) -> None: f = str(tmp_path / "temp.blp") diff --git a/src/PIL/BlpImagePlugin.py b/src/PIL/BlpImagePlugin.py index 2d03af9d7dd..c932b3b9c8a 100644 --- a/src/PIL/BlpImagePlugin.py +++ b/src/PIL/BlpImagePlugin.py @@ -259,21 +259,36 @@ class BlpImageFile(ImageFile.ImageFile): def _open(self) -> None: self.magic = self.fp.read(4) + if not _accept(self.magic): + msg = f"Bad BLP magic {repr(self.magic)}" + raise BLPFormatError(msg) - self.fp.seek(5, os.SEEK_CUR) - (self._blp_alpha_depth,) = struct.unpack(" tuple[int, int]: try: - self._read_blp_header() + self._read_header() self._load() except struct.error as e: msg = "Truncated BLP file" @@ -292,25 +307,9 @@ def decode(self, buffer: bytes | Image.SupportsArrayInterface) -> tuple[int, int def _load(self) -> None: pass - def _read_blp_header(self) -> None: - assert self.fd is not None - self.fd.seek(4) - (self._blp_compression,) = struct.unpack(" None: + self._offsets = struct.unpack("<16I", self._safe_read(16 * 4)) + self._lengths = struct.unpack("<16I", self._safe_read(16 * 4)) def _safe_read(self, length: int) -> bytes: assert self.fd is not None @@ -326,9 +325,11 @@ def _read_palette(self) -> list[tuple[int, int, int, int]]: ret.append((b, g, r, a)) return ret - def _read_bgra(self, palette: list[tuple[int, int, int, int]]) -> bytearray: + def _read_bgra( + self, palette: list[tuple[int, int, int, int]], alpha: bool + ) -> bytearray: data = bytearray() - _data = BytesIO(self._safe_read(self._blp_lengths[0])) + _data = BytesIO(self._safe_read(self._lengths[0])) while True: try: (offset,) = struct.unpack(" bytearray: break b, g, r, a = palette[offset] d: tuple[int, ...] = (r, g, b) - if self._blp_alpha_depth: + if alpha: d += (a,) data.extend(d) return data @@ -344,19 +345,21 @@ def _read_bgra(self, palette: list[tuple[int, int, int, int]]) -> bytearray: class BLP1Decoder(_BLPBaseDecoder): def _load(self) -> None: - if self._blp_compression == Format.JPEG: + self._compression, self._encoding, alpha = self.args + + if self._compression == Format.JPEG: self._decode_jpeg_stream() - elif self._blp_compression == 1: - if self._blp_encoding in (4, 5): + elif self._compression == 1: + if self._encoding in (4, 5): palette = self._read_palette() - data = self._read_bgra(palette) + data = self._read_bgra(palette, alpha) self.set_as_raw(data) else: - msg = f"Unsupported BLP encoding {repr(self._blp_encoding)}" + msg = f"Unsupported BLP encoding {repr(self._encoding)}" raise BLPFormatError(msg) else: - msg = f"Unsupported BLP compression {repr(self._blp_encoding)}" + msg = f"Unsupported BLP compression {repr(self._encoding)}" raise BLPFormatError(msg) def _decode_jpeg_stream(self) -> None: @@ -365,8 +368,8 @@ def _decode_jpeg_stream(self) -> None: (jpeg_header_size,) = struct.unpack(" None: class BLP2Decoder(_BLPBaseDecoder): def _load(self) -> None: + self._compression, self._encoding, alpha, self._alpha_encoding = self.args + palette = self._read_palette() assert self.fd is not None - self.fd.seek(self._blp_offsets[0]) + self.fd.seek(self._offsets[0]) - if self._blp_compression == 1: + if self._compression == 1: # Uncompressed or DirectX compression - if self._blp_encoding == Encoding.UNCOMPRESSED: - data = self._read_bgra(palette) + if self._encoding == Encoding.UNCOMPRESSED: + data = self._read_bgra(palette, alpha) - elif self._blp_encoding == Encoding.DXT: + elif self._encoding == Encoding.DXT: data = bytearray() - if self._blp_alpha_encoding == AlphaEncoding.DXT1: - linesize = (self.size[0] + 3) // 4 * 8 - for yb in range((self.size[1] + 3) // 4): - for d in decode_dxt1( - self._safe_read(linesize), alpha=bool(self._blp_alpha_depth) - ): + if self._alpha_encoding == AlphaEncoding.DXT1: + linesize = (self.state.xsize + 3) // 4 * 8 + for yb in range((self.state.ysize + 3) // 4): + for d in decode_dxt1(self._safe_read(linesize), alpha): data += d - elif self._blp_alpha_encoding == AlphaEncoding.DXT3: - linesize = (self.size[0] + 3) // 4 * 16 - for yb in range((self.size[1] + 3) // 4): + elif self._alpha_encoding == AlphaEncoding.DXT3: + linesize = (self.state.xsize + 3) // 4 * 16 + for yb in range((self.state.ysize + 3) // 4): for d in decode_dxt3(self._safe_read(linesize)): data += d - elif self._blp_alpha_encoding == AlphaEncoding.DXT5: - linesize = (self.size[0] + 3) // 4 * 16 - for yb in range((self.size[1] + 3) // 4): + elif self._alpha_encoding == AlphaEncoding.DXT5: + linesize = (self.state.xsize + 3) // 4 * 16 + for yb in range((self.state.ysize + 3) // 4): for d in decode_dxt5(self._safe_read(linesize)): data += d else: - msg = f"Unsupported alpha encoding {repr(self._blp_alpha_encoding)}" + msg = f"Unsupported alpha encoding {repr(self._alpha_encoding)}" raise BLPFormatError(msg) else: - msg = f"Unknown BLP encoding {repr(self._blp_encoding)}" + msg = f"Unknown BLP encoding {repr(self._encoding)}" raise BLPFormatError(msg) else: - msg = f"Unknown BLP compression {repr(self._blp_compression)}" + msg = f"Unknown BLP compression {repr(self._compression)}" raise BLPFormatError(msg) self.set_as_raw(data) @@ -472,10 +475,15 @@ def _save(im: Image.Image, fp: IO[bytes], filename: str | bytes) -> None: assert im.palette is not None fp.write(struct.pack("