spore_ep1stegano.py

#!/usr/bin/env python3

'''Encodes and decodes Spore EP1 stegano chunk (spOr) on PNG images.'''

__author__ = '0KepOnline' # https://github.com/0KepOnline
__license__ = 'Unlicense' # http://unlicense.org
__maintainer__ = '0KepOnline'
__email__ = 'i.am@scenariopla.net'




import argparse
import os
import struct
import traceback
import zlib
import sys
from io import BufferedReader, BufferedWriter
from typing import Callable
from crcmod import mkCrcFun as make_crc32
from crcmod.predefined import PredefinedCrc
from png import Reader as PngReader
from zopfli import zlib as zlib_zopfli




FILE_MAGIC: bytes = b'\x89PNG\r\n\x1a\n'
STEGANO_DATA_MAGIC: bytes = b'\x5c\xa3\xf1\xe6' # converted to be
UNKNOWN_NULL32: bytes = b'\0' * 4

# default/mandatory png chunks
CHUNK_HEADER: bytes = b'IHDR'
CHUNK_IMAGE_DATA: bytes = b'IDAT'
CHUNK_END: bytes = b'IEND'
# spore steganography m-i-p data chunk
CHUNK_STEGANO_DATA: bytes = b'spOr'

def _make_crc32_mpeg2(initial: int, invert: bool = False) -> Callable[[bytes], int]:
    '''Returns a customized CRC32/MPEG-2 checksum generating function.

    Arguments:
        initial (int): The initial CRC32 value.
        invert (bool): Should the value be inverted after it's generated or not.
                       Default: False.
    
    Returns:
        out (Callable[[bytes], int]): The customized CRC32/MPEG-2 function.
    '''
    crc32_mpeg2_predefined: PredefinedCrc = PredefinedCrc('crc-32-mpeg')
    crc32_mpeg2 = make_crc32(crc32_mpeg2_predefined.poly,
                             initial,
                             crc32_mpeg2_predefined.reverse,
                             crc32_mpeg2_predefined.xorOut)
    if invert:
        return lambda data: (~crc32_mpeg2(data) & 0xffffffff)
    return crc32_mpeg2

CRC32_BGRA32_INITIAL: int = 1581049650
CRC32_BGRA32: Callable[[bytes], int] = _make_crc32_mpeg2(CRC32_BGRA32_INITIAL)

CRC32_COMPLEX: Callable[[bytes, int], int] = lambda data, crc32_bgra32: (
    _make_crc32_mpeg2(crc32_bgra32, invert=True)(data)
)

CRC32_STEGANO_DATA_INITIAL: int = 786340527 # crc32/mpeg-2 (CHUNK_STEGANO_DATA)
CRC32_STEGANO_DATA: Callable[[bytes], int] = _make_crc32_mpeg2(CRC32_STEGANO_DATA_INITIAL,
                                                               invert=True)

#
# stegano chunk structure:
#
#   length ({CHUNK_STEGANO_DATA}) [4] be
#   CHUNK_STEGANO_DATA [4] {
#       STEGANO_DATA_MAGIC [4]
#       UNKNOWN_NULL32 [4]
#       zlib-compressed data [~]
#       crc32/mpeg-2 ({zlib-compressed data},
#                     initial = crc32/mpeg-2 ({BGRA32}, initial = CRC32_BGRA32_INITIAL),
#                     invert) [4] le
#       length (zlib-compressed data) [4] le
#       STEGANO_DATA_MAGIC [4]
#   } crc32/mpeg-2 ({CHUNK_STEGANO_DATA}, initial = CRC32_STEGANO_DATA_INITIAL, invert) [4] be
#

ZOPFLI_KWARGS: dict = {}

_ESC_ERROR: str = '\x1b[1;31m'
_ESC_REFRESH: str = '\x1b[0m'
_message_error: str = 'ERROR:'
if hasattr(sys.stdout, 'isatty') and sys.stdout.isatty():
    _message_error = _ESC_ERROR + _message_error + _ESC_REFRESH

def _print_error(message: str) -> None:
    '''Prints an error with ANSI color codes if supported.
    
    Arguments:
        message (str): The error message to print.
    '''
    print(_message_error + ' ' + message)



