LDD/reader/asset_reader.py

"""
    LIF Extractor - LEGO Digital Designer LIF extractor.

    Copyright (C) 2012 JrMasterModelBuilder

    You accept full responsibility for how you use this program.

    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 3 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, see <http://www.gnu.org/licenses/>.
"""

import os
import logging

logger = logging.getLogger(__name__)
logger.setLevel(logging.DEBUG)

LDD_INSTALL_PATH = r'C:\Program Files (x86)\LEGO Company\LEGO Digital Designer'
LDD_ASSET_PATH = os.path.join(LDD_INSTALL_PATH, 'Assets.lif')

fileData = ''


def asset_extract(design_ids=None, material_only=False, debug=True):
    global fileData

    def uint32(offset):
        if bytesAre == "str":
            return (ord(fileData[offset]) * 16777216) + (ord(fileData[offset + 1]) * 65536) + (
                ord(fileData[offset + 2]) * 256) + ord(fileData[offset + 3])
        else:
            return (fileData[offset] * 16777216) + (fileData[offset + 1] * 65536) + (fileData[offset + 2] * 256) + \
                   fileData[offset + 3]

    def uint16(offset):
        if bytesAre == "str":
            return (ord(fileData[offset]) * 256) + ord(fileData[offset + 1])
        else:
            return (fileData[offset] * 256) + fileData[offset + 1]

    def recurse(prefix, offset):
        if prefix == "":
            offset += 36
        else:
            folderList.extend([prefix])
            offset += 4
        for i in range(uint32(offset)):
            offset += 4
            entryType = uint16(offset)  # 1 = directory, 2 = file
            print('xxx', entryType, offset)
            offset += 6
            # Build the name.
            entryName = os.sep
            # Check both integer and byte for different versions of Python.
            while fileData[offset + 1] != 0 and fileData[offset + 1] != b"\x00":
                print('---', offset, fileData[offset + 1])
                if bytesAre == "str":
                    entryName += fileData[offset + 1]
                else:
                    entryName += chr(fileData[offset + 1])
                offset += 2
            offset += 6

            if entryType == 1:
                # Recurse through the director, and update the offset.
                packedFilesOffset[0] += 20
                offset = recurse(prefix + entryName, offset)
            elif entryType == 2:
                # File offset.
                packedFilesOffset[0] += 20
                fileOffset = packedFilesOffset[0]
                # File size.
                fileSize = uint32(offset) - 20
                offset += 24
                packedFilesOffset[0] += fileSize
                fileList.extend([[prefix + entryName, fileOffset, fileSize]])
        # Return the offset at the end of this run to update parent runnings.
        return offset
    if debug:
        logger.debug('PROCESSING: %s' % LDD_ASSET_PATH)
    # Open the file if valid.
    try:
        with open(LDD_ASSET_PATH, "rb") as f:
            fileData = f.read()

    except IOError as e:
        logger.error('ERROR: Failed to read file.')
        return

    if len(fileData) < 4 or fileData[0:4] != b"LIFF":
        logger.error('ERROR: Not a LIF file.')
        return

    if debug:
        logger.debug('EXTRACTING: Asset.lif: Please wait.')

    # Recurse through, creating a list of all the files.
    packedFilesOffset = [84]  # Array for non-global function-modifiable variable.
    folderList = []
    fileList = []

    # Check if this version of Python treats bytes as int or str
    bytesAre = type(b'a'[0]).__name__

    # Recuse through the archive.
    recurse("", (uint32(72) + 64))

    for i in fileList:
        if 'db.lif' in i[0]:
            fileData = fileData[i[1]:i[1] + i[2]]
            break
    if debug:
        logger.debug('PROCESSING: db.lif')
    if len(fileData) < 4 or fileData[0:4] != b"LIFF":
        logger.error("\tERROR: Not a LIF file.")
        return
    if debug:
        logger.debug('EXTRACTING db.lif: Please wait.')
    # Recurse through, creating a list of all the files.
    packedFilesOffset = [84]  # Array for non-global function-modifiable variable.
    folderList = []
    fileList = []

    # Check if this version of Python treats bytes as int or str
    bytesAre = type(b'a'[0]).__name__

    # Recuse through the archive.
    recurse("", (uint32(72) + 64))

    if material_only:
        for i in fileList:
            material_file = r'\Materials.xml'
            if i[0] in [material_file]:
                m_data = fileData[i[1]:i[1] + i[2]]
                return m_data
        return None

    if design_ids:
        g_data_map = {}
        g_file_list = [r'\Primitives\LOD0\%s.g' % d for d in design_ids]
        for i in fileList:
            if i[0] not in g_file_list:
                continue
            g_data = fileData[i[1]:i[1] + i[2]]
            design_id = design_ids[g_file_list.index(i[0])]
            g_data_map[design_id] = g_data

        return g_data_map
    else:
        design_id_list = []
        for i in fileList:
            if r'\Primitives\LOD0' not in i[0]:
                continue
            design_id = os.path.splitext(os.path.basename(i[0]))[0]
            design_id_list.append(design_id)

        return design_id_list
if __name__ == '__main__':
    asset_extract('99818', material_only=True)