from pngparser import PngParser, chunktypes
import zlib


def calc_crc(chunk):
    return zlib.crc32(chunk.type + chunk.data).to_bytes(chunktypes.CHUNK_CRC_SIZE, 'big')


def reconstruct_image():
    parser = PngParser("./intro-forensics-3")
    header = parser.get_header()
    header.crc = calc_crc(header)

    data_chunks = sorted([chunk  for chunk in parser.get_all() if chunk.type not in  [b"IHDR", b"IEND"]], key=lambda x: x.crc)

    for chunk in data_chunks:
        chunk.crc = calc_crc(chunk)

    end_chunk = parser.get_by_type(chunktypes.TYPE_IEND)
    end_chunk[0].crc = calc_crc(end_chunk[0])

    chunks = [header] + data_chunks + end_chunk

    png_file = bytes.fromhex("89 50 4E 47 0D 0A 1A 0A".replace(" ", ""))

    for chunk in chunks:
        png_file += chunk.to_bytes()

    with open("reconstructed.png", "wb") as f:
        f.write(png_file)

if __name__ == "__main__":
    reconstruct_image()