From 9b8ba848f6b212e71cbb6c5061983e80fb149cd5 Mon Sep 17 00:00:00 2001
From: Manuel Thalmann <m@nuth.ch>
Date: Wed, 7 Aug 2024 18:29:20 +0200
Subject: [PATCH] Refactor app association script

---
 scripts/Windows/Scripts/AppAssociations.ps1 | 136 +++++++++++++-------
 1 file changed, 88 insertions(+), 48 deletions(-)

diff --git a/scripts/Windows/Scripts/AppAssociations.ps1 b/scripts/Windows/Scripts/AppAssociations.ps1
index 3321f165..c630292a 100644
--- a/scripts/Windows/Scripts/AppAssociations.ps1
+++ b/scripts/Windows/Scripts/AppAssociations.ps1
@@ -1,62 +1,102 @@
+using namespace System.Xml;
+
 $null = New-Module {
-    class AppAssociations {
-        static $defaultAssociationsSelector = "/DefaultAssociations";
-        static $associationSelector = "./Association";
+    $rootSelector = "/DefaultAssociations";
+    $associationSelector = "./Association";
 
-        static [string] GetSelector([string] $identifier) {
-            return "$([AppAssociations]::associationSelector)[@Identifier='$identifier']";
+    <#
+        .SYNOPSIS
+        Generates a selector for getting the association of the specified identifier.
+
+        .PARAMETER Identifier
+        The identifier to get the association of.
+    #>
+    $getSelector = {
+        param(
+            [string] $Identifier
+        )
+
+        [string] $filter = "";
+
+        if ($Identifier) {
+            $filter = "[@Identifier='$Identifier']";
         }
 
-        static [xml] GetAppAssociations() {
-            [xml] $associations = [xml]::new();
-            $associations.LoadXml(((DISM /Online /Get-DefaultAppAssociations) | Select-Object -Skip 6 | Select-Object -SkipLast 2 | Out-String));
-            return $associations;
-        }
-
-        static SetAssociation([string] $identifier, [string] $progId, [string] $applicationName) {
-            $document = [AppAssociations]::GetAppAssociations();
-            $associations = $document.SelectSingleNode([AppAssociations]::defaultAssociationsSelector);
-
-            [System.Xml.XmlNode] $association = $null;
-            $candidates = $associations.SelectNodes([AppAssociations]::GetSelector($identifier));
-
-            if ($candidates.Count -eq 1) {
-                $association = $candidates[0];
-            } else {
-                $association = ($associations.SelectNodes([AppAssociations]::associationSelector) | Select-Object -Last 1).CloneNode($true);
-                $association.Identifier = $identifier;
-                $association = $associations.AppendChild($association);
-            }
-
-            $association.ProgId = $progId;
-            $association.ApplicationName = $applicationName;
-            [AppAssociations]::SaveAssociations($document);
-        }
-
-        static SaveAssociations([xml] $document) {
-            $associations = $document.SelectSingleNode([AppAssociations]::defaultAssociationsSelector);
-            $defaultAssociations = $associations.SelectNodes([AppAssociations]::associationSelector);
-            $defaultAssociations | ForEach-Object { $associations.RemoveChild($_); } | Sort-Object -Property "Identifier" | ForEach-Object { $associations.AppendChild($_); };
-
-            $configFile = New-TemporaryFile;
-            $writerSettings = [System.Xml.XmlWriterSettings]::new();
-            $writerSettings.Indent = $true;
-            $writerSettings.Encoding = [System.Text.UTF8Encoding]::new();
-            $writer = [System.Xml.XmlWriter]::Create($configFile.FullName, $writerSettings);
-            $document.Save($writer);
-            $writer.Dispose();
-            DISM /Online "/Import-DefaultAppAssociations:$($configFile.FullName)";
-            Remove-Item $configFile;
-        }
+        "$rootSelector/$associationSelector$filter";
     }
 
+    <#
+        .SYNOPSIS
+        Gets the default app associations of the running Windows system.
+    #>
+    function Get-DefaultAppAssociations {
+        [OutputType([xml])]
+        param()
+        [xml]::new().LoadXml(((DISM /Online /Get-DefaultAppAssociations) | Select-Object -Skip 6 | Select-Object -SkipLast 2 | Out-String));
+    }
+
+    <#
+        .SYNOPSIS
+        Sets a default app association.
+
+        .PARAMETER Identifier
+        The identifier of the association to set.
+
+        .PARAMETER ProgId
+        The ProgId to set the association to.
+
+        .PARAMETER ApplicationName
+        The ApplicationName to set the association to.
+    #>
     function Set-DefaultAppAssociation {
-        param (
+        param(
             [string] $Identifier,
             [string] $ProgId,
             [string] $ApplicationName
         )
 
-        [AppAssociations]::SetAssociation($Identifier, $ProgId, $ApplicationName);
+        [System.Xml.XmlNode] $association = $null;
+        $document = Get-DefaultAppAssociations;
+        $candidates = $document.SelectNodes((& $getSelector $Identifier));
+
+        if ($candidates.Count -eq 1) {
+            $association = $candidates[0];
+        } else {
+            $association = ([System.Xml.XmlNode]($document.SelectNodes((& $getSelector)) | Select-Object -Last 1)[0]).CloneNode($true);
+            $association.Identifier = $Identifier;
+        }
+
+        $association.ProgId = $ProgId;
+        $association.ApplicationName = $ApplicationName;
+        Save-DefaultAppAssociations $document;
+    }
+
+    <#
+        .SYNOPSIS
+        Saves the default app associations.
+
+        .PARAMETER Document
+        The xml document containing the declarations of the app associations.
+    #>
+    function Save-DefaultAppAssociations {
+        param(
+            [xml] $Document
+        )
+
+        $root = $document.SelectSingleNode($rootSelector);
+        $associations = $root.SelectNodes((& $getSelector));
+
+        # Reorder associations by their Identifier
+        $associations | ForEach-Object { $root.RemoveChild($_); } | Sort-Object -Property "Identifier" | ForEach-Object { $root.AppendChild($_); };
+
+        $configFile = New-TemporaryFile;
+        $writerSettings = [XmlWriterSettings]::new();
+        $writerSettings.Indent = $true;
+        $writerSettings.Encoding = [System.Text.UTF8Encoding]::new();
+        $writer = [XmlWriter]::Create($configFile.FullName, $writerSettings);
+        $document.Save($writer);
+        $writer.Dispose();
+        DISM /Online "/Import-DefaultAppAssociations:$($configFile.FullName)";
+        Remove-Item $configFile;
     }
 }