Move Sessions Hosts to New Host Pools Based on VM Location.

# Import the Azure module if needed
#Import-Module Az.Accounts
#Import-Module Az.DesktopVirtualization
#Import-Module Az.AVD
#Install-Module AZ.AVD

# Authenticate to Azure
Connect-AzAccount

#Old HostPool Name and Resource Group Name
$hostPoolName = "AVD-WEU-01"
$resourceGroupName = "RG"

#New HostPool Names and Resource Group Names. (Replace HostPool Names and Resource Group Names)
$Newhostpoolnameweu = "AVD-West-Europe"
$Newhostpoolnamewus = "AVD-West-US"
$Newhostpoolnamecin = "AVD-Central-India"
$NewResourceGroupNameWEU = "RG-AVD-WEU"
$NewResourceGroupNameWUS = "RG-AVD-WUS"
$NewResourceGroupNameCIN = "RG-AVD-CIN"


# Update registration tokens for the new host pools with + 4 hours
Update-AvdRegistrationToken -HostpoolName $Newhostpoolnameweu -ResourceGroupName $NewResourceGroupNameWEU -HoursActive 4
Update-AvdRegistrationToken -HostpoolName $Newhostpoolnamewus -ResourceGroupName $NewResourceGroupNameWUS -HoursActive 4
Update-AvdRegistrationToken -HostpoolName $Newhostpoolnamecin -ResourceGroupName $NewResourceGroupNameCIN -HoursActive 4

# Get and store registration tokens to change RegistrationToken registry value with new host pool token.
$weutoken = (Get-AzWvdHostPoolRegistrationToken -ResourceGroupName $NewResourceGroupNameWEU -HostPoolName $Newhostpoolnameweu).Token
$wustoken = (Get-AzWvdHostPoolRegistrationToken -ResourceGroupName $NewResourceGroupNameWUS -HostPoolName $Newhostpoolnamewus).Token
$cintoken = (Get-AzWvdHostPoolRegistrationToken -ResourceGroupName $NewResourceGroupNameCIN -HostPoolName $Newhostpoolnamecin).Token

# Get all session hosts in the host pool
$sessionHosts = Get-AzWvdSessionHost -ResourceGroupName $resourceGroupName -HostPoolName $hostPoolName
Write-Host "Found $($sessionHosts.Count) session hosts in the host pool."

#Add Start VMs so they are all powered on.

# Array to store the information
$vmInfoArray = @()

# Get all resource groups
$resourceGroups = Get-AzResourceGroup

foreach ($sessionHost in $sessionHosts) {
    # Get the assigned user for each session host
    $assignedUser = $sessionHost.AssignedUser

    # Extract the host name from the session host name
    $hostname = $sessionHost.Name.Split("/")[-1].Split(".")[0]
    $fqdn = $sessionHost.Name.Split("/")[-1]

    # Retrieve the VM to get the location
    $vm = $null
    foreach ($resourceGroup in $resourceGroups) {
        $vm = Get-AzVM -Name $hostname -ResourceGroupName $resourceGroup.ResourceGroupName -ErrorAction SilentlyContinue
        if ($null -ne $vm) {
            break
        }
    }

    # Check if VM was found
    if ($null -eq $vm) {
        Write-Host "Could not find VM for session host $($sessionHost.Name)."
        continue
    }

    # Store the information in the array
    $vmInfo = @{
        "HostName"     = $hostname
        "FQDN"         = $fqdn
        "Location"     = $vm.Location
        "AssignedUser" = $assignedUser
    }

    $vmInfoArray += $vmInfo

    Write-Host "Stored info for session host $hostname in the array."

    # Remove the session host from the old host pool
    Remove-AzWvdSessionHost -HostPoolName $hostPoolName -ResourceGroupName $resourceGroupName -Name $fqdn -Force
    Write-Output "Removed SessionHost $($sessionHost.Name) from $HostpoolName"
}

# Display the contents of the array
$vmInfoArray | Format-Table

