#!/usr/bin/python

__author__    = "Anoop Menon <codelogic at gmail dot com>"
__copyright__ = "Copyright (c) 2007"
__license__   = "GPL Version 2.0"
__version__   = "0.12"
__date__      = "12/12/07"


#########################################################################
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; version 2 of the License.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# If you have any comments or would like to suggest improvements, please
# email me at 'codelogic at gmail dot com'
#
#########################################################################


#########################################################################
#
# Changelog
#
# 12/12/07 - added support to save images to PPM if Python Imaging
#            Library is not detected.
#
# 12/07/07 - added support to extract audio files from the theme. They
#            are in the VAG format.
#
# 12/07/07 - fixed the skewed icons issue where if the icon's width
#            was not a multiple of 4, it woud extract a skewed version.
#            (thx to FluBBa for pointing out the cause).
#
# 12/06/07 - dragging and dropping a p3t file onto the executable now
#            works in Windows (thx to Matt Endersby for reporting the
#            issue).
#
#          - default extracted path is now 'extracted.<theme filename>'
#            in the current working directory.
#
# 12/03/07 - fixed bug in output XML filename
#
# 11/20/07 - fixed bug in path seperator (thx to Kamil Sarniak for
#            pointing it out).
#
# 11/19/07 - added support to recreate an XML which can be fed to the
#            theme compiler to recreate original p3t file.
#
# 11/19/07 - multiple bgimages are correctly extracted instead of being
#            overwritten.
#
# 11/19/07 - removed requirement of having p3tcompiler.exe in the same
#            directory.
#
# 11/16/07 - added support to save GIM files as PNGs
#
# 11/16/07 - added a basic GIM file reader
#
# 11/16/07 - non 'bgimages' are zlib compressed. Added support for that
#            so that icons are decompressed to GIM files.
#
# 11/15/07 - First release
#
#########################################################################

import zipimport
import sys
import struct
import os
import os.path
import traceback
import zlib
import xml.dom.minidom

g_strings = []
g_PIL = 0

try:
    from PIL import Image
    g_PIL = 1
except:
    pass

class GimFile:
    '''
    GIM files have a 128 byte header followed by raw RGBA data
    width and height are stored at offsets 72 and 74 as short ints (big endian)
    '''
    def __init__(self, imagedata):
        self.buf = imagedata
        if self.buf[1:4]!="GIM":
            raise Exception("This data is not in GIM format")
        self.header = self.buf[0:128]
        (self.width, self.height) = struct.unpack(">hh", self.header[72:76])
        self.buf = self.buf[128:]

    def save(self, dest):
        global g_PIL
        if g_PIL==0:
            print "Python Imaging Library is required to save as PNG, saving to PPM instead ",
            stride = self.width
            if stride%4 != 0:
                stride = self.width+(4-self.width%4)
            header = "P6\n%d\n%d\n255\n" % (self.width, self.height)
            f = open(dest+".ppm", "wb")
            f.write(header)
            for row in range(0, self.height):
                offset = stride*row*4
                for col in range(0, self.width):
                    f.write(self.buf[offset+col*4:offset+col*4+3])
            f.close()
            return        
        try:
            # stride is padded to a multiple of 4 (thx FluBBa)
            stride = self.width
            if stride%4 != 0:
                stride = self.width+(4-self.width%4)
            img = Image.frombuffer("RGBA", (self.width, self.height), self.buf, 'raw', 'RGBA', stride*4, 1)
            img.save(dest)
        except Exception, e:
            raise e


class P3THeader:
    def __init(self):
        self.magic = ""
        self.version = 0

        self.tree_offset = 0
        self.tree_size = 0

        self.idtable_offset = 0
        self.idtable_size = 0

        self.stringtable_offset = 0
        self.stringtable_size = 0

        self.intarray_offset = 0
        self.intarray_size = 0

        self.floatarray_offset = 0
        self.floatarray_size = 0

        self.filetable_offset = 0
        self.filetable_size = 0


    def __str__(self):
        return '''
        magic:               "%s"
        version:             %d
        tree_offset:         %d
        tree_size:           %d
        idtable_offset:      %d
        idtable_size:        %d
        stringtable_offset:  %d
        stringtable_size:    %d
        intarray_offset:     %d
        intarray_size:       %d
        floatarray_offset:   %d
        floatarray_size:     %d
        filetable_offset:    %d
        filetable_size:      %d

        ''' % (self.magic, self.version, self.tree_offset,
               self.tree_size, self.idtable_offset, self.idtable_size, self.stringtable_offset,
               self.stringtable_size, self.intarray_offset, self.intarray_size,
               self.floatarray_offset, self.floatarray_size, self.filetable_offset,
               self.filetable_size)


