Install-VMInsights

1.10.4

Install-VMInsights.ps1

<#PSScriptInfo
 
.VERSION 1.10.4
 
.GUID 76a487ef-47bf-4537-8942-600a66a547b1
 
.AUTHOR mithomps@microsoft.com
 
.COMPANYNAME Microsoft Corporation
 
.COPYRIGHT Microsoft Corporation. All rights reserved.
 
.TAGS
 
.LICENSEURI
 
.PROJECTURI
 
.ICONURI
 
.EXTERNALMODULEDEPENDENCIES
 
.REQUIREDSCRIPTS
 
.EXTERNALSCRIPTDEPENDENCIES
 
.RELEASENOTES
Remove Legacy Agent (LA/MMA/OMS) support. AMA is now the only supported onboarding path.
Remove Dependency Agent (DA) support. -ProcessAndDependencies now throws a retirement error.
 
#> 

#Requires -Module @{ ModuleName = 'Az.Compute'; ModuleVersion = '7.1.0' }
#Requires -Module @{ ModuleName = 'Az.Resources'; ModuleVersion = '6.12.1' }
#Requires -Module @{ ModuleName = 'Az.Accounts'; ModuleVersion = '2.13.2' }
#Requires -Module @{ ModuleName = 'Az.PolicyInsights'; ModuleVersion = '1.6.4' }
#Requires -Module @{ ModuleName = 'Az.Monitor'; ModuleVersion = '5.0.0' }
#Requires -Module @{ ModuleName = 'Az.ManagedServiceIdentity'; ModuleVersion = '1.2.0' }



<#
.SYNOPSIS
This script installs VM extensions for Azure Monitor Agent (AMA) for VM Insights.
A Data Collection Rule (DCR) and a User Assigned Managed Identity (UAMI) are associated with the VM's and VMSS.
 
.DESCRIPTION
This script installs or re-configures the following on VM's and VMSS under a Subscription.
Azure Monitor Agent along with Data Collection Rule association and User Assigned Managed Identity.
 
 
Scope can further narrowed down to :
- Resource Group
- Specific VM/VMSS
- Policy Assignment Scope
 
Script will show you list of VM's/VMSS that will apply to and let you confirm to continue.
Use -Approve Switch to run without prompting, if all required parameters are provided.
 
Use -WhatIf to get info about expected effect of the commands in the script.
 
.PARAMETER SubscriptionId
Subscription identifier for the VMs or VMSS
If using PolicyAssignmentName parameter, VMs part of the parameter SubscriptionId are considered.
 
.PARAMETER ResourceGroup
<Optional> Name of the Resource Group of VMs or VMSS. It accepts wildcard characters. The default behavior is to match all.
 
.PARAMETER PolicyAssignmentName
<Optional> Only include VMs associated with this policy.
 
 
.PARAMETER Name
<Optional> Name qualifier to match on VM/VMSS's name in the scope. It accepts wildcard characters. The default behavior is to match all.
 
.PARAMETER TriggerVmssManualVMUpdate
<Optional> Set this flag to trigger update of VM instances in a scale set whose upgrade policy is set to Manual.
 
.PARAMETER Approve
<Optional> Set this flag to provide the approval for the installation to start with no confirmation prompt for the listed VM's/VMSS.
 
.PARAMETER Whatif
<Optional> Set this flag to get info about expected effect of the commands in the script.
 
.PARAMETER Confirm
<Optional> Set this flag to confirm each VM / VMSS.
 
 
 
.PARAMETER DcrResourceId
Data Collection Rule (DCR) azure resource ID identifier.
 
.PARAMETER UserAssignedManagedIdentityResourceGroup
Name of User Assigned Managed Identity (UAMI) resource group.
 
.PARAMETER UserAssignedManagedIdentityName
Name of User Assigned Managed Identity (UAMI).
 
.PARAMETER ProcessAndDependencies
<Retired> This flag is no longer functional. VM Insights Map and the Dependency Agent will be retired on
June 30th, 2028. Using this flag results in an error.
See: https://aka.ms/DependencyAgentRetirement
 
.PARAMETER WorkspaceId
<Retired> The Log Analytics agent was retired on August 31, 2024. Using this parameter results in an error.
See: https://aka.ms/MigrateToAMA
 
.PARAMETER WorkspaceKey
<Retired> The Log Analytics agent was retired on August 31, 2024. Using this parameter results in an error.
See: https://aka.ms/MigrateToAMA
 
.PARAMETER ReInstall
<Retired> The Log Analytics agent was retired on August 31, 2024. Using this parameter results in an error.
See: https://aka.ms/MigrateToAMA
 
.EXAMPLE
Install-VMInsights.ps1 -SubscriptionId <SubscriptionId> -ResourceGroup <ResourceGroup> -DcrResourceId <DataCollectionRuleResourceId> -UserAssignedManagedIdentityName <UserAssignedIdentityName> -UserAssignedManagedIdentityResourceGroup <UserAssignedIdentityResourceGroup>
Onboard VMI for all VM's or VMSS in a Resource Group in a subscription with Azure Monitoring Agent (AMA).
Specify DcrResourceId to associate a data collection rule with VM or VMSS
Specify UserAssignedManagedIdentityName and UserAssignedManagedIdentityResourceGroup for AMA authentication purposes.
 
.EXAMPLE
Install-VMInsights.ps1 -SubscriptionId <SubscriptionId> -ResourceGroup <ResourceGroup> -ProcessAndDependencies -DcrResourceId <DataCollectionRuleResourceId> -UserAssignedManagedIdentityName <UserAssignedIdentityName> -UserAssignedManagedIdentityResourceGroup <UserAssignedIdentityResourceGroup>
NOTE: -ProcessAndDependencies is no longer supported. Dependency Agent onboarding has been retired.
See: https://aka.ms/DependencyAgentRetirement
 
.EXAMPLE
Install-VMInsights.ps1 -SubscriptionId <SubscriptionId> -PolicyAssignmentName a4f79f8ce891455198c08736 -DcrResourceId <DataCollectionRuleResourceId> -UserAssignedManagedIdentityName <UserAssignedIdentityName> -UserAssignedManagedIdentityResourceGroup <UserAssignedIdentityResourceGroup>
Specify PolicyAssignmentName to onboard VMI for VMs part of the policy assignment scope with Azure Monitoring Agent (AMA).
 
.LINK
This script is posted to and further documented at the following location :
http://aka.ms/OnBoardVMInsights
#>

<#CmdletBinding ConfirmImpact level Info : High - Customer data-flow disruptive action.
                                           Medium - Everything else that changes the system configuration.
#>
#Assumption - The script assumes the entity running the script has access to all VMs/VMSS in the script.
[CmdletBinding(SupportsShouldProcess = $True, ConfirmImpact = 'High', DefaultParameterSetName = 'AzureMonitoringAgent')]
param(
    [Parameter(mandatory = $True)][String]$SubscriptionId,
    [Parameter(mandatory = $False)][Switch]$TriggerVmssManualVMUpdate,
    [Parameter(mandatory = $False)][Switch]$Approve,
    [Parameter(mandatory = $False)][String]$ResourceGroup = "*",
    [Parameter(mandatory = $False)][String]$Name = "*",
    [Parameter(mandatory = $False)][String]$PolicyAssignmentName,

    [Parameter(mandatory = $False, ParameterSetName = 'AzureMonitoringAgent')][Switch]$ProcessAndDependencies,
    [Parameter(mandatory = $True,  ParameterSetName = 'AzureMonitoringAgent')][String]$DcrResourceId,
    [Parameter(mandatory = $True,  ParameterSetName = 'AzureMonitoringAgent')][String]$UserAssignedManagedIdentityResourceGroup,
    [Parameter(mandatory = $True,  ParameterSetName = 'AzureMonitoringAgent')][String]$UserAssignedManagedIdentityName,

    [Parameter(mandatory = $True,  ParameterSetName = 'LogAnalyticsAgent')][String]$WorkspaceId,
    [Parameter(mandatory = $True,  ParameterSetName = 'LogAnalyticsAgent')][String]$WorkspaceKey,
    [Parameter(mandatory = $False, ParameterSetName = 'LogAnalyticsAgent')][Switch]$ReInstall
)

$ErrorActionPreference = "Stop"

class CustomerSkip : System.Exception {
    CustomerSkip() : base () {}
}

class FatalException : System.Exception {
    FatalException([String]$errorMessage, [Exception]$innerException) : base($errorMessage, $innerException) {}
}

class VirtualMachineException : System.Exception {
    VirtualMachineException(<#[Microsoft.Azure.Commands.Compute.Models.PSVirtualMachine]#>$vmObject,
                            [String]$errorMessage,
                            [Exception]$innerException) : base("$(FormatVmIdentifier -VMObject $vmObject) : $errorMessage", $innerException) {}
}

class VirtualMachineScaleSetException : System.Exception {
    VirtualMachineScaleSetException(<#[Microsoft.Azure.Commands.Compute.Automation.Models.PSVirtualMachineScaleSet]#>$vmssObject,
                                    [String]$errorMessage,
                                    [Exception]$innerException) : base("$(FormatVmssIdentifier -VMssObject $vmssObject) : $errorMessage", $innerException) {}
}

class ResourceGroupDoesNotExist : System.Exception {
    ResourceGroupDoesNotExist ([String]$rgName, [Exception]$innerException) : base("$rgName : Resource Group does not exist", $innerException) {}
}

class VirtualMachineUnknownException : VirtualMachineException {
    VirtualMachineUnknownException(<#[Microsoft.Azure.Commands.Compute.Models.PSVirtualMachine]#>$vmObject,
                                   [String]$errorMessage,
                                   [Exception]$innerException) : base($vmObject, $errorMessage, $innerException) {}
}

class VirtualMachineDoesNotExist : VirtualMachineException {
    VirtualMachineDoesNotExist (<#[Microsoft.Azure.Commands.Compute.Models.PSVirtualMachine]#>$vmObject,
                                [Exception]$innerException) : base($vmObject, "VM does not exist", $innerException) {}
}

class VirtualMachineOperationFailed : VirtualMachineException {
    VirtualMachineOperationFailed(<#[Microsoft.Azure.Commands.Compute.Models.PSVirtualMachine]#>$vmObject,
                                  [String]$errorMessage) : base($vmObject, $errorMessage, $null) {}
}

class VirtualMachinePoweredDown : VirtualMachineException {
    VirtualMachinePoweredDown(<#[Microsoft.Azure.Commands.Compute.Models.PSVirtualMachine]#>$vmObject,
                              [Exception]$innerException) : base($vmObject, "VM is powered down", $innerException) {}
}

class VirtualMachineExtensionError : VirtualMachineException {
    VirtualMachineExtensionError (<#[Microsoft.Azure.Commands.Compute.Models.PSVirtualMachine]#>$vmObject,
                                  [String]$extension,
                                  [Exception]$innerException) : base($vmObject, "VM extension error for $extension`n$($innerException.Message)", $innerException) {}
}

