bz2wat.py - BZ2 Water Parser

from struct import pack, unpack
from math import sqrt

def from_file(file_path):
    return WAT().from_file(file_path)

class WAT:
    glow =       0x80 # 10000000 On/Off
    direction =  0x03 # 00000011 First 2 bits to represnt 0, 1, 2, or 3
    repeat =     0x7C # 01111100 5 bits to represnt 0 to 31 (Appears as 1 to 32 in editor)
    rep_offset = 0x02 # How much to shift right by

    def __init__(self, tiles = None):
        self.layers = []
        self.index = []

        if tiles:
            self.blank(tiles)

    def __str__(self):
        return "WAT tiles=%d layers=%d" % (len(self.index), len(self.layers))

    def blank(self, tiles):
        self.layers = []
        for layer in range(15):
            self.layers += [Layer()]
        self.index = [0]*tiles

    def from_file(self, file_path):
        with open(file_path, "rb") as fl:
            # [Header:16 Bytes]
            magic_number = fl.read(4)
            version = unpack("i", fl.read(4))[0]
            if magic_number != b"WATR" or version != 2:
                raise Exception("Invalid WAT file.")

            # Width(inverted), Height(inverted), Width, Height
            size = unpack("HHHH", fl.read(8))
            water_chunks = (size[2]//2)*(size[3]//2)

            # [Layers:92] *15 Total: 1380
            self.layers = []
            for layer in range(15):
                height, damage = 0.0, 0.0
                velocity, rgba, texture = [0, 0], [0, 0], [0, 0]

                height = unpack("f", fl.read(4))[0]
                damage = unpack("f", fl.read(4))[0]
                for tex in range(2):
                    velocity[tex] = unpack("f", fl.read(4))[0]
                    rgba[tex] = unpack("BBBB", fl.read(4))
                    texture[tex] = fl.read(32).decode("ascii", "ignore")
                    texture[tex] = texture[tex][0:texture[tex].index("\x00")]

                flags = unpack("BB", fl.read(2))
                fl.seek(2, 1)

                glow = [bool(x & WAT.glow) for x in flags]
                direction = [(x & WAT.direction) for x in flags]
                repeat = [((x & WAT.repeat) >> WAT.rep_offset) for x in flags]

                self.layers += [Layer(height, damage, texture, rgba, velocity, glow, direction, repeat)]

            # [Indices:1] Total: (water_chunks)
            self.index = list(unpack("B"*water_chunks, fl.read(water_chunks)))
        return self

    def to_file(self, file_path):
        with open(file_path, "wb") as fl:
            fl.write(b"WATR")
            fl.write(pack("i", 2))
            size = int(sqrt(len(self.index)))
            fl.write(pack("HHHH", 0x10000-size*2, 0x10000-size*2, size*2, size*2))

            for layer in self.layers:
                fl.write(layer.pack())

            fl.write(pack("B"*len(self.index), *self.index))

class Layer:
    def __init__(self,
    height    = 0.0,
    damage    = 0.0,
    texture   = None,
    rgba      = None,
    velocity  = None,
    glow      = None,
    direction = None,
    repeat    = None):
        if texture is None:
            texture = ["", ""]
        if rgba is None:
            rgba = [[0, 0, 0, 0], [0, 0, 0, 0]]
        if velocity is None:
            velocity = [0.0, 0.0]
        if glow is None:
            glow = [False, False]
        if direction is None:
            direction = [0, 0]
        if repeat is None:
            repeat = [31, 31]

        self.h = height
        self.dmg = damage
        self.t = [texture[0][0:31], texture[1][0:31]]
        self.rgba = rgba
        self.vel = velocity
        self.g = glow
        self.d = direction
        self.rep = repeat

    def __str__(self):
        return "h=%f dmg=%f %s" % (
            self.h,
            self.dmg,
            str([("%r rgba%s glow=%d dir=%d veloc=%f rep=%d" % (
                self.t[i],
                str(self.rgba[i]),
                self.g[i],
                self.d[i],
                self.vel[i],
                self.rep[i])) for i in range(2)]))

    def pack(self):
        """Pack data and return as bytes."""
        data = bytearray(92)
        data[0:8] = pack("ff", self.h, self.dmg)

        for tex in range(2):
            i = tex*40+8
            data[i+0:i+4] = pack("f", self.vel[tex])
            data[i+4:i+8] = pack("BBBB", *self.rgba[tex])
            data[i+8:i+40] = self.t[tex].ljust(32, "\x00").encode("ascii", "ignore")

        for tex in range(2):
            data[88+tex] = self.d[tex] | self.g[tex]*WAT.glow | (self.rep[tex] << WAT.rep_offset)

        data[90:92] = b"\x00\x00" # Struct padding

        return data

if __name__ == "__main__":
    file = "./Subject.WAT"
    wat = WAT().from_file(file)
    print(wat)