modules/BackendPoolMigration/BackendPoolMigration.psm1

# Load Modules
Import-Module ((Split-Path $PSScriptRoot -Parent) + "/Log/Log.psd1")
Import-Module ((Split-Path $PSScriptRoot -Parent) + "/UpdateVmss/UpdateVmss.psd1")
Import-Module ((Split-Path $PSScriptRoot -Parent) + "/UpdateVmssInstances/UpdateVmssInstances.psd1")
Import-Module ((Split-Path $PSScriptRoot -Parent) + "/GetVmssFromBasicLoadBalancer/GetVmssFromBasicLoadBalancer.psd1")

function _HardCopyObject {
    [CmdletBinding()]
    param (
        [Parameter(Mandatory = $True)][System.Collections.Generic.List[Microsoft.Azure.Management.Compute.Models.SubResource]] $listSubResource
    )
    $options = [System.Text.Json.JsonSerializerOptions]::new()
    $options.WriteIndented = $true
    $options.IgnoreReadOnlyProperties = $true
    $cgenericListSubResource = [System.Text.Json.JsonSerializer]::Serialize($listSubResource, "System.Collections.Generic.List[Microsoft.Azure.Management.Compute.Models.SubResource]", $options)
    $cgenericListSubResource = [System.Text.Json.JsonSerializer]::Deserialize($cgenericListSubResource, "System.Collections.Generic.List[Microsoft.Azure.Management.Compute.Models.SubResource]")
    # To preserve the original object type in the return we must use a , before the object to be returned
    return , [System.Collections.Generic.List[Microsoft.Azure.Management.Compute.Models.SubResource]]$cgenericListSubResource
}

function _MigrateHealthProbe {
    [CmdletBinding()]
    param (
        [Parameter(Mandatory = $True)][Microsoft.Azure.Commands.Network.Models.PSLoadBalancer] $StdLoadBalancer,
        [Parameter(Mandatory = $True)][Microsoft.Azure.Commands.Compute.Automation.Models.PSVirtualMachineScaleSet] $vmss,
        [Parameter(Mandatory = $True)][Microsoft.Azure.Commands.Compute.Automation.Models.PSVirtualMachineScaleSet] $refVmss
    )

    log -Message "[_MigrateHealthProbe] Migrating Health Probes"
    try {
        $refHealthProbe = $refVmss.VirtualMachineProfile.NetworkProfile.HealthProbe
        if (![string]::IsNullOrEmpty($refHealthProbe)) {
            $refProbeName = $refHealthProbe.Id.Split('/')[-1]
            log -Message "[_MigrateHealthProbe] Health Probes found in reference VMSS $($refProbeName)"
            $refProbeId = ($StdLoadBalancer.Probes | Where-Object { $_.Name -eq $refProbeName }).id
            $vmss.VirtualMachineProfile.NetworkProfile.HealthProbe = $refProbeId
        }
        else {
            log -Message "[_MigrateHealthProbe] Health Probes not found in reference VMSS $($refVmss.Name)"
        }
    }
    catch {
        $message = "[_MigrateHealthProbe] An error occured migrating a health probe to the VMSS. To recover, address the cause of the following error, then follow the steps at https://aka.ms/basiclbupgradefailure to retry the migration. Error: $_"
        log 'Error' $message -terminateOnError
    }
    log -Message "[_MigrateHealthProbe] Migrating Health Probes completed"
}

function _RestoreUpgradePolicyMode {
    param (
        [Parameter(Mandatory = $True)][Microsoft.Azure.Commands.Compute.Automation.Models.PSVirtualMachineScaleSet] $vmss,
        [Parameter(Mandatory = $True)][Microsoft.Azure.Commands.Compute.Automation.Models.PSVirtualMachineScaleSet] $refVmss
    )

    log -Message "[_RestoreUpgradePolicyMode] Restoring VMSS Upgrade Policy Mode"
    if ($vmss.UpgradePolicy.Mode -ne $refVmss.UpgradePolicy.Mode) {
        log -Message "[_RestoreUpgradePolicyMode] Restoring VMSS Upgrade Policy Mode to $($refVmss.UpgradePolicy.Mode)"
        $vmss.UpgradePolicy.Mode = $refVmss.UpgradePolicy.Mode
    }
    else {
        log -Message "[_RestoreUpgradePolicyMode] VMSS Upgrade Policy Mode not changed"
    }

    log -Message "[_RestoreUpgradePolicyMode] Restoring VMSS Upgrade Policy Mode completed"
}

