path: root/src/tools/heat/UtilMutator.cs

// Copyright (c) .NET Foundation and contributors. All rights reserved. Licensed under the Microsoft Reciprocal License. See LICENSE.TXT file in the project root for full license information.

namespace WixToolset.Harvesters
{
    using System;
    using System.Collections;
    using System.Diagnostics;
    using System.Globalization;
    using System.IO;
    using WixToolset.Harvesters.Extensibility;
    using Wix = WixToolset.Harvesters.Serialize;

    /// <summary>
    /// The template type.
    /// </summary>
    public enum TemplateType
    {
        /// <summary>
        /// A fragment template.
        /// </summary>
        Fragment,

        /// <summary>
        /// A module template.
        /// </summary>
        Module,

        /// <summary>
        /// A product template.
        /// </summary>
        Package
    }

    /// <summary>
    /// The mutator for the WiX Toolset Internet Information Services Extension.
    /// </summary>
    public sealed class UtilMutator : BaseMutatorExtension
    {
        private ArrayList components;
        private ArrayList componentGroups;
        private string componentGroupName;
        private bool createFragments;
        private ArrayList directories;
        private ArrayList directoryRefs;
        private ArrayList files;
        private ArrayList features;
        private SortedList fragments;
        private bool autogenerateGuids;
        private bool generateGuids;
        private string guidFormat = "B"; // Defaults to guid in {}
        private Wix.IParentElement rootElement;
        private bool setUniqueIdentifiers;
        private TemplateType templateType;

        /// <summary>
        /// Instantiate a new UtilMutator.
        /// </summary>
        public UtilMutator()
        {
            this.components = new ArrayList();
            this.componentGroups = new ArrayList();
            this.directories = new ArrayList();
            this.directoryRefs = new ArrayList();
            this.features = new ArrayList();
            this.files = new ArrayList();
            this.fragments = new SortedList();
        }

        /// <summary>
        /// Gets or sets the value of the component group name.
        /// </summary>
        /// <value>The component group name.</value>
        public string ComponentGroupName
        {
            get { return this.componentGroupName; }
            set { this.componentGroupName = value; }
        }

        /// <summary>
        /// Gets or sets the option to create fragments.
        /// </summary>
        /// <value>The option to create fragments.</value>
        public bool CreateFragments
        {
            get { return this.createFragments; }
            set { this.createFragments = value; }
        }

        /// <summary>
        /// Gets or sets the option to autogenerate component guids at compile time.
        /// </summary>
        /// <value>The option to autogenerate component guids.</value>
        public bool AutogenerateGuids
        {
            get { return this.autogenerateGuids; }
            set { this.autogenerateGuids = value; }
        }

        /// <summary>
        /// Gets or sets the option to generate missing guids.
        /// </summary>
        /// <value>The option to generate missing guids.</value>
        public bool GenerateGuids
        {
            get { return this.generateGuids; }
            set { this.generateGuids = value; }
        }

        /// <summary>
        /// Gets or sets the option to set the format of guids.
        /// D - 32 digits separated by hyphens:
        ///     xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx
        /// B - 32 digits separated by hyphens, enclosed in brackets:
        ///     {xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx}
        /// </summary>
        /// <value>Guid format either B or D.</value>
        public string GuidFormat
        {
            get { return this.guidFormat; }
            set { this.guidFormat = value; }
        }

        /// <summary>
        /// Gets the sequence of the extension.
        /// </summary>
        /// <value>The sequence of the extension.</value>
        public override int Sequence
        {
            get { return 1000; }
        }

        /// <summary>
        /// Gets of sets the option to set unique identifiers.
        /// </summary>
        /// <value>The option to set unique identifiers.</value>
        public bool SetUniqueIdentifiers
        {
            get { return this.setUniqueIdentifiers; }
            set { this.setUniqueIdentifiers = value; }
        }

        /// <summary>
        /// Gets or sets the template type.
        /// </summary>
        /// <value>The template type.</value>
        public TemplateType TemplateType
        {
            get { return this.templateType; }
            set { this.templateType = value; }
        }

