WTS/game/scripts/utility/punk.rpy

#####################################
##     Created by briandeheus      ##
## https://github.com/briandeheus  ##
##    Implementation and changes   ##
##           LoafyLemon            ##
#####################################
init python in image_payload:
    import binascii
    import struct
    import zlib

    CHUNK_TYPE_END = "IEND"
    CHUNK_TYPE_PUNK = "wtSi"
    MAX_BYTES = 2147483647
    SIGNATURE_BYTES = 8
    BYTES_IN_KB = 2014

    def bytes_to_hex(b):
        return b.hex()

    def bytes_to_utf(b):
        return b.decode()

    def bytes_to_int(b):
        return int(bytes_to_hex(b), 16)

    def read_bytes(f, byte_count: int):
        return f.read(byte_count)

    def rewind_bytes(f, byte_count):
        f.seek(f.tell() - byte_count)

    def get_file_length(f):
        f.seek(0, os.SEEK_END)
        file_length = f.tell()
        f.seek(0)

        return file_length

    def read_chunk(f):
        chunk_size = read_bytes(f, 4)
        chunk_type = read_bytes(f, 4)
        chunk_content = read_bytes(f, bytes_to_int(chunk_size))
        chunk_crc = read_bytes(f, 4)

        return [chunk_size, chunk_type, chunk_content, chunk_crc]

    def inject_punk_chunk(f, content):
        chunk_size = len(content)

        if chunk_size > MAX_BYTES:
            raise ValueError(f"Cannot inject more than {MAX_BYTES} bytes")

        print(f"Injecting {CHUNK_TYPE_PUNK} chunk {chunk_size / BYTES_IN_KB} kb")

        # Create a byte array to store our chunk data in.
        tmp_bytes = bytearray()
        # First write the chunk type
        tmp_bytes.extend(CHUNK_TYPE_PUNK.encode())
        # Now write the bytes of whatever we're trying to hide
        tmp_bytes.extend(content)

        # Write the chunk size
        f.write(bytearray(struct.pack("!i", chunk_size)))

        # And the content
        f.write(tmp_bytes)

        crc = binascii.crc32(tmp_bytes)
        crc_bytes = crc.to_bytes(4, "big")
        print("Chunk CRC", bytes_to_hex(crc_bytes))
        f.write(crc_bytes)

        print("Chunk injected!")

    def list(input):
        path = os.path.join(config.gamedir, "outfits", input)

        with open(path, "rb") as input_file:

            input_file_length = get_file_length(input_file)
            input_file.read(SIGNATURE_BYTES)

            while True:
                chunk_size, chunk_type, chunk_content, chunk_crc = read_chunk(input_file)
                chunk_type_str = bytes_to_utf(chunk_type)
                print(f"Chunk {chunk_type_str}, {bytes_to_int(chunk_size)} bytes")

                if input_file.tell() >= input_file_length:
                    return

    def inject(input, output, content):
        input_path = os.path.join(config.gamedir, "outfits", input)
        output_path = os.path.join(config.gamedir, "outfits", output)
        content = zlib.compress(str(content).encode())

        with open(input_path, "rb") as input_file, open(output_path, "wb") as output_file:

            input_file_length = get_file_length(input_file)
            output_file.write(input_file.read(SIGNATURE_BYTES))

            while True:
                chunk_size, chunk_type, chunk_content, chunk_crc = read_chunk(input_file)
                chunk_type_str = bytes_to_utf(chunk_type)
                print(f"Chunk {chunk_type_str}, {bytes_to_int(chunk_size)} bytes")

                if chunk_type_str == CHUNK_TYPE_END:
                    inject_punk_chunk(output_file, content)

                output_file.write(chunk_size)
                output_file.write(chunk_type)
                output_file.write(chunk_content)
                output_file.write(chunk_crc)

                if input_file.tell() >= input_file_length:
                    return

    def extract(input):
        print("Attempting to extract punked data from", input)
        path = os.path.join(config.gamedir, "outfits", input)

        with open(path, "rb") as input_file:

            input_file_length = get_file_length(input_file)
            input_file.read(SIGNATURE_BYTES)

            while True:
                chunk_size, chunk_type, chunk_content, chunk_crc = read_chunk(input_file)
                chunk_type_str = bytes_to_utf(chunk_type)

                if chunk_type_str == CHUNK_TYPE_PUNK:
                    print("Found a punk chunk worth", bytes_to_int(chunk_size), "bytes")
                    return zlib.decompress(chunk_content).decode()

                if input_file.tell() >= input_file_length:
                    print("No punked data found")
                    return