function _RestoreAutomaticOSUpgradePolicy {
    param (
        [Parameter(Mandatory = $True)][Microsoft.Azure.Commands.Compute.Automation.Models.PSVirtualMachineScaleSet] $vmss,
        [Parameter(Mandatory = $True)][Microsoft.Azure.Commands.Compute.Automation.Models.PSVirtualMachineScaleSet] $refVmss
    )

    log -Message "[_RestoreAutomaticOSUpgradePolicy] Restoring VMSS Upgrade Policy Mode"
    if ($vmss.upgradePolicy.AutomaticOSUpgradePolicy.enableAutomaticOSUpgrade -ne $refVmss.upgradePolicy.AutomaticOSUpgradePolicy.enableAutomaticOSUpgrade) {
        log -Message "[_RestoreAutomaticOSUpgradePolicy] Restoring VMSS Upgrade Policy Mode to '$($refVmss.upgradePolicy.AutomaticOSUpgradePolicy.enableAutomaticOSUpgrade)'"
        $vmss.upgradePolicy.AutomaticOSUpgradePolicy.enableAutomaticOSUpgrade = $refVmss.upgradePolicy.AutomaticOSUpgradePolicy.enableAutomaticOSUpgrade
    }
    else {
        log -Message "[_RestoreAutomaticOSUpgradePolicy] VMSS Upgrade Policy Mode not changed"
    }

    log -Message "[_RestoreAutomaticOSUpgradePolicy] Restoring VMSS Upgrade Policy Mode completed"
}

function _MigrateNetworkInterfaceConfigurationsVmss {
    param (
        [Parameter(Mandatory = $True)][Microsoft.Azure.Commands.Network.Models.PSLoadBalancer] $BasicLoadBalancer,
        [Parameter(Mandatory = $True)][Microsoft.Azure.Commands.Network.Models.PSLoadBalancer] $StdLoadBalancer,
        [Parameter(Mandatory = $True)][Microsoft.Azure.Commands.Compute.Automation.Models.PSVirtualMachineScaleSet] $vmss
    )

    log -Message "[_MigrateNetworkInterfaceConfigurationsVmss] Adding BackendAddressPool to VMSS $($vmss.Name)"
    foreach ($networkInterfaceConfiguration in $vmss.VirtualMachineProfile.NetworkProfile.NetworkInterfaceConfigurations) {
        $genericListSubResource = New-Object System.Collections.Generic.List[Microsoft.Azure.Management.Compute.Models.SubResource]
        foreach ($ipConfiguration in $networkInterfaceConfiguration.IpConfigurations) {
            If (![string]::IsNullOrEmpty($ipConfiguration.LoadBalancerBackendAddressPools)) {
                $genericListSubResource.AddRange($ipConfiguration.LoadBalancerBackendAddressPools)
            }

            foreach ($BackendAddressPool in $BasicLoadBalancer.BackendAddressPools) {
                foreach ($BackendIpConfiguration in $BackendAddressPool.BackendIpConfigurations) {
                    $lbBeNicName = $BackendIpConfiguration.Id.Split('/')[-3]
                    $lbBeipConfigName = $BackendIpConfiguration.Id.Split('/')[-1]
                    if ($lbBeNicName -eq $networkInterfaceConfiguration.Name -and $lbBeipConfigName -eq $ipConfiguration.Name) {
                        try {
                            $subResource = New-Object Microsoft.Azure.Management.Compute.Models.SubResource
                            $subResource.Id = ($StdLoadBalancer.BackendAddressPools | Where-Object { $_.Name -eq $BackendAddressPool.Name }).Id
                            log -Message "[_MigrateNetworkInterfaceConfigurationsVmss] Adding BackendAddressPool $($subResource.Id.Split('/')[-1]) to VMSS Nic: $lbBeNicName ipConfig: $lbBeipConfigName"
                            $genericListSubResource.Add($subResource)
                        }
                        catch {
                            $message = "[_MigrateNetworkInterfaceConfigurationsVmss] An error occured creating a new VMSS IP Config. To recover, address the cause of the following error, then follow the steps at https://aka.ms/basiclbupgradefailure to retry the migration. Error: $_"
                            log 'Error' $message -terminateOnError
                        }
                    }
                }
            }
            # Taking a hard copy of the object and assigning, it's important because the object was passed by reference
            If ($genericListSubResource.Count -gt 0) {
                $ipConfiguration.LoadBalancerBackendAddressPools = _HardCopyObject -listSubResource $genericListSubResource
            }
            $genericListSubResource.Clear()
        }
    }
    log -Message "[_MigrateNetworkInterfaceConfigurationsVmss] Migrate NetworkInterface Configurations completed"
}