        /// <summary>
        /// Mutate a WiX document.
        /// </summary>
        /// <param name="wix">The Wix document element.</param>
        public override void Mutate(Wix.Wix wix)
        {
            this.components.Clear();
            this.directories.Clear();
            this.directoryRefs.Clear();
            this.features.Clear();
            this.files.Clear();
            this.fragments.Clear();
            this.rootElement = null;

            // index elements in this wix document
            this.IndexElement(wix);

            this.MutateWix(wix);

            this.MutateFiles();

            this.MutateDirectories();

            this.MutateComponents();

            if (null != this.componentGroupName)
            {
                this.CreateComponentGroup(wix);
            }

            // add the components to the product feature after all the identifiers have been set
            if (TemplateType.Package  == this.templateType)
            {
                var feature = (Wix.Feature)this.features[0];

                foreach (Wix.ComponentGroup group in this.componentGroups)
                {
                    var componentGroupRef = new Wix.ComponentGroupRef();
                    componentGroupRef.Id = group.Id;

                    feature.AddChild(componentGroupRef);
                }
            }
            else if (TemplateType.Module == this.templateType)
            {
                foreach (Wix.ISchemaElement element in wix.Children)
                {
                    if (element is Wix.Module)
                    {
                        foreach (Wix.ComponentGroup group in this.componentGroups)
                        {
                            var componentGroupRef = new Wix.ComponentGroupRef();
                            componentGroupRef.Id = group.Id;

                            ((Wix.IParentElement)element).AddChild(componentGroupRef);
                        }
                        break;
                    }
                }
            }

            //if(!this.createFragments && TemplateType.Package
            foreach (Wix.Fragment fragment in this.fragments.Values)
            {
                wix.AddChild(fragment);
            }
        }

        /// <summary>
        /// Creates a component group with a given name.
        /// </summary>
        /// <param name="wix">The Wix document element.</param>
        private void CreateComponentGroup(Wix.Wix wix)
        {
            var componentGroup = new Wix.ComponentGroup();
            componentGroup.Id = this.componentGroupName;
            this.componentGroups.Add(componentGroup);

            var cgFragment = new Wix.Fragment();
            cgFragment.AddChild(componentGroup);
            wix.AddChild(cgFragment);

            var componentCount = 0;
            for (; componentCount < this.components.Count; componentCount++)
            {
                var c = this.components[componentCount] as Wix.Component;

                if (this.createFragments)
                {
                    if (c.ParentElement is Wix.Directory)
                    {
                        var parentDirectory = c.ParentElement as Wix.Directory;

                        componentGroup.AddChild(c);
                        c.Directory = parentDirectory.Id;
                        parentDirectory.RemoveChild(c);
                    }
                    else if (c.ParentElement is Wix.DirectoryRef || c.ParentElement is Wix.StandardDirectory)
                    {
                        var parentDirectory = c.ParentElement as Wix.DirectoryBase;

                        componentGroup.AddChild(c);
                        c.Directory = parentDirectory.Id;
                        parentDirectory.RemoveChild(c);

                        // Remove whole fragment if moving the component to the component group just leaves an empty DirectoryRef
                        if (0 < this.fragments.Count && parentDirectory.ParentElement is Wix.Fragment)
                        {
                            var parentFragment = parentDirectory.ParentElement as Wix.Fragment;
                            var childCount = 0;
                            foreach (Wix.ISchemaElement element in parentFragment.Children)
                            {
                                childCount++;
                            }

                            // Component should always have an Id but the SortedList creation allows for null and bases the name on the fragment count which we cannot reverse engineer here.
                            if (1 == childCount && !String.IsNullOrEmpty(c.Id))
                            {
                                var removeIndex = this.fragments.IndexOfKey(String.Concat("Component:", c.Id));
                                if (0 <= removeIndex)
                                {
                                    this.fragments.RemoveAt(removeIndex);
                                }
                            }
                        }
                    }
                }
                else
                {
                    var componentRef = new Wix.ComponentRef();
                    componentRef.Id = c.Id;
                    componentGroup.AddChild(componentRef);
                }
            }
        }

        /// <summary>
        /// Index an element.
        /// </summary>
        /// <param name="element">The element to index.</param>
        private void IndexElement(Wix.ISchemaElement element)
        {
            if (element is Wix.Component)
            {
                this.components.Add(element);
            }
            else if (element is Wix.ComponentGroup)
            {
                this.componentGroups.Add(element);
            }
            else if (element is Wix.Directory)
            {
                this.directories.Add(element);
            }
            else if (element is Wix.DirectoryRef)
            {
                this.directoryRefs.Add(element);
            }
            else if (element is Wix.Feature)
            {
                this.features.Add(element);
            }
            else if (element is Wix.File)
            {
                this.files.Add(element);
            }
            else if (element is Wix.Module || element is Wix.PatchCreation || element is Wix.Package)
            {
                Debug.Assert(null == this.rootElement);
                this.rootElement = (Wix.IParentElement)element;
            }

            // index the child elements
            if (element is Wix.IParentElement parentElement)
            {
                foreach (Wix.ISchemaElement childElement in parentElement.Children)
                {
                    this.IndexElement(childElement);
                }
            }
        }