class VirtualMachineScaleSetUnknownException : VirtualMachineScaleSetException {
    VirtualMachineScaleSetUnknownException(<#[Microsoft.Azure.Commands.Compute.Automation.Models.PSVirtualMachineScaleSet]#>$vmssObject,
                                           [String]$errorMessage,
                                           [Exception]$innerException) : base($vmssObject, $errorMessage, $innerException) {}
}

class VirtualMachineScaleSetDoesNotExist : VirtualMachineScaleSetException {
    VirtualMachineScaleSetDoesNotExist (<#[Microsoft.Azure.Commands.Compute.Automation.Models.PSVirtualMachineScaleSet]#>$vmssObject,
                                        [Exception]$innerException) : base($vmssObject, "VMSS does not exist", $innerException) {}
}

class VirtualMachineScaleSetOperationFailed : VirtualMachineScaleSetException {
    VirtualMachineScaleSetOperationFailed(<#[Microsoft.Azure.Commands.Compute.Automation.Models.PSVirtualMachineScaleSet]#>$vmssObject,
                                          [String]$errorMessage) : base($vmssObject, $errorMessage, $null) {}
}

class VirtualMachineScaleSetExtensionError : VirtualMachineScaleSetException {
    VirtualMachineScaleSetExtensionError (<#[Microsoft.Azure.Commands.Compute.Automation.Models.PSVirtualMachineScaleSet]#>$vmssObject,
                                          [Exception]$innerException) : base($vmssObject, "VMSS extension error`n$($innerException.Message)", $innerException) {}
}

class DataCollectionRuleForbidden : FatalException {
    DataCollectionRuleForbidden([String]$dcrResourceId,
                                [Exception]$innerException) : base("$dcrResourceId : Access to Data Collection Rule is forbidden", $innerException) {}
}

class DataCollectionRuleDoesNotExist : FatalException {
    DataCollectionRuleDoesNotExist([String]$dcrResourceId,
                                   [Exception]$innerException) : base("$dcrResourceId : Data Collection Rule does not exist.", $innerException) {}
}

class DataCollectionRuleIncorrect : FatalException {
    DataCollectionRuleIncorrect([String]$dcrResourceId,
                                [Exception]$innerException) : base("$dcrResourceId : Data Collection Rule incorrect format.", $innerException) {}
}

class UserAssignedManagedIdentityDoesNotExist : FatalException {
    UserAssignedManagedIdentityDoesNotExist($uamiName,
                                            [Exception]$innerException) : base("$uamiName : User Assigned Managed Identity does not exist.", $innerException) {}
}

class UserAssignedManagedIdentityResourceGroupDoesNotExist : FatalException {
    UserAssignedManagedIdentityResourceGroupDoesNotExist($uamiResourceGroup,
                                                         [Exception]$innerException) : base("$uamiResourceGroup : User Assigned Managed Identity Resource Group does not exist.", $innerException) {}
}

class UserAssignedManagedIdentityUnknownException : FatalException {
    UserAssignedManagedIdentityUnknownException([String]$errorMessage,
                                                [Exception]$innerException) : base($errorMessage, $innerException) {}
}

class ResourceGroupTableElement {
    [System.Collections.ArrayList] $VirtualMachineList = [System.Collections.ArrayList]::new()
    [System.Collections.ArrayList] $VirtualMachineScaleSetList = [System.Collections.ArrayList]::new()
}

class OnboardingCounters {
    [Decimal]$Succeeded = 0
    [Decimal]$Total = 0
    [Decimal]$Skipped = 0
    [Decimal]$VMSSInstanceUpgradeFailure = 0
}

# Azure Monitoring Agent Extension constants
Set-Variable -WhatIf:$False -Confirm:$False -Name amaExtensionConstantMap -Option Constant -Value @{ 
       "Windows" = @{  ExtensionType = "AzureMonitorWindowsAgent"
                       TypeHandlerVersion = "1.16"
                       Publisher = "Microsoft.Azure.Monitor" 
                    }
       "Linux" =   @{  ExtensionType = "AzureMonitorLinuxAgent" 
                       TypeHandlerVersion = "1.16"
                       Publisher = "Microsoft.Azure.Monitor"
                    }
}
Set-Variable -WhatIf:$False -Confirm:$False -Name amaDefaultExtensionName -Option Constant -Value "AzureMonitoringAgent"

$extensionVmDefaultUpgradeSettings = @{
    EnableAutomaticUpgrade = $True
    DisableAutoUpgradeMinorVersion = $False
}

Set-Variable -WhatIf:$False -Confirm:$False -Name unknownExceptionVirtualMachineConsequentCounterLimit -Option Constant -Value 3
Set-Variable -WhatIf:$False -Confirm:$False -Name unknownExceptionVirtualMachineScaleSetConsequentCounterLimit -Option Constant -Value 3
Set-Variable -WhatIf:$False -Confirm:$False -Name unknownExceptionTotalCounterLimit -Option Constant -Value 6

#
# FUNCTIONS
#
function PrintSummaryMessage {
    <#
    .SYNOPSIS
    Print the Total of eligible VM/VMSS, number of succeeded and failed.
    #>
    param (
        [Parameter(mandatory = $True)]
        [OnboardingCounters]
        $OnboardingCounters
    )
    Write-Host ""
    Write-Host "Summary :"
    Write-Host "Total VM/VMSS to be processed : $($OnboardingCounters.Total)"
    Write-Host "Succeeded : $($OnboardingCounters.Succeeded)"
    Write-Host "Skipped : $($OnboardingCounters.Skipped)"
    Write-Host "Failed : $($OnboardingCounters.Total - $OnboardingCounters.Skipped - $OnboardingCounters.Succeeded)"
    Write-Host "VMSS Instance Upgrade Failures : $($OnboardingCounters.VMSSInstanceUpgradeFailure)"
}

function ExtractCloudExceptionErrorMessage {
    <#
    .SYNOPSIS
    Extract error code from the Cloud Exception.
    #>
    param
    (
        [Parameter(Mandatory=$True)]
        [System.Management.Automation.ErrorRecord]
        $ErrorRecord
    )

    if ($ErrorRecord.Exception.Message -match 'ErrorMessage *: *(.*)') {
        return $matches[1]
    }

    return $null
}

function ExtractExceptionErrorCode {
    <#
    .SYNOPSIS
    Extract error code from the Cloud Exception.
    #>
    param
    (
        [Parameter(Mandatory=$True)]
        [System.Management.Automation.ErrorRecord]
        $ErrorRecord
    )

    if ($ErrorRecord.Exception.Message -match 'ErrorCode *: *([^\s]+)') {
        return $matches[1]
    }

    return $null
}

function ExtractExceptionPrefixErrCode {
    <#
    .SYNOPSIS
    Extract error code from the Cloud Exception.
    #>
    param
    (
        [Parameter(Mandatory=$True)]
        [System.Management.Automation.ErrorRecord]
        $ErrorRecord
    )

    if ($ErrorRecord.Exception.Message -match '^[[]([^]]+)[]] *:') {
        return $matches[1]
    }

    return $null
}

function FormatVmIdentifier {
    <#
    .SYNOPSIS
    Format VM Information for messages
    #>
    param (
        [Parameter(Mandatory=$True)] 
        [Microsoft.Azure.Commands.Compute.Models.PSVirtualMachine]
        $VMObject
    )
    return "($($VMObject.ResourceGroupName)) $($VMObject.Name)"
}

function FormatVmssIdentifier {
    <#
    .SYNOPSIS
    Format VMSS Information for messages
    #>
    param (
        [Parameter(Mandatory=$True)] 
        [Microsoft.Azure.Commands.Compute.Automation.Models.PSVirtualMachineScaleSet]
        $VMssObject
    )
    
    return "($($VMssObject.ResourceGroupName)) $($VMssObject.Name)"
}

function IsVmAndGuestAgentUpAndHealthy {
    <#
    .SYNOPSIS
    Checks whether a VM is up and Guest agent is in a healthy state.
    #>
    param (
        [Parameter(Mandatory=$True)]
        [Microsoft.Azure.Commands.Compute.Models.PSVirtualMachine]
        $VMObject
    )
  
    try {
        $vmResourceGroupName = $VMObject.ResourceGroupName
        $vmlogheader = FormatVmIdentifier -VMObject $VMObject
        
        $vmWithStatus = Get-AzVm -ResourceGroupName $vmResourceGroupName -Name $VMObject.Name -Status
        
        $vmhealthStatusCode = $vmWithStatus.Statuses.Code
        if (!($vmhealthStatusCode -contains 'ProvisioningState/succeeded' -and $vmhealthStatusCode -contains 'PowerState/running')) {
            Write-Host "$vmlogheader : Virtual Machine is not operational - $($vmWithStatus.Statuses.DisplayStatus)"
            return $False
        }

        $guestAgentHealthCode = $vmWithStatus.VMAgent.Statuses.Code
        if (!($guestAgentHealthCode -contains 'ProvisioningState/succeeded')) {
            Write-Host "$vmlogheader : Guest Agent is not healthy - $($vmWithStatus.VMAgent.Statuses.DisplayStatus)"
            return $False
        }

        return $True
    } catch [Microsoft.Azure.Commands.Compute.Common.ComputeCloudException] {
        $errorCode = ExtractExceptionErrorCode -ErrorRecord $_
        if ($errorCode -eq "ResourceNotFound") { 
            throw [VirtualMachineDoesNotExist]::new($VMObject, $_.Exception)
        }

        if ($errorCode -eq "ResourceGroupNotFound") {
            throw [ResourceGroupDoesNotExist]::new($vmResourceGroupName, $_.Exception)  
        }

        throw [VirtualMachineUnknownException]::new($VMObject, "Failed to get status of virtual machine", $_.Exception)
    }
}

function VerboseDisplayException {
    <#
    .SYNOPSIS
    Renders customer actionable message, inner exception and stack trace to verbose stream.
    Does not throw any exceptions.
    #>
    param (
        [Parameter(Mandatory=$True)]
        [System.Management.Automation.ErrorRecord]
        $ErrorRecord
    )

    try {
        $ex = $ErrorRecord.Exception
        Write-Host "ExceptionMessage : $($ex.Message)"
        Write-Verbose "StackTrace :"
        Write-Verbose "$($ex.StackTrace)"
        while ($ex = $ex.InnerException) {
            Write-Verbose "InnerExceptionMessage : $($ex.Message)"
            Write-Verbose "StackTrace :"
            Write-Verbose "$($ex.StackTrace)"
        }
        Write-Verbose "ScriptStackTrace :" $ErrorRecord.ScriptStackTrace
    }
    catch {
        # silently ignore
    }
}

