templates/CapacityMonitor/Utils.ps1

function Connect-CBSArray {
    param (
        [String] $KeyVaultName,
        [int] $DefaultUtilizationThreshold = 80
    )

    $CBSArrays = @()
    Import-Module -Name PureStoragePowerShellSDK2
    $Secrets = Get-AzKeyVaultSecret -VaultName $KeyVaultName | where-object {$_.Name -like "*-$($KeyVaultName)-capacity-username"}
    foreach ($Secret in $Secrets) {
        $ArrayName = ($Secret.name -Split "-$($KeyVaultName)-capacity-username")[0]
        $CBSUsername = Get-AzKeyVaultSecret -VaultName $KeyVaultName -Name $Secret.Name -AsPlainText
        #TODO: Password we got here is plain text. Any security concern?
        $CBSPassword = Get-AzKeyVaultSecret -VaultName $KeyVaultName -Name "$($ArrayName)-$($KeyVaultName)-capacity-password" -AsPlainText
        $SecureCBSPassword = ConvertTo-SecureString -String $CBSPassword -AsPlainText -Force
        $Credential = New-Object System.Management.Automation.PSCredential -ArgumentList ($CBSUsername, $SecureCBSPassword)

        $Location = Get-AzKeyVaultSecret -VaultName $KeyVaultName -Name "$($ArrayName)-$($KeyVaultName)-capacity-location" -AsPlainText
        $Threshold = $DefaultUtilizationThreshold
        $ThresholdString = Get-AzKeyVaultSecret -VaultName $KeyVaultName -Name "$($ArrayName)-$($KeyVaultName)-capacity-threshold" -AsPlainText
        if ($ThresholdString) {
            $Threshold = [int]$ThresholdString
        }
        # If the string matches the format like "172-168-1-0", the array ip addressed was processed because secret name does not allow "."
        if ($ArrayName -match "^\d+-\d+-\d+-\d+$") {
            $ArrayName = $ArrayName.Replace("-", ".")
        }

        $Array = Connect-Pfa2Array -EndPoint $ArrayName -Credential $Credential -IgnoreCertificateError
        if (-not $Array) {
            Write-Warning "Failed to connect to array ($ArrayName)"
        }

        $CBSArrays += @{ArrayConnection= $Array; Threshold = $Threshold; Location= $Location}
    }
    return $CBSArrays
}

function Get-AzureAuthHeader {
    param (
        $AzContext
    )
    $AzProfile = [Microsoft.Azure.Commands.Common.Authentication.Abstractions.AzureRmProfileProvider]::Instance.Profile
    $ProfileClient = New-Object -TypeName Microsoft.Azure.Commands.ResourceManager.Common.RMProfileClient -ArgumentList ($AzProfile)
    $Token = $profileClient.AcquireAccessToken($azContext.Subscription.TenantId)
    $AuthHeader = @{
        'Content-Type'='application/json'
        'Authorization'='Bearer ' + $token.AccessToken
    }

    return $AuthHeader
}

function Check-AzureQuota {
    param (
        $CBSArray,
        $CBSLocation,
        $CurrentCapacity,
        $DesiredCapacity
    )
    $ArrayModel = (Get-Pfa2Controller -Array $CBSArray).Model | Select-Object -First 1
    $QuotaDevice = "PremiumV2TotalDiskSizeinGB"
    if ($ArrayModel -eq "CBS-V10MUR1" -or $ArrayModel -eq "CBS-V20MUR2") {
        $QuotaDevice = "UltraSSDDiskSizeinGB"
    }
    elseif ($ArrayModel -eq "CBS-V20MP2R2") {
        $QuotaDevice = "PremiumV2TotalDiskSizeinGB"
    }
    else {
        throw "Unsupported array model: $ArrayModel for array ($($CBSArray.ArrayName))"
    }

    $QuotaInfo =  Get-AzVMUsage -Location $CBSLocation | Where-Object { $_.Name.LocalizedValue -eq $QuotaDevice }
    $AvailableQuota = ($QuotaInfo.Limit - $QuotaInfo.CurrentValue) * 1GB
    if ($AvailableQuota -ge ($DesiredCapacity - $CurrentCapacity)) {
        return $true
    }
    return $false
}

function Get-ArrayUtilization {
    param (
        $CBSArray
    )
    Import-Module -Name PureStoragePowerShellSDK2
    $ArraySpace = Get-Pfa2ArraySpace -Array $CBSArray
    $Utilization = $ArraySpace.Space.TotalPhysical / $ArraySpace.Capacity
    return $Utilization
}

function Get-ArrayNextUpgradeStep {
    param (
        $CBSArray,
        $CBSLocation
    )

    $CloudCapacity =  Invoke-Pfa2RestCommand -Array $CBSArray -Method GET -RelativeUri arrays/cloud-capacity | ConvertFrom-Json
    if ($CloudCapacity.Items.Status -ne "idle") {
        throw "Array ($($CBSArray.ArrayName)) is not ready for upgrade. Status: $($CloudCapacity.Items.Status)"
    }
    $CurrentCapacity = $CloudCapacity.Items.current_capacity
    $SupportedStepsData =  Invoke-Pfa2RestCommand -Array $CBSArray -Method GET -RelativeUri arrays/cloud-capacity/supported-steps | ConvertFrom-Json
    # Steps should be returned in sorted order from the REST API but we sort it here to make sure
    $SupportedSteps = $SupportedStepsData.Items.supported_capacity | Sort-Object
    foreach ($Step in $SupportedSteps) {
        if ($Step -gt $CurrentCapacity) {
            return $Step
        }
    }
    throw "No upgrade step found for array ($($CBSArray.ArrayName))"
}

function Start-ArrayStorageUpgrade {
    param (
        $CBSArray,
        $NewStorageSize
    )
    Invoke-Pfa2RestCommand -Array $CBSArray -Method PATCH -RelativeUri arrays/cloud-capacity -Body "{""requested_capacity"":$NewStorageSize}"
}

function Start-UpgradeCheck {
    param (
        $CBSArrays,
        $AuthHeader
    )
    Write-Host "Starting upgrade check ..."
    foreach ($CBSArray in $CBSArrays) {
        try {
            $ArrayUtilization = Get-ArrayUtilization -CBSArray $CBSArray.ArrayConnection
            Write-Host "Array ($($CBSArray.ArrayConnection.ArrayName)) utilization: $($ArrayUtilization)"
            $UtilizationThreshold = $CBSArray.Threshold/100
            Write-Host "Array ($($CBSArray.ArrayConnection.ArrayName)) utilization threshold: $($UtilizationThreshold)"
            if ($ArrayUtilization -ge $UtilizationThreshold) {
                Write-Host "Array ($($CBSArray.ArrayName)) utilization has reeached the utilization threshold. Checking if array can be upgraded..."
                # Check if array can be upgraded
                $NextUpgradeStep = Get-ArrayNextUpgradeStep -CBSArray $CBSArray.ArrayConnection -CBSLocation $CBSArray.Location
                if ($NextUpgradeStep -gt 0) {
                    # TODO: Check Azure Quota
                    Write-Host "Starting upgrade for array ($($CBSArray.ArrayConnection.ArrayName)) ...."
                    Start-ArrayStorageUpgrade -CBSArray $CBSArray.ArrayConnection -NewStorageSize $NextUpgradeStep
                }
            }
            else {
                Write-Host "Array ($($CBSArray.ArrayConnection.ArrayName)) utilization is under threshold. No need to upgrade."
            }
        }
        catch {
            Write-Host "Failed to check capacity upgrade for array $($CBSArray.ArrayName). Error: $_"
        }
    }
}