class Stegano:
    '''A class that represents Spore EP1 stegano chunk data in a modifiable form.
    
    Attributes:
        data (bytes): The data stored in a stegano container.
        image_checksum (int): The CRC32/MPEG-2 checksum of the image
                              (in BGRA32; initial = 1581049650).
        use_zopfli (bool): Determines if Zopfli should be used for compression instead of pure Zlib.
                           Default: True.
        zlib_data (bytes): The data stored in a stegano container, in a compressed form (read-only).
    '''
    def _compile_chunk_data(self) -> bytes:
        '''Returns the compiled data to be stored in a PNG chunk.'''
        return (STEGANO_DATA_MAGIC
                + UNKNOWN_NULL32
                + self.zlib_data
                + struct.pack('<I', CRC32_COMPLEX(self.zlib_data, self.image_checksum))
                + struct.pack('<I', len(self.zlib_data) + 4)
                + STEGANO_DATA_MAGIC)

    def __init__(self,
                 image_checksum: int,
                 data: bytes | BufferedReader = b'',
                 use_zopfli: bool = True,
                 compressed: bool = False) -> None:
        '''Initializes a new Spore EP1 stegano container based on given data and configuration.
        
        Arguments:
            image_checksum (int): The CRC32/MPEG-2 checksum of the image
                                  (in BGRA32; initial = 1581049650).
            data (bytes): The data stored in a stegano container.
            use_zopfli (bool): Determines if Zopfli should be used to improve compression.
                               Default: True.
            compressed (bool): Determines if the data should be treated as already Zlib-compressed.

        Returns:
            out (Stegano): The initialized Spore EP1 stegano container.
        '''
        if data == b'':
            compressed = True
        self.image_checksum: int = image_checksum
        self.use_zopfli: bool = use_zopfli
        if isinstance(data, BufferedReader):
            data = data.read()
        if compressed:
            self._zlib_data: bytes = data
        else:
            self.data: bytes = data

    @property
    def zlib_data(self) -> bytes:
        '''The data stored in a stegano container, in a compressed form (read-only).'''
        return self._zlib_data

    @property
    def data(self) -> bytes:
        '''The data stored in a stegano container.'''
        return zlib.decompress(self.zlib_data)
    @data.setter
    def data(self, data_modified: bytes | BufferedReader) -> None:
        if isinstance(data_modified, BufferedReader):
            data_modified = data_modified.read()
        self._zlib_data: bytes = _compress_zlib(data_modified, self.use_zopfli)

    def __len__(self) -> int:
        '''Returns the length of the compiled stegano data in a PNG chunk.'''
        # 2xSTEGANO_DATA_MAGIC + UNKNOWN_NULL32 + zlib data length + complex checksum
        return len(self._zlib_data) + 20

    def calculate_checksum(self) -> int:
        '''Returns the CRC32/MPEG-2 checksum of the compiled stegano data
        (initial = 786340527, inverted).
        '''
        return CRC32_STEGANO_DATA(self._compile_chunk_data())

    def compile(self) -> bytes:
        '''Compiles the full PNG chunk of the stegano container.'''
        return (struct.pack('>I', len(self))
                + CHUNK_STEGANO_DATA
                + self._compile_chunk_data()
                + struct.pack('>I', self.calculate_checksum()))



def _compress_zlib(data: bytes, use_zopfli: bool = True) -> bytes:
    '''Compresses the data with Zlib algorithm.
        
    Arguments:
        data (bytes): The data to compress.
        use_zopfli (bool): Determines if Zopfli should be used to improve compression.
                           Default: True.

    Returns:
        out (bytes): The compressed data.
    '''
    if use_zopfli:
        return zlib_zopfli.compress(data, **ZOPFLI_KWARGS)
    return zlib.compress(data, 9)

def convert_rgba32_to_bgra32(rgba32_pixels: bytearray) -> bytearray:
    '''Converts RGBA32 pixels to BGRA32 image format.
        
    Arguments:
        rgba32_pixels (bytearray): The pixels in RGBA32 format.

    Returns:
        out (bytearray): The pixels in BGRA32 format.
    '''
    bgra32_pixels: bytearray = bytearray(len(rgba32_pixels))
    bgra32_pixels[0::4] = rgba32_pixels[2::4]
    bgra32_pixels[1::4] = rgba32_pixels[1::4]
    bgra32_pixels[2::4] = rgba32_pixels[0::4]
    bgra32_pixels[3::4] = rgba32_pixels[3::4]
    return bgra32_pixels

