o3de/Tools/build/JenkinsScripts/distribution/Installer/Heat.py

#
# All or portions of this file Copyright (c) Amazon.com, Inc. or its affiliates or
# its licensors.
#
# For complete copyright and license terms please see the LICENSE at the root of this
# distribution (the "License"). All use of this software is governed by the License,
# or, if provided, by the license below or the license accompanying this file. Do not
# remove or modify any license notices. This file is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#

import os
import sys
import traceback
import xml.etree.ElementTree as ET

from BuildInstallerWixUtils import *

# HEAT COMMANDLINE TEMPLATES
heatCommandBase = 'heat.exe {harvestType} "{harvestSource}" -nologo -sreg -gg'
heatCommandRequiredArgs = " -dr {directoryRef} -cg {componentGroup} -out {outputPath}"


def convertToForwardSlashAndLower(pathString):
    """
    Used to ensure path strings are formatted in a consistent manner.
    """
    return pathString.replace('\\', '/').lower()


def lowerDictValues(dictToLower):
    """
    Used to lower-case all values in the given dictionary.
    """
    for key in dictToLower.keys():
        if isinstance(dictToLower[key], list):
            dictToLower[key][:] = [convertToForwardSlashAndLower(valueStr) for valueStr in dictToLower[key]]


def removeWxsEntriesWithJsonDirFilelist(wxsFilepath, jsonDirFilelist, directoryKey):
    """
    Removes entries from the given *.wxs file with the given jsonDirFilelist.
    @param wxsFilepath - Path to *.wxs file to operate the given jsonDirFilelist against
    @param jsonDirFilelist - Dictionary populated from given "dirFilelist" JSON. Files
        in the *.wxs files that don't have corresponding entries in dirFilelist will be
        removed from the *.wxs XML content.
    """
    # Without registering the namespace, the ET XML serialized output
    # is pretty funky, and WiX will not be happy.
    namespaceStr = 'http://schemas.microsoft.com/wix/2006/wi'
    ET.register_namespace('', namespaceStr)
    tree = ET.ElementTree()
    tree.parse(wxsFilepath)
    xmlRoot = tree.getroot()

    # Find all Component tags that have File tags
    ns = {'wixns': namespaceStr}
    componentGroupList = xmlRoot.findall('.//wixns:ComponentGroup', ns)
    for componentGroup in componentGroupList:
        componentList = componentGroup.findall('.//wixns:Component[wixns:File]', ns)

        for component in componentList:
            for childFileTag in component:
                # Entries are typically prefixed with an unnecessary 'SourceDir\' string
                sourceValue = convertToForwardSlashAndLower(childFileTag.attrib['Source']).replace('sourcedir/', '')
                if sourceValue not in jsonDirFilelist[directoryKey]:
                    component.remove(childFileTag)

            # WXS files generated by our installer typically only have one File tag per
            # ComponentGroup, but we'll check that it's empty, just in case.
            if len(component.getchildren()) < 1:
                componentGroup.remove(component)

    tree.write(wxsFilepath)


def heatDirectory(wxsName, sourceDirectory, outputPath, componentGroup, directoryRefName, verbose):
    """
    Generate a .wxs file for a given directory, including all subdirectories,
        and place it at the given outputPath.
    @remarks - Intentionally hardcoding harvest type to "dir" here. For other
        types of harvests, a new function should be created. Any logic not
        directly related to executing the heat command should exist outside this
        function.
    """
    verboseCmd = getVerboseCommand(verbose)

    # Intentionally hardcoding harvest type to "dir" here. For other types of harvests, a new
    #   function should be created. Any logic not directly related to executing
    #   the heat command should exist outside this function.
    heatCommand = heatCommandBase.format(harvestType="dir",
                                         harvestSource=sourceDirectory)
    heatCommand += verboseCmd
    heatCommand += heatCommandRequiredArgs.format(directoryRef=directoryRefName,
                                                  componentGroup=componentGroup, outputPath=outputPath)

    verbose_print(verbose, '\n{}\n'.format(heatCommand))
    return os.system(heatCommand)


