import unittest import zlib from io import BytesIO from PIL import Image, ImageFile, PngImagePlugin from .helper import PillowTestCase TEST_FILE = "Tests/images/png_decompression_dos.png" class TestPngDos(PillowTestCase): def test_ignore_dos_text(self): ImageFile.LOAD_TRUNCATED_IMAGES = True try: im = Image.open(TEST_FILE) im.load() finally: ImageFile.LOAD_TRUNCATED_IMAGES = False for s in im.text.values(): assert len(s) < 1024 * 1024, "Text chunk larger than 1M" for s in im.info.values(): assert len(s) < 1024 * 1024, "Text chunk larger than 1M" def test_dos_text(self): try: im = Image.open(TEST_FILE) im.load() except ValueError as msg: assert msg, "Decompressed Data Too Large" return for s in im.text.values(): assert len(s) < 1024 * 1024, "Text chunk larger than 1M" def test_dos_total_memory(self): im = Image.new("L", (1, 1)) compressed_data = zlib.compress(b"a" * 1024 * 1023) info = PngImagePlugin.PngInfo() for x in range(64): info.add_text("t%s" % x, compressed_data, zip=True) info.add_itxt("i%s" % x, compressed_data, zip=True) b = BytesIO() im.save(b, "PNG", pnginfo=info) b.seek(0) try: im2 = Image.open(b) except ValueError as msg: assert "Too much memory" in msg return total_len = 0 for txt in im2.text.values(): total_len += len(txt) assert total_len < 64 * 1024 * 1024, "Total text chunks greater than 64M" if __name__ == "__main__": unittest.main()