function BackendPoolMigrationVmss {
    [CmdletBinding()]
    param (
        [Parameter(Mandatory = $True)][Microsoft.Azure.Commands.Network.Models.PSLoadBalancer] $BasicLoadBalancer,
        [Parameter(Mandatory = $True)][Microsoft.Azure.Commands.Network.Models.PSLoadBalancer] $StdLoadBalancer,
        [Parameter(Mandatory = $True)][Microsoft.Azure.Commands.Compute.Automation.Models.PSVirtualMachineScaleSet] $refVmss
    )
    log -Message "[BackendPoolMigrationVmss] Initiating Backend Pool Migration"

    log -Message "[BackendPoolMigrationVmss] Adding Standard Load Balancer back to the VMSS"
    log -Message "[BackendPoolMigrationVmss] Building VMSS object from Basic Load Balancer $($BasicLoadBalancer.Name)"
    $vmss = GetVmssFromBasicLoadBalancer -BasicLoadBalancer $BasicLoadBalancer

    # Migrating Health Probe in case it exist
    _MigrateHealthProbe -StdLoadBalancer $StdLoadBalancer -vmss $vmss -refVmss $refVmss

    # Migrating Network Interface Configurations back to the VMSS
    _MigrateNetworkInterfaceConfigurationsVmss -BasicLoadBalancer $BasicLoadBalancer -StdLoadBalancer $StdLoadBalancer -vmss $vmss

    # Update VMSS on Azure
    try {
        Update-Vmss -vmss $vmss
    }
    catch {
        $message = "[BackendPoolMigrationVmss] An error occured while updating the VMSS to associate it with the standard load balancer's backend pool. To recover address the following error, and try again specifying the -FailedMigrationRetryFilePath parameter and Basic Load Balancer backup State file located either in this directory or the directory specified with -RecoveryBackupPath. To manually complete the migration, add the VMSS to the appropriate backend pools and check that the VMSS Upgrade Policy mode is correct. `nError message: $_"
        log 'Error' $message -terminateOnError
    }

    # Update Instances
    UpdateVmssInstances -vmss $vmss

    # Restore VMSS Upgrade Policy Mode
    _RestoreUpgradePolicyMode -vmss $vmss -refVmss $refVmss

    # Restore Automatic OS Upgrade Policy
    _RestoreAutomaticOSUpgradePolicy -vmss $vmss -refVmss $refVmss

    # Update VMSS on Azure
    try {
        Update-Vmss -vmss $vmss
    }
    catch {
        $message = "[BackendPoolMigrationVmss] An error occured while restoring the VMSS Upgrade Policy Mode. To recover address the following error, and try again specifying the -FailedMigrationRetryFilePath parameter and Basic Load Balancer backup State file located either in this directory or the directory specified with -RecoveryBackupPath. To manually complete the migration, the VMSS Upgrade Policy mode matches the mode in the VMSS state file export. `nError message: $_"
        log 'Error' $message -terminateOnError
    }

    #log -Message "[BackendPoolMigrationVmss] Updating VMSS Instances $($vmss.Name)"
    #UpdateVmssInstances -vmss $vmss

    #log -Message "[BackendPoolMigrationVmss] StackTrace $($StackTrace)" -Severity "Debug"
    log -Message "[BackendPoolMigrationVmss] Backend Pool Migration Completed"
}

