Refactor ImportRegistryFile function to simplify error handling and improve fallback logic

Scripts/Features/ImportRegistryFile.ps1

@@ -8,12 +8,13 @@ function ImportRegistryFile {
     Write-Host $message
 
     $usesOfflineHive = $script:Params.ContainsKey("Sysprep") -or $script:Params.ContainsKey("User")
-    $regFilePath = if ($usesOfflineHive) {
-        "$script:RegfilesPath\Sysprep\$path"
+    $regFileDirectory = if ($usesOfflineHive) {
+        Join-Path $script:RegfilesPath "Sysprep"
     }
     else {
-        "$script:RegfilesPath\$path"
+        $script:RegfilesPath
     }
+    $regFilePath = Join-Path $regFileDirectory $path
 
     if (-not (Test-Path $regFilePath)) {
         $errorMessage = "Unable to find registry file: $path ($regFilePath)"
@@ -22,19 +23,28 @@ function ImportRegistryFile {
         throw $errorMessage
     }
 
-    # Reset exit code before running reg.exe for reliable success detection
-    $global:LASTEXITCODE = 0
+    $regResult = $null
+    $offlineHiveLoaded = $false
 
+    try {
         if ($usesOfflineHive) {
             # Sysprep targets Default user, User targets the specified user
-        $hiveDatPath = if ($script:Params.ContainsKey("Sysprep")) {
-            GetUserDirectory -userName "Default" -fileName "NTUSER.DAT"
-        } else {
-            GetUserDirectory -userName $script:Params.Item("User") -fileName "NTUSER.DAT"
+            $targetUserName = if ($script:Params.ContainsKey("Sysprep")) { "Default" } else { $script:Params.Item("User") }
+            $hiveDatPath = GetUserDirectory -userName $targetUserName -fileName "NTUSER.DAT"
+
+            $global:LASTEXITCODE = 0
+            reg load "HKU\Default" $hiveDatPath | Out-Null
+            $loadExitCode = $LASTEXITCODE
+
+            if ($loadExitCode -ne 0) {
+                throw "Failed importing registry file '$path'. Offline hive load failed: Failed to load user hive at '$hiveDatPath' (exit code: $loadExitCode)"
+            }
+
+            $offlineHiveLoaded = $true
         }
 
         $regResult = Invoke-NonBlocking -ScriptBlock {
-            param($hivePath, $targetRegFilePath)
+            param($targetRegFilePath)
             $result = @{
                 Output = @()
                 ExitCode = 0
@@ -43,13 +53,6 @@ function ImportRegistryFile {
 
             try {
                 $global:LASTEXITCODE = 0
-                reg load "HKU\Default" $hivePath | Out-Null
-                $loadExitCode = $LASTEXITCODE
-
-                if ($loadExitCode -ne 0) {
-                    throw "Failed to load user hive at '$hivePath' (exit code: $loadExitCode)"
-                }
-
                 $output = reg import $targetRegFilePath 2>&1
                 $importExitCode = $LASTEXITCODE
 
@@ -66,27 +69,9 @@ function ImportRegistryFile {
                 $result.Error = $_.Exception.Message
                 $result.ExitCode = if ($LASTEXITCODE -ne 0) { $LASTEXITCODE } else { 1 }
             }
-            finally {
-                $global:LASTEXITCODE = 0
-                reg unload "HKU\Default" | Out-Null
-                $unloadExitCode = $LASTEXITCODE
-                if ($unloadExitCode -ne 0 -and -not $result.Error) {
-                    $result.Error = "Failed to unload registry hive HKU\Default (exit code: $unloadExitCode)"
-                    $result.ExitCode = $unloadExitCode
-                }
-            }
 
             return $result
-        } -ArgumentList @($hiveDatPath, $regFilePath)
-    }
-    else {
-        $regResult = Invoke-NonBlocking -ScriptBlock {
-            param($targetRegFilePath)
-            $global:LASTEXITCODE = 0
-            $output = reg import $targetRegFilePath 2>&1
-            return @{ Output = @($output); ExitCode = $LASTEXITCODE; Error = $null }
         } -ArgumentList $regFilePath
-    }
 
         $regOutput = @($regResult.Output)
         $hasSuccess = ($regResult.ExitCode -eq 0) -and -not $regResult.Error
@@ -107,11 +92,27 @@ function ImportRegistryFile {
 
         if (-not $hasSuccess) {
             $details = if ($regResult.Error) { $regResult.Error } else { "Exit code: $($regResult.ExitCode)" }
-        $errorMessage = "Failed importing registry file '$path'. $details"
-        Write-Host $errorMessage -ForegroundColor Red
-        Write-Host ""
-        throw $errorMessage
+            Write-Warning "reg import failed for '$path'. Falling back to PowerShell registry writer. Details: $details"
+            Invoke-RegistryOperationsFromRegFile -RegFilePath $regFilePath
+            Write-Host "Fallback import succeeded for '$path'." -ForegroundColor Yellow
         }
 
         Write-Host ""
     }
+    catch {
+        Write-Host $_.Exception.Message -ForegroundColor Red
+        Write-Host ""
+        throw
+    }
+    finally {
+        if ($offlineHiveLoaded) {
+            $global:LASTEXITCODE = 0
+            reg unload "HKU\Default" | Out-Null
+            $unloadExitCode = $LASTEXITCODE
+
+            if ($unloadExitCode -ne 0) {
+                throw "Failed to unload registry hive HKU\Default after importing '$path' (exit code: $unloadExitCode)"
+            }
+        }
+    }
+}

Scripts/Helpers/ApplyRegistryRegFile.ps1

@@ -0,0 +1,223 @@
+function Get-NormalizedRegistryValueName {
+    param(
+        [AllowNull()]
+        $ValueName
+    )
+
+    if ([string]::IsNullOrEmpty([string]$ValueName)) {
+        return ''
+    }
+
+    return [string]$ValueName
+}
+
+function Convert-RegOperationToValueKind {
+    param(
+        [Parameter(Mandatory)]
+        $Operation
+    )
+
+    $valueName = if ([string]::IsNullOrEmpty([string]$Operation.ValueName)) { '' } else { [string]$Operation.ValueName }
+    $valueType = [string]$Operation.ValueType
+    $operationKeyPath = [string]$Operation.KeyPath
+
+    switch ($valueType) {
+        'DWord' {
+            $unsigned = [uint32]$Operation.ValueData
+            $value = [BitConverter]::ToInt32([BitConverter]::GetBytes($unsigned), 0)
+            return @{ Name = $valueName; Kind = [Microsoft.Win32.RegistryValueKind]::DWord; Value = $value }
+        }
+        'String' {
+            return @{ Name = $valueName; Kind = [Microsoft.Win32.RegistryValueKind]::String; Value = [string]$Operation.ValueData }
+        }
+        'Binary' {
+            return @{ Name = $valueName; Kind = [Microsoft.Win32.RegistryValueKind]::Binary; Value = [byte[]]$Operation.ValueData }
+        }
+        default {
+            throw "Unsupported value type '$valueType' while applying reg operation for '$operationKeyPath'"
+        }
+    }
+}
+
+function Remove-RegistrySubKeyTreeIfExists {
+    param(
+        [Parameter(Mandatory)]
+        [Microsoft.Win32.RegistryKey]$RootKey,
+        [Parameter(Mandatory)]
+        [string]$SubKeyPath
+    )
+
+    try {
+        $RootKey.DeleteSubKeyTree($SubKeyPath, $false)
+    }
+    catch [System.UnauthorizedAccessException], [System.Security.SecurityException] {
+        throw
+    }
+    catch {
+        # Best-effort cleanup only; missing keys are fine.
+    }
+}
+
+function Get-RegistryKeyForOperation {
+    param(
+        [Parameter(Mandatory)]
+        [string]$RegistryPath,
+        [switch]$CreateIfMissing,
+        [bool]$OpenKey = $true
+    )
+
+    $parts = Split-RegistryPath -path $RegistryPath
+    if (-not $parts) {
+        throw "Unsupported registry path: $RegistryPath"
+    }
+
+    $rootKey = Get-RegistryRootKey -hiveName $parts.Hive
+    if (-not $rootKey) {
+        throw "Unsupported registry hive '$($parts.Hive)' in path '$RegistryPath'"
+    }
+
+    $subKeyPath = $parts.SubKey
+    if ([string]::IsNullOrWhiteSpace($subKeyPath)) {
+        return [PSCustomObject]@{ RootKey = $rootKey; SubKeyPath = $null; Key = $rootKey }
+    }
+
+    if (-not $OpenKey) {
+        return [PSCustomObject]@{ RootKey = $rootKey; SubKeyPath = $subKeyPath; Key = $null }
+    }
+
+    $key = if ($CreateIfMissing) {
+        $rootKey.CreateSubKey($subKeyPath)
+    }
+    else {
+        $rootKey.OpenSubKey($subKeyPath, $true)
+    }
+
+    return [PSCustomObject]@{ RootKey = $rootKey; SubKeyPath = $subKeyPath; Key = $key }
+}
+
+function Invoke-RegistryDeleteValueOperation {
+    param(
+        [Parameter(Mandatory)]
+        $Operation,
+        [Parameter(Mandatory)]
+        $KeyInfo
+    )
+
+    if ($null -eq $KeyInfo.Key) {
+        $valueName = Get-NormalizedRegistryValueName -ValueName $Operation.ValueName
+        $displayValueName = if ([string]::IsNullOrEmpty($valueName)) { '(Default)' } else { $valueName }
+        Write-Verbose "Unable to find or open key '$($Operation.KeyPath)' and value '$displayValueName'"
+        return
+    }
+
+    try {
+        $valueName = Get-NormalizedRegistryValueName -ValueName $Operation.ValueName
+        $KeyInfo.Key.DeleteValue($valueName, $false)
+    }
+    finally {
+        $KeyInfo.Key.Close()
+    }
+}
+
+function Invoke-RegistrySetValueOperation {
+    param(
+        [Parameter(Mandatory)]
+        $Operation,
+        [Parameter(Mandatory)]
+        $KeyInfo
+    )
+
+    if ($null -eq $KeyInfo.Key) {
+        throw [System.UnauthorizedAccessException]::new("Unable to open or create registry key '$($Operation.KeyPath)'")
+    }
+
+    try {
+        $setArgs = Convert-RegOperationToValueKind -Operation $Operation
+        $KeyInfo.Key.SetValue($setArgs.Name, $setArgs.Value, $setArgs.Kind)
+    }
+    finally {
+        $KeyInfo.Key.Close()
+    }
+}
+
+function Write-RegistryOperationAccessDeniedWarning {
+    param(
+        [Parameter(Mandatory)]
+        $Operation,
+        [Parameter(Mandatory)]
+        [string]$ExceptionMessage
+    )
+
+    $keyPath = [string]$Operation.KeyPath
+    $operationType = [string]$Operation.OperationType
+
+    if ($operationType -eq 'SetValue' -or $operationType -eq 'DeleteValue') {
+        $valueName = Get-NormalizedRegistryValueName -ValueName $Operation.ValueName
+        $displayValueName = if ([string]::IsNullOrEmpty($valueName)) { '(Default)' } else { $valueName }
+        Write-Warning "Skipping operation '$operationType' on key '$keyPath' value '$displayValueName' due to access restrictions: $ExceptionMessage"
+        return
+    }
+
+    Write-Warning "Skipping operation '$operationType' on key '$keyPath' due to access restrictions: $ExceptionMessage"
+}
+
+function Invoke-RegistryOperation {
+    param(
+        [Parameter(Mandatory)]
+        $Operation,
+        [Parameter(Mandatory)]
+        [string]$RegFilePath
+    )
+
+    $operationType = [string]$Operation.OperationType
+    $isSetValueOperation = $operationType -eq 'SetValue'
+    $isDeleteKeyOperation = $operationType -eq 'DeleteKey'
+
+    $keyInfo = Get-RegistryKeyForOperation -RegistryPath $Operation.KeyPath -CreateIfMissing:$isSetValueOperation -OpenKey:(-not $isDeleteKeyOperation)
+
+    switch ($operationType) {
+        'DeleteKey' {
+            if ($null -ne $keyInfo.SubKeyPath) {
+                Remove-RegistrySubKeyTreeIfExists -RootKey $keyInfo.RootKey -SubKeyPath $keyInfo.SubKeyPath
+            }
+        }
+        'DeleteValue' {
+            Invoke-RegistryDeleteValueOperation -Operation $Operation -KeyInfo $keyInfo
+        }
+        'SetValue' {
+            Invoke-RegistrySetValueOperation -Operation $Operation -KeyInfo $keyInfo
+        }
+        default {
+            throw "Unsupported reg operation type '$($Operation.OperationType)' in '$RegFilePath'"
+        }
+    }
+}
+
+function Invoke-RegistryOperationsFromRegFile {
+    param(
+        [Parameter(Mandatory)]
+        [string]$RegFilePath
+    )
+
+    $accessDeniedCount = 0
+    $operations = @(Get-RegFileOperations -regFilePath $RegFilePath)
+    $totalOperations = $operations.Count
+
+    foreach ($operation in $operations) {
+        try {
+            Invoke-RegistryOperation -Operation $operation -RegFilePath $RegFilePath
+        }
+        catch [System.UnauthorizedAccessException], [System.Security.SecurityException] {
+            $accessDeniedCount++
+            Write-RegistryOperationAccessDeniedWarning -Operation $operation -ExceptionMessage $_.Exception.Message
+        }
+    }
+
+    if ($totalOperations -gt 0 -and $accessDeniedCount -eq $totalOperations) {
+        throw "Registry fallback import could not apply any operations in '$RegFilePath' because all $accessDeniedCount operation(s) were blocked by access restrictions."
+    }
+
+    if ($accessDeniedCount -gt 0) {
+        Write-Warning "Registry fallback import completed with $accessDeniedCount access-restricted operation(s) skipped in '$RegFilePath'."
+    }
+}

Win11Debloat.ps1

@@ -349,6 +349,7 @@ if (-not $script:WingetInstalled -and -not $Silent) {
 . "$PSScriptRoot/Scripts/Helpers/GetUserDirectory.ps1"
 . "$PSScriptRoot/Scripts/Helpers/GetUserName.ps1"
 . "$PSScriptRoot/Scripts/Helpers/RegistryPathHelpers.ps1"
+. "$PSScriptRoot/Scripts/Helpers/ApplyRegistryRegFile.ps1"
 . "$PSScriptRoot/Scripts/Helpers/TestIfUserIsLoggedIn.ps1"
 
 # Threading functions
@@ -401,7 +402,7 @@ else {
 }
 
 if ($script:Params.ContainsKey("Sysprep")) {
-    $defaultUserPath = GetUserDirectory -userName "Default"
+    GetUserDirectory -userName "Default" | Out-Null
 
     # Exit script if run in Sysprep mode on Windows 10
     if ($WinVersion -lt 22000) {
@@ -412,10 +413,10 @@ if ($script:Params.ContainsKey("Sysprep")) {
 
 # Ensure that target user exists, if User or AppRemovalTarget parameter was provided
 if ($script:Params.ContainsKey("User")) {
-    $userPath = GetUserDirectory -userName $script:Params.Item("User")
+    GetUserDirectory -userName $script:Params.Item("User") | Out-Null
 }
 if ($script:Params.ContainsKey("AppRemovalTarget")) {
-    $userPath = GetUserDirectory -userName $script:Params.Item("AppRemovalTarget")
+    GetUserDirectory -userName $script:Params.Item("AppRemovalTarget") | Out-Null
 }
 
 # Remove LastUsedSettings.json file if it exists and is empty