        /// <summary>
        /// Mutate the components.
        /// </summary>
        private void MutateComponents()
        {
            var identifierGenerator = new IdentifierGenerator("Component", this.Core);
            if (TemplateType.Module == this.templateType)
            {
                identifierGenerator.MaxIdentifierLength = IdentifierGenerator.MaxModuleIdentifierLength;
            }

            foreach (Wix.Component component in this.components)
            {
                if (null == component.Id)
                {
                    var firstFileId = String.Empty;

                    // attempt to create a possible identifier from the first file identifier in the component
                    foreach (Wix.File file in component[typeof(Wix.File)])
                    {
                        firstFileId = file.Id;
                        break;
                    }

                    if (String.IsNullOrEmpty(firstFileId))
                    {
                        firstFileId = this.GetGuid();
                    }

                    component.Id = identifierGenerator.GetIdentifier(firstFileId);
                }

                if (null == component.Guid)
                {
                    if (this.AutogenerateGuids)
                    {
                        component.Guid = "*";
                    }
                    else
                    {
                        component.Guid = this.GetGuid();
                    }
                }

                if (this.createFragments && component.ParentElement is Wix.Directory directory)
                {
                    // parent directory must have an identifier to create a reference to it
                    if (null == directory.Id)
                    {
                        break;
                    }

                    if (this.rootElement is Wix.Module)
                    {
                        // add a ComponentRef for the Component
                        var componentRef = new Wix.ComponentRef();
                        componentRef.Id = component.Id;
                        this.rootElement.AddChild(componentRef);
                    }

                    // create a new Fragment
                    var fragment = new Wix.Fragment();
                    this.fragments.Add(String.Concat("Component:", component.Id ?? this.fragments.Count.ToString()), fragment);

                    // create a new DirectoryRef
                    var directoryRef = DirectoryHelper.CreateDirectoryReference(directory.Id);
                    fragment.AddChild(directoryRef);

                    // move the Component from the the Directory to the DirectoryRef
                    directory.RemoveChild(component);
                    directoryRef.AddChild(component);
                }
            }
        }

        /// <summary>
        /// Mutate the directories.
        /// </summary>
        private void MutateDirectories()
        {
            if (!this.setUniqueIdentifiers)
            {
                // assign all identifiers before fragmenting (because fragmenting requires them all to be present)
                var identifierGenerator = new IdentifierGenerator("Directory", this.Core);
                if (TemplateType.Module == this.templateType)
                {
                    identifierGenerator.MaxIdentifierLength = IdentifierGenerator.MaxModuleIdentifierLength;
                }

                foreach (Wix.Directory directory in this.directories)
                {
                    if (null == directory.Id)
                    {
                        directory.Id = identifierGenerator.GetIdentifier(directory.Name);
                    }
                }
            }

            if (this.createFragments)
            {
                foreach (Wix.Directory directory in this.directories)
                {
                    if (directory.ParentElement is Wix.Directory)
                    {
                        var parentDirectory = (Wix.DirectoryBase)directory.ParentElement;

                        // parent directory must have an identifier to create a reference to it
                        if (null == parentDirectory.Id)
                        {
                            return;
                        }

                        // create a new Fragment
                        var fragment = new Wix.Fragment();
                        this.fragments.Add(String.Concat("Directory:", ("TARGETDIR" == directory.Id ? null : (null != directory.Id ? directory.Id : this.fragments.Count.ToString()))), fragment);

                        // create a new DirectoryRef
                        var directoryRef = DirectoryHelper.CreateDirectoryReference(parentDirectory.Id);
                        fragment.AddChild(directoryRef);

                        // move the Directory from the parent Directory to DirectoryRef
                        parentDirectory.RemoveChild(directory);
                        directoryRef.AddChild(directory);
                    }
                    else if (directory.ParentElement is Wix.Fragment parent)
                    {
                        // When creating fragments, remove any top-level Directory elements;
                        // the fragments should be pulled in by their DirectoryRefs instead.
                        parent.RemoveChild(directory);

                        // Remove the fragment if it is empty.
                        if (parent.Children.GetEnumerator().Current == null && parent.ParentElement != null)
                        {
                            ((Wix.IParentElement)parent.ParentElement).RemoveChild(parent);
                        }
                    }
                    else if (directory.ParentElement == this.rootElement)
                    {
                        // create a new Fragment
                        var fragment = new Wix.Fragment();
                        this.fragments.Add(String.Concat("Directory:", ("TARGETDIR" == directory.Id ? null : (null != directory.Id ? directory.Id : this.fragments.Count.ToString()))), fragment);

                        // move the Directory from the root element to the Fragment
                        this.rootElement.RemoveChild(directory);
                        fragment.AddChild(directory);
                    }
                }
            }
        }

