Pillow/Tests/check_png_dos.py

69 lines
1.8 KiB
Python
Raw Normal View History

2014-12-30 04:10:27 +03:00
import zlib
from io import BytesIO
from PIL import Image, ImageFile, PngImagePlugin
from .helper import PillowTestCase, unittest
2014-12-31 03:57:24 +03:00
TEST_FILE = "Tests/images/png_decompression_dos.png"
2015-04-24 02:26:52 +03:00
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():
2019-06-13 18:53:42 +03:00
self.assertLess(len(s), 1024 * 1024, "Text chunk larger than 1M")
for s in im.info.values():
2019-06-13 18:53:42 +03:00
self.assertLess(len(s), 1024 * 1024, "Text chunk larger than 1M")
def test_dos_text(self):
try:
2014-12-31 03:57:24 +03:00
im = Image.open(TEST_FILE)
im.load()
except ValueError as msg:
2014-12-31 03:57:24 +03:00
self.assertTrue(msg, "Decompressed Data Too Large")
return
for s in im.text.values():
2019-06-13 18:53:42 +03:00
self.assertLess(len(s), 1024 * 1024, "Text chunk larger than 1M")
2014-12-30 04:10:27 +03:00
def test_dos_total_memory(self):
2019-06-13 18:53:42 +03:00
im = Image.new("L", (1, 1))
compressed_data = zlib.compress(b"a" * 1024 * 1023)
2014-12-30 04:10:27 +03:00
info = PngImagePlugin.PngInfo()
for x in range(64):
2019-06-13 18:53:42 +03:00
info.add_text("t%s" % x, compressed_data, zip=True)
info.add_itxt("i%s" % x, compressed_data, zip=True)
2014-12-30 04:10:27 +03:00
b = BytesIO()
2019-06-13 18:53:42 +03:00
im.save(b, "PNG", pnginfo=info)
2014-12-30 04:10:27 +03:00
b.seek(0)
2015-04-24 02:26:52 +03:00
2014-12-30 04:10:27 +03:00
try:
im2 = Image.open(b)
except ValueError as msg:
2014-12-31 03:57:24 +03:00
self.assertIn("Too much memory", msg)
2014-12-30 04:10:27 +03:00
return
total_len = 0
for txt in im2.text.values():
total_len += len(txt)
2019-06-13 18:53:42 +03:00
self.assertLess(
total_len, 64 * 1024 * 1024, "Total text chunks greater than 64M"
)
2015-04-24 02:26:52 +03:00
2018-03-03 12:54:00 +03:00
2019-06-13 18:53:42 +03:00
if __name__ == "__main__":
unittest.main()