function DisplayException {
    <#
    .SYNOPSIS
    Renders customer actionable message, inner exception and stack trace to console.
    Does not throw any exceptions.
    #>
    param (
        [Parameter(Mandatory=$True)]
        [System.Management.Automation.ErrorRecord]
        $ErrorRecord
    )

    try {
        $ex = $ErrorRecord.Exception
        Write-Host "ExceptionMessage : $($ex.Message)"
        Write-Host "StackTrace :"
        Write-Host "$($ex.StackTrace)"
        while ($ex = $ex.InnerException) {
            Write-Host "InnerExceptionMessage : $($ex.Message)"
            Write-Host "StackTrace :"
            Write-Host "$($ex.StackTrace)"
        }
        Write-Host "ScriptStackTrace :" $ErrorRecord.ScriptStackTrace
    }
    catch {
        # silently ignore
    }
}

function GetRgObject {

    <#
    .SYNOPSIS
    PopulateRgHashTableVMs/VMss util function
    #>
    param(
        [Parameter(Mandatory=$True)]
        [Hashtable]
        $Rghashtable,
        [Parameter(Mandatory=$True)]
        [String]
        $ResourceGroupName
    )

    $rgTableElemObject = $Rghashtable[$ResourceGroupName]
    if ($null -eq $rgTableElemObject) {
        $rgTableElemObject = [ResourceGroupTableElement]::new()
        $Rghashtable.Add($ResourceGroupName,$rgTableElemObject)
    }

    return $rgTableElemObject
}

function PopulateRgHashTableVm {
    <#
    .SYNOPSIS
    Populate Resource Group hash table for VMs
    #>
    param(
        [Parameter(Mandatory=$True)]
        [Hashtable]
        $Rghashtable,
        [Parameter(Mandatory=$True)]
        [Microsoft.Azure.Commands.Compute.Models.PSVirtualMachine]
        $VMObject
    )

    $rgTableElemObject = GetRgObject -Rghashtable $Rghashtable -ResourceGroupName $VMObject.ResourceGroupName
    $rgTableElemObject.VirtualMachineList.Add($VMObject)  > $null
}

function PopulateRgHashTableVmss {
    <#
    .SYNOPSIS
    Populate Resource Group hash table for VMSS
    #>
    param(
        [Parameter(Mandatory=$True)]
        [Hashtable]
        $Rghashtable,
        [Parameter(Mandatory=$True)]
        [Microsoft.Azure.Commands.Compute.Automation.Models.PSVirtualMachineScaleSet]
        $VMssObject
    )

    $rgTableElemObject = GetRgObject -Rghashtable $Rghashtable -ResourceGroupName $VMssObject.ResourceGroupName
    $rgTableElemObject.VirtualMachineScaleSetList.Add($VMssObject)  > $null
}

function GetVMExtension {
    <#
    .SYNOPSIS
    Return the VM extension of specified Type and Publisher.
    #>
    param
    (
        [Parameter(mandatory = $True)]
        [Microsoft.Azure.Commands.Compute.Models.PSVirtualMachine]
        $VMObject,
        [Parameter(mandatory = $True)]
        [Hashtable]
        $ExtensionProperties
    )

    $extensionType = $ExtensionProperties.ExtensionType
    $publisher = $ExtensionProperties.Publisher

    try {
        $extensions = Get-AzVMExtension -ResourceGroupName $VMObject.ResourceGroupName -VMName $VMObject.Name
        foreach ($extension in $extensions) {
            if ($extension.ExtensionType -eq $extensionType -and $extension.Publisher -eq $publisher) {
                return $extension
            }
        }
    } catch [Microsoft.Azure.Commands.Compute.Common.ComputeCloudException] {
        $errorCode = ExtractExceptionErrorCode -ErrorRecord $_
        if ($errorCode -eq "ParentResourceNotFound") {
            throw [VirtualMachineDoesNotExist]::new($VMObject, $_.Exception)
        }
        if ($errorCode -eq "ResourceGroupNotFound") {
            throw [ResourceGroupDoesNotExist]::new($($VMObject.ResourceGroupName),$_.Exception)   
        }    
        throw [VirtualMachineUnknownException]::new($VMObject, "Failed to locate extension with type = $($publisher).$($extensionType)", $_.Exception)
    }
    
    return $null
}

function GetVMssExtension {
    <#
    .SYNOPSIS
    Return the VMSS extension of specified ExtensionType and Publisher.
    #>
    param
    (
        [Parameter(Mandatory = $True)]
        [Microsoft.Azure.Commands.Compute.Automation.Models.PSVirtualMachineScaleSet]
        $VMssObject,
        [Parameter(mandatory = $True)]
        [String]
        $ExtensionType,
        [Parameter(mandatory = $True)]
        [String]
        $Publisher
    )

    $extensions = $VMssObject.VirtualMachineProfile.ExtensionProfile.Extensions

    foreach ($extension in $extensions) {
        if ($extension.Type -eq $ExtensionType -and $extension.Publisher -eq $Publisher) {
            return $extension
        }
    }

    return $null
}

function RemoveVMExtension {
    <#
    .SYNOPSIS
    Remove a VM Extension.
    #>
    [CmdletBinding(SupportsShouldProcess = $True, ConfirmImpact = 'High')]
    param
    (
        [Parameter(mandatory = $True)]
        [Microsoft.Azure.Commands.Compute.Models.PSVirtualMachine]
        $VMObject,
        [Parameter(mandatory = $True)]
        [String]
        $ExtensionName,
        [Parameter(mandatory = $True)]
        [hashtable]
        $ExtensionProperties
    )

    $vmlogheader = FormatVmIdentifier -VMObject $VMObject
    $extensionType = $ExtensionProperties.ExtensionType
    $extensionPublisher = $ExtensionProperties.Publisher
    
    if (!$PSCmdlet.ShouldProcess($vmlogheader, "Remove $ExtensionName, type $($extensionPublisher).$($extensionType)")) {
        #-WhatIf skip processing here, return to the caller as we have completed our work.
        if ($WhatIfPreference) {
            return
        }
        throw [CustomerSkip]::new()
    }

    try {
        #Remove operation on non existent VM, extension still return a success
        $removeResult = Remove-AzVMExtension -ResourceGroupName $VMObject.ResourceGroupName `
                                             -VMName $VMObject.Name `
                                             -Name $ExtensionName -Confirm:$false -Force
    } catch [Microsoft.Azure.Commands.Compute.Common.ComputeCloudException] {
        $errorCode = ExtractExceptionErrorCode -ErrorRecord $_
        if ($errorCode -eq "ResourceGroupNotFound") {
            throw [ResourceGroupDoesNotExist]::new($VMObject.ResourceGroupName, $_.Exception)       
        } 
        
        throw [VirtualMachineUnknownException]::new($VMObject, "Failed to remove extension $ExtensionName, type $($extensionPublisher).$($extensionType)", $_.Exception)
    }
    
    if ($removeResult.IsSuccessStatusCode) {
        Write-Host "$vmlogheader : Successfully removed extension $ExtensionName, type $($extensionPublisher).$($extensionType)"
        return
    }

    throw [VirtualMachineOperationFailed]::new($VMObject, 
            "Failed to remove extension $ExtensionName, type $($extensionPublisher).$($extensionType). StatusCode = $($removeResult.StatusCode). ReasonPhrase = $($removeResult.ReasonPhrase)")
}

