sld-analysis.py

#!/usr/bin/env python3

# This file is licensed under LGPL 3.0+ (see https://www.gnu.org/licenses/lgpl-3.0.en.html
# for the license text).
# In short this means: - you can freely share, edit and publish this code
#                      - if you publish the script or changes to it somewhere else,
#                        please link back to the original file

import argparse
from pathlib import Path
from struct import Struct

from PIL import Image

LAYER_LOOKUP = {
    0x01: "Main",
    0x02: "Shadow",
    0x04: "???",
    0x08: "Damage Mask",
    0x10: "Player Color",
}


def parse_args():
    parser = argparse.ArgumentParser()

    parser.add_argument("file")

    return parser.parse_args()


def decode_frame_header(sld_bytes, offset):
    # header?
    header = Struct("< H H H H B B H")
    print("  Header:")
    print("    OFFSET:", hex(offset))

    content = header.unpack_from(sld_bytes, offset)
    offset += header.size

    print("    ???:", hex(content[0]))
    print("    ???:", hex(content[1]))
    print("    ???:", hex(content[2]))
    print("    ???:", hex(content[3]))
    print("    Layer types:", hex(content[4]))
    print("    ???:", content[5])
    print("    Frame index:", content[6])

    return content, offset


def decode_layer(sld_bytes, offset, layer_type, width=0, height=0):
    if layer_type == 0x01:
        return decode_layer0x01(sld_bytes, offset)

    if layer_type == 0x02:
        return decode_layer0x02(sld_bytes, offset)

    if layer_type == 0x04:
        return decode_layer0x04(sld_bytes, offset)

    if layer_type == 0x08:
        return decode_layer0x08(sld_bytes, offset, width, height)

    if layer_type == 0x10:
        return decode_layer0x10(sld_bytes, offset, width, height)

    # else:
    layer_header = Struct("< I")
    return layer_header.unpack_from(sld_bytes, offset)


def decode_layer0x01(sld_bytes, offset):
    layer_header = Struct("< I H H H H H H")
    header = layer_header.unpack_from(sld_bytes, offset)

    layer_len = header[0]
    content1_len = 2 * header[6]
    content2_len = layer_len - layer_header.size - content1_len
    content2_chunks = content2_len / 8
    width = header[3] - header[1]
    height = header[4] - header[2]

    # Chunk sizes
    content1_start = offset + layer_header.size
    chunk_sizes = get_chunk_offsets(sld_bytes, content1_start, content1_len)

    # Image
    img = Image.new('RGBA', (width, height))
    chunk_offset = offset + layer_header.size + content1_len
    offset_x = 0
    offset_y = 0
    for line in chunk_sizes:
        chunks_count = line[1]
        offset_x += line[0] * 4
        while offset_x >= img.width:
            # Wrap around if line end reached
            offset_y += 4
            offset_x -= img.width

        for _ in range(chunks_count):
            quad_bytes = sld_bytes[chunk_offset:chunk_offset + 8]
            quad_px = get_bc1_quad(quad_bytes)
            draw_quad(quad_px, img, offset_x, offset_y)

            offset_x += 4
            if offset_x >= img.width:
                offset_x = 0
                offset_y += 4

            chunk_offset += 8

    img.save("layer0x01.png")
    # img.show()
    return layer_len, content1_len, content2_len, content2_chunks, header


def decode_layer0x02(sld_bytes, offset):
    layer_header = Struct("< I H H H H H H")
    header = layer_header.unpack_from(sld_bytes, offset)

    layer_len = header[0]
    content1_len = 2 * header[6]
    content2_len = layer_len - layer_header.size - content1_len
    content2_chunks = content2_len / 8
    width = header[3] - header[1]
    height = header[4] - header[2]

    # Chunk sizes
    content1_start = offset + layer_header.size
    chunk_sizes = get_chunk_offsets(sld_bytes, content1_start, content1_len)

    # Image
    img = Image.new('RGBA', (width, height))
    chunk_offset = offset + layer_header.size + content1_len
    offset_x = 0
    offset_y = 0
    for line in chunk_sizes:
        chunks_count = line[1]
        offset_x += line[0] * 4
        while offset_x >= img.width:
            # Wrap around if line end reached
            offset_y += 4
            offset_x -= img.width

        for _ in range(chunks_count):
            quad_bytes = sld_bytes[chunk_offset:chunk_offset + 8]
            quad_px = get_bc4_quad(quad_bytes)
            draw_quad(quad_px, img, offset_x, offset_y)

            offset_x += 4
            if offset_x >= img.width:
                offset_x = 0
                offset_y += 4

            chunk_offset += 8

    img.save("layer0x02.png")
    # img.show()
    return layer_len, content1_len, content2_len, content2_chunks, header


