Enable RDP ShortPath for AVD on an image version automated

param(
    [parameter(mandatory = $true, ValueFromPipelineByPropertyName)]$virtualMachineName,
    [parameter(mandatory = $true, ValueFromPipelineByPropertyName)]$resourceGroupName,
    [parameter(mandatory = $true, ValueFromPipelineByPropertyName)]$password,
    [parameter(mandatory = $true, ValueFromPipelineByPropertyName)]$username
)

import-module az.desktopvirtualization
import-module az.network
import-module az.compute


function add-firewallRule() {
    param(
    [parameter(mandatory = $true, ValueFromPipelineByPropertyName)]$NSG,
    [parameter(mandatory = $true, ValueFromPipelineByPropertyName)]$Name,
    [parameter(mandatory = $true, ValueFromPipelineByPropertyName)]$Port,
    [parameter(mandatory = $true, ValueFromPipelineByPropertyName)]$Protocol,
    [parameter(mandatory = $false, ValueFromPipelineByPropertyName)]$Source

)
    if ($null -eq $Source){$Source = "*";$SourceReadable="Any"}
    # Pick random number for setting priority. It will exclude current priorities.
    $InputRange = 100..200
    $Exclude = ($NSG | Get-AzNetworkSecurityRuleConfig | select Priority).priority
    $RandomRange = $InputRange | Where-Object { $Exclude -notcontains $_ }
    $priority = Get-Random -InputObject $RandomRange
    $nsgParameters = @{
        Name                     = $Name
        Description              = $Name
        Access                   = 'Allow'
        Protocol                 = $Protocol 
        Direction                = "Inbound" 
        Priority                 = $priority 
        SourceAddressPrefix      = $Source 
        SourcePortRange          = "*"
        DestinationAddressPrefix = "*" 
        DestinationPortRange     = $Port
    }
    $NSG | Add-AzNetworkSecurityRuleConfig @NSGParameters  | Set-AzNetworkSecurityGroup 
}

$vm = Get-AzVM -name $virtualMachineName -ResourceGroupName $resourceGroupName
$virtualNetworkSubnet = (Get-AzNetworkInterface -ResourceId $vm.NetworkProfile.NetworkInterfaces.id).IpConfigurations.subnet.id
$NSG = Get-AzNetworkSecurityGroup | ? { $_.subnets.id -eq $virtualNetworkSubnet }# Enabling Powershell Remote Extention
$virtualMachinePublicIp = (Get-AzPublicIpAddress | where { $_.name -match $VirtualMachineName }).IpAddress

$connectionUri = "https://" + $virtualMachinePublicIp + ":5986" 
[securestring]$secStringPassword = ConvertTo-SecureString $password -AsPlainText -Force
[pscredential]$creds = New-Object System.Management.Automation.PSCredential ($userName, $secStringPassword)
$session = $null


# Adding Remote PowerShell Extention and NSG configuration
Write-Output "Enabling Powershell Remote Extention"
Invoke-AzVMRunCommand -CommandId "EnableRemotePS" -VM $vm
add-firewallRule -NSG $NSG -port 5986 -Source 127.0.0.1
 -Protocol "TCP" -Name "Allow-port-5986"

while (!($session)) {
    $session = New-PSSession -ConnectionUri $connectionUri -Credential $creds -SessionOption (New-PSSessionOption -SkipCACheck -SkipCNCheck -SkipRevocationCheck)
    Write-Output "Creating Remote Powershell session"
    $session
}

# Configure base needs
Invoke-Command -Session $session -ScriptBlock { Set-ExecutionPolicy -ExecutionPolicy Unrestricted }
Invoke-Command -Session $session -ScriptBlock { if(!(Test-Path C:\Scripts)){New-Item -ItemType "directory" -Path "C:\Scripts" } }
# RDP Short path
Invoke-Command -Session $session -ScriptBlock { Invoke-WebRequest -uri https://raw.githubusercontent.com/srozemuller/Windows-Virtual-Desktop/master/SessionHost-Management/Prepare-ForRDPShortPath.ps1 -OutFile C:\Scripts\Perpare-ForRDPShortPath.ps1} 
Invoke-Command -Session $session -ScriptBlock { C:\Scripts\Perpare-ForRDPShortPath.ps1 } 

Disconnect-PSSession -Session $session
Remove-PSSession -Session $session

$WinstationsKey = 'HKLM:\SYSTEM\CurrentControlSet\Control\Terminal Server\WinStations'
if(Test-Path $WinstationsKey){
    New-ItemProperty -Path $WinstationsKey -Name 'fUseUdpPortRedirector' -ErrorAction:SilentlyContinue -PropertyType:dword -Value 1 -Force
    New-ItemProperty -Path $WinstationsKey -Name 'UdpPortNumber' -ErrorAction:SilentlyContinue -PropertyType:dword -Value 3390 -Force
}

New-NetFirewallRule -DisplayName 'Remote Desktop - Shortpath (UDP-In)'  -Action Allow -Description 'Inbound rule for the Remote Desktop service to allow RDP traffic. [UDP 3390]' -Group '@FirewallAPI.dll,-28752' -Name 'RemoteDesktop-UserMode-In-Shortpath-UDP'  -PolicyStore PersistentStore -Profile Domain, Private -Service TermService -Protocol udp -LocalPort 3390 -Program '%SystemRoot%\system32\svchost.exe' -Enabled:True
Restart-Computer $env:computername -Force

# Adding the role for RDP Shortpath
add-firewallRule -NSG $NSG -port 3390 -Protocol "UDP" -Name "Allow-3390"