#VMI supports Customers onboarding DCR from different subscription to which it has access to.
#Cannot validate DCRResourceId parameter Get-AzDataCollectionRule -ResourceId is getting deprecated.
function NewDCRAssociationVm {
    <#
    .SYNOPSIS
    Create a new DCRAssociation with VMs.
    #>
    [CmdletBinding(SupportsShouldProcess = $True, ConfirmImpact = 'Medium')]
    param
    (
        [Parameter(mandatory = $True)]
        [Microsoft.Azure.Commands.Compute.Models.PSVirtualMachine]
        $VMObject
    )

    $vmId = $VMObject.Id
    $vmlogheader = FormatVmIdentifier -VMObject $VMObject
    
    try {
        # A VM may have zero or more Data Collection Rule Associations.
        $dcrAssociationList = Get-AzDataCollectionRuleAssociation -TargetResourceId $vmId 
        foreach ($dcra in $dcrAssociationList) {
            if ($dcra.DataCollectionRuleId -eq $DcrResourceId) {
                Write-Host "$vmlogheader : Data Collection Rule Id $DcrResourceId already associated with the VM."
                return
            }
        }
    } catch [System.Management.Automation.ParameterBindingException] {
        throw [VirtualMachineDoesNotExist]::new($VMObject, $_.Exception)
    }

    if (!($PSCmdlet.ShouldProcess($vmlogheader, "Install Data Collection Rule Association"))) {
        #-WhatIf skip processing here, return to the caller as we have completed our work.
        if ($WhatIfPreference) {
            return
        }
        throw [CustomerSkip]::new()
    }

    #Customer can associate multiple DCRs to a VM.
    $dcrAssociationName = "VM-Insights-DCR-Association-$(New-Guid)"
    Write-Host "$vmlogheader : Deploying Data Collection Rule Association $dcrAssociationName"
    try {
        #BUG! To be fixed by AMCS team - New-AzDataCollectionRuleAssociation allows creating multiple DCRAs for the same {DCR,VMSS} combination.
        $dcrAssociation = New-AzDataCollectionRuleAssociation -TargetResourceId $vmId `
                                                              -AssociationName $dcrAssociationName `
                                                              -RuleId $DcrResourceId `
                                                              -Confirm:$false
    } catch [System.Management.Automation.PSInvalidOperationException] {
        $exceptionMessage = $_.Exception.Message
        
        if ($exceptionMessage.Contains('Invalid format of the resource identifier')) {
            throw [DataCollectionRuleIncorrect]::new($DcrResourceId, $_.Exception)
        } 

        if (!($exceptionMessage -match "status code '([^\s]+)'")) {
            throw [VirtualMachineUnknownException]::new($VMObject, "Failed to create Data Collection Rule Association with $DcrResourceId", $_.Exception)
        }

        $statusCode = $matches[1]
        if ($statusCode -eq 'BadRequest') {
            throw [DataCollectionRuleDoesNotExist]::new($DcrResourceId, $_.Exception)
        }
        if ($statusCode -eq 'NotFound') {
            throw [VirtualMachineDoesNotExist]::new($VMObject, $_.Exception)
        }
        if ($statusCode -eq 'Forbidden') {
            throw [DataCollectionRuleForbidden]::new($DcrResourceId, $_.Exception)     
        }

        throw [VirtualMachineUnknownException]::new($VMObject, "Failed to create Data Collection Rule Association with with $DcrResourceId. StatusCode = $statusCode", $_.Exception)
    }

    #Tmp fix task :- 21191002
    if (($null -eq $dcrAssociation) -or ($dcrAssociation -is [Microsoft.Azure.Management.Monitor.Models.ErrorResponseCommonV2Exception])) {
        throw [VirtualMachineUnknownException]::new($VMObject, "Unable to create Data Collection Rule Association with $DcrResourceId", $_.Exception)
    }

    Write-Host "$vmlogheader : Successfully created Data Collection Rule Association"
}

function NewDCRAssociationVmss {
    <#
    .SYNOPSIS
    Create a new DCR Association with VMSS.
    #>
    [CmdletBinding(SupportsShouldProcess = $True, ConfirmImpact = 'Medium')]
    param
    (
        [Parameter(mandatory = $True)]
        [Microsoft.Azure.Commands.Compute.Automation.Models.PSVirtualMachineScaleSet]
        $VMssObject
    )

    $vmsslogheader = FormatVmssIdentifier -VMssObject $VMssObject
    $vmssId = $VMssObject.Id

    try {
        # A VMSS may have zero or more Data Collection Rule Associations.
        $dcrAssociationList = Get-AzDataCollectionRuleAssociation -TargetResourceId $vmssId 
        foreach ($dcra in $dcrAssociationList) {
            if ($dcra.DataCollectionRuleId -eq $DcrResourceId) {
                Write-Host "$vmsslogheader : Data Collection Rule Id $DcrResourceId already associated with the VMSS."
                return
            }
        }
    } catch [System.Management.Automation.ParameterBindingException] {
        throw [VirtualMachineScaleSetDoesNotExist]::new($VMssObject, $_.Exception)
    }

    if (!($PSCmdlet.ShouldProcess($vmsslogheader, "Install Data Collection Rule Association"))) {
        #-WhatIf skip processing here, return to the caller as we have completed our work.
        if ($WhatIfPreference) {
            return
        }
        throw [CustomerSkip]::new()
    }

    #Customer can associate multiple DCRs to a VMSS.
    $dcrAssociationName = "VM-Insights-DCR-Association-$(New-Guid)"
    Write-Host "$vmsslogheader : Deploying Data Collection Rule Association $dcrAssociationName"
    try {
        #BUG! To be fixed by AMCS team - New-AzDataCollectionRuleAssociation allows creating multiple DCRAs for the same {DCR,VMSS} combination.
        $dcrAssociation = New-AzDataCollectionRuleAssociation -TargetResourceId $vmssId `
                                                              -AssociationName $dcrAssociationName `
                                                              -RuleId $DcrResourceId `
                                                              -Confirm:$false
    } catch [System.Management.Automation.PSInvalidOperationException] {
        $exceptionMessage = $_.Exception.Message
        
        if ($exceptionMessage.Contains('Invalid format of the resource identifier')) {
            throw [DataCollectionRuleIncorrect]::new($DcrResourceId, $_.Exception)
        } 
        
        if (!($exceptionMessage -match "status code '([^\s]+)'")) {
            throw [VirtualMachineScaleSetUnknownException]::new($VMssObject, "Failed to create Data Collection Rule Association with $DcrResourceId", $_.Exception)
        }

        $statusCode = $matches[1]
        if ($statusCode -eq 'BadRequest') {
            throw [DataCollectionRuleDoesNotExist]::new($DcrResourceId, $_.Exception)
        }
        if ($statusCode -eq 'NotFound') {
            throw [VirtualMachineScaleSetDoesNotExist]::new($VMssObject, $_.Exception)
        }
        if ($statusCode -eq 'Forbidden') {
            throw [DataCollectionRuleForbidden]::new($DcrResourceId, $_.Exception)     
        }

        throw [VirtualMachineUnknownException]::new($VMssObject, "Failed to create Data Collection Rule Association with with $DcrResourceId. StatusCode = $statusCode", $_.Exception)
    }
    #Tmp fix task :- 21191002
    if (($null -eq $dcrAssociation) -or ($dcrAssociation -is [Microsoft.Azure.Management.Monitor.Models.ErrorResponseCommonV2Exception])) {
        throw [VirtualMachineScaleSetDoesNotExist]::new($VMssObject, "Unable to create Data Collection Rule Association with $DcrResourceId", $_.Exception)
    }
}

function RetainExtensionUpgradeSettings {
    <#
    .SYNOPSIS
    Retain extension upgrade settings
    #>
    param
    (
        [Parameter(mandatory = $True)]
        [Microsoft.Azure.Commands.Compute.Models.PSVirtualMachineExtension]
        $Extension,
        [Parameter(mandatory = $True)]
        [hashtable]
        $ExtUpgradeSettings
    )

    $ExtUpgradeSettings.DisableAutoUpgradeMinorVersion = !($Extension.AutoUpgradeMinorVersion)
    $ExtUpgradeSettings.EnableAutomaticUpgrade = $Extension.EnableAutomaticUpgrade
}

function OnboardAmaVm {
    <#
    .SYNOPSIS
    Onboard AMA on VM, handling if already installed.
    #>
    param
    (
        [Parameter(mandatory = $True)]
        [Microsoft.Azure.Commands.Compute.Models.PSVirtualMachine]
        $VMObject,
        [Parameter(mandatory = $True)]
        [hashtable]
        $ExtensionSettings
    )
    
    $amaExtensionConstantProperties = $amaExtensionConstantMap[$VMObject.StorageProfile.OsDisk.OsType.ToString()]
    $extensionName = $amaDefaultExtensionName
    $extUpgradeSettings = $extensionVmDefaultUpgradeSettings.clone()
    $extension = GetVMExtension -VMObject $VMObject -ExtensionProperties $amaExtensionConstantProperties
    # Use supplied name unless already deployed, use same name.
    if ($extension) {
        $extensionName = $extension.Name
        RetainExtensionUpgradeSettings -Extension $extension -ExtUpgradeSettings $extUpgradeSettings
        Write-Host "$(FormatVmIdentifier -VMObject $VMObject) : Extension $extensionName, type = $($amaExtensionConstantProperties.Publisher).$($amaExtensionConstantProperties.ExtensionType) already installed. Provisioning State : $($extension.ProvisioningState)"
    }
    
    return SetVMExtension -VMObject $VMObject `
                          -ExtensionName $extensionName `
                          @amaExtensionConstantProperties `
                          -ExtensionUpgradeSettings $extUpgradeSettings `
                          -ExtensionSettings $ExtensionSettings
}

function OnboardVmiWithAmaVm {
    <#
    .SYNOPSIS
    Onboard VMI with AMA on VM.
    #>
    param
    (
        [Parameter(mandatory = $True)]
        [Microsoft.Azure.Commands.Compute.Models.PSVirtualMachine]
        $VMObject,
        [Parameter(mandatory = $True)]
        [hashtable]
        $ExtensionSettings
    )

    $VMObject = AssignVmUserManagedIdentity -VMObject $VMObject
    NewDCRAssociationVm -VMObject $VMObject
    return OnboardAmaVm -VMObject $VMObject -ExtensionSettings $ExtensionSettings
}

function OnboardVmiWithAmaVmss {
    <#
    .SYNOPSIS
    Onboard VMI with AMA on VMSS
    #>
    param
    (
        [Parameter(mandatory = $True)]
        [Microsoft.Azure.Commands.Compute.Automation.Models.PSVirtualMachineScaleSet]
        $VMssObject,
        [Parameter(mandatory = $True)]
        [hashtable]
        $ExtensionSettings
    )
            
    $VMssObject = AssignVmssUserManagedIdentity -VMssObject $VMssObject
    NewDCRAssociationVmss -VMssObject $VMssObject
    return OnboardVMssExtension -VMssObject $VMssObject `
                                -ExtensionName $amaDefaultExtensionName `
                                -ExtensionConstantMap $amaExtensionConstantMap `
                                -ExtensionSettings $ExtensionSettings
}

function SetManagedIdentityRoles {
    <#
    .SYNOPSIS
    Set roles to User Assigned Managed Identity.
    #>
    [CmdletBinding(SupportsShouldProcess = $True, ConfirmImpact = 'Medium')]
    param
    (
        [Parameter(Mandatory = $True)]
        [String]
        $ResourceGroupId,
        [Parameter(Mandatory = $True)]
        [object]
        $UserAssignedManagedIdentity,
        [Parameter(Mandatory = $True)]
        [String[]]
        $Roles
    )

    $uamiName = $UserAssignedManagedIdentity.Name
    if (!($PSCmdlet.ShouldProcess($ResourceGroupId, "Assign $Roles to User Assigned Managed Identity $uamiName"))) {
        #-WhatIf skip processing here, return to the caller as we have completed our work.
        if ($WhatIfPreference) {
            return
        }
        throw [CustomerSkip]::new()
    }

    foreach ($role in $Roles) {
        Write-Verbose "Assigning role $role"
        try {
            New-AzRoleAssignment -ObjectId $($UserAssignedManagedIdentity.principalId) `
                                 -RoleDefinitionName $role `
                                 -Scope $ResourceGroupId `
                                 -Confirm:$false
            Write-Verbose "$ResourceGroupId : $role has been successfully assigned to $uamiName"
        } catch {
            $excepMessage = $_.Exception.Message
            if ($excepMessage.Contains('Conflict')) {
                Write-Verbose "$uamiName : Role $role has been assigned already"
            }
            if ($excepMessage.Contains('BadRequest')) {
                throw [FatalException]::new("$uamiName : User Assigned Managed Identity doesn't exist", $_.Exception)
            }
            if ($excepMessage.Contains('NotFound')) {
                throw [ResourceGroupDoesNotExist]::new($($VMObject.ResourceGroupName), $_.Exception)
            }
        }
        
    }
}

