#!/usr/bin/env python
# -*- encoding: iso-8859-1 -*-

###############################################################################
## treinfo.py - Print and extract file information from SWG TRE files.
## Copyright (C) 2004 por D. Rodríguez
## Fecha: 2004-09-12
## Autores: Daniel Rodríguez Troitiño <notzcoolx@yahoo.es>
##
##  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; either version 2 of the License, or
##  (at your option) any later version.
##
##  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.
##
##  You should have received a copy of the GNU General Public License
##  along with this program; if not, write to the Free Software
##  Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
##############################################################################/


import os
import sys
import zlib
import struct
import string
from optparse import OptionParser


fii_filename = 0
fii_unk1 = 1
fii_reallength = 2
fii_offset = 3
fii_compression = 4
fii_internallength = 5
fii_nameoffset = 6


def printTREinfo(fileData):
	totalSize = 0
	for fileItem in fileData:
		print "%9d %s" % (fileItem[fii_reallength], fileItem[fii_filename])
		totalSize += fileItem[fii_reallength]
	print "%9d All files" % totalSize

def extractTREFileItem(fileItemInfo):
	destdir = os.path.abspath(os.curdir)
	(dirtree, filename) = os.path.split(fileItemInfo[fii_filename])
	dirtree = dirtree.split('/') # files are stored Unix style, so not os.path.sep here
	for nextdir in dirtree:
		destdir = os.path.join(destdir, nextdir)
		if not os.path.exists(destdir):
			os.mkdir(destdir)
	newfile = open(os.path.join(destdir, filename), 'wb')
	
	trefile.seek(fileItemInfo[fii_offset])
	
	if fileItemInfo[fii_compression] == 2: size = fileItemInfo[fii_internallength]
	else: size = fileItemInfo[fii_reallength]
	rawData = trefile.read(size)
	
	if fileItemInfo[fii_compression] == 2:
		realData = zlib.decompress(rawData)
	else:
		realData = rawData
		
	newfile.write(realData)
	
	newfile.close()

parser = OptionParser(usage='%prog [options] filename.tre', version='%prog 0.02')
parser.add_option('-i', '--info', action='store_true', dest='onlyinfo',
                  default=False, help='Print file info and exit')
                  
(option, args) = parser.parse_args()

if len(args) != 1:
    parser.print_help()
    sys.exit(-1)

trefile = open(args[0], 'rb')

treident = trefile.read(8)
if treident != "EERT5000":
    print "%s is not a valid TRE file"
    sys.exit(-2)

(unk1,
 tableOffset,
 tableCompression,
 tableLength,
 nameCompression,
 nameLength,
 unk4) = struct.unpack("LLLLLLL", trefile.read(7*4))
 
trefile.seek(tableOffset)
tableCompressedData = trefile.read(tableLength)
nameCompressedData = trefile.read(nameLength)

if option.onlyinfo: # if we print we close the TRE file, we don't need it anymore
	trefile.close()

if tableCompression == 2: # if table is zlib compressed
	tableData = zlib.decompress(tableCompressedData)
else:
	tableData = tableCompressedData

if nameCompression == 2: # if names are zlib compressed
	nameData = string.split(zlib.decompress(nameCompressedData), '\x00')
else:
	nameData = string.split(nameCompressedData, '\x00')
	
fileData = []
for i in range(0, len(tableData), 6*4):
  fileInfo = struct.unpack("LLLLLL", tableData[i:i+6*4])
  fileInfo = (nameData[i/(6*4)],) + fileInfo # add the real name to the tuple
  fileData.append(fileInfo)

if option.onlyinfo:
	printTREinfo(fileData)
else:
	print 'Extracting...'
	for item in fileData:
		print item[0] # progress indicator
		extractTREFileItem(item)
	trefile.close()