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()