Connecting Nerdio Manager to Neo

#description: Neo Agent: logs off one user's Azure Virtual Desktop sessions in a host pool of this account.
#tags: Neo Agent

<# Variables:
{
  "HostPoolName": {"Description": "Name of the host pool.", "IsRequired": true},
  "UserPrincipalName": {"Description": "The user to log off, for example jane@contoso.com.", "IsRequired": true},
  "SessionHostName": {"Description": "Only the session on this session host. Leave empty for every session of the user.", "IsRequired": false, "DefaultValue": ""}
}
#>

$ErrorActionPreference = 'Stop'
# The subscription is the one Nerdio Manager runs this account's scripted actions in, never one the caller names.
Set-AzContext -SubscriptionId $AzureSubscriptionID | Out-Null

$pools = @(Get-AzWvdHostPool | Where-Object { $_.Name -eq $HostPoolName })
if ($pools.Count -ne 1) {
    throw "Found $($pools.Count) host pools named $HostPoolName in this account's subscription."
}
$resourceGroupName = ($pools[0].Id -split '/')[4]

$sessions = Get-AzWvdUserSession -HostPoolName $HostPoolName -ResourceGroupName $resourceGroupName |
    Where-Object { $_.UserPrincipalName -eq $UserPrincipalName }
if ($SessionHostName) {
    # Compare the short host name, so "avd-1" never matches "avd-10" and a full name matches too.
    $shortName = ($SessionHostName -split '\.')[0]
    $sessions = $sessions | Where-Object { ((($_.Name -split '/')[1]) -split '\.')[0] -eq $shortName }
}
if (-not $sessions) {
    Write-Output "No session of $UserPrincipalName in $HostPoolName."
    return
}
foreach ($session in $sessions) {
    $parts = $session.Name -split '/'
    Remove-AzWvdUserSession -HostPoolName $HostPoolName -ResourceGroupName $resourceGroupName `
        -SessionHostName $parts[1] -Id $parts[2] -Force
    Write-Output "Logged off session $($parts[2]) of $UserPrincipalName on $($parts[1])."
}