def read_png_as_bgra32(file: bytes | BufferedReader) -> bytearray | None:
    '''Reads pixels of the PNG image in BGRA32 image format.
        
    Arguments:
        file (bytes | BufferedReader): The PNG file to read.

    Returns:
        out (bytearray): The pixels in BGRA32 format.
    '''
    reader: PngReader
    try:
        if isinstance(file, bytes):
            reader = PngReader(bytes=file)
        else:
            reader = PngReader(file)
        _, _, rows, _ = reader.asRGBA8()
        pixel_data: bytearray = bytearray()
        for row in rows:
            pixel_data.extend(row)
        return convert_rgba32_to_bgra32(pixel_data)
    except Exception:
        return None

def read_stegano_from_chunk(stegano_chunk: bytes,
                            image_checksum: int,
                            validate: bool = True) -> Stegano | None:
    '''Reads the stegano container from the PNG chunk.
        
    Arguments:
        stegano_chunk (bytes): The PNG chunk to read.
        image_checksum (int): The CRC32/MPEG-2 checksum of the image
                              (in BGRA32; initial = 1581049650).
        validate (bool): Determines if the checksums should be double-checked.
                         Default: True.

    Returns:
        out (Stegano | None): The stegano container, or None if failed.
    '''
    if stegano_chunk[4:8] != CHUNK_STEGANO_DATA or stegano_chunk[8:12] != STEGANO_DATA_MAGIC:
        return None
    stegano_chunk_length: int = struct.unpack('>I', stegano_chunk[:4])[0]
    stegano_data: bytes = stegano_chunk[8 : stegano_chunk_length + 8]
    if (stegano_data[-4:] != STEGANO_DATA_MAGIC or
        struct.unpack('<I', stegano_data[-8:-4])[0] != len(stegano_data[8:-8])):
        return None
    zlib_data: bytes = stegano_data[8:-12]
    stegano: Stegano = Stegano(image_checksum, zlib_data, compressed=True)
    if validate:
        stegano_chunk_checksum: int = struct.unpack('>I', stegano_chunk[-4:])[0]
        complex_checksum: int = struct.unpack('<I', stegano_data[-12:-8])[0]
        if (stegano_chunk_checksum != CRC32_STEGANO_DATA(stegano_data) or
            complex_checksum != CRC32_COMPLEX(zlib_data, image_checksum)):
            return None
    return stegano

def read_stegano_from_png(file: bytes | BufferedReader, validate: bool = True) -> Stegano | None:
    '''Reads the stegano container from the PNG image.
        
    Arguments:
        file (bytes | BufferedReader): The PNG file to read.
        validate (bool): Determines if the checksums should be double-checked.
                         Default: True.

    Returns:
        out (Stegano | None): The stegano container, or None if failed.
    '''
    chunk_length_bytes: bytes
    chunk_length: int
    chunk_type: bytes
    stegano_chunk: bytes
    try:
        if isinstance(file, BufferedReader):
            while True:
                chunk_length_bytes = file.read(4)
                if not chunk_length_bytes:
                    return None
                chunk_length = struct.unpack('>I', chunk_length_bytes)[0]
                chunk_type = file.read(4)
                if chunk_type == CHUNK_STEGANO_DATA:
                    stegano_chunk = (chunk_length_bytes
                                    + chunk_type
                                    + file.read(chunk_length + 4))
                else:
                    file.seek(chunk_length + 4, 1)
        else:
            input_cursor: int = 8
            while input_cursor < len(file):
                chunk_length_bytes = file[input_cursor : input_cursor + 4]
                input_cursor += 4
                chunk_length = struct.unpack('>I', chunk_length_bytes)[0]
                chunk_type = file[input_cursor : input_cursor + 4]
                input_cursor += 4
                if chunk_type == CHUNK_STEGANO_DATA:
                    stegano_chunk = (chunk_length_bytes
                                    + chunk_type
                                    + file[input_cursor : input_cursor + chunk_length + 4])
                else:
                    input_cursor += chunk_length + 4
        image_checksum: int = CRC32_BGRA32(read_png_as_bgra32(file))
        return read_stegano_from_chunk(stegano_chunk, image_checksum, validate)
    except Exception:
        return None