# Set the registry keys and restart services, then reassign users back to their original session hosts after they are added to a new host pool.
foreach ($sessionHost in $vmInfoArray) {

    Write-Output "Setting registry keys for session host $($sessionHost.HostName)..."
    
    #You can Uncomment to use specific credentials that have local admin rights.
    #$credential = Get-Credential -UserName 'Domain\localAdmin' -Message 'Please enter your password'
    
    # Create a new PSSession to the target VM
    #$session = New-PSSession -ComputerName $sessionHost.FQDN -Credential $credential
    $session = New-PSSession -ComputerName $sessionHost.FQDN

    if ($sessionHost.Location -eq 'centralindia') {
        # Invoke the scriptblock on the remote session
        Invoke-Command -Session $session -ScriptBlock {
            # Set registry keys
            Set-ItemProperty -Path 'Registry::HKEY_LOCAL_MACHINE\SOFTWARE\Microsoft\RDInfraAgent' -Name 'RegistrationToken' -Value $args[0]
            Set-ItemProperty -Path 'Registry::HKEY_LOCAL_MACHINE\SOFTWARE\Microsoft\RDInfraAgent' -Name 'IsRegistered' -Value 0
            # Restart service
            Restart-Service -Name 'RDAgentBootLoader'
        } -ArgumentList $cintoken
    }
    elseif ($sessionHost.Location -eq 'westus' -or $sessionHost.Location -eq 'eastus') {
        # Invoke the scriptblock on the remote session
        Invoke-Command -Session $session -ScriptBlock {
            # Set registry keys
            Set-ItemProperty -Path 'Registry::HKEY_LOCAL_MACHINE\SOFTWARE\Microsoft\RDInfraAgent' -Name 'RegistrationToken' -Value $args[0]
            Set-ItemProperty -Path 'Registry::HKEY_LOCAL_MACHINE\SOFTWARE\Microsoft\RDInfraAgent' -Name 'IsRegistered' -Value 0
            #Restart service
            Restart-Service -Name 'RDAgentBootLoader'
        } -ArgumentList $wustoken
    }
    elseif ($sessionHost.Location -eq 'westeurope') {
        # Invoke the scriptblock on the remote session
        Invoke-Command -Session $session -ScriptBlock {
            # Set registry keys
            Set-ItemProperty -Path 'Registry::HKEY_LOCAL_MACHINE\SOFTWARE\Microsoft\RDInfraAgent' -Name 'RegistrationToken' -Value $args[0]
            Set-ItemProperty -Path 'Registry::HKEY_LOCAL_MACHINE\SOFTWARE\Microsoft\RDInfraAgent' -Name 'IsRegistered' -Value 0
            # Restart service
            Restart-Service -Name 'RDAgentBootLoader'
        } -ArgumentList $weutoken
    }
    else {
        Write-Output "Unexpected location: $($sessionHost.Location)"
        continue
    }
    # Close the session
    Remove-PSSession -Session $session
    # Pause for 30 seconds for SessionHost to Join New Pool. Can decrease if you think its to long.
    Start-Sleep -Seconds 30
    Write-Output "Registry keys set and service restarted for session host $($sessionHost.HostName)."

    Write-Output "Assigning user back to the session host $($sessionHost.HostName)..."

    if ($sessionHost.Location -eq 'centralindia') {
        Update-AzWvdSessionHost -HostPoolName $Newhostpoolnamecin -Name $sessionHost.FQDN -ResourceGroupName $NewResourceGroupNameCIN -AssignedUser $sessionHost.AssignedUser 
    }
    elseif ($sessionHost.Location -eq 'westus' -or $sessionHost.Location -eq 'eastus') {
        Update-AzWvdSessionHost -HostPoolName $Newhostpoolnamewus -Name $sessionHost.FQDN -ResourceGroupName $NewResourceGroupNameWUS -AssignedUser $sessionHost.AssignedUser
    }
    elseif ($sessionHost.Location -eq 'westeurope') {
        Update-AzWvdSessionHost -HostPoolName $Newhostpoolnameweu -Name $sessionHost.FQDN -ResourceGroupName $NewResourceGroupNameWEU -AssignedUser $sessionHost.AssignedUser
    }
    Write-Output "User assigned back to the session host $($sessionHost.HostName)."
}