def decode_layer0x04(sld_bytes, offset):
    layer_header = Struct("< I I 14H")
    header = layer_header.unpack_from(sld_bytes, offset)

    layer_len = header[0]
    content1_len = 2 * header[2]
    content2_len = layer_len - layer_header.size - 0x3A
    content2_chunks = content2_len / 8
    content2_chunks = 0x82

    # Chunk sizes
    content1_start = offset + layer_header.size
    chunk_sizes = get_chunk_offsets(sld_bytes, content1_start, 0x3A)
    # chunk_sizes = [(29, 3), (66, 2), (1, 2), (62, 3), (4, 1), (62, 1), (6, 1), (2, 6), (54, 1), (6, 4), (4, 2), (53, 1), (15, 3), (50, 2), (17, 3), (48, 2), (19, 3), (46, 2), (21, 3), (44, 1), (24, 3), (41, 2), (26, 3), (39, 2), (28, 3), (36, 2), (31, 3), (31, 3), (35, 2), (29, 2), (39, 3), (21, 1), (2, 2), (43, 3), (18, 1), (50, 3), (69, 2), (69, 1), (69, 2), (14, 1), (53, 2), (13, 1), (55, 2), (11, 3), (55, 1), (12, 2), (55, 2), (12, 1), (55, 2), (11, 2), (5, 2), (48, 2), (9, 4), (55, 1), (8, 6), (55, 1), (7, 4), (58, 1), (4, 2), (1, 4), (58, 1), (4, 6), (63, 8), (52, 2), (8, 4), (3, 1), (23, 2), (33, 2), (3, 3), (3, 2), (9, 1), (41, 1), (5, 8), (14, 3), (11, 1), (25, 4), (2, 3), (2, 5), (15, 1), (9, 1), (1, 5), (3, 2), (15, 11), (5, 3), (2, 1), (13, 1), (5, 16), (12, 5), (13, 7), (7, 2), (2, 9), (3, 11), (11, 5), (18, 4), (1, 5), (4, 13), (4, 6), (1, 1), (8, 1), (2, 1), (19, 7), (1, 1), (6, 12), (7, 6), (6, 2), (24, 7), (6, 12), (9, 5), (2, 1), (1, 3), (27, 4), (6, 2), (2, 7), (11, 10), (30, 3), (3, 5), (5, 4), (12, 1), (38, 10), (5, 4), (53, 2), (1, 4), (43, 0)]

    # Image
    img = Image.new('RGBA', (280, 180))
    chunk_offset = 0x45BE
    offset_x = 0
    offset_y = 0
    for line in chunk_sizes:
        chunks_count = line[1]
        offset_x += line[0] * 4
        while offset_x >= img.width:
            # Wrap around if line end reached
            offset_y += 4
            offset_x -= img.width

        for _ in range(int(content2_chunks)):
            # for _ in range(int(chunks_count)):
            quad_bytes = sld_bytes[chunk_offset:chunk_offset + 8]
            quad_px = get_bc4_quad(quad_bytes)
            draw_quad(quad_px, img, offset_x, offset_y)

            offset_x += 4
            if offset_x >= img.width:
                offset_x = 0
                offset_y += 4

            chunk_offset += 8

            if chunk_offset > chunk_offset + content2_len:
                break

        if chunk_sizes.index(line) == 0:
            break

    img.save("layer0x04.png")
    # img.show()
    return layer_len, content1_len, content2_len, content2_chunks, header


def decode_layer0x08(sld_bytes, offset, width, height):
    layer_header = Struct("< I H H")
    header = layer_header.unpack_from(sld_bytes, offset)

    layer_len = header[0]
    content1_len = 2 * header[2]
    content2_len = layer_len - layer_header.size - content1_len
    content2_chunks = content2_len / 8

    # Chunk sizes
    content1_start = offset + layer_header.size
    chunk_sizes = get_chunk_offsets(sld_bytes, content1_start, content1_len)

    # Image
    img = Image.new('RGBA', (width, height))
    chunk_offset = offset + layer_header.size + content1_len
    offset_x = 0
    offset_y = 0
    for line in chunk_sizes:
        chunks_count = line[1]
        offset_x += line[0] * 4
        while offset_x >= img.width:
            # Wrap around if line end reached
            offset_y += 4
            offset_x -= img.width

        for _ in range(chunks_count):
            quad_bytes = sld_bytes[chunk_offset:chunk_offset + 8]
            quad_px = get_bc1_quad(quad_bytes)
            draw_quad(quad_px, img, offset_x, offset_y)

            offset_x += 4
            if offset_x >= img.width:
                offset_x = 0
                offset_y += 4

            chunk_offset += 8

    img.save("layer0x08.png")
    # img.show()
    return layer_len, content1_len, content2_len, content2_chunks, header


