Python PNG parser

png_parser.py

# PNG spec http://www.libpng.org/pub/png/spec/1.2/png-1.2-pdg.html

import argparse
import binascii
import struct
from typing import cast

parser = argparse.ArgumentParser(prog="png_parse",
                                 description="Parse a PNG file and show the file information",
                                 epilog="See http://www.libpng.org/pub/png/spec/1.2/png-1.2-pdg.html for the PNG spec")

parser.add_argument("filenames",
                    type=str,
                    help="PNG file(s) to parse",
                    nargs='+')

args = parser.parse_args()
image_names = args.filenames

class ParsingError(Exception):
    pass

def propcol(s: str) -> str:
    return f"\033[31;1m{s}\033[0m"

def err(s: str) -> str:
    print(f"\033[97;41;1m{s}\033[0m")
    raise ParsingError

def pop_first(n: int, b: bytearray) -> bytearray:
    if len(b) < n:
        err("Not enough bytes remaining in file")
    ret = b[:n]
    del b[:n]
    return ret

def uint32(b) -> int:
    assert len(b) == 4
    return struct.unpack('!I', b)[0]

def parseIHDR(chunkData):
    ihdrData = chunkData.copy()
    width = uint32(pop_first(4, ihdrData))
    height = uint32(pop_first(4, ihdrData))
    bitDepth = pop_first(1, ihdrData)[0]
    colorType = pop_first(1, ihdrData)[0]
    if not colorType in [0, 2, 3, 4, 6]:
        err(f"Invalid color type {colorType}, only 0, 2, 3, 4, 6 are allowed")
    if (colorType == 0 and not bitDepth in [1, 2, 4, 8, 16]
        or colorType == 2 and not bitDepth in [8, 16]
        or colorType == 3 and not bitDepth in [1, 2, 4, 8]
        or colorType == 4 and not bitDepth in [8, 16]
        or colorType == 6 and not bitDepth in [8, 16]):
            err(f"Invalid bit depth {bitDepth} for color type {colorType}. See http://www.libpng.org/pub/png/spec/1.2/png-1.2-pdg.html#C.IHDR")
    
    compressionMethod = pop_first(1, ihdrData)[0]
    filterMethod = pop_first(1, ihdrData)[0]
    interlaceMethod = pop_first(1, ihdrData)[0]
    print("    " + propcol("Width: ") + str(width),
            propcol("Height: ") + str(height), propcol("Bit depth: ") + str(bitDepth),
            propcol("Color type: ") + str(colorType) + ' \033[90m' + ' '.join(("PALETTE" if colorType & 1 else '',
                                                                                "COLOR" if colorType & 2 else '',
                                                                                "ALPHA" if colorType & 4 else '')) + '\033[0m',
            propcol("Compression method: ") + str(compressionMethod), propcol("Filter method: ") + str(filterMethod),
            sep='\t')

for image_name in image_names:
    try:
        with open(image_name, 'rb') as f:
            data = bytearray(f.read())

        print(f"\033[1mImage \033[32m\"{image_name}\"\033[0m")

        PNGmagick = pop_first(8, data)
        print(propcol("PNGmagick:"), bytes(PNGmagick), sep='\t')
        if PNGmagick != b"\x89PNG\r\n\x1a\n":
            err("File magic does not match b'\x89PNG\r\n\x1a\n'")

        print(propcol("Chucks:"))

        chunk_i = 0
        seen_iend = False  # the chunk of type IEND must appear last
        while data:
            print(f"\033[34;1m  Chunk nr. {chunk_i}\033[0m")

            chunkLen = uint32(pop_first(4, data))
            print(propcol("  chunkLen:"), chunkLen, sep='\t')

            chunkType = pop_first(4, data)
            print(propcol("  chunkType:"), bytes(chunkType), sep='\t', end=' \033[90m')
            for c in chunkType:
                if c < 65 or c > 90 and c < 97 or c > 122:
                    err(f"    Byte {c} not allowed in Chunk Type")
            print("ANCILLARY" if chunkType[0] & 32 else "CRITICAL",
                    "PRIVATE" if chunkType[1] & 32 else "PUBLIC",
                    "NOT CONFORMING" if chunkType[1] & 32 else "1.2-CONFORMING",
                    "SAFE-TO-COPY" if chunkType[1] & 32 else "UNSAFE-TO-COPY",
                    end='\033[0m\n')
            
            if chunk_i == 0 and chunkType != b"IHDR":
                err(f"    Chunk type of chunk nr. 0 must be b'IHDR' and not {chunkType}")
            if seen_iend:
                err(f"    There are more chunks after the b'IEND' chunk")
            if chunkType == b"IEND":
                seen_iend = True
                if chunkLen != 0:
                    err("Chunk with type b'IEND' must be empty")
            
            chunkData = pop_first(chunkLen, data)
            if chunkType == b"IHDR":        # Parse IHDR chunk
                parseIHDR(chunkData)
            else:
                print(propcol("  chunkData:"), str(bytes(chunkData[:8])) + ("..." if chunkLen > 8 else ''), sep='\t')

            chunkCRC = uint32(pop_first(4, data))
            print(propcol("  chunkCRC:"), "0x" + format(chunkCRC, 'X'), sep='\t', end=' ')
            if (calcCRC := binascii.crc32(chunkType + chunkData)) != chunkCRC:
                print()
                err(f"ChunkCRC (0x{chunkCRC:X}) does not match calculated crc (0x{calcCRC:X})")
            else:
                print("\033[32mCORRECT\033[0m")

            print()
            chunk_i += 1
        print()
    except ParsingError:
        pass