function OnboardVMssExtension {
    <#
    .SYNOPSIS
    Install/Update Extension VMSS, handling if already installed
    #>
    #check if this supposed to be all the stack. that is functioing call it as well.
    [CmdletBinding(SupportsShouldProcess = $True, ConfirmImpact = 'Medium')]
    param
    (
        [Parameter(Mandatory = $True)]
        [Microsoft.Azure.Commands.Compute.Automation.Models.PSVirtualMachineScaleSet]
        $VMssObject,
        [Parameter(mandatory = $True)]
        [String]
        $ExtensionName,
        [Parameter(mandatory = $True)]
        [Hashtable]
        $ExtensionConstantMap,
        [Parameter(mandatory = $True)]
        [Hashtable]
        $ExtensionSettings
    )

    $extensionConstantProperties = $ExtensionConstantMap[$VMssObject.VirtualMachineProfile.StorageProfile.OsDisk.OsType.ToString()]
    $extensionType = $extensionConstantProperties.ExtensionType
    $publisher = $extensionConstantProperties.Publisher
    $typeHandlerVersion = $extensionConstantProperties.TypeHandlerVersion
    $extensionUpgradeSettings = @{
        AutoUpgradeMinorVersion = $True
        EnableAutomaticUpgrade = $True 
    }
    $extension = GetVMssExtension -VMssObject $VMssObject -ExtensionType $extensionType -Publisher $publisher
    $vmsslogheader = FormatVmssIdentifier -VMssObject $VMssObject

    if ($extension) {
        if (!($PSCmdlet.ShouldProcess($vmsslogheader, "Update extension $($extension.Name), type = $($publisher).$($extensionType)"))) {
            #-WhatIf skip processing here, return to the caller as we have completed our work.
            if ($WhatIfPreference) {
                return $VMssObject
            }
            throw [CustomerSkip]::new()
        }
        Write-Host "$vmsslogheader : Extension $($extension.Name), type = $($publisher).$($extensionType) already installed."
        #VMSS extension has field "Settings" but Add-AzVmssExtension supports 'Setting'
        if ($ExtensionSettings.ContainsKey("Setting")) {
            $extension.Settings = $ExtensionSettings.Setting
        }
        #VMSS extension has field "ProtectedSettings" but Add-AzVmssExtension supports 'ProtectedSetting'
        if ($ExtensionSettings.ContainsKey("ProtectedSetting")) {
            $extension.ProtectedSettings = $ExtensionSettings.ProtectedSetting
        }
        $extensionUpgradeSettings.GetEnumerator() | ForEach-Object { $extension.($_.Key) = $_.Value }
        $extension.TypeHandlerVersion = $typeHandlerVersion
        return $VMssObject
    }

    if (!($PSCmdlet.ShouldProcess($vmsslogheader, "Install extension $ExtensionName, type = $($publisher).$($extensionType)"))) {
        #-WhatIf skip processing here, return to the caller as we have completed our work.
        if ($WhatIfPreference) {
            return $VMssObject
        }
        throw [CustomerSkip]::new()
    }
    
    $VMssObject = Add-AzVmssExtension -VirtualMachineScaleSet $VMssObject `
                                      -Name $ExtensionName `
                                      -Type $extensionType `
                                      -Publisher $publisher `
                                      -TypeHandlerVersion $typeHandlerVersion `
                                      @extensionUpgradeSettings `
                                      @ExtensionSettings `
                                      -Confirm:$false


    Write-Host "$vmsslogheader : Extension $ExtensionName, type = $($publisher).$($extensionType) added."
    return $VMssObject
}

function SetVMExtension {
    <#
    .SYNOPSIS
    Install/Update VM Extension, handling if already installed
    #>
    [CmdletBinding(SupportsShouldProcess = $True, ConfirmImpact = 'Medium')]
    param
    (
        [Parameter(mandatory = $True)]
        [Microsoft.Azure.Commands.Compute.Models.PSVirtualMachine]
        $VMObject,
        [Parameter(mandatory = $True)]
        [String]
        $ExtensionName,
        [Parameter(mandatory = $True)]
        [String]
        $ExtensionType,
        [Parameter(mandatory = $True)]
        [String]
        $Publisher,
        [Parameter(mandatory = $True)]
        [String]
        $TypeHandlerVersion,
        [Parameter(mandatory = $True)]
        [Hashtable]
        $ExtensionUpgradeSettings,
        [Parameter(mandatory = $True)]
        [Hashtable]
        $ExtensionSettings
    )
    
    $vmlogheader = $(FormatVmIdentifier -VMObject $VMObject)
    
    if (!($PSCmdlet.ShouldProcess($vmlogheader, "Install/Update extension $ExtensionName, type = $($Publisher).$($ExtensionType)"))) {
        #-WhatIf skip processing here, return to the caller as we have completed our work.
        if ($WhatIfPreference) {
            return $VMObject
        }
        throw [CustomerSkip]::new()
    }

    Write-Host "$vmlogheader : Installing/Updating extension $ExtensionName, type = $($Publisher).$($ExtensionType)"
    
    try {
        $result = Set-AzVMExtension -ResourceGroupName $($VMObject.ResourceGroupName) `
                                    -VMName $($VMObject.Name) `
                                    -Name $ExtensionName `
                                    -ExtensionType $ExtensionType `
                                    -Publisher $Publisher `
                                    -TypeHandlerVersion $TypeHandlerVersion `
                                    @ExtensionUpgradeSettings `
                                    @ExtensionSettings -ForceRerun $True `
                                    -Confirm:$false

        if (!$result.IsSuccessStatusCode) {
            throw [VirtualMachineOperationFailed]::new($VMObject, "Failed to update extension. StatusCode = $($result.StatusCode). ReasonPhrase = $($result.ReasonPhrase)")
        }
    
        Write-Host "$vmlogheader : Successfully installed/updated extension $ExtensionName, type = $($Publisher).$($ExtensionType)"
        return $VMObject
    } catch [Microsoft.Azure.Commands.Compute.Common.ComputeCloudException] {
        $errorMessage = ExtractCloudExceptionErrorMessage -ErrorRecord $_
        $errorCode = ExtractExceptionErrorCode -ErrorRecord $_
        #Found by experimentation.
        if ($errorCode -eq "OperationNotAllowed" -and $errorMessage -eq "Cannot modify extensions in the VM when the VM is not running.") {
            throw [VirtualMachinePoweredDown]::new($VMObject, $_.Exception)
        }
        
        if ($errorCode -eq "ParentResourceNotFound") {
            throw [VirtualMachineDoesNotExist]::new($VMObject, $_.Exception)
        } 
        
        if ($errorCode -eq "ResourceGroupNotFound") {
            throw [ResourceGroupDoesNotExist]::new($VMObject.ResourceGroupName, $_.Exception)       
        } 
        
        if ($errorCode -eq "VMExtensionHandlerNonTransientError") {
            throw [VirtualMachineExtensionError]::new($VMObject, "$($Publisher).$($ExtensionType)", $_.Exception)       
        }

        throw [VirtualMachineUnknownException]::new($VMObject, "Failed to install/update extension $ExtensionName, type = $($Publisher).$($ExtensionType)", $_.Exception)
    }
}

function UpgradeVmssExtensionManualUpdateEnabled {
    <#
    .SYNOPSIS
    Upgrade VMss Extension with manual update.
    #>
    [CmdletBinding(SupportsShouldProcess = $True, ConfirmImpact = 'Medium')]
    param
    (
        [Parameter(mandatory = $True)]
        [Microsoft.Azure.Commands.Compute.Automation.Models.PSVirtualMachineScaleSet]
        $VMssObject,
        [Parameter(mandatory = $True)]
        [ref]
        $InstanceUpgradeFailCounter
    )
    
    $vmssResourceGroupName = $VMssObject.ResourceGroupName
    $vmssName = $VMssObject.Name
    $vmsslogheader = FormatVmssIdentifier -VMssObject $VMssObject
    
    $scaleSetInstances = @()
    try {
        $scaleSetInstances = Get-AzVmssVm -ResourceGroupName $vmssResourceGroupName -VMScaleSetName $vmssName -InstanceView
    } catch [Microsoft.Azure.Commands.Compute.Common.ComputeCloudException] {
        $errorCode = ExtractExceptionErrorCode -ErrorRecord $_
        if ($errorCode -eq "ParentResourceNotFound") {
            throw [VirtualMachineScaleSetDoesNotExist]::new($VMssObject, $_.Exception)
        }
        if ($errorCode -eq "ResourceGroupNotFound") {
            throw [ResourceGroupDoesNotExist]::new($VMssObject.ResourceGroupName, $_.Exception)       
        } 
        
        throw [VirtualMachineScaleSetUnknownException]::new($VMssObject, "Failed to upgrade VMSS", $_.Exception)
    }

    $i = 0
    $instanceCount = $scaleSetInstances.Count
    $unexpectedUpgradeExceptionCounter = 0
    $unexpectedUpgradeExceptionLimit = 5 
    
    Foreach ($scaleSetInstance in $scaleSetInstances) {
        $i++
        $healthstatus = $scaleSetInstance.InstanceView.Statuses.Code

        if (!($healthstatus -contains 'ProvisioningState/succeeded' -and $healthstatus -contains 'PowerState/running')) {
            Write-Host "VMSS instance $scaleSetInstanceName, $i of $instanceCount is not operational. State = $($scaleSetInstance.InstanceView.Statuses.DisplayStatus)"
            Write-Host "Continuing ..."
            $InstanceUpgradeFailCounter.Value += 1
            continue
        }

        $scaleSetInstanceName = $($scaleSetInstance.Name)
        Write-Host "$vmsslogheader : Upgrading $scaleSetInstanceName, $i of $instanceCount"

        if ($scaleSetInstance.LatestModelApplied) {
            Write-Verbose "$vmsslogheader : Latest model already applied for $scaleSetInstanceName, $i of $instanceCount"
            continue
        }
        
        Write-Verbose "$vmsslogheader : Upgrading VMSS instance name $scaleSetInstanceName, $i of $instanceCount"
        try {
            $result = Update-AzVmssInstance -ResourceGroupName $vmssResourceGroupName `
                                            -VMScaleSetName $vmssName `
                                            -InstanceId $scaleSetInstance.InstanceId `
                                            -Confirm:$false
            if ($result.Status -ne "Succeeded") {
                Write-Host "$vmsslogheader : Failed to upgrade VMSS instance name $scaleSetInstanceName, $i of $instanceCount. $($result.Status)"
                $InstanceUpgradeFailCounter.Value += 1
            } else {
                Write-Verbose "$vmsslogheader : Upgrade VMSS instance name $scaleSetInstanceName, $i of $instanceCount"
            }
        } catch [Microsoft.Azure.Commands.Compute.Common.ComputeCloudException] {
            $errorCode = ExtractExceptionErrorCode -ErrorRecord $_
            if ($errorCode -eq "ResourceNotFound") { 
                throw [VirtualMachineScaleSetDoesNotExist]::new($VMssObject, $_.Exception)
            }

            if ($errorCode -eq "ResourceGroupNotFound") {
                throw [ResourceGroupDoesNotExist]::new($VMssObject.ResourceGroupName ,$_.Exception)  
            }

            if ($unexpectedUpgradeExceptionCounter -gt $unexpectedUpgradeExceptionLimit) {
                throw [VirtualMachineScaleSetUnknownException]::new($VMssObject, "More than $unexpectedUpgradeExceptionLimit unexpected exceptions encountered", $_.Exception)
            }

            if ($errorCode -eq "OperationNotAllowed") {
                Write-Host "$vmsslogheader : Failed to locate VMSS instance name $scaleSetInstanceName. Continuing ..."
            } elseif ($errorCode -eq "VMExtensionHandlerNonTransientError") {
                Write-Host "$vmsslogheader : Failed to upgrade extension for VMSS instance name $scaleSetInstanceName, $i of $instanceCount."
                VerboseDisplayException -ErrorRecord $_
                Write-Host "Continuing ..."
            } else {
                Write-Host "$vmsslogheader : Failed to upgrade VMSS instance name $scaleSetInstanceName, $i of $instanceCount. ErrorCode $errorCode." 
                DisplayException -ErrorRecord $_
                Write-Host "Continuing ..."
                $unexpectedUpgradeExceptionCounter += 1
            }
            
            $InstanceUpgradeFailCounter.Value += 1
        }
    }
}

