unfinished nintendo particle format (.ptcl) parser

ptcl.py

# .ptcl parser for miitomo
# heavily unfinished, i couldn't work out how to find subtexture sizes

import struct
import numpy as np 
from PIL import Image

class PTCLSection:
  def __init__(self, magic=b'', offset=0, size=0, data_offset=0):
    self.magic = magic
    self.offset = offset
    self.size = size
    self.data_offset = data_offset
    self.subsections = []

  def get_data(self, buffer):
    orig_offset = buffer.tell()
    buffer.seek(self.offset + self.data_offset)
    data = buffer.read(self.size - self.data_offset)
    buffer.seek(orig_offset)
    return data

class PTCLTexture:
  def __init__(self):
    self.width = 0
    self.height = 0
    self.imagedata = []

  def get_image(self, index):
    width = self.width
    height = self.height
    pixels = np.frombuffer(self.imagedata[index], dtype=np.uint32)
    pixels = pixels.reshape((width, height))
    return Image.fromarray(pixels, mode="RGBA")

class PTCL:
  def __init__(self, buffer):
    self.textures = []
    if buffer: 
      self.buffer = buffer
      self.read()

  @classmethod
  def Open(cls, path):
    with open(path, 'rb') as f:
      instance = cls(f)
      return instance

  def read(self):
    self.buffer.seek(0, 2)
    buffer_size = self.buffer.tell()
    self.buffer.seek(0)
    # read file magic
    magic = self.buffer.read(4)
    if not magic == b'EFTB': exit("The %s format type isn't supported" % magic.decode('ascii'))
    self.buffer.seek(48)

    # read file sections
    for section in self.read_sections():
      
      if section.magic == b'ESTA':
        self.read_emitters(section)

      elif section.magic == b'TEXA':
        self.read_textures(section)

      # elif section.magic == b'PRMA':
      #   self.read_primitives(section)

      elif section.magic == b'SHCA':
        self.read_shaders(section)

      elif section.magic == b'SHDA':
        # this section seems to be empty
        pass

  def read_emitters(self, section):
    for emitter in section.subsections:
      print(emitter.magic)
  
  def read_primitives(self, section):
    # not implemented, seems to be model data though
    pass
  
  def read_textures(self, section):
    for index, texture in enumerate(section.subsections):
      tex = PTCLTexture()
      self.buffer.seek(texture.offset + texture.data_offset)
      tex.width, tex.height = struct.unpack("<HH", self.buffer.read(4))
      tilemode, swizzle = struct.unpack("<II", self.buffer.read(8))

      # TODO:
      # can't figure out how to get subtexture sizes?

      for subindex, subtexture in enumerate(texture.subsections):
        self.buffer.seek(subtexture.offset + subtexture.data_offset)
        data = subtexture.get_data(self.buffer)
        # dump data for debugging
        with open('./tex%d_%d.bin' % (index, subindex), 'wb') as f:
          f.write(data)

        tex.imagedata.append(data)
      
      # write first image to file
      self.textures.append(tex)
      img = tex.get_image(0)
      img.save('./tex%d.png' % index)

  def read_shaders(self, section):
    for shader in section.subsections:
      print(shader.magic)
  
  def read_string(self, length):
    return self.buffer.read(length).decode("utf-8").strip("\x00")

  def read_sections(self):
    while True:
      # get section offset
      offset = self.buffer.tell()
      # unpack section header
      magic, size, subsection_offset, next_offset, _, data_offset, _, subsection_count, _ = struct.unpack('<4s6I2H', self.buffer.read(32))
      section = PTCLSection(magic=magic, offset=offset, size=size, data_offset=data_offset)
      # read subsections
      if subsection_count > 0:
        self.buffer.seek(offset + subsection_offset)
        section.subsections = list(self.read_sections())
      # yeild section 
      yield section
      # 0xffffffff is a null offset, so we stop looking for next sections here
      if next_offset == 0xFFFFFFFF:
        break
      # otherwise seek to the next section
      else:
        self.buffer.seek(offset + next_offset)

PTCL.Open('./effect.ptcl')