Add support for undo features in registry backup capture and validation processes

Scripts/Features/BackupRegistrySnapshotCapture.ps1

@@ -1,10 +1,12 @@
 function Get-RegistryBackupCapturePlans {
     param(
         [object[]]$SelectedRegistryFeatures = @(),
+        [object[]]$UndoRegistryFeatures = @(),
         [switch]$UseSysprepRegFiles
     )
 
     $planMap = @{}
+
     foreach ($feature in $SelectedRegistryFeatures) {
         $regFilePath = Get-RegistryFilePathForFeature -RegistryKey $feature.RegistryKey -UseSysprepRegFiles:$UseSysprepRegFiles
         if (-not (Test-Path $regFilePath)) {
@@ -13,34 +15,30 @@ function Get-RegistryBackupCapturePlans {
 
         foreach ($operation in @(Get-RegFileOperations -regFilePath $regFilePath)) {
             if (-not $operation.KeyPath) { continue }
+            Add-RegistryPlanOperation -PlanMap $planMap -Operation $operation
+        }
+    }
 
-            $mapKey = $operation.KeyPath.ToLowerInvariant()
-            if (-not $planMap.ContainsKey($mapKey)) {
-                $planMap[$mapKey] = [PSCustomObject]@{
-                    Path = $operation.KeyPath
-                    IncludeSubKeys = $false
-                    CaptureAllValues = $false
-                    ValueNames = New-Object 'System.Collections.Generic.HashSet[string]' ([System.StringComparer]::OrdinalIgnoreCase)
-                }
+    foreach ($feature in $UndoRegistryFeatures) {
+        $regFilePath = Resolve-RegistryBackupUndoFilePath -Feature $feature
+        if ([string]::IsNullOrWhiteSpace($regFilePath)) {
+            continue
+        }
+
+        if (-not (Test-Path $regFilePath)) {
+            $undoKeyDescription = if (-not [string]::IsNullOrWhiteSpace([string]$feature.RegistryUndoKey)) {
+                [string]$feature.RegistryUndoKey
+            }
+            else {
+                [string]$feature.RegistryKey
             }
 
-            $plan = $planMap[$mapKey]
-            switch ($operation.OperationType) {
-                'DeleteKey' {
-                    $plan.IncludeSubKeys = $true
-                    $plan.CaptureAllValues = $true
-                }
-                'SetValue' {
-                    if (-not $plan.CaptureAllValues) {
-                        $null = $plan.ValueNames.Add([string]$operation.ValueName)
-                    }
-                }
-                'DeleteValue' {
-                    if (-not $plan.CaptureAllValues) {
-                        $null = $plan.ValueNames.Add([string]$operation.ValueName)
-                    }
-                }
-            }
+            throw "Unable to find registry undo file for backup: $undoKeyDescription ($regFilePath)"
+        }
+
+        foreach ($operation in @(Get-RegFileOperations -regFilePath $regFilePath)) {
+            if (-not $operation.KeyPath) { continue }
+            Add-RegistryPlanOperation -PlanMap $planMap -Operation $operation
         }
     }
 
@@ -56,6 +54,65 @@ function Get-RegistryBackupCapturePlans {
     )
 }
 
+function Add-RegistryPlanOperation {
+    param(
+        [hashtable]$PlanMap,
+        [PSCustomObject]$Operation
+    )
+
+    $mapKey = $Operation.KeyPath.ToLowerInvariant()
+    if (-not $PlanMap.ContainsKey($mapKey)) {
+        $PlanMap[$mapKey] = [PSCustomObject]@{
+            Path = $Operation.KeyPath
+            IncludeSubKeys = $false
+            CaptureAllValues = $false
+            ValueNames = New-Object 'System.Collections.Generic.HashSet[string]' ([System.StringComparer]::OrdinalIgnoreCase)
+        }
+    }
+
+    $plan = $PlanMap[$mapKey]
+    switch ($Operation.OperationType) {
+        'DeleteKey' {
+            $plan.IncludeSubKeys = $true
+            $plan.CaptureAllValues = $true
+        }
+        'SetValue' {
+            if (-not $plan.CaptureAllValues) {
+                $null = $plan.ValueNames.Add([string]$Operation.ValueName)
+            }
+        }
+        'DeleteValue' {
+            if (-not $plan.CaptureAllValues) {
+                $null = $plan.ValueNames.Add([string]$Operation.ValueName)
+            }
+        }
+    }
+}
+
+function Resolve-RegistryBackupUndoFilePath {
+    param(
+        [Parameter(Mandatory)]
+        $Feature
+    )
+
+    $undoRegistryKey = [string]$Feature.RegistryUndoKey
+    if (-not [string]::IsNullOrWhiteSpace($undoRegistryKey)) {
+        $resolvedUndoPath = Resolve-UndoRegFilePath -FileName $undoRegistryKey
+        return Join-Path $script:RegfilesPath $resolvedUndoPath
+    }
+
+    $resolvedRegistryKey = [string]$Feature.RegistryKey
+    if ([string]::IsNullOrWhiteSpace($resolvedRegistryKey)) {
+        return $null
+    }
+
+    if ([System.IO.Path]::IsPathRooted($resolvedRegistryKey)) {
+        return $resolvedRegistryKey
+    }
+
+    return Join-Path $script:RegfilesPath $resolvedRegistryKey
+}
+
 function Get-RegistrySnapshotsForBackup {
     param(
         [object[]]$CapturePlans = @()

Scripts/Features/BackupRegistryState.ps1

@@ -84,9 +84,11 @@ function Get-RegistryBackupPayload {
         }
     }
 
-    $allCapturableFeatures = @($SelectedFeatures) + @($UndoFeatures)
-    $selectedRegistryFeatures = @(Get-RegistryBackedFeatures -Features $allCapturableFeatures)
-    $capturePlans = @(Get-RegistryBackupCapturePlans -SelectedRegistryFeatures $selectedRegistryFeatures)
+    $selectedRegistryFeatures = @(Get-RegistryBackedFeatures -Features $SelectedFeatures)
+    $undoRegistryFeatures = @($UndoFeatures | Where-Object { 
+        -not [string]::IsNullOrWhiteSpace([string]$_.RegistryUndoKey) -or -not [string]::IsNullOrWhiteSpace([string]$_.RegistryKey)
+    })
+    $capturePlans = @(Get-RegistryBackupCapturePlans -SelectedRegistryFeatures $selectedRegistryFeatures -UndoRegistryFeatures $undoRegistryFeatures)
     $registryKeys = @(Get-RegistrySnapshotsForBackup -CapturePlans $capturePlans)
 
     $backupPayload = @{

Scripts/Features/RegistryBackupValidation.ps1

@@ -133,6 +133,9 @@ function Test-RegistryBackupMatchesSelectedFeatures {
         [AllowEmptyCollection()]
         [string[]]$SelectedFeatureIds,
         [Parameter(Mandatory)]
+        [AllowEmptyCollection()]
+        [string[]]$SelectedUndoFeatureIds,
+        [Parameter(Mandatory)]
         [string]$Target,
         [Parameter(Mandatory)]
         [AllowEmptyCollection()]
@@ -146,12 +149,13 @@ function Test-RegistryBackupMatchesSelectedFeatures {
         return $errors.ToArray()
     }
 
-    $selectedRegistryFeatures = @(Get-SelectedRegistryFeaturesForBackupValidation -SelectedFeatureIds @($SelectedFeatureIds) -Errors $errors)
+    $selectedRegistryFeatures = @(Get-SelectedRegistryFeaturesForBackupValidation -SelectedFeatureIds @($SelectedFeatureIds) -IsUndoFeature:$false -Errors $errors)
+    $undoRegistryFeatures = @(Get-SelectedRegistryFeaturesForBackupValidation -SelectedFeatureIds @($SelectedUndoFeatureIds) -IsUndoFeature:$true -Errors $errors)
     $useSysprepRegFiles = ($Target -eq 'DefaultUserProfile') -or ($Target -like 'User:*')
 
     $capturePlans = @()
-    if ($errors.Count -eq 0 -and $selectedRegistryFeatures.Count -gt 0) {
-        $capturePlans = @(Get-RegistryBackupCapturePlans -SelectedRegistryFeatures @($selectedRegistryFeatures) -UseSysprepRegFiles:$useSysprepRegFiles)
+    if ($errors.Count -eq 0 -and ($selectedRegistryFeatures.Count -gt 0 -or $undoRegistryFeatures.Count -gt 0)) {
+        $capturePlans = @(Get-RegistryBackupCapturePlans -SelectedRegistryFeatures @($selectedRegistryFeatures) -UndoRegistryFeatures @($undoRegistryFeatures) -UseSysprepRegFiles:$useSysprepRegFiles)
     }
 
     $planMap = New-RegistryBackupAllowListPlanMap -CapturePlans @($capturePlans)
@@ -173,6 +177,8 @@ function Get-SelectedRegistryFeaturesForBackupValidation {
         [AllowEmptyCollection()]
         [string[]]$SelectedFeatureIds,
         [Parameter(Mandatory)]
+        [bool]$IsUndoFeature,
+        [Parameter(Mandatory)]
         [AllowEmptyCollection()]
         $Errors
     )
@@ -189,7 +195,26 @@ function Get-SelectedRegistryFeaturesForBackupValidation {
         }
 
         $feature = $script:Features[$featureId]
-        if ($feature -and -not [string]::IsNullOrWhiteSpace([string]$feature.RegistryKey)) {
+        if (-not $feature) {
+            continue
+        }
+
+        # For undo features, check RegistryUndoKey if present (real features)
+        # Otherwise check RegistryKey (for synthetic features from backup capture)
+        $registryKeyToUse = if ($IsUndoFeature) {
+            $key = [string]$feature.RegistryUndoKey
+            if (-not [string]::IsNullOrWhiteSpace($key)) {
+                $key
+            }
+            else {
+                [string]$feature.RegistryKey
+            }
+        }
+        else {
+            [string]$feature.RegistryKey
+        }
+
+        if (-not [string]::IsNullOrWhiteSpace($registryKeyToUse)) {
             $selectedRegistryFeatures.Add($feature)
         }
     }

Scripts/Features/RestoreRegistryBackup.ps1

@@ -97,7 +97,7 @@ function Normalize-RegistryBackup {
     if ($allSelectedFeatures.Count -eq 0) {
         $errors.Add('Backup must contain at least one feature ID in SelectedFeatures or SelectedUndoFeatures.')
     }
-    $allowListValidationErrors = @(Test-RegistryBackupMatchesSelectedFeatures -SelectedFeatureIds @($allSelectedFeatures) -Target $normalizedTarget -RegistryKeys @($normalizedKeys))
+    $allowListValidationErrors = @(Test-RegistryBackupMatchesSelectedFeatures -SelectedFeatureIds @($selectedFeatures) -SelectedUndoFeatureIds @($selectedUndoFeatures) -Target $normalizedTarget -RegistryKeys @($normalizedKeys))
     foreach ($allowListValidationError in $allowListValidationErrors) {
         $errors.Add([string]$allowListValidationError)
     }