Exercise 3 solution

    import os
    import struct
    import tempfile
    import zlib

    SIG = b"\x89PNG\r\n\x1a\n"


    def chunk(ctype, data=b""):
        body = ctype + data
        return struct.pack(">I", len(data)) + body + struct.pack(">I", zlib.crc32(body))


    def write_gray_png(path, rows):
        h, w = len(rows), len(rows[0])
        if any(len(r) != w for r in rows):
            raise ValueError("rows must all be the same length")
        ihdr = struct.pack(">IIBBBBB", w, h, 8, 0, 0, 0, 0)   # 8-bit greyscale
        raw = b"".join(b"\x00" + bytes(r) for r in rows)     # filter type 0 per row
        with open(path, "wb") as f:
            f.write(SIG + chunk(b"IHDR", ihdr)
                    + chunk(b"IDAT", zlib.compress(raw))
                    + chunk(b"IEND"))


    def read_gray_png(path):
        with open(path, "rb") as f:
            raw = f.read()
        if raw[:8] != SIG:
            raise ValueError("not a PNG")
        pos, idat, header = 8, b"", None
        while pos + 12 <= len(raw):
            length, ctype = struct.unpack_from(">I4s", raw, pos)
            data = raw[pos + 8:pos + 8 + length]
            (stored,) = struct.unpack_from(">I", raw, pos + 8 + length)
            if stored != zlib.crc32(ctype + data):
                raise ValueError(f"CRC mismatch in {ctype!r}")
            if ctype == b"IHDR":
                header = struct.unpack(">IIBBBBB", data)
            elif ctype == b"IDAT":
                idat += data
            elif ctype == b"IEND":
                break
            pos += 12 + length
        w, h, depth, ctype_, comp, filt, interlace = header
        if (depth, ctype_, interlace) != (8, 0, 0):
            raise ValueError("only 8-bit greyscale, non-interlaced")
        pixels = zlib.decompress(idat)
        stride = 1 + w
        rows = []
        for y in range(h):
            row = pixels[y * stride:(y + 1) * stride]
            if row[0] != 0:
                raise ValueError(f"row {y} uses filter {row[0]}, only 0 supported")
            rows.append(list(row[1:]))
        return rows


    folder = tempfile.mkdtemp()
    path = os.path.join(folder, "t.png")
    grid = [[0, 50, 100, 150], [200, 250, 255, 0], [10, 20, 30, 40]]
    write_gray_png(path, grid)
    print(os.path.getsize(path) > 8, open(path, "rb").read(8).hex(" "))
    back = read_gray_png(path)
    print(back)
    print(back == grid)

    with open(path, "r+b") as f:
        f.seek(17)
        f.write(b"\x09")
    try:
        read_gray_png(path)
    except ValueError as e:
        print("ValueError:", e)

    os.remove(path)
    os.rmdir(folder)
    print(os.path.exists(folder))

    # True 89 50 4e 47 0d 0a 1a 0a
    # [[0, 50, 100, 150], [200, 250, 255, 0], [10, 20, 30, 40]]
    # True
    # ValueError: CRC mismatch in b'IHDR'
    # False

(Run as written; the output matched the comments.)

Explanation: The writer puts a zero filter byte in front of every row, compresses the rows with
zlib and wraps them in IHDR, IDAT and IEND. The reader checks each CRC, joins the
IDAT data, decompresses it and strips the filter bytes, refusing any row that uses a
filter it does not implement. Changing one byte inside IHDR is caught by the CRC
before the header is ever interpreted. The folder is removed afterwards.