base64_handler.py

import os
from pathlib import Path
import hashlib
import zlib
import struct
import base64


try:
    from .definitions import APP_ROOT
    from .helpers import get_bytes
except ImportError as e:
    print(e)
    import traceback


    print(traceback.format_exc())
    from definitions import APP_ROOT
    from helpers import get_bytes


# http://coreygoldberg.blogspot.com/2013/01/python-verify-png-file-and-get-image.html
def get_image_info(data):
    if is_png(data):
        w, h = struct.unpack(b'>LL', data[16:24])
        width = int(w)
        height = int(h)
    else:
        raise Exception('not a png image')
    return width, height


def is_png(data):
    return data[:8] == b'\211PNG\r\n\032\n' and (data[12:16] == b'IHDR')


# https://blender.stackexchange.com/questions/240137/is-it-possible-to-create-image-data-from-a-base64-encoded-png
def image_from_data(name, data, height=1, width=1):
    import bpy

    img = bpy.data.images.get(name)
    if img is not None:
        return img

    # Create image, width and height are dummy values
    img = bpy.data.images.new(name, height, width)
    img.use_fake_user = True  # otherwise it won't save to the file

    # Set packed file data
    img.pack(data=data, data_len=len(data))
    # img.reload()

    # Switch to file source so it uses the packed file
    img.source = 'FILE'

    return img


# TODO: will be used for stud.io parts that have textures
# TexMap.base64_to_png(filename, img_data)
def base64_to_png_data(string):
    new_string = get_bytes(string)
    return base64.decodebytes(new_string)


def image_from_base64_str(filename, base64_str):
    img_data = base64_to_png_data(base64_str)
    return image_from_data(filename, img_data)


def named_png_from_base64_str(filename, base64_str):
    filename = f"{Path(filename).stem}.png"
    return image_from_base64_str(filename, base64_str)


def sha_named_png_from_base64_str(base64_str):
    img_data = base64_to_png_data(base64_str)
    hash_object = hashlib.sha256(img_data)
    basename = hash_object.hexdigest()[:16]
    filename = f"{basename}.png"
    return image_from_base64_str(filename, base64_str)


# basename prevents writing to any place but APP_ROOT
def write_png_data(app_root, filename, data):
    filepath = os.path.join(app_root, f"{os.path.basename(filename)}.png")
    with open(filepath, 'wb') as file:
        file.write(data)


def png_pack(png_tag, data):
    chunk_head = png_tag + data
    return (struct.pack("!I", len(data)) +
            chunk_head +
            struct.pack("!I", 0xFFFFFFFF & zlib.crc32(chunk_head)))


def blender_image_to_base64(image):
    width = image.size[0]
    height = image.size[1]

    buf = bytearray([int(p * 255) for p in image.pixels])

    # Reverse vertical order and prepend null bytes for PNG scanline filter
    width_byte_4 = width * 4
    raw_data = b''.join(b'\x00' + buf[span:span + width_byte_4] for span in range((height - 1) * width_byte_4, -1, -width_byte_4))

    png_bytes = b''.join([
        b'\x89PNG\r\n\x1a\n',
        png_pack(b'IHDR', struct.pack("!2I5B", width, height, 8, 6, 0, 0, 0)),
        png_pack(b'IDAT', zlib.compress(raw_data, 9)),
        png_pack(b'IEND', b'')
    ])

    return base64.b64encode(png_bytes).decode()