#Update $RegistryPath Value for your Environment. 
$RegistryPath = "HKLM:\SOFTWARE\WaaS"
$RegistryRollBack = "HKLM:\SYSTEM\Setup\Rollback"
$LogFile = "C:\ProgramData\WaaS\WaaS_RollbackRecovery.log" #Default if one isn't specified in Task Sequence
$LastOSUpgradeFrom = Get-ItemPropertyValue -Path "$RegistryPath" -Name LastOSUpgradeFrom
$LastOSUpgradeTo = Get-ItemPropertyValue -Path "$RegistryPath" -Name LastOSUpgradeTo
$RegistryPathFull = "$RegistryPath\$LastOSUpgradeTo"
$WaaSStage = Get-ItemPropertyValue "$RegistryPathFull" 'WaaS_Stage' -ErrorAction SilentlyContinue
$ScriptName = $MyInvocation.MyCommand.Name

function Test-RegistryValue
{
    <# 
    .SYNOPSIS 
    Tests if a registry value exists. 
     
    .DESCRIPTION 
    The usual ways for checking if a registry value exists don't handle when a value simply has an empty or null value. This function actually checks if a key has a value with a given name. 
     
    .EXAMPLE 
    Test-RegistryKeyValue -Path 'hklm:\Software\Carbon\Test' -Name 'Title' 
     
    Returns `True` if `hklm:\Software\Carbon\Test` contains a value named 'Title'. `False` otherwise. 
    #>
    [CmdletBinding()]
    param(
        [Parameter(Mandatory=$true)]
        [string]
        # The path to the registry key where the value should be set. Will be created if it doesn't exist.
        $Path,
        
        [Parameter(Mandatory=$true)]
        [string]
        # The name of the value being set.
        $Name
    )
    
    Set-StrictMode -Version 'Latest'

    #Use-CallerPreference -Cmdlet $PSCmdlet -Session $ExecutionContext.SessionState

    if( -not (Test-Path -Path $Path -PathType Container) )
    {
        return $false
    }
    
    $properties = Get-ItemProperty -Path $Path 
    if( -not $properties )
    {
        return $false
    }
    
    $member = Get-Member -InputObject $properties -Name $Name
    if( $member )
    {
        return $true
    }
    else
    {
        return $false
    }
}


if (((Test-RegistryValue -Path "$RegistryPath" -Name "CustomLogPreFix") -eq $true) -and ((Test-RegistryValue -Path "$RegistryPath" -Name "CustomLogPath") -eq $true))
    {
    $CustomLogPath = Get-ItemPropertyValue -Path $RegistryPath -Name "CustomLogPath" -ErrorAction SilentlyContinue
    $CustomLogPreFix = Get-ItemPropertyValue -Path $RegistryPath -Name "CustomLogPreFix" -ErrorAction SilentlyContinue
    $LogFile = "$CustomLogPath\$($CustomLogPreFix)_RollbackRecovery.log"
    }


if ((Test-RegistryValue -Path "$RegistryPath" -Name "LegalText_DefaultIsEnabled") -eq $true)
    {$LegalText_DefaultIsEnabled = Get-ItemPropertyValue -Path $RegistryPath -Name "LegalText_DefaultIsEnabled"}

if ((Test-RegistryValue -Path "$RegistryPath" -Name "LegalText_EnableMods") -eq $true)
    {
    $LegalText_EnableMods = Get-ItemPropertyValue -Path $RegistryPath -Name "LegalText_EnableMods"
    if ($LegalText_EnableMods -eq $true)
        {
        if ($LegalText_DefaultIsEnabled -eq $True)
            {
            $legalnoticecaption = Get-ItemPropertyValue -Path $RegistryPath -Name "LegalText_DefaultCaption" -ErrorAction SilentlyContinue
            $legalnoticetext = Get-ItemPropertyValue -Path $RegistryPath -Name "LegalText_DefaultText" -ErrorAction SilentlyContinue
            }
        }
    }

# Schedule Task Names
$DefaultSTN = "WaaS_SetLegalNoticeText-Default"
$OSUpgradeCompleteSTN = "WaaS_SetLegalNoticeText-OSUpgradeComplete"
$OSUninstallSuccessfulSTN = "WaaS_SetLegalNoticeText-OSUninstallSuccessful"
$OSUninstallSTN = "WaaS_SetLegalNoticeText-OSUninstall"
$OSUninstallDefaultSTN = "WaaS_SetLegalNoticeText-OSUninstallDefault"
$ForceLockScreenCleanup2HoursSTN = "WaaS_ForceLockScreenCleanup2Hours"
$LockScreenCleanUpSTN = "WaaS_LockScreenCleanUp"
$OSUninstallCleanUpSTN = "WaaS_OSUninstallCleanUp"
$Post_OSUninstall_ProcessesSTN = "WaaS_Post_OSUninstall_Processes"
$RollBackRecoverySTN = "WaaS_SetLegalNoticeText-RollBackRecovery"
$RollBackRecoveryDefaultSTN = "WaaS_SetLegalNoticeText-RollBackRecoveryDefault"
$RollBackRecovery_Process_DownlevelSTN = "WaaS_RollBackRecovery_Process_Downlevel"
$RollBackRecovery_Process_Phase2STN = "WaaS_RollBackRecovery_Process_Phase2"
$RollBackRecovery_CleanUpSTN = "WaaS_RollBackRecovery_CleanUp"


