#!/usr/bin/env python3

from binascii import crc32
from pathlib import Path

def overwrite_at_pos(original: bytes, offset: int, replacement: bytes) -> bytes:
    return original[:offset] + replacement + original[offset + len(replacement):]

def num_to_4b(i: int) -> bytes:
    return i.to_bytes(4, byteorder='big')

def main():
    with (Path(__file__).parent / "mystery").open('rb') as file:
        content = file.read()

    # From https://www.libpng.org/pub/png/spec/1.2/PNG-Structure.html:
    # > The first eight bytes of a PNG file always contain the following (decimal) values:
    # > 137 80 78 71 13 10 26 10

    PNG_MAGIC_BYTES = bytes([137, 80, 78, 71, 13, 10, 26, 10])

    # https://www.libpng.org/pub/png/spec/1.2/PNG-Chunks.html:
    #
    # > 3.2. Chunk layout
    # > Each chunk consists of four parts:
    # >
    # > Length
    # >     A 4-byte unsigned integer giving the number of bytes in the chunk's data field.
    # >     The length counts only the data field, not itself, the chunk type code, or the CRC.
    # >     Zero is a valid length. Although encoders and decoders should treat the length as unsigned,
    # >     its value must not exceed 231 bytes.
    # >
    # > Chunk Type
    # >     A 4-byte chunk type code. For convenience in description and in examining PNG files,
    # >     type codes are restricted to consist of uppercase and lowercase ASCII letters
    # >     (A-Z and a-z, or 65-90 and 97-122 decimal). However, encoders and decoders must treat the
    # >     codes as fixed binary values, not character strings. For example, it would not be correct
    # >     to represent the type code IDAT by the EBCDIC equivalents of those letters.
    # >     Additional naming conventions for chunk types are discussed in the next section.
    # >
    # > Chunk Data
    # >     The data bytes appropriate to the chunk type, if any. This field can be of zero length.
    # >
    # > CRC
    # >     A 4-byte CRC (Cyclic Redundancy Check) calculated on the preceding bytes in the chunk,
    # >     including the chunk type code and chunk data fields, but not including the length field.
    # >     The CRC is always present, even for chunks containing no data. See CRC algorithm.

    S_LENGTH = 4
    S_TYPE = 4
    S_CRC = 4

    # https://www.libpng.org/pub/png/spec/1.2/PNG-Chunks.html:
    # > IHDR length is static:
    # > Width:              4 bytes
    # > Height:             4 bytes
    # > Bit depth:          1 byte
    # > Color type:         1 byte
    # > Compression method: 1 byte
    # > Filter method:      1 byte
    # > Interlace method:   1 byte

    IHDR_LENGTH = 4 + 4 + 1 * 5
    IHDR_CHUNK_TYPE = b"IHDR"

    PREFIX = PNG_MAGIC_BYTES + num_to_4b(IHDR_LENGTH) + IHDR_CHUNK_TYPE
    new_content = overwrite_at_pos(content, 0 , PREFIX)

    # $ pngcheck mystery.png
    # mystery.png  CRC error in chunk pHYs (computed 38d82c82, expected 495224f0)

    checksum_location = new_content.find(b'\x49\x52\x24\xf0')
    new_content = overwrite_at_pos(new_content, checksum_location, b'\x38\xd8\x2c\x82')

    # $ pngcheck -v mystery.png
    # File: mystery.png (202940 bytes)
    #   chunk IHDR at offset 0x0000c, length 13
    #     1642 x 1095 image, 24-bit RGB, non-interlaced
    #   chunk sRGB at offset 0x00025, length 1
    #     rendering intent = perceptual
    #   chunk gAMA at offset 0x00032, length 4: 0.45455
    #   chunk pHYs at offset 0x00042, length 9: 2852132389x5669 pixels/meter
    # :  invalid chunk length (too large)
    # ERRORS DETECTED in mystery.png

    # https://www.libpng.org/pub/png/spec/1.2/PNG-Chunks.html:
    # > pHYs length is static:
    # > Pixels per unit, X axis: 4 bytes (unsigned integer)
    # > Pixels per unit, Y axis: 4 bytes (unsigned integer)
    # > Unit specifier:          1 byte

    PHYS_SIZE = 4 + 4 + 1

    # Next chunk seems broken, both the length and the type, maybe type is ?DET -> IDAT?

    phys_type_location = new_content.find(b'pHYs')
    next_chunk_length_pos = phys_type_location + S_TYPE + PHYS_SIZE + S_CRC
    new_content = overwrite_at_pos(new_content, next_chunk_length_pos + S_LENGTH, b'IDAT')

    # https://www.libpng.org/pub/png/spec/1.2/PNG-Chunks.html:
    # > There can be multiple IDAT chunks; if so, they must appear consecutively with no other intervening chunks.
    # So we can calculate the expected size by finding the next idat chunk

    chunk_data_length = new_content[next_chunk_length_pos + S_LENGTH + S_TYPE:].find(b'IDAT') - S_LENGTH - S_CRC
    new_content = overwrite_at_pos(new_content, next_chunk_length_pos, num_to_4b(chunk_data_length))


    with (Path(__file__).parent / "mystery.png").open('wb') as file:
        file.write(new_content)

if __name__ == '__main__':
    main()