class P3TElement(xml.dom.minidom.Element):
    def __init__(self, header):
        xml.dom.minidom.Element.__init__(self, "")
        self.numattr  = 0
        self.stringoffset = 0
        self.parentoffset = 0
        self.name = ""
        self.offset = 0
        self.t = []
        self.attributes = {}
        self.has_file = 0
        self.filename = ""
        self.header = header
        self.elemap = {}

    def __str__(self):
        return ("Name: '%s', no. of attributes: %d" % (self.name, self.numattr)) #+ "\n"+str(self.t)

    def add_subelement(self, ele):
        self.elemap[ele.offset] = ele
        self.appendChild(ele)

    def add_attribute(self, attr):
        if attr.type==6:
            self.has_file = 1
        self.attributes[attr.name] = attr
        if attr.name=="size":
            return
        if attr.name and attr.value:
            self.setAttribute(str(attr.name), str(attr.value))

    def dump_files(self, rootdir, hfile):
        for subele in self.elemap.keys():
            self.elemap[subele].dump_files(rootdir, hfile)

        if not self.has_file:
            return
        try:
            os.mkdir(rootdir)
        except:
            pass
        if not os.path.exists(rootdir):
            raise Exception("Unable to create directory: %s" % rootdir)

        fileroot = rootdir + os.path.sep

        pos = hfile.tell()
        fileoffset = self.header.filetable_offset

        for k,v in self.attributes.iteritems():

            if v.type==6:
                if self.name=="se":
                    ext = ".vag"
                else:
                    ext = ".gim"
                suffix = 1
                if self.name=="bgimage":
                    ext = ".jpg"
                while True:
                    try:
                        filename = fileroot + self.attributes['id'].value + "_" + str(suffix) + ext
                    except:
                        filename = fileroot + v.name + "_" + str(suffix) + ext
                    if os.path.exists(filename):
                        suffix+=1
                    else:
                        break
                hfile.seek(fileoffset + v.fileoffset)
                outfile = open(filename, "wb")
                print "Writing original : '%s'" % filename
                cbuf = hfile.read(v.filesize)
                self.filename = filename
                if ext==".gim":
                    # zlib compressed image

                    try:
                        dbuf = zlib.decompress(cbuf)
                    except Exception, e:
                        dbuf = cbuf
                        print str(e)
                        pass
                    try:
                        gimfile = GimFile(dbuf)
                        print "Coverting to PNG : '%s' : " % (filename+".png"),
                        gimfile.save(filename+".png")
                        self.filename = filename+".png"
                        print "OK"
                    except Exception, e:
                        print "Failed"
                        print str(e)
                    outfile.write(dbuf)
                else:
                    outfile.write(cbuf)
                outfile.close()

                # remove
                self.setAttribute(str(v.name), os.path.basename(str(self.filename)))

        hfile.seek(pos)

    def parse(self, data, format, hfile):
        global g_strings

        f = hfile

        self.offset = hfile.tell() - struct.calcsize(format) - self.header.tree_offset

        t = struct.unpack(format, data)
        self.t = t
        # 1st element is offset into string table
        self.stringoffset = t[0]

        # 2nd element is the no. of attributes
        self.numattr = t[1]

        # 3rd element is the offset of the parent element
        self.parentoffset = t[2]

        pos = f.tell()
        f.seek(self.header.stringtable_offset+self.stringoffset)
        self.name = ""
        a = f.read(1)
        while a!='\x00':
            self.name = self.name + str(a)
            a = f.read(1)
        f.seek(pos)
        self.tagName = self.name

class P3TAttribute:
    def __init__(self):
        self.type  =0
        self.handle = 0
        self.offset = 0
        self.size = 0
        self.value = 0
        self.name = ""
        self.fileoffset = 0
        self.filesize = 0
        self.id = 0
        self.t = []

    def __str__(self):
        #return "type: %d, name: '%s', value: '%s', offset: %d, handle: %d, size: %d" % (self.type, str(self.name), str(self.value), self.offset, self.handle, self.size)
        return ("type: %d, name: '%s', value: '%s', id: %d" % (self.type, str(self.name), str(self.value), self.id)) #+ "\n    "+str(self.t)

    def parse(self, data, header, hfile):
        global g_strings

        f = hfile
        pos = f.tell()

        t = struct.unpack(">iiii", data)
        self.type = t[1]
        atype = self.type
        self.handle = t[0]

        if atype==1: #int
            self.value = t[2]

        elif atype==2: #float
            t = struct.unpack(">iif4x", data)
            self.value = t[2]

        elif atype==3: #string
            t = struct.unpack(">iiii", data)
            self.offset = t[2]
            self.size = t[3]
            f.seek(header.stringtable_offset+self.offset)
            #self.value = unicode(f.read(self.size), 'utf-8')
            self.value = f.read(self.size)

        elif atype==6: #filename
            t = struct.unpack(">iiii", data)
            self.fileoffset = t[2]
            self.filesize = t[3]

        elif atype==7: #id
            t = struct.unpack(">iii4x", data)
            self.offset = t[2]
            f.seek(header.idtable_offset+self.offset)
            id_bin = f.read(4)
            (self.id,) = struct.unpack('>i', id_bin)
            self.value = ""
            a = f.read(1)
            while a!='\x00':
                self.value = self.value + str(a)
                a = f.read(1)

        f.seek(header.stringtable_offset+self.handle)
        self.name = ""
        a = f.read(1)
        while a!='\x00':
            self.name = self.name + str(a)
            a = f.read(1)
        f.seek(pos)

        self.t = t