def decode_layer0x10(sld_bytes, offset, width, height):
    layer_header = Struct("< I H H")
    header = layer_header.unpack_from(sld_bytes, offset)

    layer_len = header[0]
    content1_len = 2 * header[2]
    content2_len = layer_len - layer_header.size - content1_len
    content2_chunks = content2_len / 8

    # Chunk sizes
    content1_start = offset + layer_header.size
    chunk_sizes = get_chunk_offsets(sld_bytes, content1_start, content1_len)

    # Image
    img = Image.new('RGBA', (width, height))
    chunk_offset = offset + layer_header.size + content1_len
    offset_x = 0
    offset_y = 0
    for line in chunk_sizes:
        chunks_count = line[1]
        offset_x += line[0] * 4
        while offset_x >= img.width:
            # Wrap around if line end reached
            offset_y += 4
            offset_x -= img.width

        for _ in range(chunks_count):
            quad_bytes = sld_bytes[chunk_offset:chunk_offset + 8]
            quad_px = get_bc4_quad(quad_bytes)
            draw_quad(quad_px, img, offset_x, offset_y)

            offset_x += 4
            if offset_x >= img.width:
                offset_x = 0
                offset_y += 4

            chunk_offset += 8

    img.save("layer0x10.png")
    # img.show()
    return layer_len, content1_len, content2_len, content2_chunks, header


def get_chunk_offsets(sld_bytes, offset, length):
    chunk_offsets = []
    for chunk_offset in range(offset, offset + length, 2):
        skipped_chunks = sld_bytes[chunk_offset]
        read_chunks = sld_bytes[chunk_offset + 1]

        chunk_offsets.append((skipped_chunks, read_chunks))

    count_skipped = 0
    count_read = 0
    for ch in chunk_offsets:
        count_skipped += ch[0]
        count_read += ch[1]

    count = count_skipped + count_read
    print("Skipped", count_skipped, hex(count_skipped))
    print("Read", count_read, hex(count_read))
    print("Total", count, hex(count))

    return chunk_offsets


def draw_quad(quad_px: tuple, img: Image, offset_x: int, offset_y: int):
    quad_offset_x = 0
    quad_offset_y = 0
    for px in quad_px:
        img.putpixel((offset_x + quad_offset_x, offset_y + quad_offset_y), px)

        quad_offset_x += 1
        if quad_offset_x >= 4:
            quad_offset_x = 0
            quad_offset_y += 1


def get_bc1_quad(quad_bytes):
    quad = []
    color0 = quad_bytes[0:2]
    color1 = quad_bytes[2:4]

    c0, c1, c2, c3 = colors_bc1(color0, color1)
    # for byte_start in range(4):
    #     byte_val = quad_bytes[4 + byte_start]
    #     for shift in range(4):
    #         mask = 0b00000011

    #         bc1_val = byte_val >> (shift * 2) & mask
    #         if bc1_val == 0b00:
    #             quad.append(c0)

    #         elif bc1_val == 0b01:
    #             quad.append(c1)

    #         elif bc1_val == 0b10:
    #             quad.append(c2)

    #         elif bc1_val == 0b11:
    #             quad.append(c3)
    px_vals = int.from_bytes(quad_bytes[4:8], byteorder='little')
    for shift in range(16):
        mask = 0b00000011

        bc1_val = px_vals >> (shift * 2) & mask
        if bc1_val == 0b00:
            quad.append(c0)

        elif bc1_val == 0b01:
            quad.append(c1)

        elif bc1_val == 0b10:
            quad.append(c2)

        elif bc1_val == 0b11:
            quad.append(c3)

    return tuple(quad)


def get_bc4_quad(quad_bytes):
    quad = []
    red0 = quad_bytes[0]
    red1 = quad_bytes[1]

    c0, c1, c2, c3, c4, c5, c6, c7 = colors_bc4(red0, red1)
    for byte_start in range(2):
        start = 2 + byte_start * 3
        end = 2 + byte_start * 3 + 3
        red_vals = int.from_bytes(quad_bytes[start:end], byteorder='little')
        print(hex(red_vals))
        for shift in range(8):
            mask = 0b00000111

            bc1_val = red_vals >> (shift * 3) & mask
            if bc1_val == 0b000:
                quad.append(c0)

            elif bc1_val == 0b001:
                quad.append(c1)

            elif bc1_val == 0b010:
                quad.append(c2)

            elif bc1_val == 0b011:
                quad.append(c3)

            elif bc1_val == 0b100:
                quad.append(c4)

            elif bc1_val == 0b101:
                quad.append(c5)

            elif bc1_val == 0b110:
                quad.append(c6)

            elif bc1_val == 0b111:
                quad.append(c7)

    return tuple(quad)