#region: CMTraceLog Function formats logging in CMTrace style
        function CMTraceLog {
         [CmdletBinding()]
    Param (
		    [Parameter(Mandatory=$false)]
		    $Message,
 
		    [Parameter(Mandatory=$false)]
		    $ErrorMessage,
 
		    [Parameter(Mandatory=$false)]
		    $Component = $env:computername,
 
		    [Parameter(Mandatory=$false)]
		    [int]$Type,
		
		    [Parameter(Mandatory=$true)]
		    $LogFile
	    )
    <#
    Type: 1 = Normal, 2 = Warning (yellow), 3 = Error (red)
    #>
	    $Time = Get-Date -Format "HH:mm:ss.ffffff"
	    $Date = Get-Date -Format "MM-dd-yyyy"
 
	    if ($ErrorMessage -ne $null) {$Type = 3}
	    if ($Component -eq $null) {$Component = " "}
	    if ($Type -eq $null) {$Type = 1}
 
	    $LogMessage = "<![LOG[$Message $ErrorMessage" + "]LOG]!><time=`"$Time`" date=`"$Date`" component=`"$Component`" context=`"`" type=`"$Type`" thread=`"`" file=`"`">"
	    $LogMessage | Out-File -Append -Encoding UTF8 -FilePath $LogFile
    }

if (get-ScheduledTask -TaskName $RollBackRecovery_Process_DownlevelSTN -ErrorAction SilentlyContinue)
    {CMTraceLog -Message  " $RollBackRecovery_Process_DownlevelSTN Exist, Contining Process" -Type 2 -LogFile $LogFile}
else 
    {
    if ($WaaSStage -eq "Deployment_Error")
        {       
        CMTraceLog -Message  " $RollBackRecovery_Process_DownlevelSTN Doesn't Exist" -Type 1 -LogFile $LogFile
        Set-ItemProperty -Path $RegistryPathFull -Name "WaaS_Stage" -Value "Deployment_RollBack" -Force
        }
    }

if (get-ScheduledTask -TaskName $RollBackRecovery_Process_Phase2STN -ErrorAction SilentlyContinue)
    {CMTraceLog -Message  " $RollBackRecovery_Process_Phase2STN Exist, Contining Process" -Type 2 -LogFile $LogFile}
else 
    {
    if ($WaaSStage -eq "Deployment_Error")
        {       
        CMTraceLog -Message  " $RollBackRecovery_Process_Phase2STN Doesn't Exist" -Type 1 -LogFile $LogFile
        Set-ItemProperty -Path $RegistryPathFull -Name "WaaS_Stage" -Value "Deployment_RollBack" -Force
        }
    }

$WaaSStage = Get-ItemPropertyValue "$RegistryPathFull" 'WaaS_Stage' -ErrorAction SilentlyContinue


if ($WaaSStage -eq "Deployment_RollBack") #OSUninstall is set using a OSUninstall TS post Upgrade, to allow being reverted, so this will never run if this key isn't set by the OSUninstall TS
    {
    CMTraceLog -Message  "---Starting $ScriptName Script---" -Type 1 -LogFile $LogFile -Component LegalText-Default
    $CurrentLegalCaption = Get-ItemPropertyValue 'HKLM:\Software\Microsoft\Windows\CurrentVersion\Policies\System' 'legalnoticecaption' -ErrorAction SilentlyContinue
    CMTraceLog -Message  "LegalCaption Data Currently: $CurrentLegalCaption" -Type 1 -LogFile $LogFile -Component LegalText-Default
    CMTraceLog -Message  "Updated Values: legalnoticecaption & legalnoticetext" -Type 1 -LogFile $LogFile -Component LegalText-Default
    if ($LegalText_DefaultIsEnabled -eq $True)
        {
        Set-ItemProperty -Path 'HKLM:\Software\Microsoft\Windows\CurrentVersion\Policies\System' -Name legalnoticecaption -Value $legalnoticecaption
        Set-ItemProperty -Path 'HKLM:\Software\Microsoft\Windows\CurrentVersion\Policies\System' -Name legalnoticetext -Value $legalnoticetext
        $CurrentLegalCaption = Get-ItemPropertyValue 'HKLM:\Software\Microsoft\Windows\CurrentVersion\Policies\System' 'legalnoticecaption' -ErrorAction SilentlyContinue
        CMTraceLog -Message  "LegalCaption Data Currently: $CurrentLegalCaption" -Type 1 -LogFile $LogFile -Component LegalText-Default
        }
    else
        {
        Remove-ItemProperty -Path 'HKLM:\Software\Microsoft\Windows\CurrentVersion\Policies\System' -Name legalnoticecaption
        Remove-ItemProperty -Path 'HKLM:\Software\Microsoft\Windows\CurrentVersion\Policies\System' -Name legalnoticetext
        CMTraceLog -Message  "LegalCaption Was Deleted" -Type 1 -LogFile $LogFile -Component LegalText-Default
        }
    Unregister-ScheduledTask -TaskName $OSUninstallSTN -Confirm:$false -ErrorAction SilentlyContinue
    Unregister-ScheduledTask -TaskName $OSUninstallDefaultSTN -Confirm:$false -ErrorAction SilentlyContinue
    Unregister-ScheduledTask -TaskName $Post_OSUninstall_ProcessesSTN -Confirm:$false -ErrorAction SilentlyContinue
    Unregister-ScheduledTask -TaskName $OSUninstallCleanUpSTN -Confirm:$false -ErrorAction SilentlyContinue
    CMTraceLog -Message  "Removed Scheduled Tasks for OSUninstall" -Type 1 -LogFile $LogFile -Component LegalText-Default
    Unregister-ScheduledTask -TaskName $RollBackRecoverySTN -Confirm:$false -ErrorAction SilentlyContinue
    Unregister-ScheduledTask -TaskName $RollBackRecoveryDefaultSTN -Confirm:$false -ErrorAction SilentlyContinue
    Unregister-ScheduledTask -TaskName $RollBackRecovery_Process_DownlevelSTN -Confirm:$false -ErrorAction SilentlyContinue
    Unregister-ScheduledTask -TaskName $RollBackRecovery_Process_Phase2STN -Confirm:$false -ErrorAction SilentlyContinue
    CMTraceLog -Message  "Removed Scheduled Tasks for RollBackRecovery" -Type 1 -LogFile $LogFile -Component LegalText-Default
                    
    Start-ScheduledTask -TaskName $RollBackRecovery_CleanUpSTN
    CMTraceLog -Message  "Triggering $RollBackRecovery_CleanUpSTN Scheduled Task" -Type 1 -LogFile $LogFile
    
    CMTraceLog -Message  "---Exiting $ScriptName Script---" -Type 1 -LogFile $LogFile -Component LegalText-Default
    }

if ($WaaSStage -eq "OSUninstallComplete")
    {
    CMTraceLog -Message  "---Starting $ScriptName Script---" -Type 1 -LogFile $LogFile -Component LegalText-Default
    Unregister-ScheduledTask -TaskName $OSUninstallSTN -Confirm:$false -ErrorAction SilentlyContinue
    Unregister-ScheduledTask -TaskName $OSUninstallDefaultSTN -Confirm:$false -ErrorAction SilentlyContinue
    Unregister-ScheduledTask -TaskName $Post_OSUninstall_ProcessesSTN -Confirm:$false -ErrorAction SilentlyContinue
    Unregister-ScheduledTask -TaskName $OSUninstallCleanUpSTN -Confirm:$false -ErrorAction SilentlyContinue
    CMTraceLog -Message  "Removed Scheduled Tasks for OSUninstall" -Type 1 -LogFile $LogFile -Component LegalText-Default
    Unregister-ScheduledTask -TaskName $RollBackRecoverySTN -Confirm:$false -ErrorAction SilentlyContinue
    Unregister-ScheduledTask -TaskName $RollBackRecoveryDefaultSTN -Confirm:$false -ErrorAction SilentlyContinue
    Unregister-ScheduledTask -TaskName $RollBackRecovery_Process_DownlevelSTN -Confirm:$false -ErrorAction SilentlyContinue
    Unregister-ScheduledTask -TaskName $RollBackRecovery_Process_Phase2STN -Confirm:$false -ErrorAction SilentlyContinue
    CMTraceLog -Message  "Removed Scheduled Tasks for RollBackRecovery" -Type 1 -LogFile $LogFile -Component LegalText-Default
    CMTraceLog -Message  "---Exiting $ScriptName Script---" -Type 1 -LogFile $LogFile -Component LegalText-Default
    }