Refactor registry operation functions

Scripts/Helpers/ApplyRegistryRegFile.ps1

@@ -1,3 +1,16 @@
+function Get-NormalizedRegistryValueName {
+    param(
+        [AllowNull()]
+        $ValueName
+    )
+
+    if ([string]::IsNullOrEmpty([string]$ValueName)) {
+        return ''
+    }
+
+    return [string]$ValueName
+}
+
 function Convert-RegOperationToValueKind {
     param(
         [Parameter(Mandatory)]
@@ -6,6 +19,7 @@ function Convert-RegOperationToValueKind {
 
     $valueName = if ([string]::IsNullOrEmpty([string]$Operation.ValueName)) { '' } else { [string]$Operation.ValueName }
     $valueType = [string]$Operation.ValueType
+    $operationKeyPath = [string]$Operation.KeyPath
 
     switch ($valueType) {
         'DWord' {
@@ -49,7 +63,7 @@ function Convert-RegOperationToValueKind {
         'HexB' {
             $qwordBytes = [byte[]]$Operation.ValueData
             if ($qwordBytes.Count -gt 8) {
-                throw "Unsupported hex value type '$valueType' with invalid byte count '$($qwordBytes.Count)' while applying reg operation for '$($Operation.KeyPath)'."
+                throw "Unsupported hex value type '$valueType' with invalid byte count '$($qwordBytes.Count)' while applying reg operation for '$operationKeyPath'."
             }
 
             if ($qwordBytes.Count -lt 8) {
@@ -64,10 +78,10 @@ function Convert-RegOperationToValueKind {
         }
         default {
             if ($valueType -like 'Hex*') {
-                throw "Unsupported hex value type '$valueType' while applying reg operation for '$($Operation.KeyPath)'."
+                throw "Unsupported hex value type '$valueType' while applying reg operation for '$operationKeyPath'."
             }
 
-            throw "Unsupported value type '$valueType' while applying reg operation for '$($Operation.KeyPath)'"
+            throw "Unsupported value type '$valueType' while applying reg operation for '$operationKeyPath'"
         }
     }
 }
@@ -128,6 +142,101 @@ function Get-RegistryKeyForOperation {
     return [PSCustomObject]@{ RootKey = $rootKey; SubKeyPath = $subKeyPath; Key = $key }
 }
 
+function Invoke-RegistryDeleteValueOperation {
+    param(
+        [Parameter(Mandatory)]
+        $Operation,
+        [Parameter(Mandatory)]
+        $KeyInfo
+    )
+
+    if ($null -eq $KeyInfo.Key) {
+        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)]
@@ -140,58 +249,11 @@ function Invoke-RegistryOperationsFromRegFile {
 
     foreach ($operation in $operations) {
         try {
-            $keyInfo = Get-RegistryKeyForOperation -RegistryPath $operation.KeyPath -CreateIfMissing:($operation.OperationType -eq 'SetValue') -OpenKey:($operation.OperationType -ne 'DeleteKey')
+            Invoke-RegistryOperation -Operation $operation -RegFilePath $RegFilePath
-
-            switch ($operation.OperationType) {
-                'DeleteKey' {
-                    if ($null -ne $keyInfo.SubKeyPath) {
-                        Remove-RegistrySubKeyTreeIfExists -RootKey $keyInfo.RootKey -SubKeyPath $keyInfo.SubKeyPath
-                    }
-                }
-                'DeleteValue' {
-                    if ($null -ne $keyInfo.Key) {
-                        try {
-                            $valueName = if ([string]::IsNullOrEmpty([string]$operation.ValueName)) { '' } else { [string]$operation.ValueName }
-                            $keyInfo.Key.DeleteValue($valueName, $false)
-                        }
-                        finally {
-                            $keyInfo.Key.Close()
-                        }
-                    }
-                }
-                'SetValue' {
-                    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()
-                    }
-                }
-                default {
-                    throw "Unsupported reg operation type '$($operation.OperationType)' in '$RegFilePath'"
-                }
-            }
         }
         catch [System.UnauthorizedAccessException], [System.Security.SecurityException] {
             $accessDeniedCount++
-            $keyPath = [string]$operation.KeyPath
+            Write-RegistryOperationAccessDeniedWarning -Operation $operation -ExceptionMessage $_.Exception.Message
-            $opType = [string]$operation.OperationType
-            $valueName = [string]$operation.ValueName
-
-            if ($opType -eq 'SetValue' -or $opType -eq 'DeleteValue') {
-                $display = if ([string]::IsNullOrEmpty($valueName)) { '(Default)' } else { $valueName }
-                Write-Warning "Skipping operation '$opType' on key '$keyPath' value '$display' due to access restrictions: $($_.Exception.Message)"
-            }
-            else {
-                Write-Warning "Skipping operation '$opType' on key '$keyPath' due to access restrictions: $($_.Exception.Message)"
-            }
-
-            continue
         }
     }
 
@@ -218,8 +280,10 @@ function Invoke-RegistryImportViaPowerShell {
         Invoke-RegistryOperationsFromRegFile -RegFilePath $targetRegFilePath
     }
 
+    $hiveAlreadyLoaded = $HiveAlreadyLoaded.IsPresent
+
     if ($UseOfflineHive) {
-        if ((Get-Command -Name Invoke-WithLoadedBackupHive -ErrorAction SilentlyContinue) -and -not $HiveAlreadyLoaded.IsPresent) {
+        if ((Get-Command -Name Invoke-WithLoadedBackupHive -ErrorAction SilentlyContinue) -and -not $hiveAlreadyLoaded) {
             return Invoke-WithLoadedBackupHive -ScriptBlock $applyScript -ArgumentObject $RegFilePath
         }
 
@@ -227,7 +291,7 @@ function Invoke-RegistryImportViaPowerShell {
             throw "Offline hive path was not provided for fallback import of '$RegFilePath'"
         }
 
-        if (-not $HiveAlreadyLoaded.IsPresent) {
+        if (-not $hiveAlreadyLoaded) {
             $global:LASTEXITCODE = 0
             reg load "HKU\Default" $OfflineHiveDatPath | Out-Null
             $loadExitCode = $LASTEXITCODE