function UpdateVMssExtension {
    <#
    .SYNOPSIS
    Update VMss Extension
    #>
    [CmdletBinding(SupportsShouldProcess = $True, ConfirmImpact = 'Medium')]
    param
    (
        [Parameter(mandatory = $True)]
        [Microsoft.Azure.Commands.Compute.Automation.Models.PSVirtualMachineScaleSet]
        $VMssObject
    )

    $vmsslogheader = FormatVMssIdentifier -VMssObject $VMssObject

    if (!($PSCmdlet.ShouldProcess($vmsslogheader, "Update VMSS"))) {
        #-WhatIf skip processing here, return to the caller as we have completed our work.
        if ($WhatIfPreference) {
            return $VMssObject
        }
        throw [CustomerSkip]::new()
    }

    Write-Host "$vmsslogheader : Updating VMSS"
    try {
        $VMssObject = Update-AzVmss -VMScaleSetName $VMssObject.Name `
                                    -ResourceGroupName $VMssObject.ResourceGroupName `
                                    -VirtualMachineScaleSet $VMssObject `
                                    -Confirm:$false
        if ($VMssObject.ProvisioningState -ne "Succeeded") {
            throw [VirtualMachineScaleSetOperationFailed]::new($VMssObject, "Failed to update VMSS")
        }
        Write-Host "$vmsslogheader : Successfully updated scale set with extension"
        return $VMssObject                           
    } catch [Microsoft.Azure.Commands.Compute.Common.ComputeCloudException] {
        $errorCode = ExtractExceptionErrorCode -ErrorRecord $_
        if ($errorCode -eq "ParentResourceNotFound") {
            throw [VirtualMachineScaleSetDoesNotExist]::new($VMssObject, $_.Exception)
        } 
        
        if ($errorCode -eq "ResourceGroupNotFound") {
            throw [ResourceGroupDoesNotExist]::new($($VMssObject.ResourceGroupName), $_.Exception)       
        }
        
        if ($errorCode -eq "VMExtensionHandlerNonTransientError") {
            throw [VirtualMachineScaleSetExtensionError]::new($VMssObject, $_.Exception)   
        }

        throw [VirtualMachineScaleSetUnknownException]::new($VMssObject, "Failed to update VMSS", $_.Exception)
    }
}

function Get-MissingUserAssignedIdentities {
    <#
    .SYNOPSIS
    Validates that all User Assigned Managed Identities in the list still exist.
    Returns resource IDs of UAMIs that are referenced but have been deleted.
    #>

    [CmdletBinding()]

    param(
        [Parameter(Mandatory = $True)]
        [string[]]$IdentityIds
    )
    
    $missingUamis = @()

    foreach ($uamiId in $IdentityIds) {
        try {
            $uamiParts = $uamiId -split '/'
            if ($uamiParts.Count -lt 9) {
                Write-Verbose "Invalid UAMI resource ID format: $uamiId"
                continue
            }
            
            $uamiResourceGroup = $uamiParts[4]
            $uamiName = $uamiParts[8]
            
            # Try to get the UAMI - if it doesn't exist, this will throw
            $null = Get-AzUserAssignedIdentity -ResourceGroupName $uamiResourceGroup -Name $uamiName -ErrorAction Stop
        } catch {
            $errorCode = ExtractExceptionPrefixErrCode -ErrorRecord $_
            if ($errorCode -eq "ResourceNotFound") {
                Write-Verbose "UAMI no longer exists: $uamiId"
                $missingUamis += $uamiId
            } else {
                # For other errors (throttling, network, permission), log and skip validation for this UAMI
                # Re-throwing would fail the entire operation; skipping allows the Update call to fail with its own error
                Write-Verbose "Unable to validate UAMI $uamiId (error: $errorCode), skipping validation"
            }
        }

    }
    
    return $missingUamis
}

function Format-UamiForDisplay {
    <#
    .SYNOPSIS
    Formats a UAMI resource ID for user-friendly display.
    Returns format: "name (RG: resourceGroup)" which is unique and readable.
    #>
    [CmdletBinding()]
    param(
        [Parameter(Mandatory = $True)]
        [string]$UamiId
    )
    
    try {
        $parts = $UamiId -split '/'
        if ($parts.Count -ge 9) {
            $resourceGroup = $parts[4]
            $uamiName = $parts[8]
            return "$uamiName (RG: $resourceGroup)"
        }
    } catch {
        Write-Verbose "Failed to parse UAMI ID for display: $UamiId"
    }
    
    # Fallback: return full ID if parsing fails
    return $UamiId
}

function Handle-FailedIdentityOperation {
    <#
    .SYNOPSIS
    Handles FailedIdentityOperation errors by validating which UAMIs are missing.
    Returns user-friendly error details for VMSS or VM operations.
    #>
    [CmdletBinding()]
    param(
        [Parameter(Mandatory = $True)]
        [string]$ResourceType,  # "VMSS" or "VM"
        
        [Parameter(Mandatory = $True)]
        [string]$LogHeader,
        
        [Parameter(Mandatory = $True)]
        [string]$NewUamiId,
        
        [Parameter(Mandatory = $True)]
        [string]$NewUamiName,
        
        [Parameter(Mandatory = $False)]
        [System.Collections.Generic.Dictionary[string,Microsoft.Azure.Management.Compute.Models.UserAssignedIdentitiesValue]]$ExistingUamis,
        
        [Parameter(Mandatory = $True)]
        [System.Management.Automation.ErrorRecord]$ErrorRecord
    )
    
    # First, check if the NEW UAMI being added has been deleted
    $missingNewUami = Get-MissingUserAssignedIdentities -IdentityIds @($NewUamiId)
    if ($missingNewUami.Count -gt 0) {
        Write-Host "$LogHeader : Cannot assign User Assigned Managed Identity '$NewUamiName' because it no longer exists. The UAMI may have been deleted after the script started. Please verify the UAMI exists and try again."
        throw [UserAssignedManagedIdentityDoesNotExist]::new($NewUamiName, $ErrorRecord.Exception)
    }
    
    # Then, check if any EXISTING UAMIs on the resource have been deleted
    if ($ExistingUamis) {
        $missingExistingUamis = Get-MissingUserAssignedIdentities -IdentityIds $ExistingUamis.Keys
        
        if ($missingExistingUamis.Count -gt 0) {
            # Format UAMIs for readable display: "name (RG: resourceGroup)"
            $formattedUamis = $missingExistingUamis | ForEach-Object { Format-UamiForDisplay $_ }
            $missingUamisList = $formattedUamis -join ', '
            
            # Resource-specific documentation links
            $docLink = if ($ResourceType -eq "VMSS") {
                "https://learn.microsoft.com/entra/identity/managed-identities-azure-resources/how-to-configure-managed-identities-scale-sets"
            } else {
                "https://learn.microsoft.com/entra/identity/managed-identities-azure-resources/how-manage-user-assigned-managed-identities"
            }
            
            Write-Host "$LogHeader : Cannot assign User Assigned Managed Identity because the $ResourceType has invalid UAMI associations. The following UAMIs are associated with this $ResourceType but no longer exist: $missingUamisList. Please disassociate the deleted UAMIs from the $ResourceType before onboarding. See: $docLink"
            
            # Format exception message for plural UAMIs to avoid awkward "uami1, uami2 : User Assigned Managed Identity does not exist."
            $exceptionUamiName = if ($missingExistingUamis.Count -eq 1) { $missingExistingUamis[0] } else { "multiple UAMIs ($missingUamisList)" }
            throw [UserAssignedManagedIdentityDoesNotExist]::new($exceptionUamiName, $ErrorRecord.Exception)
        }
    }
    
    # If we get here, the FailedIdentityOperation error occurred for an unknown reason
    # This could be a transient issue or a permission problem we couldn't diagnose
    throw [UserAssignedManagedIdentityUnknownException]::new("$NewUamiName : FailedIdentityOperation error occurred for an unknown reason. This may be a transient issue or a permission problem. Please try again or verify that the script has sufficient permissions to assign identities on this $ResourceType.", $ErrorRecord.Exception)
}

function AssignVmssUserManagedIdentity {
    <#
    .SYNOPSIS
    Checking if User Assigned Managed Identity is already assigned to VMSS, if not assigning it
    #>
    [CmdletBinding(SupportsShouldProcess = $True, ConfirmImpact = 'Medium')]
    param
    (
        [Parameter(Mandatory = $True)]
        [Microsoft.Azure.Commands.Compute.Automation.Models.PSVirtualMachineScaleSet]
        $VMssObject
    )

    $userAssignedManagedIdentityName = $UserAssignedManagedIdentityObject.Name
    $vmsslogheader = FormatVmssIdentifier -VMssObject $VmssObject
    $vmssResourceGroupName = $VMssObject.ResourceGroupName
                                  
    if ($VMssObject.Identity.UserAssignedIdentities.Keys -contains $UserAssignedManagedIdentityObject.Id) {
        Write-Host "$vmsslogheader : User Assigned Managed Identity $userAssignedManagedIdentityName already assigned."
        return $VMssObject
    }

    if (!($PSCmdlet.ShouldProcess($vmsslogheader, "Assign User Assigned Managed Identity $userAssignedManagedIdentityName"))) {
        #-WhatIf skip processing here, return to the caller as we have completed our work.
        if ($WhatIfPreference) {
            return $VMssObject
        }
        throw [CustomerSkip]::new()
    }

    $identityType = if ("$($VMssObject.Identity.Type)".StartsWith("SystemAssigned")) { "SystemAssignedUserAssigned" } else { "UserAssigned" }
    $identityList = [System.Collections.Generic.List[string]]@($UserAssignedManagedIdentityObject.Id)
    if ($VMssObject.Identity.UserAssignedIdentities) {
        $identityList += $VMssObject.Identity.UserAssignedIdentities.Keys
    }

    Write-Host "$vmsslogheader : Assigning User Assigned Managed Identity $userAssignedManagedIdentityName"
    try {
        $VMssObject = Update-AzVmss -VMScaleSetName $VMssObject.Name `
                                -ResourceGroupName  $vmssResourceGroupName `
                                -VirtualMachineScaleSet $VMssObject `
                                -IdentityType $identityType `
                                -IdentityID $identityList `
                                -Confirm:$false
        if ($VMssObject.ProvisioningState -ne "Succeeded") {
            throw [VirtualMachineOperationScaleSetFailed]::new($VMssObject, "Failed to assign User Assigned Managed Identity $userAssignedManagedIdentityName")
        }
        
        Write-Host "$vmsslogheader : Successfully assigned User Assigned Managed Identity $userAssignedManagedIdentityName"
        return $VMssObject
    } catch [Microsoft.Azure.Commands.Compute.Common.ComputeCloudException] {
        $errorCode = ExtractExceptionErrorCode -ErrorRecord $_
        if ($errorCode -eq "FailedIdentityOperation") {
            Handle-FailedIdentityOperation `
                -ResourceType "VMSS" `
                -LogHeader $vmsslogheader `
                -NewUamiId $UserAssignedManagedIdentityObject.Id `
                -NewUamiName $userAssignedManagedIdentityName `
                -ExistingUamis $VMssObject.Identity.UserAssignedIdentities `
                -ErrorRecord $_
        }
        if ($errorCode -eq "ResourceGroupNotFound") {
            throw [ResourceGroupDoesNotExist]::new($vmssResourceGroupName, $_.Exception)       
        } 
        if ($errorCode -eq "InvalidParameter") {
            throw [VirtualMachineScaleSetDoesNotExist]::new($VMssObject, $_.Exception) 
        } 
            
        throw [VirtualMachineScaleSetUnknownException]::new($VMssObject, "Failed to assign User Assigned Managed Identity $userAssignedManagedIdentityName", $_.Exception)
    }
}

