Poll for an Azure Build Agent to become available and when available copy the capabilities

Copy-UserCapablityToAgent.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. 'java = JAVA_HOME, jdk = JAVA_HOME'")]
    $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 Get-AgentUserCapabilites {
    param
    (
        $pat,
        $url ,
        $poolid,
        $agentid
    )

    $wc = Get-WebClient($pat)

    # get the current capabilities, we the update does a replace
    $uri = "$url/_apis/distributedtask/pools/$poolid/agents/$($agentID)?includeCapabilities=true&api-version=5.1"
    $agentSettings = $wc.DownloadString($uri) | ConvertFrom-Json

    $capabilitiesObject = @{}
    # bit ugly, deals with PSCustomObject
    $currentCapabilities = $agentSettings.userCapabilities | Get-Member -MemberType NoteProperty | % Name | %{$_ + '=' + $($agentSettings.userCapabilities).$_}
    foreach ($capability in $currentCapabilities.Split(";")){
        $item = $capability.Split("=")
        $capabilitiesObject.Add($item[0].Trim(), $(if ($item[1] -eq $null) {""} else {$item[1].Trim()}))
    }

    $capabilitiesObject
}

function Get-AgentSystemCapabilites {
    param
    (
        $pat,
        $url ,
        $poolid,
        $agentid
    )

    $wc = Get-WebClient($pat)

    # get the current capabilities, we the update does a replace
    $uri = "$url/_apis/distributedtask/pools/$poolid/agents/$($agentID)?includeCapabilities=true&api-version=5.1"
    $agentSettings = $wc.DownloadString($uri) | ConvertFrom-Json

    $agentSettings.systemCapabilities
}

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


    $systemCapabilities = Get-AgentSystemCapabilites -pat $pat -url $url -poolid $poolid -agentid $agentid
    $capabilitiesObject = Get-AgentUserCapabilites -pat $pat -url $url -poolid $poolid -agentid $agentid

    # Now add the new ones
    foreach ($capability in $capabilities.Split(",")){
       $item = $capability.Split("=")
       if ($capabilitiesObject.contains($item[0].Trim()) -eq $false) {
	  write-host "$($item[0].trim()) = $($systemCapabilities.$($item[1]))"
	  $capabilitiesObject.Add($item[0].Trim(), $($systemCapabilities.$($item[1])))
       }
    }

    # 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/
    $wc = Get-WebClient($pat)
    $uri = "$url/_apis/distributedtask/pools/$poolid/agents/$($agentID)/usercapabilities?api-version=5.1"
    $data = $capabilitiesObject | ConvertTo-Json
    $data
    $jsondata = $wc.UploadString($uri,"PUT", $data) | ConvertFrom-Json
}


if ([string]::IsNullOrEmpty($capabilities)) {
    write-host "No user capability passed to copy 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) duplication of these pairs $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"
    }
}