class P3TExtractor:
    '''
    Basic format of a p3t theme file:

      - header (64)
      - tree
      - idtable
      - stringtable
      - intarraytable
      - floatarraytable
      - filetable

    '''
    def __init__(self, filename):
        self.themefile = filename
        self.header = P3THeader()
        self.elemap = {}

        try:
            p3tcompiler = zipimport.zipimporter('p3tcompiler.exe')
            self.cxml = p3tcompiler.load_module('cxml')
            self.header_size = self.cxml.header_bin_size
            self.header_fmt = ">"+self.cxml.header_bin_fmt
            self.element_size = self.cxml.element_bin_size
            self.element_fmt = ">"+self.cxml.element_bin_fmt
        except:
            self.header_size = 64
            self.header_fmt = ">4siiiiiiiiiiiii8x"
            self.element_size = 28
            self.element_fmt = ">iiiiiii"
            pass

        self.attr_fmt = ">iii4x"
        self.attr_size = struct.calcsize(self.attr_fmt)

    def parse(self):
        global g_strings

        h = self.header
        f = open(self.themefile, "rb")

        # parse header
        header_bin = f.read(self.header_size)
        (h.magic, h.version, h.tree_offset, h.tree_size,
         h.idtable_offset, h.idtable_size, h.stringtable_offset,
         h.stringtable_size, h.intarray_offset, h.intarray_size,
         h.floatarray_offset, h.floatarray_size, h.filetable_offset,
         h.filetable_size) = struct.unpack(self.header_fmt, header_bin)

        # load string table
        f.seek(h.stringtable_offset)
        stringtable_bin = f.read(h.stringtable_size)
        temp_strings = stringtable_bin.split('\x00')
        for a in temp_strings:
            g_strings.append(unicode(a, 'utf8'))

        # parse element tree
        f.seek(h.tree_offset)
        ele_size = self.element_size

        # read and parse the tree
        parentele = None
        while (f.tell()-h.tree_offset) < h.tree_size:
            element_bin = f.read(ele_size)
            ele = P3TElement(h)
            ele.parse(element_bin, self.element_fmt, f)

            # print ele

            # parse and add attributes to the element
            for y in range(0, ele.numattr):
                attr_bin = f.read(self.attr_size)
                attr = P3TAttribute()
                attr.parse(attr_bin, h, f)
                ele.add_attribute(attr)
                # print "    "+str(attr)

            # add to element map
            self.elemap[ele.offset] = ele

        # assign parents
        for k in self.elemap.keys():
            v = self.elemap[k]
            if k==0:
                continue
            if v.parentoffset==0:
                self.elemap[0].add_subelement(v)
            else:
                self.elemap[v.parentoffset].add_subelement(v)

        # restructure root element
        eleroot = self.elemap[0]
        self.elemap.clear()
        self.elemap[0] = eleroot

        # close file
        f.close()

    def dump_files(self, rootdir="extracted"):
        f = open(self.themefile, "rb")
        self.elemap[0].dump_files(rootdir, f)
        f.close()

    def dump_xml(self, rootdir="extracted"):
        document = xml.dom.minidom.Document()
        document.appendChild(self.elemap[0])
        f = open(rootdir+os.path.sep+os.path.basename(self.themefile)+".xml", "wb")
        f.write(document.toprettyxml())
        f.close()

def main():
    destination = None
    try:
        srcfile = sys.argv[1]
        try:
            destination = sys.argv[2]
        except:
            pass
    except:
        usage()
        return

    if destination==None:
        destination = "extracted."+os.path.basename(srcfile)

    print "Output directory: "+os.getcwd()+os.path.sep+destination

    try:
        extr = P3TExtractor(srcfile)
        extr.parse()
        extr.dump_files(destination)
        extr.dump_xml(destination)
        print extr.header
    except Exception, e:
        usage()
        print str(e)
        traceback.print_exc()
        raw_input("Press Enter to continue.")

def usage():
    print '''
P3T Unpacker %s (%s)
Copyright (c) 2007. Anoop Menon <codelogic at gmail dot com>

This program unpacks Playstation 3 Theme files (.p3t). By
default, it will extract the contents of the theme file
to the directory 'extracted.<theme filename>' in the current
working directory.

It will also output an XML file (named <themename>.xml) which
can be used to recreate the P3T file using the theme compiler,
with the extracted images. The XML file will have the correct
filenames, so no renaming will need to be done.

This program is still in alpha stage and probably has dozens
of bugs. Feel free to fix them if you want.

If you have any suggestions, feel free to email me.

Usage:
   p3textractor <input theme file> [destination path]
''' % (__version__,__date__)

if __name__=="__main__":
    main()