Poll for an Azure Build Agent to become available and when available then add capabilities

Add-UserCapablitiesToAgent.ps1

param
(
    [parameter(Mandatory = $true, HelpMessage = "Azure DevOps PAT token")]
    $pat,
    [parameter(Mandatory = $true, HelpMessage = "URL of Azure DevOps instance e.g. https://dev.aure.com/myinstance")]
    $url,
    [parameter(Mandatory = $true, HelpMessage = "Azure DevOps Agent Pool name")]
    $pool ,
    [parameter(Mandatory = $true, HelpMessage = "Based name for agent to search for e.g MyAgent as part of B1MyAgent-123")]
    $agentBaseName ,
    [parameter(Mandatory = $true, HelpMessage = "Prefix for agent number e.g B as part of B1MyAgent-123")]
    $prefix,
    [parameter(Mandatory = $true, HelpMessage = "Agent number e.g 1 as part of B1MyAgent-123")]
    $index,
    [parameter(Mandatory = $true, HelpMessage = "List of capabilies to added e.g. 'aaa, bbb=123'")]
    $capabilities,
    [parameter(Mandatory = $false, HelpMessage = "Polling interval")]
    $interval = 180,
    [parameter(Mandatory = $false, HelpMessage = "Retry count")]
    $maxRetryCount = 10
)


function Get-WebClient {
    param
    (
        [string]$pat,
        [string]$ContentType = "application/json"
    )

    $wc = New-Object System.Net.WebClient
    $wc.Headers["Content-Type"] = $ContentType
    $pair = ":${pat}"
    $bytes = [System.Text.Encoding]::ASCII.GetBytes($pair)
    $base64 = [System.Convert]::ToBase64String($bytes)
    $wc.Headers.Add("Authorization", "Basic $base64");
    $wc
}

function Get-AgentPool {
    param
    (
        $pat,
        $url ,
        $pool
    )

    $wc = Get-WebClient($pat)

    # get the pool ID from the name
    $uri = "$url/_apis/distributedtask/pools?poolName=$pool&api-version=5.1"
    $jsondata = $wc.DownloadString($uri) | ConvertFrom-Json
    $jsondata.value
}

function Get-Agent {
    param
    (
        $pat,
        $url ,
        $poolid,
        $agentBaseName,
        $prefix,
        $index
    )

    $wc = Get-WebClient($pat)

    # get the agent, we can't use the url filter as the name has a random number
    $uri = "$url/_apis/distributedtask/pools/$poolid/agents?api-version=5.1"
    $jsondata = $wc.DownloadString($uri) | ConvertFrom-Json
    $jsondata.value | where { $_.name -like "$prefix$index$agentBaseName-*" }

}


function Update-Agent {
    param
    (
        $pat,
        $url ,
        $poolid,
        $agentid,
        $capabilities
    )

    $wc = Get-WebClient($pat)

    # The documented call in the api does not work, but I found this https://blogs.blackmarble.co.uk/rfennell/2018/12/06/programmatically-adding-user-capabilities-to-azure-devops-agents/

    $uri = "$url/_apis/distributedtask/pools/$poolid/agents/$($agentID)/usercapabilities?api-version=5.1"
    $uri
    # have to pass the ID too else there is a 404 error
    # this is the expected format
    # $capabilities = "AAA , BBB=123"
    $capabilitiesObject = @{}
    foreach ($capability in $capabilities.Split(",")){
        $item = $capability.Split("=")
        $capabilitiesObject.Add($item[0].Trim(), $(if ($item[1] -eq $null) {""} else {$item[1].Trim()}))
    }

    $data = $capabilitiesObject | ConvertTo-Json
    $jsondata = $wc.UploadString($uri,"PUT", $data) | ConvertFrom-Json

}


if ([string]::IsNullOrEmpty($capabilities)) {
    write-host "No user capabilities passing so skipping"
} else {
    Write-Host "The capabilities to set are $capabilities"
    write-host "Finding Agent Pool $pool"
    $agentPool = Get-AgentPool -url $url -pat $pat -pool $pool
    $retryCount = 0
    do {
        write-host "Waiting for Agent $prefix$index$agentbasename-XXX to become available. Waiting $($retryCount*$interval) of $($maxRetryCount*$interval) seconds"
        sleep -Seconds $interval
        $retryCount +=1
        $agent = Get-Agent -url $url -pat $pat -poolid $agentPool.id -agentBaseName $agentBaseName -prefix $prefix -index $index
     } while ( $agent -eq $null -and $retryCount -lt $maxRetryCount)

    if ($agent -ne $null)
    {
        write-host "Found agent $($agent.name) adding user capabilities $capabilities"
        Update-Agent -url $url -pat $pat -poolid $agentPool.id -agentid $agent.id -capabilities $capabilities
    } else {
        write-error "No agent with a name in form $prefix$index$agentbasename-XXX is registered with timeout period"
    }
}