Pillow/Tests/check_png_dos.py
Jon Dufresne d50445ff30 Introduce isort to automate import ordering and formatting
Similar to the recent adoption of Black. isort is a Python utility to
sort imports alphabetically and automatically separate into sections. By
using isort, contributors can quickly and automatically conform to the
projects style without thinking. Just let the tool do it.

Uses the configuration recommended by the Black to avoid conflicts of
style.

Rewrite TestImageQt.test_deprecated to no rely on import order.
2019-07-06 16:11:35 -07:00

69 lines
1.8 KiB
Python

import zlib
from io import BytesIO
from PIL import Image, ImageFile, PngImagePlugin
from .helper import PillowTestCase, unittest
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():
self.assertLess(len(s), 1024 * 1024, "Text chunk larger than 1M")
for s in im.info.values():
self.assertLess(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:
self.assertTrue(msg, "Decompressed Data Too Large")
return
for s in im.text.values():
self.assertLess(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:
self.assertIn("Too much memory", msg)
return
total_len = 0
for txt in im2.text.values():
total_len += len(txt)
self.assertLess(
total_len, 64 * 1024 * 1024, "Total text chunks greater than 64M"
)
if __name__ == "__main__":
unittest.main()