def heatDirectories(rootDirectory, outputDirectory, directoryRefName, verbose, outputPrefix="", dirFilelist=None):
    """
    Generates package info for every directory in the given rootDirectory, and
        gathers each package's content information into a .wxs file.
    @param rootDirectory - The directory containing the source of many packages.
    @param outputDirectory - The directory to put all generated .wxs files.
    @param directoryRefName - The ID of the directory for these source files to
        be placed when installed. (Must match HeatPackageBase directory ID.)
    @param verbose - Running in verbose mode?
    @param outputPrefix - A string to append to the beginning of the name of the
        package when creating the output file. (Default = "").
    @param dirFilelist - A string that gives a path to a JSON file containing a
        list of directories, and for each directory, a list of files. Only the 
        files listed in the JSON file will be included in the installer output
        for the directory specified. This can be used to remove unnecessary
        files that aren't needed on a customer's machine, but are included in a
        packaged build created by Jenkins, for example.
    @return - A dictionary of package names to their associated package information (dictionaries).
    """

    # Attempt to parse jsonDirFilelist JSON.
    jsonDirFilelist = {}
    if dirFilelist is not None:
        assert (os.path.exists(dirFilelist)), 'The "dirFilelist" argument was provided but the JSON file at {} does not exist.'.format(dirFilelist)
        with open(dirFilelist, 'r') as source:
            try:
                jsonDirFilelist = json.load(source)
            except ValueError as e:
                print(traceback.format_exc())
                print('Error parsing the given JSON file at {} with exception: {}'.format(dirFilelist, e))
                sys.exit()
            except:
                print(traceback.format_exc())
                print('Unexpected error parsing the given JSON file at {}. Please verify that the file is correctly formatted.'.format(dirFilelist))
                sys.exit()

    combinedWxsResults = {}
    jsonValuesLowered = False

    for directoryName in get_immediate_subdirectories(rootDirectory):
        packageInfo = createPackageInfo(directoryName, rootDirectory, outputDirectory, outputPrefix)

        moduleName = packageInfo['name']
        sourceDirectory = packageInfo['sourcePath']
        wxsName = packageInfo['wxsName']
        outputPath = packageInfo['wxsPath']
        # There will only be one component group in the reference list at this point.
        componentGroup = packageInfo['componentGroupRefs']

        # check for existence of name collision
        if moduleName in combinedWxsResults:
            print('ERROR when creating module "{0}" from "{1}". A module with the name "{0}" already exists.'.format(moduleName, sourceDirectory))
            # Passing let us rapidly iterate on this tool, feel free to upgrade to raising an exception.
            pass

        combinedWxsResults[moduleName] = packageInfo
        success = heatDirectory(wxsName, sourceDirectory, outputPath, componentGroup, directoryRefName, verbose)
        assert (success == 0), 'Failed to generate WXS file for {}.'.format(moduleName)

        sourceDirFormatted = convertToForwardSlashAndLower(sourceDirectory)
        for key in jsonDirFilelist.keys():
            if sourceDirFormatted.endswith(convertToForwardSlashAndLower(key)):

                # Lower-case all entries to allow case-insensitive compare
                if not jsonValuesLowered:
                    lowerDictValues(jsonDirFilelist)
                    jsonValuesLowered = True

                # Alter XML contents of WXS file by filtering it against the JSON
                # directory list of files.
                removeWxsEntriesWithJsonDirFilelist(outputPath, jsonDirFilelist, key)

    return combinedWxsResults


def heatFile(file,
             directoryRefName,
             rootDirectory,
             verbose,
             componentGroup,
             outputPath):
    """
    Generates a WXS file for an individual file.
    @param file: Full path to the file to heat.
    @param directoryRefName: The ID of the directory for these source files to
        be placed when installed. (Must match HeatPackageBase directory ID.)
    @param verbose: Running in verbose mode?
    @param componentGroup: The component group to set in the file.
    @param outputPath: Where to output the generated WXS file.
    @return:
    """
    heatCommand = heatCommandBase.format(harvestType="file", harvestSource=file)
    heatCommand += getVerboseCommand(verbose)
    # According to Wix's docs ( http://wixtoolset.org/documentation/manual/v3/overview/heat.html )
    # SRD's description implies that it suppresses root directory harvesting, which seems to imply it works only
    # in directory harvesting mode. It also implies that it takes in no parameters.
    # This is either poor documentation, or incorrect (I'm assuming poor documentation).
    # The actual behavior is:
    # Normally when harvesting with Heat (directory or file), directories and directory references are generated
    # based on the path to the root directory. By calling suppress root directory harvesting and passing in a directory,
    # then the generated nested directory path in the generated wxs file will not include pathing based on this.
    # This is necessary when harvesting the loose files in a folder's root: Heat's harvesting does not support
    # whitelist / blacklist functionality, and all subfolders of the package have been harvested in other ways.
    # This means that all of the loose files in the directory roots that weren't included elsewhere need to be
    # included in some other installers. If the root directory is not suppressed, then they generate a directory
    # hierarchy that collides with each other.
    heatCommand += " -srd " + rootDirectory
    heatCommand += heatCommandRequiredArgs.format(directoryRef=directoryRefName,
                                                  componentGroup=componentGroup,
                                                  outputPath=outputPath)

    verbose_print(verbose, '\n{}\n'.format(heatCommand))
    return os.system(heatCommand)


def heatFiles(fileList,
              rootDirectory,
              outputDirectory,
              directoryRefName,
              verbose,
              outputPrefix = ""):
    """
    Generates WXS files for every file in the file list, and returns a mapping containing the associated
    WXS files and component groups.
    @param fileList: The list of files to generate WXS files for.
    @param outputDirectory: The directory to put generated WXS files.
    @param directoryRefName: The ID of the directory for these source files to
        be placed when installed. (Must match HeatPackageBase directory ID.)
    @param verbose: Running in verbose mode?
    @param outputPrefix: A string to append to the beginning of the name of the
        package when creating the output file. (Default = "").
    @return: A dictionary of input files to their associated WXS information (component group, WXS file location).
    """
    combinedWxsResults = {}
    for file in fileList:
        moduleName = strip_special_characters(file)
        wxsName = '{}{}'.format(outputPrefix, moduleName)
        outputPath = os.path.join(outputDirectory, '{}.wxs'.format(wxsName))
        componentGroup = '{}CG'.format(replace_leading_numbers(strip_special_characters(file)))

        wxsInfo = {
            'name': moduleName,
            'wxsName': wxsName,
            'wxsPath': outputPath,
            'componentGroupRefs': componentGroup
        }
        combinedWxsResults[file] = wxsInfo

        success = heatFile(os.path.join(rootDirectory,file),
                 directoryRefName,
                 rootDirectory,
                 verbose,
                 componentGroup,
                 outputPath)
        assert (success == 0), 'Failed to generate WXS file for {}.'.format(moduleName)

    return combinedWxsResults