function BackendPoolMigrationVM {
    [CmdletBinding()]
    param (
        [Parameter(Mandatory = $True)][Microsoft.Azure.Commands.Network.Models.PSLoadBalancer] $BasicLoadBalancer,
        [Parameter(Mandatory = $True)][Microsoft.Azure.Commands.Network.Models.PSLoadBalancer] $StdLoadBalancer
    )
    log -Message "[BackendPoolMigrationVM] Initiating Backend Pool Migration"

    log -Message "[BackendPoolMigrationVM] Adding original VMs to the new Standard Load Balancer backend pools"

    # build table of NICs and their ipconfigs and associate them to backend pools
    $backendPoolNicsAssociationArray = @()
    ForEach ($backendAddressPool in $BasicLoadBalancer.BackendAddressPools) {
        ForEach ($BackendIpConfiguration in $backendAddressPool.BackendIpConfigurations) {
            $lbBeNicId = ($BackendIpConfiguration.Id -split '/ipConfigurations/')[0]
            $ipConfigName = ($BackendIpConfiguration.Id -split '/ipConfigurations/')[1]

            $nicRow = [PSCustomObject]@{
                nicId         = $lbBeNicId
                ipConfigName  = $ipConfigName
                backendPoolId = $StdLoadBalancer.BackendAddressPools | Where-Object { $_.Name -eq $backendAddressPool.Name } | Select-Object -ExpandProperty Id
            }

            $backendPoolNicsAssociationArray += $nicRow
        }
    }

    # loop though nics and associate ipconfigs to backend pools
    $jobs = @()
    ForEach ($nicRecordGroup in ($backendPoolNicsAssociationArray | Group-Object -Property nicId)) {
        $nicName = ($nicRecordGroup.Name -split '/networkInterfaces/')[1]

        log -Message "[BackendPoolMigrationVM] Adding ipconfigs on NIC '$($nicName)' to backend pools"
        try {
            log -Message "[BackendPoolMigrationVM] Getting NIC '$($nicRecordGroup.Name)'"
            $nic = Get-AzNetworkInterface -ResourceId $nicRecordGroup.Name
        }
        catch {
            $message = "[BackendPoolMigrationVM] An error occured getting the Network Interface '$($nicRecordGroup.Name)'. Check that the NIC exists. To recover, address the cause of the following error, then follow the steps at https://aka.ms/basiclbupgradefailure to retry the migration. Error: $_"
            log 'Error' $message -terminateOnError
        }

        ForEach ($ipconfigGroup in ($nicRecordGroup.Group | Group-Object -Property ipConfigName)) {
            $ipConfigName = $ipconfigGroup.Name

            try {
                log -Message "[BackendPoolMigrationVM] Getting IP Config '$($ipConfigName)' on NIC '$($nicName)'"
                $ipConfig = Get-AzNetworkInterfaceIpConfig -NetworkInterface $nic -Name $ipConfigName

                ForEach ($backendPoolEntry in $ipconfigGroup.Group) {
                    $backendPoolObj = new-object Microsoft.Azure.Commands.Network.Models.PSBackendAddressPool 

                    log -Message "[BackendPoolMigrationVM] Getting Backend Pool '$($backendPoolEntry.backendPoolId.split('/')[-1])' on Standard Load Balancer and adding to IP Config '$($ipConfigName)'" -Severity Debug
                    $backendPoolObj.id = $StdLoadBalancer.BackendAddressPools | Where-Object { $_.Name -eq $backendPoolEntry.backendPoolId.split('/')[-1] } | Select-Object -ExpandProperty Id

                    log -Message "[BackendPoolMigrationVM] Adding backend pool '$($backendPoolEntry.backendPoolId)' to NIC '$($nicName)' IP Config '$($ipConfigName)'" -Severity "Debug"
                    $ipConfig.LoadBalancerBackendAddressPools.Add($backendPoolObj)
                }
            }
            catch {
                $message = "[BackendPoolMigrationVM] An error occured adding the backend pools to the Network Interface '$($nic.Name)' IP Configuration '$($ipConfigName)'. To recover, address the cause of the following error, then follow the steps at https://aka.ms/basiclbupgradefailure to retry the migration. Error: $_"
                log 'Error' $message -terminateOnError
            }
        }

        log -Message "[BackendPoolMigrationVM] Setting NIC '$($nic.Name)' with updated IP Configurations as a job" -Severity "Debug"
        $jobs += Set-AzNetworkInterface -NetworkInterface $nic -AsJob
    }

    log -Message "[BackendPoolMigrationVM] Waiting for all '$($jobs.count)' NIC backend pool association jobs to complete"
    $jobs | Wait-Job -Timeout $defaultJobWaitTimeout | Foreach-Object {
        $job = $_
        If ($job.Error -or $job.State -eq 'Failed') {
            log -Severity Error -Message "Associating NIC with LB Backend Pool failed with the following errors: $($job.error; $job | Receive-Job). Migration will continue--to recover, manually associate NICs with the backend pool after the script completes. See association table: `n $($backendPoolNicTable | ConvertTo-Json -depth 10)"
        }
    }

    log -Message "[BackendPoolMigrationVM] Backend Pool Migration Completed"
}

Export-ModuleMember -Function BackendPoolMigrationVmss
Export-ModuleMember -Function BackendPoolMigrationVM