einstein/einsteingame/resources.py

from os import listdir
import struct
import io
import zlib

from .buffer import Buffer
from .i18n import splitFileName, getScore, locale

class ResourceError(Exception):
    def __init__(self, code=None):
        self.code = code

class ResVariant(object):
    def __init__(self, resFile, i18nScore, dirEntry):
        self.__file = resFile
        self.__i18nScore = i18nScore
        self.__offset = dirEntry['offset']
        self.__unpackedSize = dirEntry['unpackedSize']
        self.__packedSize = dirEntry['packedSize']
        self.__level = dirEntry['level']
        self.__refCnt = 0
        self.__data = None

    @property
    def i18nScore(self):
        return self.__i18nScore

    def getData(self):
        buf = Buffer()

        buf.setSize(self.__unpackedSize)

        if self.__refCnt == 0:
            self.__data = self.__file.load(self.__offset, self.__packedSize, self.__unpackedSize, self.__level)

        return self.__data

class Resource(object):
    def __init__(self, file, i18nScore, entry, name):
        self.__variants = []
        self.__name = name

        self.addVariant(file, i18nScore, entry)

    def addVariant(self, file, i18nScore, entry):
        if len(self.__variants) == 0:
            self.__variants.append(ResVariant(file, i18nScore, entry))

            return

        try:
            best_variant = [idx for idx, variant in enumerate(self.__variants) \
                                    if variant.i18nScore == i18nScore][0]

            self.__variants[best_variant] = ResVariant(file, i18nScore, entry)
        except KeyError:
            self.__variants.append(ResVariant(file, i18nScore, entry))
            self.__variants.sort(key=lambda variant: variant.getI18nScore())

    def getVariantsCount(self):
        return len(self.__variants)

    @property
    def variants(self):
        return self.__variants

    def getName(self):
        return self.__name

class ResourceFile(object):
    def __init__(self, fileName, buf=None):
        self.__name = fileName

        if buf is not None:
            self.__buffer = buf
            self.__ownBuffer = False
        else:
            self.__buffer = Buffer()
            self.__ownBuffer = True

        try:
            self.__stream = open(self.__name, 'rb')
            sign = self.__stream.read(4)
            if sign != b'CRF\x00':
                raise ResourceError(code='invalid')

            major = self.readInt()
            minor = self.readInt()
            priority = self.readInt()

            if major != 2 or minor < 0:
                raise ResourceError(code='incompatible')
        except ResourceError as e:
            raise Exception("Invalid resource file: {}".format(e.code))
        except IOError as e:
            raise Exception("Error loading resource file '{}': {}".format(self.__name, e.message))

    def load(self, offset, packedSize, unpackedSize, level):
        try:
            if level == 0:
                self.__stream.seek(offset, io.SEEK_SET)
                self.__stream.read()

                return self.__stream.read(unpackedSize)

            self.__stream.seek(offset, io.SEEK_SET)
            data = self.__stream.read(packedSize)

            return self.unpack(data)
        except Exception as e:
            print("Eror loading resource: {}".format(e))

            raise

    # void ResourceFile::unpack(char *in, int inSize, char *out, int outSize)
    def unpack(self, data):
        return zlib.decompress(data)

    def readInt(self):
        return struct.unpack('i', self.__stream.read(4))[0]

    def readString(self):
        string = b''

        while True:
            b = self.__stream.read(1)

            if b == b'\x00' or b == b'':
                break

            string += b

        return string.decode('ascii')

    def getDirectory(self):
        try:
            self.__stream.seek(-8, io.SEEK_END)
        except IOError:
            raise Exception("Error reading {} directory".format(self.__name))

        start = self.readInt()
        count = self.readInt()

        try:
            self.__stream.seek(start, io.SEEK_SET)
        except IOError:
            raise Exception("Error reading {} directory".format(self.__name))

        directory = []

        try:
            for i in range(0, count):
                directory.append({
                    "name": self.readString(),
                    "unpackedSize": self.readInt(),
                    "offset": self.readInt(),
                    "packedSize": self.readInt(),
                    "level": self.readInt(),
                    "group": self.readString(),
                })
        except IOError:
            raise Exception("Error reading {} directory".format(self.__name))

        return directory

class ResourcesCollection:
    def __init__(self, directories):
        self.__files = []
        self.__buffer = None
        self.__resources = {}
        self.__groups = {}
        self.loadResourceFiles(directories)
        self.processFiles()

    def __iter__(self):
        return iter(self.__resources.values())

    def loadResourceFiles(self, directories):
        for directory in directories:
            try:
                for name in listdir(directory):
                    if not name.startswith('.') and name.endswith('.res'):
                        self.__files.append(
                            ResourceFile(directory + '/' + name,
                                         buf=self.__buffer))
            except OSError:
                pass

    def processFiles(self):
        for res_file in self.__files:
            directory = res_file.getDirectory()

            for de in directory:
                name, ext, language, country = splitFileName(de['name'])
                score = getScore(language, country, locale)

                if score > 0:
                    resName = name + '.' + ext
                    res = self.__resources.get(resName, None)

                    if not res:
                        res = Resource(res_file, score, de, resName)
                        self.__resources[resName] = res

                        if len(de['group']) > 0:
                            if de['group'] not in self.__groups:
                                self.__groups[de['group']] = []

                            self.__groups[de['group']].append(res)
                    else:
                        res.addVariant(res_file, score, de)

            directory.clear()

    def forEachInGroup(self, group, func):
        for resource in self.__groups[group]:
            func(resource)

    def dump(self):
        for name, resource in self.__resources.items():
            for variant in resource.variants:
                with open('resources/' + name, 'wb') as f:
                    f.write(variant.getData())

resources = None