#Requires -RunAsAdministrator

<#
.SYNOPSIS
    Enables vTPM and state/migration traffic encryption on all VMs of the local Hyper-V host.

.DESCRIPTION
    For every local VM:
      - Sets a local key protector (required for standalone hosts without a Host Guardian Service),
        unless the VM already has one.
      - Enables the virtual Trusted Platform Module.
      - Enables encryption of saved state and live migration traffic.

    Security settings can only be changed while a VM is off, so running VMs are shut down
    gracefully first and started again after the changes are applied. VMs that are already
    fully configured are skipped without a shutdown. Generation 1 VMs do not support vTPM
    and are skipped with a warning.

.PARAMETER ShutdownTimeoutSeconds
    How long to wait for a graceful guest shutdown before powering the VM off.
#>

[CmdletBinding()]
param(
    [int]$ShutdownTimeoutSeconds = 300
)

$ErrorActionPreference = 'Stop'

function Stop-VMWithTimeout {
    param($VM, [int]$TimeoutSeconds)

    # -Force suppresses the unsaved-data confirmation prompt; the shutdown itself
    #   is still graceful (performed by the guest via integration services).
    $job = Stop-VM -VM $VM -Force -AsJob
    Wait-Job -Job $job -Timeout $TimeoutSeconds | Out-Null
    Remove-Job -Job $job -Force

    if ((Get-VM -Id $VM.Id).State -ne 'Off') {
        Write-Warning "[$($VM.Name)] Graceful shutdown did not finish within $TimeoutSeconds seconds. Powering off."
        Stop-VM -VM $VM -TurnOff -Force
    }
}

$failed = @()

foreach ($vm in Get-VM) {
    Write-Host "Processing '$($vm.Name)' (state: $($vm.State))..."

    if ($vm.Generation -ne 2) {
        Write-Warning "[$($vm.Name)] Generation $($vm.Generation) VM does not support vTPM. Skipping."
        continue
    }

    $security = $vm | Get-VMSecurity
    if ($security.TpmEnabled -and $security.EncryptStateAndVmMigrationTraffic) {
        Write-Host "[$($vm.Name)] Already configured. Skipping."
        continue
    }

    if ($vm.State -eq 'Saved') {
        # Changing the key protector invalidates a saved state, so deleting it here is no
        #   worse -- but it equals a hard power-off, so leave the decision to the operator.
        Write-Warning "[$($vm.Name)] VM has a saved state. Start and shut it down, then re-run this script. Skipping."
        continue
    }

    if ($vm.State -notin 'Off', 'Running', 'Paused') {
        Write-Warning "[$($vm.Name)] Unexpected state '$($vm.State)'. Skipping."
        continue
    }

    $wasRunning = $vm.State -in 'Running', 'Paused'

    try {
        if ($vm.State -eq 'Paused') {
            Write-Host "[$($vm.Name)] Resuming and shutting down..."
            Resume-VM -VM $vm
            Stop-VMWithTimeout -VM $vm -TimeoutSeconds $ShutdownTimeoutSeconds
        } elseif ($vm.State -eq 'Running') {
            Write-Host "[$($vm.Name)] Shutting down..."
            Stop-VMWithTimeout -VM $vm -TimeoutSeconds $ShutdownTimeoutSeconds
        }

        # Enable Trusted Platform Module and encrypt state/migration traffic.
        #   A local key protector is required for standalone hosts without a Host Guardian Service.
        #   Don't replace an existing key protector: vTPM state is encrypted with it, so a new
        #   protector would make any existing vTPM data unrecoverable. An unset protector is a
        #   4-byte (or empty) placeholder.
        if ((Get-VMKeyProtector -VM $vm).Length -le 4) {
            $vm | Set-VMKeyProtector -NewLocalKeyProtector
        }
        $vm | Enable-VMTPM
        $vm | Set-VMSecurity -EncryptStateAndVmMigrationTraffic $true

        if ($wasRunning) {
            Write-Host "[$($vm.Name)] Starting..."
            Start-VM -VM $vm
        }

        Write-Host "[$($vm.Name)] Done."
    } catch {
        Write-Warning "[$($vm.Name)] Failed: $_"
        $failed += $vm.Name

        # Don't leave a previously running VM down because of a failed settings change.
        if ($wasRunning -and (Get-VM -Id $vm.Id).State -eq 'Off') {
            try {
                Write-Host "[$($vm.Name)] Starting back up after failure..."
                Start-VM -VM $vm
            } catch {
                Write-Warning "[$($vm.Name)] Could not start VM back up: $_"
            }
        }
    }
}

if ($failed) {
    Write-Warning "Failed VMs: $($failed -join ', ')"
    exit 1
}

Write-Host 'All done!'