function AssignVmUserManagedIdentity {
     <#
    .SYNOPSIS
    Checking if User Assigned Managed Identity is already assigned to VM, if not assigning it
    #>
    [CmdletBinding(SupportsShouldProcess = $True, ConfirmImpact = 'Medium')]
    param
    (
        [Parameter(mandatory = $True)]
        [Microsoft.Azure.Commands.Compute.Models.PSVirtualMachine]
        $VMObject
    )

    $userAssignedManagedIdentityName = $UserAssignedManagedIdentityObject.Name
    $vmlogheader = FormatVmIdentifier -VMObject $VMObject
    $vmResourceGroupName = $VMObject.ResourceGroupName 
    
    if ($VMObject.Identity.UserAssignedIdentities.Keys -contains $UserAssignedManagedIdentityObject.Id) {
        Write-Host "$vmlogheader : User Assigned Managed Identity $userAssignedManagedIdentityName already assigned."
        return $VMObject
    }

    if (!($PSCmdlet.ShouldProcess($vmlogheader, "Assign User Assigned Managed Identity $userAssignedManagedIdentityName"))) {
        #-WhatIf skip processing here, return to the caller as we have completed our work.
        if ($WhatIfPreference) {
            return $VMObject
        } 
        throw [CustomerSkip]::new()
    }

    $identityType = if ("$($VMObject.Identity.Type)".StartsWith("SystemAssigned")) { "SystemAssignedUserAssigned" } else { "UserAssigned" }
    $identityList = [System.Collections.Generic.List[string]]@($UserAssignedManagedIdentityObject.Id)
    if ($VMObject.Identity.UserAssignedIdentities) {
        $identityList += $VMObject.Identity.UserAssignedIdentities.Keys
    }

    Write-Host "$vmlogheader : Assigning User Assigned Managed Identity $userAssignedManagedIdentityName"

    try {
        $result = Update-AzVM -VM $VMObject `
                                -ResourceGroupName $vmResourceGroupName `
                                -IdentityType $identityType `
                                -IdentityID $identityList `
                                -Confirm:$false
    } catch [Microsoft.Azure.Commands.Compute.Common.ComputeCloudException] {
        $errorCode = ExtractExceptionErrorCode -ErrorRecord $_
        if ($errorCode -eq "FailedIdentityOperation") {
            Handle-FailedIdentityOperation `
                -ResourceType "VM" `
                -LogHeader $vmlogheader `
                -NewUamiId $UserAssignedManagedIdentityObject.Id `
                -NewUamiName $userAssignedManagedIdentityName `
                -ExistingUamis $VMObject.Identity.UserAssignedIdentities `
                -ErrorRecord $_
        }
        if ($errorCode -eq "ResourceGroupNotFound") {
            throw [ResourceGroupDoesNotExist]::new($vmResourceGroupName, $_.Exception)       
        }
        if ($errorCode -eq "InvalidParameter") {
            throw [VirtualMachineDoesNotExist]::new($VMObject, $_.Exception)
        }
        
        throw [VirtualMachineUnknownException]::new($VMObject, "Failed to update VM with $userAssignedManagedIdentityName", $_.Exception)
    }

    if (!($result.IsSuccessStatusCode)) {
        throw [VirtualMachineOperationFailed]::new($VMObject, "Failed to assign User Assigned Managed Identity $userAssignedManagedIdentityName. StatusCode : $($result.StatusCode). ReasonPhrase : $($result.ReasonPhrase)")
    }
    
    Write-Host "$vmlogheader : Successfully assigned User Assigned Managed Identity $userAssignedManagedIdentityName"
    return $VMObject
}

function SetManagedIdentityRolesAma {
    <#
    .SYNOPSIS
    Set Roles to a User Assigned Managed Identity
    #>
    param(
        [Parameter(Mandatory = $True)]
        [String]
        $ResourceGroupName
    )
    
    try { 
        $rgObj = Get-AzResourceGroup -Name $ResourceGroupName
    } catch { 
        $errorCode = ExtractExceptionErrorCode -ErrorRecord $_
        if ($errorCode -eq "ResourceGroupNotFound") {
            throw [FatalException]::new("$ResourceGroupName : Not found.", $_.Exception)
        }
         
        throw [FatalException]::new("Unable to lookup Resource Group.", $_.Exception)
    }

    $roles = @("Virtual Machine Contributor", "Azure Connected Machine Resource Administrator", "Log Analytics Contributor")
    Write-Host "($ResourceGroupName) : Assigning roles"
    SetManagedIdentityRoles -ResourceGroupId $rgObj.ResourceId `
                            -UserAssignedManagedIdentity $UserAssignedManagedIdentityObject `
                            -Roles $roles
}