def colors_bc1(color0: bytes, color1: bytes):
    c0 = int.from_bytes(color0, byteorder='little')
    c1 = int.from_bytes(color1, byteorder='little')

    r0 = (c0 & 0xF800) >> 11
    g0 = (c0 & 0x07E0) >> 5
    b0 = c0 & 0x001F

    r1 = (c1 & 0xF800) >> 11
    g1 = (c1 & 0x07E0) >> 5
    b1 = c1 & 0x001F

    r0, g0, b0 = int(r0 * 8), int(g0 * 4), int(b0 * 8)
    r1, g1, b1 = int(r1 * 8), int(g1 * 4), int(b1 * 8)

    if c0 > c1:
        r2 = int((2 * r0 + r1 + 1) / 3)
        g2 = int((2 * g0 + g1 + 1) / 3)
        b2 = int((2 * b0 + b1 + 1) / 3)
        r3 = int((r0 + 2 * r1 + 1) / 3)
        g3 = int((g0 + 2 * g1 + 1) / 3)
        b3 = int((b0 + 2 * b1 + 1) / 3)
        a = 255

    else:
        r2 = int((r0 + r1) / 2)
        g2 = int((g0 + g1) / 2)
        b2 = int((b0 + b1) / 2)
        r3 = 0
        g3 = 0
        b3 = 0
        a = 0

    return (r0, g0, b0), (r1, g1, b1), (r2, g2, b2), (r3, g3, b3, a)


def colors_bc4(red0: int, red1: int):
    colors = []

    colors.append((red0, 0, 0))
    colors.append((red1, 0, 0))

    if red0 > red1:
        for i in range(1, 7):
            red = int(((7 - i) * red0 + i * red1) / 7.0)
            colors.append((red, 0, 0))

    else:
        for i in range(1, 5):
            red = int(((5 - i) * red0 + i * red1) / 5.0)
            colors.append((red, 0, 0))

        colors.append((0, 0, 0, 0))
        colors.append((255, 0, 0))

    return tuple(colors)


def decode_frame(sld_bytes, offset):
    # header
    header, offset = decode_frame_header(sld_bytes, offset)

    layers = {}

    layer_mask = header[4]
    layer_bit = 1
    layer_count = 0
    while layer_bit < 0x100:
        if layer_mask & layer_bit:
            layer_count += 1
            layers[layer_bit] = None

        layer_bit = layer_bit << 1

    main_width = 0
    main_height = 0
    for layer_idx, layer_type in enumerate(layers.keys()):
        layer = decode_layer(sld_bytes, offset, layer_type, main_width, main_height)
        layer_len = layer[0]

        print("  Layer:", LAYER_LOOKUP[layer_type], f"(type {hex(layer_type)})")
        print("    OFFSET:", hex(offset))
        print("    length:", hex(layer_len), "=", layer_len)

        if layer_len % 4 != 0:
            layer_len += 4 - (layer_len % 4)

        print("    padded length:", hex(layer_len))

        content1_len = layer[1]
        content2_len = layer[2]
        content2_chunks = layer[3]
        print("    content1 length:", hex(content1_len), "=", content1_len)
        print("    content2 length:", hex(content2_len), "=", content2_len)
        print("    content2 chunks:", hex(int(content2_chunks)), "=", content2_chunks)

        if layer_type == 0x01:
            main_width = layer[4][3] - layer[4][1]
            main_height = layer[4][4] - layer[4][2]

        layers[layer_type] = sld_bytes[offset + 4:offset + layer_len]

        offset += layer_len

    return None, offset


def decode(sld_bytes):
    offset = 0
    file_header = Struct("< 4s H H H H I")
    file_header_content = file_header.unpack_from(sld_bytes, offset)

    print("Format:", file_header_content[0])
    print("Version:", file_header_content[1])
    print("Frame Count:", file_header_content[2])
    print("???:", file_header_content[3])
    print("???:", file_header_content[4])
    print("???:", file_header_content[5])

    frame_idx = 0
    offset += file_header.size
    while offset < len(sld_bytes):
        print("Frame:", frame_idx)
        print("  OFFSET:", hex(offset))
        _, offset = decode_frame(sld_bytes, offset)
        frame_idx += 1

    if offset != len(sld_bytes):
        print("Offset != File Length:", hex(offset), "!=", hex(len(sld_bytes)))


def main():
    args = parse_args()

    sld_path = Path(args.file)
    with sld_path.open("rb") as sld_file:
        sld_bytes = sld_file.read()

    # find_lz4_start(sld_bytes)
    decode(sld_bytes)


if __name__ == "__main__":
    main()