def _action_extract_stegano_zlib(input_png: bytes | BufferedReader,
                                 output_stegano: BufferedWriter,
                                 validate: bool = True) -> int: # status code
    '''Action: extracts the data from the stegano container in the PNG image.
        
    Arguments:
        input_png (bytes | BufferedReader): The PNG file to read.
        output_stegano (BufferedWriter): The data file to write.
        validate (bool): Determines if the checksums should be double-checked.
                         Default: True.

    Returns:
        out (int): The status code.
    '''
    if not output_stegano.writable:
        _print_error('Output file is not writeable.')
        return 1
    if isinstance(input_png, BufferedReader):
        input_png = input_png.read()
    stegano: Stegano = read_stegano_from_png(input_png, validate)
    if stegano is None:
        _print_error('Stegano reading failed. Please ensure that the file is a M-i-P of Spore EP1.')
        return 1
    try:
        zlib_data_decompressed: bytes = zlib.decompress(stegano.zlib_data)
    except Exception:
        _print_error('Unable to decompress Zlib data.')
        return 1
    output_stegano.write(zlib_data_decompressed)
    return 0

def _action_skip_namesakes(file_path: str) -> str:
    '''Action: skips existing files by adding the index in brackets.
        
    Arguments:
        file_path (str): The original file path.

    Returns:
        out (str): The modified file path that avoids already existing files.
    '''
    if os.path.exists(file_path):
        file_path_split: tuple[str, str] = os.path.splitext(file_path)
        file_path_name: str = file_path_split[0]
        file_path_ext: str = file_path_split[1]
        namesake_counter: int = 1
        while os.path.exists(file_path):
            file_path = f'{file_path_name} ({namesake_counter}){file_path_ext}'
            namesake_counter += 1
    return file_path

def _action_try_remove_unfinished_file(file_path: str) -> None:
    '''Action: tries to remove the unfinished file.
        
    Arguments:
        file_path (str): The file path.
    '''
    if os.path.exists(file_path):
        try:
            os.remove(file_path)
        except Exception:
            _print_error('Unable to remove the unfinished file.')

def _action_decode(input_png_path: str,
                   output_stegano_path: str,
                   validate: bool = True) -> int: # status code
    try:
        with open(input_png_path, mode='rb') as input_png:
            output_stegano_path = _action_skip_namesakes(output_stegano_path)
            try:
                status_code_decode: int = 0
                with open(output_stegano_path, mode='wb') as output_stegano:
                    status_code_decode = _action_extract_stegano_zlib(input_png,
                                                                      output_stegano,
                                                                      validate)
                if status_code_decode != 0:
                    _action_try_remove_unfinished_file(output_stegano_path)
                return status_code_decode
            except Exception:
                traceback.print_exc()
                _action_try_remove_unfinished_file(output_stegano_path)
    except FileNotFoundError:
        _print_error(f'Input PNG ({input_png_path}) not found.')
    return 1

def _action_modify_png(input_png: bytes | BufferedReader,
                       input_stegano: bytes | BufferedReader,
                       output_png: BufferedWriter,
                       use_zopfli: bool = True) -> int: # status code
    '''Action: appends the stegano container with the data to the PNG image.
        
    Arguments:
        input_png (bytes | BufferedReader): The input PNG file to modify.
        input_stegano (bytes | BufferedReader): The data file to put inside the stegano container.
        output_stegano (BufferedWriter): The output PNG file to write.
        use_zopfli (bool): Determines if Zopfli should be used for compression instead of pure Zlib.
                           Default: True.

    Returns:
        out (int): The status code.
    '''
    if not output_png.writable:
        _print_error('Output PNG file is not writeable.')
        return 1
    if isinstance(input_png, BufferedReader):
        input_png = input_png.read()
    png_pixels: bytearray = read_png_as_bgra32(input_png)
    if png_pixels is None:
        _print_error('PNG is invalid (missing mandatory chunks or is corrupted by another way).')
        return 1
    output_png.write(FILE_MAGIC)
    output_cursor: int = 8
    while output_cursor < len(input_png):
        chunk_length_bytes: bytes = input_png[output_cursor : output_cursor + 4]
        output_png.write(chunk_length_bytes)
        output_cursor += 4
        chunk_length: int = struct.unpack('>I', chunk_length_bytes)[0]
        chunk_type: bytes = input_png[output_cursor : output_cursor + 4]
        output_png.write(chunk_type)
        output_cursor += 4
        chunk_data: bytes = input_png[output_cursor : output_cursor + chunk_length + 4]
        output_png.write(chunk_data)
        output_cursor += chunk_length + 4
        if chunk_type == CHUNK_END:
            break
    output_png.write(Stegano(CRC32_BGRA32(png_pixels), input_stegano, use_zopfli).compile())
    return 0