#
# Main Script
#
#
#
try {
    # To report on overall status
    $onboardingCounters = [OnboardingCounters]::new()
    $unknownExceptionVirtualMachineScaleSetConsequentCounter = 0
    $unknownExceptionVirtualMachineConsequentCounter = 0
    $unknownExceptionTotalCounter = 0

    if ($ProcessAndDependencies) {
        throw [FatalException]::new("The -ProcessAndDependencies flag is no longer supported. VM Insights Map and the Dependency Agent are being retired on June 30th, 2028. Please see our retirement guidance for more details: https://aka.ms/DependencyAgentRetirement.", $null)
    }

    if ($PSCmdlet.ParameterSetName -eq 'LogAnalyticsAgent') {
        throw [FatalException]::new("The -WorkspaceId, -WorkspaceKey, and -ReInstall parameters are no longer supported. The Log Analytics agent was retired on August 31, 2024. Please use Azure Monitor Agent (AMA) onboarding instead. See: https://aka.ms/MigrateToAMA", $null)
    }

    # First make sure we are authenticed and Select the subscription supplied and input parameters are valid.
    $account =  Get-AzContext
    if ($null -eq $account.Account) {
        Write-Host "Account Context not found, please login"
        Connect-AzAccount -WhatIf:$False -Confirm:$False -subscriptionid $SubscriptionId
    }
    else {
        if ($account.Subscription.Id -eq $SubscriptionId) {
            Write-Verbose "Subscription : $SubscriptionId is already selected."
            $account
        }
        else {
            Write-Host "Current Subscription :"
            $account
            Write-Host "Changing to subscription : $SubscriptionId"
            Select-AzSubscription -WhatIf:$False -Confirm:$False -SubscriptionId $SubscriptionId
        }
    }

    #script block
    Set-Variable -WhatIf:$False -Confirm:$False -Name sb_nop_block_upgrade -Option Constant -Value { `
        param([Microsoft.Azure.Commands.Compute.Automation.Models.PSVirtualMachineScaleSet]$vmssObj, [ref]$instanceUpgradeFailCounter)
    }
    
    $Rghashtable = @{}
    
    try {
        Write-Verbose "Validating ($UserAssignedManagedIdentityResourceGroup, $UserAssignedManagedIdentityName)"
        Set-Variable -WhatIf:$False -Confirm:$False -Name UserAssignedManagedIdentityObject -Option Constant -Value `
                        (Get-AzUserAssignedIdentity -Name $UserAssignedManagedIdentityName `
                                                    -ResourceGroupName $UserAssignedManagedIdentityResourceGroup `
                                                    -ErrorAction Stop)
    } catch {
        $errorCode = ExtractExceptionPrefixErrCode -ErrorRecord $_
        if ($errorCode -eq "ResourceNotFound") {
            throw [UserAssignedManagedIdentityDoesNotExist]::new($UserAssignedManagedIdentityName, $_.Exception)
        }
        
        if ($errorCode -eq "ResourceGroupNotFound") {
            throw [UserAssignedManagedIdentityResourceGroupDoesNotExist]::new($UserAssignedManagedIdentityResourceGroup, $_.Exception)
        }

        throw [UserAssignedManagedIdentityUnknownException]::new("($UserAssignedManagedIdentityResourceGroup) $UserAssignedManagedIdentityName : Failed to locate User Assigned Managed Identity.", $_.Exception)
    }
    
    $local:amaSettings = @{
        'authentication' = @{ 
            'managedIdentity' = @{
                'identifier-name' = 'mi_res_id'
                'identifier-value' = $($UserAssignedManagedIdentityObject.Id) 
            }
        }
    }
    
    Set-Variable -WhatIf:$False -Confirm:$False -Name amaExtensionSettingsVm -Option Constant -Value `
    @{
        "Settings" = $local:amaSettings
    }

    Set-Variable -WhatIf:$False -Confirm:$False -Name amaExtensionSettingsVmss -Option Constant -Value `
    @{
        "Setting" = $local:amaSettings
    }
    
    Set-Variable -WhatIf:$False -Confirm:$False -Name sb_vm -Option Constant -Value { `
        param([Microsoft.Azure.Commands.Compute.Models.PSVirtualMachine]$vmObj) `
        OnboardVmiWithAmaVm -VMObject $vmObj -ExtensionSettings $amaExtensionSettingsVm
    }
    Set-Variable -WhatIf:$False -Confirm:$False -Name sb_vmss -Option Constant -Value { `
        param([Microsoft.Azure.Commands.Compute.Automation.Models.PSVirtualMachineScaleSet]$vmssObj) `
        OnboardVmiWithAmaVmss -VMssObject $vmssObj -ExtensionSettings $amaExtensionSettingsVmss
    }
    
    Set-Variable -WhatIf:$False -Confirm:$False -Name sb_roles -Option Constant -Value { `
        param([String]$rgName) SetManagedIdentityRolesAma -ResourceGroupName $rgName
    }

    if ($TriggerVmssManualVMUpdate) {
        Set-Variable -WhatIf:$False -Confirm:$False -Name sb_upgrade -Option Constant -Value { `
            param([Microsoft.Azure.Commands.Compute.Automation.Models.PSVirtualMachineScaleSet]$vmssObj, [ref]$instanceUpgradeFailCounter) `
            UpgradeVmssExtensionManualUpdateEnabled -VMssObject $vmssObj -InstanceUpgradeFailCounter $instanceUpgradeFailCounter
        }
    } else {
        Set-Variable -WhatIf:$False -Confirm:$False -Name sb_upgrade -Option Constant -Value $sb_nop_block_upgrade
    }

    if ($PolicyAssignmentName) {
        #this section is only for VMs
        Write-Host "Looking up Virtual Machines in policy assignment $PolicyAssignmentName"

        $policyAssignmentNameResources = @{}
        Get-AzPolicyState `
            -Filter "PolicyAssignmentName eq '$PolicyAssignmentName' and ResourceType eq 'Microsoft.Compute/virtualMachines'" `
            | ForEach-Object {
                $policyAssignmentNameResources.Add($_.ResourceId, $True)
              }

        try {
            #Virtual Machines part of a VMSS will be skipped.
            Get-AzVM -ResourceGroupName $ResourceGroup -Name $Name `
                | Where-Object {!($_.VirtualMachineScaleSet) -and $policyAssignmentNameResources.ContainsKey($_.Id)} `
                | ForEach-Object {
                    $onboardingCounters.Total +=1 ;
                    PopulateRgHashTableVm -Rghashtable $Rghashtable -VMObject $_
                }
        } catch {
            $errorCode = ExtractExceptionErrorCode -ErrorRecord $_
            if (!($errorCode -eq "ResourceGroupNotFound" -or $errorCode -eq "ResourceNotFound")) {
                throw [FatalException]::new("Unable to lookup VMs", $_.Exception)
            }

            #quietly do nothing.
        } 
    } else {
        Write-Host ""
        Write-Host "Getting list of VMs or VM Scale Sets matching specified criteria."
        #Skipping heath status check as API output throttling noticed at scale and wild-cards are not accepted with Get-AzVM -Status.
        try {
            Get-AzVM -ResourceGroupName $ResourceGroup -Name $Name `
                | Where-Object {!($_.VirtualMachineScaleSet)} `
                | ForEach-Object { 
                    $onboardingCounters.Total +=1 ; 
                    PopulateRgHashTableVm -Rghashtable $Rghashtable -VMObject $_
                }
        } catch {
            $errorCode = ExtractExceptionErrorCode -ErrorRecord $_
            if (!($errorCode -eq "ResourceGroupNotFound" -or $errorCode -eq "ResourceNotFound")) {
                throw [FatalException]::new("Unable to lookup VMs", $_.Exception)
            }

            #quietly do nothing.
        }

        try {
            #VMI does not support VMSS with flexible orchestration.
            Get-AzVmss -ResourceGroupName $ResourceGroup -Name $Name `
                | Where-Object {$_.OrchestrationMode -ne 'Flexible'} `
                | ForEach-Object {
                    $onboardingCounters.Total +=1 ; 
                    PopulateRgHashTableVmss -RgHashTable $Rghashtable -VMssObject $_
                }
        } catch {
            $errorCode = ExtractExceptionErrorCode -ErrorRecord $_
            if (!($errorCode -eq "ResourceGroupNotFound" -or $errorCode -eq "ResourceNotFound")) {
                throw [FatalException]::new("Unable to lookup VMSS", $_.Exception)
            }

            #quietly do nothing.
        } 
    }

    $rgList = $Rghashtable.GetEnumerator() | Sort-Object -Property Key
    Write-Host "VM's and VMSS matching selection criteria :"
    $ManualUpgrade = 0
    Foreach ($entry in $rgList) {
        $rg = $entry.Key
        $rgTableObj = $entry.Value
        $vmList = $rgTableObj.VirtualMachineList
        $vmssList = $rgTableObj.VirtualMachineScaleSetList
        Write-Host ""
        Write-Host "ResourceGroup : $rg"


        if ($vmList.Count -gt 0) {
            $vmList = Sort-Object -Property Name -InputObject $vmList
            $vmList | ForEach-Object { Write-Host " $($_.Name)" }
            $rgTableObj.VirtualMachineList = $vmList
        }
        
        if ($vmssList.Count -gt 0) {
            $vmssList = Sort-Object -Property Name -InputObject $vmssList
            $vmssList | ForEach-Object { Write-Host " $($_.Name) - Upgrade mode $($_.UpgradePolicy.Mode)"; `
                                         if ($_.UpgradePolicy.Mode -eq "Manual") {$ManualUpgrade+=1}
                                       }
            $rgTableObj.VirtualMachineScaleSetList = $vmssList
        }
    }

    if ($ManualUpgrade -gt 0 -and !$TriggerVmssManualVMUpdate) {
        Write-Host "Found $ManualUpgrade VMSS with upgrade mode 'Manual'."
        Write-Host "'-TriggerVmssManualVMUpdate' was not provided. Upgrade will not be performed." 
    }

    Write-Host ""

    # Validate customer wants to continue
    if ($Approve -or $WhatIfPreference -or $PSCmdlet.ShouldContinue("Continue?", "")) {
        Write-Host ""
    } else {
        Write-Host "You selected No - exiting"
        $onboardingCounters.Total = 0
        exit 1
    }
    
    ForEach ($entry in $rgList) {
        $rg = $entry.Key
        $rgTableObj = $entry.Value
        try {
            &$sb_roles -rgName $rg
            
            foreach ($vm in $rgTableObj.VirtualMachineList) {
                try {
                    Write-Host ""
                    #This is on best-effort basis. State can change later.
                    if (!(IsVmAndGuestAgentUpAndHealthy -VMObject $vm)) {
                        Write-Host "Continuing to the next VM ..."
                        $onboardingCounters.Skipped +=1
                        continue
                    }
                    $vm = &$sb_vm -vmObj $vm
                    Write-Host "$(FormatVmIdentifier -VMObject $vm) : Successfully onboarded VM insights"
                    $onboardingCounters.Succeeded +=1
                    $unknownExceptionVirtualMachineConsequentCounter = 0
                } catch [VirtualMachineUnknownException] {
                    if ($unknownExceptionVirtualMachineConsequentCounter -gt $unknownExceptionVirtualMachineConsequentCounterLimit) {
                        throw [FatalException]::new("More than $unknownExceptionVirtualMachineConsequentCounterLimit unexpected exceptions encountered consequtively", $_.Exception)
                    }
                    if ($unknownExceptionTotalCounter -gt $unknownExceptionTotalCounterLimit) {
                        throw [FatalException]::new("More than $unknownExceptionTotalCounterLimit unexpected exceptions encountered", $_.Exception)
                    }
                    Write-Host "Unexpected VM Exception :"
                    $unknownExceptionTotalCounter+=1
                    $unknownExceptionVirtualMachineConsequentCounter+=1
                    DisplayException -ErrorRecord $_
                    Write-Host "Continuing to the next VM ..."
                } catch [VirtualMachineException] {
                    Write-Host "VM Exception :"
                    VerboseDisplayException -ErrorRecord $_
                    Write-Host "Continuing to the next VM ..."
                } catch [CustomerSkip] {
                    Write-Host "Onboarding operation aborted"
                    Write-Host "Continuing to the next VM ..."
                    $onboardingCounters.Skipped +=1
                }
            }

            foreach ($vmss in $rgTableObj.VirtualMachineScaleSetList) {
                try {
                    Write-Host ""
                    $vmsslogheader = FormatVmssIdentifier -VMssObject $vmss
                    $vmss = &$sb_vmss -vmssObj $vmss
                    $vmss = UpdateVMssExtension -VMssObject $vmss
                    $instanceUpgradeFailCounter = 0
                    if ($vmss.UpgradePolicy.Mode -eq 'Manual') {
                        &$sb_upgrade -vmssObj $vmss -InstanceUpgradeFailCounter ([ref]$instanceUpgradeFailCounter)
                    }
                    $onboardingCounters.VMSSInstanceUpgradeFailure += $instanceUpgradeFailCounter
                    if ($instanceUpgradeFailCounter -eq 0) {
                        Write-Host "$vmsslogheader : Successfully onboarded VM insights"
                        $onboardingCounters.Succeeded +=1
                    }
                    $unknownExceptionVirtualMachineScaleSetConsequentCounter = 0
                } catch [VirtualMachineScaleSetUnknownException] {
                    if ($unknownExceptionVirtualMachineScaleSetConsequentCounter -gt $unknownExceptionVirtualMachineScaleSetConsequentCounterLimit) {
                        throw [FatalException]::new("More than $unknownExceptionVirtualMachineScaleSetConsequentCounterLimit unexpected exceptions encountered consequtively", $_.Exception)
                    }
                    if ($unknownExceptionTotalCounter -gt $unknownExceptionTotalCounterLimit) {
                        throw [FatalException]::new("More than $unknownExceptionTotalCounterLimit unexpected exceptions encountered", $_.Exception)
                    }
                    $unknownExceptionTotalCounter+=1
                    $unknownExceptionVirtualMachineScaleSetConsequentCounter+=1
                    Write-Host "Unexpected VMSS Exception :"
                    DisplayException -ErrorRecord $_
                    Write-Host "Continuing to the next VMSS ..."
                } catch [VirtualMachineScaleSetException] {
                    Write-Host "VMSS Exception :"
                    VerboseDisplayException -ErrorRecord $_
                    Write-Host "Continuing to the next VMSS ..."
                } catch [CustomerSkip] {
                    Write-Host "Onboarding operation aborted"
                    Write-Host "Continuing to the next VMSS ..."
                    $onboardingCounters.Skipped +=1
                }
            }
        } catch [ResourceGroupDoesNotExist] {
            Write-Host "Resource Group Exception :"
            VerboseDisplayException -ErrorRecord $_
            Write-Host "Continuing to the next Resource Group ..."
        } catch [CustomerSkip] {
            Write-Host "Onboarding operation aborted"
            $onboardingCounters.Skipped += $rgTableObj.VirtualMachineList.Count + $rgTableObj.VirtualMachineScaleSetList.Count
        }
    }
}
catch [UserAssignedManagedIdentityUnknownException] {
    Write-Host "Fatal Exception :"
    DisplayException -ErrorRecord $_
    Write-Host "Exiting ..."
    exit 2
}
# Customer actionable exceptions
catch [FatalException] {
    Write-Host "Fatal Exception :"
    VerboseDisplayException -ErrorRecord $_
    Write-Host "Exiting ..."
    exit 3
}
catch {
    Write-Host "Unexpected Fatal Exception :"
    DisplayException -ErrorRecord $_
    Write-Host "Exiting ..."
    exit 4
}
finally {
    PrintSummaryMessage  $onboardingCounters
}