from ZUnpacker import CBinaryReader
from ZUnpacker import SuperStruct
from Crypto.Cipher import AES
import os
import zlib

FILE_PREFIX = 'RustyHearts' #This is the prefix placed before all files output to the hard drive

#These are options for the table outputs
DIVIDER = '|' #What kind of separation of values the tables use
SHOW_TYPES = 1 #1 = Show the types at the top of the table. 0 = Do not show types
NEW_LINES = '\r\n' #The new lines that tables use


def main():
    fileTree = GetFileTree()
    for i, (fileName, pckFileName, fileLength, fileChecksum, fileOffset) in enumerate(fileTree, 1):
        pckFile = CBinaryReader.BinaryReader(pckFileName, 'rb')
        pckFile.seek(fileOffset)
        fileText = pckFile.read(fileLength)
        pckFile.close()

        fileName = os.path.join(FILE_PREFIX, fileName)
        print('%03d:%d - %s' % (i, len(fileTree), fileName))         

        if fileName.endswith('.rh'):
            fileText = TableHandler(fileText)

        WriteFile(fileName, fileText)

def GetFileTree():
    decryptedFileTree = DecryptMIPFile('f00X.dat')
    f = CBinaryReader.BinaryReader(decryptedFileTree, 'string')

    dirTree = []
    while not f.ReachedEOF():
        fileName = GetString(f)
        pckFileName = '%s.pck' % (str(f.rU8()).zfill(3))
        fileLength = f.rU32()
        fileChecksum = f.rU32()
        fileOffset = f.rU32()
        nulls = f.rU32()
        dirTree.append([fileName, pckFileName, fileLength, fileChecksum, fileOffset])
    return dirTree

def DecryptMIPFile(fileName):
    keys = SuperStruct.Unpack(GetFile('Zbin/RustyHearts.keys'))
    mipText = SuperStruct.Unpack(GetFile(fileName))
    mipText = list(mipText)
    for i, encryptedChar in enumerate(mipText):
        mipText[i] = encryptedChar ^ keys[i & 0xFF]
    mipText = SuperStruct.Pack(mipText)
    mipText = zlib.decompress(mipText)
    return mipText

def GenerateAESCipher():
    aesKey = 'gkw3iurpamv;kj20984;asdkfjat1af\0'
    aesIV = SuperStruct.Pack([0xDB, 0x0F, 0x49, 0x40] * 4)
    mode = AES.MODE_ECB
    Cipher = AES.new(aesKey, mode, aesIV)
    return Cipher


#The tables are in encrypted, binary form
#This function decrypts the files and turns them into separated values
def TableHandler(fileText):
    Cipher = GenerateAESCipher()
    fileText = Cipher.decrypt(fileText)

    f = CBinaryReader.BinaryReader(fileText, 'string')
    numRows = f.rS32()
    numColumns = f.rS32()

    columnNames = []
    for i in range(numColumns):
        resultString = GetString(f)
        columnNames.append(resultString)

    columnTypes = []
    for i in range(numColumns):
        columnType = f.rS32()
        if columnType == 0:
            columnTypes.append('int32')
        elif columnType == 1:
            columnTypes.append('float32')
        elif columnType == 3:
            columnTypes.append('string')
        elif columnType == 4:
            columnTypes.append('int64')
        else:
            raise Exception('Unknown column type %d at offset %d' % (columnType, f.tell()))
    
    rows = []
    for i in range(numRows):
        thisRow = []
        for j in range(numColumns):
            if columnTypes[j] == 'int32':
                thisRow.append(str(f.rS32()))
            elif columnTypes[j] == 'float32':
                thisRow.append(str(f.rF32()))
            elif columnTypes[j] == 'string':
                thisRow.append(GetString(f))
            elif columnTypes[j] == 'int64':
                thisRow.append(str(f.rS64()))
            else:
                raise Exception("Unknown column type %s?" % (columnTypes[j]))
        rows.append(thisRow)
    columnTypesStr = DIVIDER.join(columnTypes)
    columnNamesStr = DIVIDER.join(columnNames)
    rows = NEW_LINES.join([DIVIDER.join(eachRow) for eachRow in rows])

    resultTable = '%s%s%s%s' % (columnNamesStr, NEW_LINES, rows, NEW_LINES)
    if SHOW_TYPES:
        resultTable = '%s%s' % (columnTypesStr, NEW_LINES) + resultTable
    return resultTable
    

def GetString(f):
    length = f.rS16()
    resultString = []
    for j in range(length):
        char = f.rU16()
        if char < 255:
            char = chr(char)
        else:
            #Windows doesn't play nicely with actual UTF-16 characters
            #However, using "&#%s;" works on Windows
            #This is HTML formatting. If there's another way to get
            #UTF-16 characters to work in Windows, let me know!
            char = '&#%s;' % (char)
        resultString.append(char)
    resultString = ''.join(resultString)
    return resultString

def GetFile(fileName):
    f = open(fileName, 'rb')
    text = f.read()
    f.close()
    return text

def WriteFile(fileName, fileText):
    if '\\' in fileName:
        fileName = fileName.replace('\\', '/')
    fileDir = '/'.join(fileName.split('/')[:-1])
    if not os.access(fileDir, os.F_OK):
        os.makedirs(fileDir)
    f = open(fileName, 'wb')
    f.write(fileText)
    f.close()

if __name__ == '__main__':
    main()