def _action_encode(input_png_path: str,
                   input_stegano_path: str,
                   output_png_path: str,
                   use_zopfli: bool = True) -> int: # status code
    try:
        with open(input_png_path, mode='rb') as input_png:
            try:
                with open(input_stegano_path, mode='rb') as input_stegano:
                    output_png_path = _action_skip_namesakes(output_png_path)
                    try:
                        status_code_encode: int = 0
                        with open(output_png_path, mode='wb') as output_png:
                            status_code_encode = _action_modify_png(input_png,
                                                                    input_stegano,
                                                                    output_png,
                                                                    use_zopfli)
                        if status_code_encode != 0:
                            _action_try_remove_unfinished_file(output_png_path)
                        return status_code_encode
                    except Exception:
                        traceback.print_exc()
                    _action_try_remove_unfinished_file(output_png_path)
            except FileNotFoundError:
                _print_error(f'Input stegano ({input_stegano_path}) not found.')
    except FileNotFoundError:
        _print_error(f'Input PNG ({input_png_path}) not found.')
    return 1

if __name__ == '__main__':
    status_code: int = 1
    parser: argparse.ArgumentParser = argparse.ArgumentParser(
        prog='spore_ep1stegano',
        description='Encodes and decodes Spore EP1 stegano chunk (spOr) on PNG images'
    )
    subparsers = parser.add_subparsers(dest='command', help='Available commands')
    # encode
    parser_encode = subparsers.add_parser('encode',
                                          help='Encode Spore EP1 stegano to PNG',
                                          aliases=['e'])
    parser_encode.add_argument('input_png_path',
                        type=str,
                        help='Input PNG image')
    parser_encode.add_argument('input_stegano_path',
                        type=str,
                        help='Input stegano')
    parser_encode.add_argument('--output_png_path',
                        type=str,
                        help='Output PNG image',
                        required=False)
    parser_encode.add_argument('--no_zopfli',
                        help='Use the default Zlib library for compression',
                        required=False,
                        dest='use_zopfli',
                        action='store_false')
    parser_encode.set_defaults(use_zopfli=True)
    # decode
    parser_decode = subparsers.add_parser('decode',
                                          help='Decode Spore EP1 stegano from PNG',
                                          aliases=['d'])
    parser_decode.add_argument('input_png_path',
                        type=str,
                        help='Input PNG image')
    parser_decode.add_argument('--output_stegano_path',
                        type=str,
                        help='Output file',
                        required=False)
    parser_decode.add_argument('--no_validation',
                        help='Avoid checksum validation',
                        required=False,
                        dest='validate',
                        action='store_false')
    parser_decode.set_defaults(validation=True)
    args: argparse.Namespace = parser.parse_args()
    # encode: action
    if args.command in ['encode', 'e']:
        if args.output_png_path is None:
            input_png_split: tuple[str, str] = os.path.splitext(args.input_png_path)
            input_png_name = input_png_split[0]
            input_png_ext = input_png_split[1]
            args.output_png_path = f'{input_png_name}_modified{input_png_ext}'
        status_code = _action_encode(args.input_png_path,
                                     args.input_stegano_path,
                                     args.output_png_path,
                                     args.use_zopfli)
    # decode: action
    elif args.command in ['decode', 'd']:
        if args.output_stegano_path is None:
            input_png_name: str = os.path.splitext(args.input_png_path)[0]
            args.output_stegano_path = input_png_name + '.stegano'
        status_code = _action_decode(args.input_png_path,
                                     args.output_stegano_path,
                                     args.validate)
    else:
        parser.print_help()
    sys.exit(status_code)