param(
    [ValidateSet('Apply','Query','Restore')][string]$Action = 'Query',
    [string]$BackupDirectory = 'C:\ProgramData\VDI\power-policy'
)
# Run elevated inside the Windows guest. No reboot, logout or Explorer restart.
$ErrorActionPreference = 'Stop'
$user = (Get-CimInstance Win32_ComputerSystem).UserName
if (-not $user) { throw 'No interactive desktop user' }
$sid = (New-Object Security.Principal.NTAccount($user)).Translate([Security.Principal.SecurityIdentifier]).Value
if (-not (Test-Path "Registry::HKEY_USERS\$sid")) { throw 'User hive is not loaded' }
$path = "Registry::HKEY_USERS\$sid\Software\Microsoft\Windows\CurrentVersion\Policies\Explorer"
$name = 'NoClose'
$backup = Join-Path $BackupDirectory 'original.clixml'
if ($Action -eq 'Apply') {
    New-Item -ItemType Directory -Force $BackupDirectory | Out-Null
    if (-not (Test-Path $backup)) {
        $key = Get-Item -LiteralPath $path -ErrorAction SilentlyContinue
        $exists = $key -and ($key.GetValueNames() -contains $name)
        [pscustomobject]@{Sid=$sid; Path=$path; Exists=[bool]$exists;
            Value=$(if($exists){$key.GetValue($name)}else{$null});
            Kind=$(if($exists){$key.GetValueKind($name).ToString()}else{$null})} | Export-Clixml $backup
    }
    if ((Import-Clixml $backup).Sid -ne $sid) { throw 'Backup belongs to another user' }
    New-Item -Path $path -Force | Out-Null
    New-ItemProperty -LiteralPath $path -Name $name -PropertyType DWord -Value 1 -Force | Out-Null
} elseif ($Action -eq 'Restore') {
    $saved = Import-Clixml $backup
    if ($saved.Sid -ne $sid -or $saved.Path -ne $path) { throw 'Backup belongs to another user' }
    if ($saved.Exists) {
        New-ItemProperty -LiteralPath $path -Name $name -PropertyType $saved.Kind -Value $saved.Value -Force | Out-Null
    } else {
        Remove-ItemProperty -LiteralPath $path -Name $name -ErrorAction SilentlyContinue
    }
}
$key = Get-Item -LiteralPath $path -ErrorAction SilentlyContinue
$value = if ($key) { $key.GetValue($name) } else { $null }
if ($Action -eq 'Apply' -and $value -ne 1) { throw 'NoClose readback failed' }
# Notify and query the shell in the actual desktop user's session. A SYSTEM
# guest-agent process sees a different HKCU and cannot verify this restriction.
$runtime = Join-Path $BackupDirectory 'runtime'
New-Item -ItemType Directory -Force $runtime | Out-Null
$acl = Get-Acl $runtime
$acl.SetAccessRule((New-Object Security.AccessControl.FileSystemAccessRule($user,'Modify','ContainerInherit,ObjectInherit','None','Allow')))
Set-Acl $runtime $acl
$resultPath = Join-Path $runtime 'restriction.txt'
Remove-Item $resultPath -ErrorAction SilentlyContinue
$probe = @'
$ErrorActionPreference='Stop'
Add-Type -TypeDefinition @"
using System;
using System.Runtime.InteropServices;
public static class VDIPowerShellPolicy {
    [DllImport("shell32.dll")] public static extern uint SHRestricted(uint restriction);
    [DllImport("user32.dll", CharSet=CharSet.Unicode)]
    public static extern IntPtr SendMessageTimeout(IntPtr window, uint message, UIntPtr wparam, string lparam, uint flags, uint timeout, out UIntPtr result);
}
"@
$sent=[UIntPtr]::Zero
[void][VDIPowerShellPolicy]::SendMessageTimeout([IntPtr]0xffff,0x1a,[UIntPtr]::Zero,'Policy',2,2000,[ref]$sent)
[VDIPowerShellPolicy]::SHRestricted(2) | Set-Content -Encoding ascii '__RESULT__'
'@
$probe = $probe.Replace('__RESULT__', $resultPath.Replace("'", "''"))
$encoded = [Convert]::ToBase64String([Text.Encoding]::Unicode.GetBytes($probe))
$task = 'VDIPowerPolicyProbe'
if (Get-ScheduledTask -TaskName $task -ErrorAction SilentlyContinue) { throw 'Probe task already exists' }
$taskAction = New-ScheduledTaskAction -Execute powershell.exe -Argument "-NoProfile -WindowStyle Hidden -EncodedCommand $encoded"
$principal = New-ScheduledTaskPrincipal -UserId $user -LogonType Interactive -RunLevel Limited
$settings = New-ScheduledTaskSettingsSet -ExecutionTimeLimit (New-TimeSpan -Minutes 1)
Register-ScheduledTask -TaskName $task -Action $taskAction -Principal $principal -Settings $settings | Out-Null
try {
    Start-ScheduledTask -TaskName $task
    for ($i=0; $i -lt 20; $i++) {
        Start-Sleep -Seconds 1
        if (Test-Path $resultPath) { break }
    }
    if (-not (Test-Path $resultPath)) { throw 'Interactive policy probe timed out' }
    $effective = [int](Get-Content $resultPath -Raw)
    if ($Action -eq 'Apply' -and $effective -ne 1) { throw 'Interactive shell restriction was not enabled' }
} finally {
    Stop-ScheduledTask -TaskName $task -ErrorAction SilentlyContinue
    Unregister-ScheduledTask -TaskName $task -Confirm:$false
}
[pscustomobject]@{Action=$Action; User=$user; Sid=$sid; NoClose=$value; Backup=$backup;
    EffectiveShellNoClose=$effective;
    LastBoot=(Get-CimInstance Win32_OperatingSystem).LastBootUpTime.ToString('o');
    Scope='Desktop power menus only; privileged administration remains available'} | ConvertTo-Json -Compress