        /// <summary>
        /// Mutate the files.
        /// </summary>
        private void MutateFiles()
        {
            var identifierGenerator = new IdentifierGenerator("File", this.Core);
            if (TemplateType.Module == this.templateType)
            {
                identifierGenerator.MaxIdentifierLength = IdentifierGenerator.MaxModuleIdentifierLength;
            }

            foreach (Wix.File file in this.files)
            {
                if (null == file.Id)
                {
                    file.Id = identifierGenerator.GetIdentifier(Path.GetFileName(file.Source));
                }
            }
        }

        /// <summary>
        /// Mutate a Wix element.
        /// </summary>
        /// <param name="wix">The Wix element to mutate.</param>
        private void MutateWix(Wix.Wix wix)
        {
            if (TemplateType.Fragment != this.templateType)
            {
                if (null != this.rootElement || 0 != this.features.Count)
                {
                    throw new Exception("The template option cannot be used with Feature, Package, or Module elements present.");
                }

                // create a package element although it won't always be used
                var package = new Wix.SummaryInformation();
                if (TemplateType.Module == this.templateType)
                {
                    package.Id = this.GetGuid();
                }
                else
                {
                    package.Compressed = Wix.YesNoType.yes;
                }

                package.InstallerVersion = 200;

                var targetDir = new Wix.Directory();
                targetDir.Id = "TARGETDIR";
                targetDir.Name = "SourceDir";

                foreach (Wix.DirectoryRef directoryRef in this.directoryRefs)
                {
                    if (String.Equals(directoryRef.Id, "TARGETDIR", StringComparison.OrdinalIgnoreCase))
                    {
                        var parent = directoryRef.ParentElement as Wix.IParentElement;

                        foreach (Wix.ISchemaElement element in directoryRef.Children)
                        {
                            targetDir.AddChild(element);
                        }

                        parent.RemoveChild(directoryRef);

                        if (null != ((Wix.ISchemaElement)parent).ParentElement)
                        {
                            var i = 0;

                            foreach (Wix.ISchemaElement element in parent.Children)
                            {
                                i++;
                            }

                            if (0 == i)
                            {
                                var supParent = (Wix.IParentElement)((Wix.ISchemaElement)parent).ParentElement;
                                supParent.RemoveChild((Wix.ISchemaElement)parent);
                            }
                        }

                        break;
                    }
                }

                if (TemplateType.Module == this.templateType)
                {
                    var module = new Wix.Module();
                    module.Id = "PUT-MODULE-NAME-HERE";
                    module.Language = "1033";
                    module.Version = "1.0.0.0";

                    package.Manufacturer = "PUT-COMPANY-NAME-HERE";
                    module.AddChild(package);
                    module.AddChild(targetDir);

                    wix.AddChild(module);
                    this.rootElement = module;
                }
                else // product
                {
                    var product = new Wix.Package();
                    product.Id = this.GetGuid();
                    product.Language = "1033";
                    product.Manufacturer = "PUT-COMPANY-NAME-HERE";
                    product.Name = "PUT-PRODUCT-NAME-HERE";
                    product.UpgradeCode = this.GetGuid();
                    product.Version = "1.0.0.0";
                    product.AddChild(package);
                    product.AddChild(targetDir);

                    var media = new Wix.Media();
                    media.Id = "1";
                    media.Cabinet = "product.cab";
                    media.EmbedCab = Wix.YesNoType.yes;
                    product.AddChild(media);

                    var feature = new Wix.Feature();
                    feature.Id = "ProductFeature";
                    feature.Title = "PUT-FEATURE-TITLE-HERE";
                    feature.Level = 1;
                    product.AddChild(feature);
                    this.features.Add(feature);

                    wix.AddChild(product);
                    this.rootElement = product;
                }
            }
        }

        /// <summary>
        /// Get a generated guid or a placeholder for a guid.
        /// </summary>
        /// <returns>A generated guid or placeholder.</returns>
        private string GetGuid()
        {
            if (this.generateGuids)
            {
                return Guid.NewGuid().ToString(this.guidFormat, CultureInfo.InvariantCulture).ToUpper(CultureInfo.InvariantCulture);
            }
            else
            {
                return "PUT-GUID-HERE";
            }
        }
    }
}