AksHciNcCompanion.psm1

$script:HostVNetNicName = "SDN_VNET"
$script:clsuterGroupName = "clustergroup"
$script:MocConfig = $null
$script:MocGroup = $null

$script:InternetUrl = "http://www.msftncsi.com/ncsi.txt"
$script:DnsResolutionName = "bing.com"
$script:NCSIResponse = "Microsoft NCSI"

$script:LinuxUserName = "clouduser"
$script:WindowsUserName = "Administrator"

$script:CloudAgentPort = "65000"
$script:CloudAgentCurlResponse = "curl: (52) Empty reply from server"

$script:VmInboundAccessState = @{}

$script:TestOutputDirectory = "C:\Windows\Tracing\AksHci"

enum AksHciNcResourceType
{
    NetworkInterface
    LoadBalancer
    VirtualNetwork
}

function Write-TestResults($results)
{
    if (-not (Test-Path $script:TestOutputDirectory )) {
        New-Item -ItemType Directory -Path $script:TestOutputDirectory  -Force
    }
    $output = $results | ConvertTo-Json -Depth 10
    $fileName = "Test-AksHciSdnNetwork-$((get-date).Ticks).Json"
    $file = Join-Path $script:TestOutputDirectory $fileName
    $output | Out-File -FilePath $file -Force
    Write-Host "Results written to file $file" -ForegroundColor Green
}

function Write-Fail {

    param (
        [string]$testName,
        [string]$msg,
        [string]$err
    )
    
    #Write-Status -moduleName $moduleName -Verbose -msg "Test Connection for $scenario, Status Succeeded"
    $out = "FAIL: $testName $msg"

    if (-not [string]::IsNullOrWhiteSpace($err))
    {
        $out += "`nERROR $err"
    }

    Write-Host $out -ForegroundColor Red
}
function Write-Pass {
    param (
        [string]$testName,
        [string]$msg
    ) 

    Write-Host "PASS: $testName $msg" -ForegroundColor Green
}
function Write-Msg {
    param (
        [string]$msg,
        [switch]$Warning
    )
    
    #Write-Status -moduleName $moduleName -Verbose -msg "Test Connection for $scenario, Status Succeeded"
    if ($Warning.IsPresent) {
        Write-Warning $msg
        return
    }

    Write-Host $msg
}

function Get-Result {
        param(
        [string]$scenario,
        [bool]$isSuccess,
        [string]$detailedStatus
    )
    
    # return @{
    # "Scenario" = $scenario;
    # "Passed" = $isSuccess;
    # "DetailedStatus" = $detailedStatus;
    # }

    $result = New-Object -TypeName psobject
    $result | Add-Member -MemberType NoteProperty -Name Scenario -Value $scenario
    $result | Add-Member -MemberType NoteProperty -Name Passed -Value $isSuccess
    $result | Add-Member -MemberType NoteProperty -Name DetailedStatus -Value $detailedStatus

    return $result
}

function Get-VMTestResult {
    param(
        [string]$vmName,
        [string]$testName,
        [string]$sshInfo,
        [bool]$isSuccess,
        [string]$detailedStatus,
        [array]$results
    )

    $vmResult = New-Object -TypeName psobject
    $vmResult | Add-Member -MemberType NoteProperty -Name VMName -Value $vmName
    $vmResult | Add-Member -MemberType NoteProperty -Name TestName -Value $testName
    $vmResult | Add-Member -MemberType NoteProperty -Name SSHInfo -Value $sshInfo
    $vmResult | Add-Member -MemberType NoteProperty -Name Passed -Value $isSuccess
    $vmResult | Add-Member -MemberType NoteProperty -Name DetailedStatus -Value $detailedStatus
    $vmResult | Add-Member -MemberType NoteProperty -Name Results -Value $results

    return $vmResult

    # return @{
    # "VMName" = $vmName;
    # "SSHInfo" = $sshInfo;
    # "Passed" = $isSuccess;
    # "DetailedStatus" = $detailedStatus;
    # "TestName" = $testName
    # "Results" = $results;
    # }
}

function Get-TestResult {
    param(
        [string]$testName,
        [bool]$isSuccess,
        [string]$detailedStatus,
        [array]$testResults
    )

    $result = New-Object -TypeName psobject
    $result | Add-Member -MemberType NoteProperty -Name TestName -Value $testName
    $result | Add-Member -MemberType NoteProperty -Name Passed -Value $isSuccess
    $result | Add-Member -MemberType NoteProperty -Name DetailedStatus -Value $detailedStatus
    $result | Add-Member -MemberType NoteProperty -Name TestResults -Value $testResults
    return $result

    # return @{
    # "TestName" = $testName;
    # "Passed" = $isSuccess;
    # "DetailedStatus" = $detailedStatus;
    # "TestResults" = $testResults;
    # }
}

function Get-ClusterTestResult {
        param(
        [string]$clusterName,
        [bool]$isSuccess,
        [string]$detailedStatus,
        [array]$testResults
    )

    $result = New-Object -TypeName psobject
    $result | Add-Member -MemberType NoteProperty -Name ClusterName -Value $clusterName
    $result | Add-Member -MemberType NoteProperty -Name Passed -Value $isSuccess
    $result | Add-Member -MemberType NoteProperty -Name DetailedStatus -Value $detailedStatus
    $result | Add-Member -MemberType NoteProperty -Name TestResults -Value $testResults
    return $result

    # return @{
    # "ClusterName" = $clusterName;
    # "Passed" = $isSuccess;
    # "TestResults" = $testResults;
    # "DetailedStatus" = $detailedStatus;
    # }
}

function Invoke-WebRequestWithRetries {
    param(
        [System.Collections.IDictionary] $Headers,
        [string] $ContentType,
        [Microsoft.PowerShell.Commands.WebRequestMethod] $Method,
        [System.Uri] $Uri,
        [object] $Body,
        [Switch] $DisableKeepAlive,
        [Switch] $UseBasicParsing,
        [Parameter(mandatory=$false)]
        [bool] $shouldRetry = $true
    )
        
    $params = @{
        'Headers'=$headers;
        'ContentType'=$content;
        'Method'=$method;
        'uri'=$uri;
        'ErrorAction'='Stop';
        }
    
    if($Body -ne $null) 
    {
        $params.Add('Body', $Body)
    }

    if($DisableKeepAlive.IsPresent) 
    {
        $params.Add('DisableKeepAlive', $true)
    }

    if($UseBasicParsing.IsPresent) 
    {
        $params.Add('UseBasicParsing', $true)
    }

    if ($script:DefaultCredParaSet -eq $true)
    {
        $params.Add('UseDefaultCredentials', $true)
    }
    elseif($script:NetworkControllerCred -ne [System.Management.Automation.PSCredential]::Empty)
    {
        $params.Add('Credential', $script:NetworkControllerCred)
    }
        
    $retryIntervalInSeconds = 30
    $maxRetry = 6
    $retryCounter = 0
    
    do {
        try {
            $result = $null
            $result = Invoke-WebRequest @params
            break
        }
        catch {

          
            if($_.Exception.Response.StatusCode.value__ -eq 404)
            {
                #Dont retry on Not Found
                break
            }

            $retryCounter++
            if($retryCounter -le $maxRetry) {
                Start-Sleep -Seconds $retryIntervalInSeconds
            }
            else {
                # last retry still fails, so throw the exception
                throw $_
            }
        }
    } while ($shouldRetry -and ($retryCounter -le $maxRetry)) 
    
    return $result       
}

function Initialize-AksHciNcCompanion
{
    $winFeatures = @("RSAT-NetworkController")
    
    foreach ($feature in $winFeatures)
    {
        Write-Verbose "Check install state of $feature"
        if ((Get-WindowsFeature -Name $feature).InstallState -ne "Installed")
        {
            Write-Verbose "Installing $feature"
            Add-WindowsFeature -Name $feature -IncludeAllSubFeature -IncludeManagementTools
        }
    }

}

function Get-NCResource
{
    param
    (
        [parameter(Mandatory=$true)]
        [string] $Uri
    )

    $result = Invoke-WebRequestWithRetries -Uri $Uri -DisableKeepAlive -UseBasicParsing -Method "Get"      

    if(-not $result)
    {
        return $null
    }            

    $toplevel = convertfrom-json $result.Content
    if ($toplevel.value -eq $null)
    {
        $jsonOut = $toplevel
    } 
    else
    {
        $jsonOut = $toplevel.value
    }

    return $jsonOut
}

function Get-NCNic
{
    param
    (
        [parameter(Mandatory=$true)]
        [string] $resourceId
    )

    Write-Verbose "Fetching nic $resourceId from NC"

    $mocConfig = Get-MocConfigCache     
    $ncRestEndPoint = $mocConfig["networkControllerFqdnOrIpAddress"]
    $uri = "https://$ncRestEndPoint/networking/v1/networkinterfaces/$resourceId"

    $nic = Get-NCResource -Uri $uri

    if ($nic -eq $null)
    {
        $nics = (Invoke-WebRequest -Uri "https://$ncRestEndPoint/networking/v1/networkinterfaces/" -UseBasicParsing).content | ConvertFrom-Json
        $nic = $nics.value | Where-Object {$_.resourceMetadata.resourceName -eq "$resourceId"}
    }

    if ($nic -eq $null)
    {
        Write-Warning "Nic with resource id $resourceId not found in NC($ncRestEndPoint)"
    }
    return $nic
}

function Get-NCNicInstanceId
{
    param
    (
        [parameter(Mandatory=$true)]
        [string] $resourceId
    )

    $nic = Get-NCNic -resourceId $resourceId

    return $nic.instanceId
}

function Get-MocConfigCache
{
    if ($null -eq $script:MocConfig)
    {
        $script:MocConfig = Get-MocConfig
    }

    return $script:MocConfig
}

function Get-MocGroupCache
{
    if ($null -eq $script:MocGroup)
    {
        $script:MocGroup = Get-MocGroup -location MocLocation
    }

    return $script:MocGroup
}

function Clear-MocCache
{
    $script:MocConfig = $null
    $script:MocGroup = $null
}

function GetClusterGroup
{
    return $script:clsuterGroupName
}

function ValidateState
{ 
   
    $isInstalled = Test-IsAksHciInstalled

    if (-not $isInstalled)
    {
        Write-Error("UnsupportedConfiguration: AKS-HCI is not installed") -ErrorAction Stop
    }

    $mocConfig = Get-MocConfigCache

    if (-not $mocConfig["UseNetworkController"])
    {
        Write-Error("UnsupportedConfiguration: SDN intergration is not present") -ErrorAction Stop
    }
}

function CleanupMoc
{
    param
    (
        [parameter(Mandatory=$true)]
        [string] $nicName
    )

    Remove-MocNetworkInterface -name $nicName -group $(GetClusterGroup) -ErrorAction SilentlyContinue
}

function CleanupHost
{
    param
    (
        [parameter(Mandatory=$true)]
        [string] $nicName
    )

    $mocConfig = Get-MocConfigCache    
    $vswitchName = $mocConfig["vswitchName"]

    $nic = Get-VMNetworkAdapter -ManagementOS -Name $nicName -SwitchName $vswitchName -ErrorAction SilentlyContinue

    if ($nic -eq $null)
    {
        return
    }

    Remove-VMNetworkAdapter -ManagementOS -Name $nicName -SwitchName $vswitchName -ErrorAction Stop
}

function CreateHostVNic
{
    param
    (
        [parameter(Mandatory=$true)]
        [string] $name,
        [parameter(Mandatory=$true)]
        [string] $mac,
        [parameter(Mandatory=$true)]
        [string] $ipAddress,
        [parameter(Mandatory=$true)]
        [string] $prefix,
        [parameter(Mandatory=$true)]
        [string] $vswitchName
    )

    Write-Verbose "Creating a new host vnic $name, MAC $mac, IP $ipAddress/$prefix on host $(hostname)"
    $nic = Get-VMNetworkAdapter -ManagementOS -Name $name -SwitchName $vswitchName -ErrorAction SilentlyContinue

    if ($nic -ne $null -and $nic.MacAddress -eq $mac)
    {
        Write-Verbose "Reusing existing VNIC "
    }
    else
    {
        CleanupHost -nicName $name
        Start-Sleep 5

        $nic = Add-VMNetworkAdapter -ManagementOS -Name $name -SwitchName $vswitchName -StaticMacAddress $mac -ErrorAction stop
        Start-Sleep 30
    }

    $ifIndex = (Get-NetAdapter -Name "*$name*").ifIndex

    $nic = Get-NetIPAddress -IPAddress $ipAddress -ErrorAction SilentlyContinue

    if ($nic -ne $null)
    {

        if ($nic.InterfaceIndex -ne $ifIndex)
        {
            Write-Error "Fatal error, Cannot connect host, $ipAddress is already present on another interface on the host" -ErrorAction Stop
        }

        return
    }

    Remove-NetIPAddress -InterfaceIndex $ifIndex -Confirm:$false -ErrorAction SilentlyContinue | Out-Null
    New-NetIPAddress -PrefixLength $prefix -IPAddress $ipAddress -InterfaceIndex $ifIndex -ErrorAction Stop | Out-Null
}


function CreateMocNic
{
    param
    (
        [parameter(Mandatory=$true)]
        [string] $nicName
    )

    $mocConfig = Get-MocConfigCache
    $vnetName = $mocConfig["vnetName"]
    $location = $mocConfig["cloudLocation"]
    $retry = $true

    while($retry)
    {
        try
        {
            New-MocNetworkInterface -name $nicName -virtualNetworkName $vnetName -group $(GetClusterGroup) -ErrorAction Stop
            break
        }
        catch
        {
            $e = $_

            $retry = $e.Exception.Message.Contains("PrivateIPAddressInUse") -eq $true

            if (-not $retry)
            {
                Write-Warning "Failed to create Nic in Moc, Error $_"
            }
            else 
            {
                Write-Verbose "IPAddress conflict, retry nic creation"
                throw
            }
        }
    }
}

function Get-MocNic
{
    
    param
    (
        [parameter(Mandatory=$true)]
        [string] $nicName
    )

    Write-Verbose "Fetching nic $nicName from MOC"
    $vnetMocNic = Get-MocNetworkInterface -name $nicName -group $(GetClusterGroup) -ErrorAction SilentlyContinue

    if ($vnetMocNic -eq $null)
    {
        Write-Verbose "Creating new nic $nicName in MOC"
        CreateMocNic -nicName $nicName
    }

    return Get-MocNetworkInterface -name $nicName -group $(GetClusterGroup) -ErrorAction SilentlyContinue
}

function SetPortProfile
{
    param
    (
        [parameter(Mandatory=$true)]
        [string] $nicName,
        [parameter(Mandatory=$true)]
        [string] $profileId
    )

    Write-Verbose "Setting port profile for $nicName on $(hostName), ProfileId $profileId"

    $vmNic = Get-VMNetworkAdapter -Name $nicName -ManagementOS -ErrorAction Stop

    $FeatureId = "9940cd46-8b06-43bb-b9d5-93d50381fd56"

    $CurrentFeature = Get-VMSwitchExtensionPortFeature -FeatureId $FeatureId -VMNetworkAdapter $vmNic

    if ($CurrentFeature -eq $null)
    {
        $Feature = Get-VMSystemSwitchExtensionPortFeature -FeatureId $FeatureId
        $Feature.SettingData.ProfileId = "{$profileId}"
        $Feature.SettingData.NetCfgInstanceId = "{56785678-a0e5-4a26-bc9b-c0cba27311a3}"
        $Feature.SettingData.CdnLabelString = "TestCdn"
        $Feature.SettingData.CdnLabelId = 1111
        $Feature.SettingData.ProfileName = "Testprofile"
        $Feature.SettingData.VendorId = "{1FA41B39-B444-4E43-B35A-E1F7985FD548}"
        $Feature.SettingData.VendorName = "NetworkController"
        $Feature.SettingData.ProfileData = 1
        Add-VMSwitchExtensionPortFeature -VMSwitchExtensionFeature $Feature -VMNetworkAdapter $vmNic
    }
    else
    {
        $CurrentFeature.SettingData.ProfileId = "{$profileId}"
        $CurrentFeature.SettingData.ProfileData = 1
        Set-VMSwitchExtensionPortFeature -VMSwitchExtensionFeature $CurrentFeature -VMNetworkAdapter $vmNic
    }

    Write-Verbose "Successfully set port profile for $nicName on $(hostName)"
}

function Remove-PortProfile
{
    param
    (
        [parameter(Mandatory=$true)]
        [string] $nicName
    )

    Write-Verbose "Removing port profile on $nicName"

    $vmNic = Get-VMNetworkAdapter -Name $nicName -ManagementOS -ErrorAction SilentlyContinue

    if ($vmNic  -eq $null)
    {
        return
    }

    $FeatureId = "9940cd46-8b06-43bb-b9d5-93d50381fd56"

    $CurrentFeature = Get-VMSwitchExtensionPortFeature -FeatureId $FeatureId -VMNetworkAdapter $vmNic

    if ($CurrentFeature -eq $null)
    {
        return
    }

    Remove-VMSwitchExtensionPortFeature -VMSwitchExtensionFeature $CurrentFeature -VMNetworkAdapterName $vmNic

}

function Test-ConnectRequired
{
    # Checks if Connect-AksHciSdnVnet is required.
    # Not required if AksHci version > 1.0.18
    $version = Get-AksHciVersion
    $verArray = $version.Split(".")

    if ($verArray[0] -ge 1) {
        return $false
    }

    if ($verArray.Count -ge 2 -and $verArray[1] -gt 0) {
        return $false
    }

    if ($verArray.Count -ge 3 -and $verArray[3] -ge 18) {
        return $false
    }

    return $true
}
function Connect-AksHciSdnVnet
{

    param
    (
        [parameter(Mandatory=$false)]
        [switch] $Force
    )

    ValidateState    

    if (-not $Force.IsPresent) {
        if (-not $(Test-ConnectRequired)) {
            Write-Warning "Connect-AksHciSdnVnet is no longer required for the current akshci version."
            Write-Warning "Log collection and certificate rotation are supported without running this command."
            Write-Warning "To ssh to the vms use Get-AksHciSdnSshEndpoint to find the ssh ip and port to use."
            Write-Warning "To override this warning and connect to the SDN virtualnetwork, rerun with -Force parameter."
            throw "CMD_DEPRECATED"
        }
    }

    Write-Verbose "Connecting host $(hostname) to the Aks-HCI control plane virtual network"
    $nicName = "$script:HostVNetNicName-$(hostname)"

    $mocNic = Get-MocNic -nicName $nicName 
    $ncNic = Get-NCNic -resourceId $nicName

    if ($ncNic  -eq $null)
    {
        Write-Error "Failed to create a VNIC" -ErrorAction Stop
    }

    $vswitchName = (Get-MocConfigCache)["vswitchName"]
    $prefix = $mocNic.properties.ipConfigurations[0].properties.prefixlength
    $ip = $mocNic.properties.ipConfigurations[0].properties.privateIPAddress
    CreateHostVNic -name $nicName -mac $ncNic.properties.privateMacAddress -vswitchName $vswitchName -ipAddress $ip -prefix $prefix

    SetPortProfile -nicName $nicName -profileId $ncNic.instanceId
}

function Disconnect-AksHciSdnVnet
{
    param([switch] $RetainResources)   

    ValidateState

    Write-Verbose "Disconnecting host $(hostname) from the Aks-HCI control plane virtual network"
    $nicName = "$script:HostVNetNicName-$(hostname)"

    if ($RetainResources.IsPresent)
    {
        Remove-PortProfile -nicName $nicName
        Remove-PortProfile -nicName $script:HostVNetNicName
        return
    }
    

    CleanupHost -nicName $nicName
    CleanupMoc -nicName $nicName

    CleanupHost -nicName $script:HostVNetNicName
    CleanupMoc -nicName $script:HostVNetNicName
}

function Remove-AksHciSdnResources
{
    <#
    .SYNOPSIS
    Remove resources created by akshci which are stale/leaked in networkcontroller.
    .DESCRIPTION
    Leaked/stale resources may be present in NetworkController if there are any failures/bugs during uninstall-akshci or removal of a workload cluster. This cmdlet helps to remove those
    leaked resources from networkcontroller.
    .PARAMETER Leaked
    Remove leaked resources. This option can be used on a working akshci installation when any leaked resources are found in NetworkController.
    .PARAMETER All
    Removes stale resources after the uninstallation of akshci. Cannot be run when the akshci installation in present.
    .PARAMETER NCRestEndPoint
    NC rest endpoint(without https:// prefix). Must be specified when using -All option.
    .EXAMPLE
    # Clean up leaked resources when akshci installation is in a working state.
    Remove-AksHciSdnResources -Leaked
    .EXAMPLE
    #Clean up any stale/leaked resources after the uninstallation of akshci.
    Remove-AksHciSdnResources -All -NCRestEndpoint nc.contoso.com
    #>


    param(
        [Parameter(Mandatory = $true, ParameterSetName = 'All')]
        [switch] $All,
        [Parameter(Mandatory = $true, ParameterSetName = 'Leaked')]
        [switch] $Leaked,
        [Parameter(Mandatory = $true, ParameterSetName = 'All')]
        [string] $NCRestEndPoint,
        [switch] $Force
    )

    Initialize-AksHciNcCompanion
    Clear-MocCache

    $isInstalled = Test-IsAksHciInstalled
    
    if ($Leaked.IsPresent -and -not $isInstalled){
        throw "AksHci is not installed, unable to enumerate leaked resources"
    }

    if ($All.IsPresent -and $isInstalled) {
        throw "AksHci is installed. Cannot use -All option. Please retry with Remove-AksHciSdnResources -Leaked"
    }
    
    if ($isInstalled) {
        $NCRestEndPoint = (Get-MocConfigCache)["networkControllerFqdnOrIpAddress"]
    }

    # Check that the $NCRestEndPoint is correct. If it is, the test will populate the uri variable
    $uri = "https://$NCRestEndPoint"
    if (-not (Test-NetworkControllerFqdnOrIpAddress $uri)) {
        throw "Cannot connect to Network Controller $NCRestEndPoint"
    }

    if ($Leaked.IsPresent) {
    
        # Check if MocGroup is configured
        if ((Get-MocGroupCache).Length -eq 0) {
            Write-Warning "No MocGroup configured, cannot enumerate leaked resources"
            return
        }

        Remove-AksHciNcResourcesLeaked -ConnectionUri $uri -Force:$Force

        return
    }

    if ($All.IsPresent) {      
       
        Remove-AksHciNcResourcesAll -ConnectionUri $uri -Force:$Force
    }

}

function Get-MocResourceName
{
    param(
        [Parameter(Mandatory)]
        [ValidateNotNullOrEmpty()]
        [AksHciNcResourceType] $Type
    )

    # Get all the moc groups
    $mocGroup = Get-MocGroupCache
    # Return an array of string with the group names
    $mocGroup = $mocGroup | Select-Object -ExpandProperty name

    $mocResources = @()
    switch ($Type)
    {
        ([AksHciNcResourceType]::NetworkInterface)
        {
            foreach ($group in $mocGroup)
            {
                $mocGroupResource = Get-MocNetworkInterface -group $group -ErrorAction Stop
                # Extract only the name property for each resource in the group and save it
                $mocGroupResource = $mocGroupResource | Select-Object -ExpandProperty name
                $mocResources += $mocGroupResource
            }
        }
        ([AksHciNcResourceType]::LoadBalancer)
        {
            foreach ($group in $mocGroup)
            {
                $mocGroupResource = Get-MocLoadBalancer -group $group -ErrorAction Stop
                # Extract only the name property for each resource in the group and save it
                $mocGroupResource = $mocGroupResource | Select-Object -ExpandProperty name
                $mocResources += $mocGroupResource
            }
        }
        ([AksHciNcResourceType]::VirtualNetwork)
        {
            foreach ($group in $mocGroup)
            {
                $mocGroupResource = Get-MocVirtualNetwork -group $group -ErrorAction Stop
                # Extract only the name property for each resource in the group and save it
                $mocGroupResource = $mocGroupResource | Select-Object -ExpandProperty name
                $mocResources += $mocGroupResource
            }
        }
        Default
        {
            $mocResources = @()
        }
    }
    
    return $mocResources
}

function Get-AksHciNcResourceNameAndId
{
    param(
        [Parameter(Mandatory)]
        [ValidateNotNullOrEmpty()]
        [AksHciNcResourceType] $Type,

        [Parameter(Mandatory)]
        [ValidateNotNullOrEmpty()]
        [string] $ConnectionUri
    )

    # Query all Nc resources from the rest api
    switch ($Type) {
        ([AksHciNcResourceType]::NetworkInterface)
        {
            $ncResponse = Get-NetworkControllerNetworkInterface -ConnectionUri $ConnectionUri
        }
        ([AksHciNcResourceType]::LoadBalancer)
        {
            $ncResponse = Get-NetworkControllerLoadBalancer -ConnectionUri $ConnectionUri
        }
        ([AksHciNcResourceType]::VirtualNetwork)
        {
            $ncResponse = Get-NetworkControllerVirtualNetwork -ConnectionUri $ConnectionUri
        }
        Default
        {
            $ncResponse = @()
        }
    }

    # Remove entries that don't have ResourceMetadata
    $ncResponse = $ncResponse | Where-Object -Property ResourceMetadata -ne nil
    # Re-format object to extract ResourceId, ResourceMetadata.Client, ResourceMetadata.ResourceName
    $ncResponse = $ncResponse | Select-Object ResourceId, @{N="Client";E={$_.ResourceMetadata.Client}}, @{N="ResourceName";E={$_.ResourceMetadata.ResourceName}}
    # Only select the AksHci Client Objects
    $ncResponse = $ncResponse | Where-Object -Property Client -eq "AksHci"

    return $ncResponse
}

function Remove-NCResource
{
    param(
        [Parameter(Mandatory)]
        [ValidateNotNullOrEmpty()]
        [AksHciNcResourceType] $Type,

        [Parameter(Mandatory)]
        [ValidateNotNullOrEmpty()]
        [string] $ConnectionUri,
        
        [Parameter(Mandatory)]
        [ValidateNotNullOrEmpty()]
        [string] $ResourceId
    )

    switch ($Type) {
        ([AksHciNcResourceType]::NetworkInterface)
        {
            Remove-NetworkControllerNetworkInterface -ConnectionUri $ConnectionUri -ResourceId $ResourceId -ErrorAction SilentlyContinue -Force
        }
        ([AksHciNcResourceType]::LoadBalancer)
        {
            Remove-NetworkControllerLoadBalancer -ConnectionUri $ConnectionUri -ResourceId $ResourceId -ErrorAction SilentlyContinue -Force
        }
        ([AksHciNcResourceType]::VirtualNetwork)
        {
            Remove-NetworkControllerVirtualNetwork -ConnectionUri $ConnectionUri -ResourceId $ResourceId -ErrorAction SilentlyContinue -Force
        }
        Default
        {
            Write-Error "Network Controller Resource type unknown, cannot delete"
        }
    }
}

function Remove-AksHciNcResourcesLeaked
{
    param(
        [switch] $Force,
        [Parameter(Mandatory = $true)]
        [ValidateNotNullOrEmpty()]
        [string] $ConnectionUri
    )

    # go over each resource type and compare Moc to NC
    foreach ($resType in [AksHciNcResourceType].GetEnumValues())
    {
        $mocResources = Get-MocResourceName -Type $resType
        if ($mocResources.Length -eq 0)
        {
            Write-Error "No $resType resources found in Moc, cannot compare with NetworkController"
            continue
        }

        $ncResources = Get-AksHciNcResourceNameAndId -ConnectionUri $ConnectionUri -Type $resType
        if ($ncResources.Length -eq 0)
        {
            Write-Error "No $resType resources found in NetworkController, nothing to remove"
            continue
        }

        # Compare the objects
        $diff = Compare-Object -ReferenceObject $mocResources -DifferenceObject ($ncResources | Select-Object -ExpandProperty ResourceName)

        # NetworkController resources not found in Moc
        $ncOnly = $diff | Where-Object -Property SideIndicator -eq "=>" | Select-Object -ExpandProperty InputObject
        if ($ncOnly.Length -gt 0)
        {
            $ncOnlyResourceName = $ncOnly
            # Expand ncOnly from ResourceName to include ResourceId. This may be a slow operation, for each resource in $ncOnly,
            # it will search through $ncNetworkInterfaces to find the full object
            $ncOnly = $ncOnly | ForEach-Object { $ncResources | Where-Object -Property ResourceName -eq $_ }

            # Prompt user to delete
            $decision = 1
            if (!$Force.IsPresent)
            {
                $title = "Delete NetworkController$($resType) Resources"
                $question = "Are you sure you want to delete the following resources:`n$($ncOnlyResourceName -join "`n")"
                $choices = @('&Yes', '&No', '&Suspend')
                $decision = $Host.UI.PromptForChoice($title, $question, $choices, 1)
            }

            # Delete the Network Controller resources
            if ($decision -eq 0 -or $Force.IsPresent)
            {
                foreach ($n in $ncOnly)
                {
                    Remove-NCResource -ConnectionUri $ConnectionUri -Type $resType -ResourceId $n.ResourceId
                }
            }
            # Suspend, exit immediatly
            elseif ($decision -eq 2)
            {
                return
            }
        }
        else
        {
            Write-Host "No Leaked NetworkController$($resType) resources found"
        }
    }
}

function Remove-AksHciNcResourcesAll
{
    param(
        [switch] $Force,
        [Parameter(Mandatory)]
        [ValidateNotNullOrEmpty()]
        [string] $ConnectionUri
    )

    foreach ($resType in [AksHciNcResourceType].GetEnumValues())
    {
        $ncResources = Get-AksHciNcResourceNameAndId -ConnectionUri $ConnectionUri -Type $resType

        # Prompt user to delete
        $decision = 1
        if (!$Force.IsPresent -and $ncResources.Length -ne 0)
        {
            $title = "Delete NetworkController$($resType) Resources"
            $question = "Are you sure you want to delete the following resources:`n$(($ncResources | Select-Object -ExpandProperty ResourceName) -join "`n")"
            $choices = @('&Yes', '&No', '&Suspend')
            $decision = $Host.UI.PromptForChoice($title, $question, $choices, 1)
        }

        # Delete the Network Controller resources
        if ($decision -eq 0 -or $Force.IsPresent)
        {
            foreach ($n in $ncResources)
            {
                Remove-NCResource -ConnectionUri $ConnectionUri -Type $resType -ResourceId $n.ResourceId
            }
        }
        # Suspend, exit immediatly
        elseif ($decision -eq 2)
        {
            return
        }
    }
}

function Test-AksHciSdnConfig
{
    # By default, we will use the values in Moc to test NC
    # If 1 parameter is manually set, all parameters must be supplied and moc will not be used
    [CmdletBinding(DefaultParameterSetName = 'Moc')]
    param(
        [Parameter(Mandatory = $true, ParameterSetName = 'Manual')]
        [ValidateNotNullOrEmpty()]
        [string] $networkControllerFqdnOrIpAddress,
        
        [Parameter(Mandatory = $true, ParameterSetName = 'Manual')]
        [ValidateNotNullOrEmpty()]
        [string] $networkControllerLbSubnetRef,
        
        [Parameter(Mandatory = $true, ParameterSetName = 'Manual')]
        [ValidateNotNullOrEmpty()]
        [string] $networkControllerLnetRef,
        
        [Parameter(Mandatory = $true, ParameterSetName = 'Manual')]
        [ValidateNotNullOrEmpty()]
        [VipPoolSettings] $vipPool
    )

    # ----- Setup -----
    $useMoc = $PSCmdlet.ParameterSetName -eq "Moc"
    $ret = $true
    Clear-MocCache

    # Read all values from Moc
    if ($useMoc)
    {
        $mocConfig = Get-MocConfigCache

        # Check that Moc has been initialized
        if ($mocConfig["installState"] -eq [InstallState]::NotInstalled)
        {
            Write-Error "Moc is not in an installed state. Cannot read NetworkController config in Moc"
            $ret = $false
            # Cannot continue any more tests without Moc
            return $ret
        }

        # Get networkControllerFqdnOrIpAddress
        $networkControllerFqdnOrIpAddress = $mocConfig["networkControllerFqdnOrIpAddress"]
        # Get networkControllerLbSubnetRef
        $networkControllerLbSubnetRef = $mocConfig["networkControllerLbSubnetRef"]
        # Get networkControllerLnetRef
        $networkControllerLnetRef = $mocConfig["networkControllerLnetRef"]
        # Get vipPool
        $vipPoolDefaultName = $mocConfig["defaultvippoolname"]
        $vipPoolDefault = Get-MocVipPool -location MocLocation -name $vipPoolDefaultName
        $vipPool = [VipPoolSettings]::new($vipPoolDefault.name, $vipPoolDefault.properties.startIp, $vipPoolDefault.properties.endIp)
    }

    # ----- Tests -----
    # Check networkControllerFqdnOrIpAddress
    Write-Verbose "Check networkControllerFqdnOrIpAddress"
    $uri = "https://$networkControllerFqdnOrIpAddress"
    $result = Test-NetworkControllerFqdnOrIpAddress $uri
    if (-not $result)
    {
        Write-Error "Cannot connect to NetworkController with the networkControllerFqdnOrIpAddress $networkControllerFqdnOrIpAddress"
        $ret = $false
        # Cannot test any more without the Network Controller Uri
        return $ret
    }

    # Check networkControllerLnetRef
    Write-Verbose "Check networkControllerLnetRef"
    $result = Test-NetworkControllerReference $uri $networkControllerLnetRef
    if (-not $result)
    {
        Write-Error "networkControllerLnetRef ($networkControllerLnetRef) cannot be found in the Network Controller"
        $ret = $false
    }

    # Check networkControllerLbSubnetRef
    Write-Verbose "Check networkControllerLbSubnetRef"
    $lbSubnetWebResult = $null
    $result = Test-NetworkControllerReference $uri $networkControllerLbSubnetRef ([ref]$lbSubnetWebResult)
    if (-not $result)
    {
        Write-Error "networkControllerLbSubnetRef ($networkControllerLbSubnetRef) cannot be found in the Network Controller"
        $ret = $false
    }

    # Check the vipPool -- uses previous networkControllerLbSubnetRef result
    Write-Verbose "Check vippool setting"
    if ($null -ne $lbSubnetWebResult)
    {
        $result = Test-NetworkControllerDefaultVipPool ($lbSubnetWebResult.Content) $vipPool
        if (-not $result)
        {
            Write-Error "AksHci Vippool ip range is not within the NetworkController Vippool range"
            $ret = $false
        }
    }
    else
    {
        Write-Warning "Skipping Vippool Test, lb Subnet Ref not found in Network Controller"
    }

    return $ret
}

function Test-NetworkControllerFqdnOrIpAddress([string]$ConnectionUri)
{
    try {
        $result = Invoke-WebRequest $ConnectionUri/networking/discovery -UseBasicParsing
        if ($result.StatusCode -eq 200) {
            return ($result.Content | ConvertFrom-Json).resourceId -eq "discovery"
        }
    }
    catch {
    }

    return $false

}

function Test-NetworkControllerReference([string]$ConnectionUri, [string]$subRef, [ref]$webResult)
{
    $uri = "$ConnectionUri/networking/v1$subRef"
    $result = Invoke-WebRequest -UseBasicParsing -Uri $uri
    if ($null -eq $result -or $result.StatusCode -ne 200)
    {
        return $false
    }

    # set the webResult if it was passed
    if ($null -ne $webResult)
    {
        $webResult.Value = $result
    }
    return $true
}

function Test-NetworkControllerDefaultVipPool([string]$lbSubnetContent, [VipPoolSettings]$mocVipPool)
{
    $ncVippool = (ConvertFrom-Json -InputObject $lbSubnetContent).Properties.ipPools.Properties
    if ([string]::IsNullOrWhiteSpace($ncVippool))
    {
        return $false
    }

    # check that the IP addresses are in the following order
    # NCVippoolStart -> MocVippoolStart -> MocVippoolEnd -> NCVippoolEnd
    return (
        [AKSHCI.IPUtilities]::ValidateRange($ncVippool.startIpAddress, $vipPool.VipPoolStart) -and
        [AKSHCI.IPUtilities]::ValidateRange($vipPool.VipPoolStart, $vipPool.VipPoolEnd) -and
        [AKSHCI.IPUtilities]::ValidateRange($vipPool.VipPoolEnd, $ncVippool.endIpAddress)
    )
}
function Get-FirstNcVmName()
{
    $networkControllerVMName = $null
    $isMultiNode = $false

    try {
        $failoverCluster = Get-FailoverCluster -ErrorAction SilentlyContinue
        $isMultiNode = $failoverCluster -ne $null
    }
    catch {}

    try 
    {
        
        if ($isMultiNode)
        {
            Get-ClusterNode -ErrorAction Stop | ForEach-Object {
                
                if ([string]::IsNullOrWhiteSpace($networkControllerVMName))
                {
                    Write-Msg "Checking $_"
                    $networkControllerVMName = Invoke-Command -ComputerName $_ -ScriptBlock {
                        (Get-VM | Get-VMNetworkAdapter | where-object { $_.IPAddresses.Contains($((Get-NetTCPConnection -RemotePort 8571).RemoteAddress)) }).VMName
                    }
                }
            }
        }
        else 
        {
            $networkControllerVMName = (Get-VM | Get-VMNetworkAdapter | where-object { $_.IPAddresses.Contains($((Get-NetTCPConnection -RemotePort 8571).RemoteAddress)) }).VMName            
        }        
    }
    catch 
    {
        Write-Error("Unable to detect NC VM name")
    }

    return $networkControllerVMName
}

function Get-NatRuleMap
{
    <#
    .DESCRIPTION
        Fetches the ssh nat lb rules for the group and caches it.
        Returns a map of all ssh nat rules for the given group
    .PARAMETER groupName
        Group name
    #>


     param (
        [Parameter()]
        [String] $groupName
    )

    if ($script:natRuleCache -ne $null)
    {
        $ruleMap = $script:natRuleCache[$groupName]
        if ($ruleMap -ne $null) 
        {
            return $ruleMap
        }
    }
    else 
    {
        $script:natRuleCache = @{}
    }

    $ruleMap = @{}
    $lbs = Get-MocLoadBalancer -group $groupName

    if ($lbs -eq $null) 
    {
        $script:natRuleCache[$groupName] = $ruleMap
        return $ruleMap
    }
    
    
    $inboundLb = $lbs | where-object { $_.properties.inboundNatRules -ne $null }
    foreach($lb in $inboundLb) 
    {
        $sshNatRules = $lb.properties.inboundNatRules | `
            Where-Object { $_.Properties.backendPort -eq 22 }
        if($sshNatRules -ne $null)
        {
            $feIP = $lb.properties.frontendIPConfigurations[0].properties.ipAddress                
            foreach($natRule in $sshNatRules)
            {
                $port = $natRule.properties.frontendPort
                $sshData = @{}
                $sshData["ip"] = $feIP
                $sshData["port"] = $port
                $ruleMap[$natRule.Name] = $sshData
            }
        }
    }

    $natRuleCache[$groupName] = $ruleMap
    return $ruleMap
}
function Get-AksHciSdnSshEndpoint
{
    <#
    .SYNOPSIS
        Fetches the ssh ip and port for the given VM.
    .DESCRIPTION
        When SDN integration is configured, ssh access to the vm is via the inbound nat rule. This cmdlet fetches the frontend nat ip address
        and port from the loadbalancer.
    .PARAMETER VMName
        Name of the VM as shown in hyper-v for which the ssh configuration should be fetched. If VMName is omitted, ssh information for all the vms in the given cluster is fetched.
    .PARAMETER ClusterName
        Name of the targetcluster where the vm belongs
    .PARAMETER Management
        Denotes that the information should be fetched for the management cluster.
    .EXAMPLE
    # Fetch ssh endpoint of the management cluster VMs
    Get-AksHciSdnSshEndpoint -Management
    .EXAMPLE
    # Fetch ssh endpoint of the given VM from the given target cluster
    Get-AksHciSdnSshEndpoint -VMName testcluster-control-plansfdsfe-c784t-d0e6fasdffe2 -clustername testcluster
    .EXAMPLE
    # Fetch ssh endpoints for all the VM from the given target cluster
    Get-AksHciSdnSshEndpoint -clustername testcluster
    #>

    
     param (
        [Parameter(Mandatory = $false)]
        [String] $VMName,
        [Parameter(Mandatory = $true, ParameterSetName = 'Cluster')]
        [String] $ClusterName,
        [Parameter(Mandatory = $true, ParameterSetName = 'Management')]
        [Switch] $Management
    )

    $groupName = "clustergroup"
    if (-not $Management.IsPresent) {
        $groupName += "-$ClusterName"
    }

    if (-not [String]::IsNullOrWhiteSpace($VMName)) {
        $vmAr = $VMName.Split("-")
        $mocVmName = ($vmAr | select -First ($vmAr.count - 1)) -join "-"
        $mocVms = Get-MocVirtualMachine -Name $mocVmName -group $groupName
    } else {
        $mocVms = Get-MocVirtualMachine -group $groupName
    }    

    $eps = @()
    $mocVms | Foreach-Object {
        $nicName = $_.virtualmachineproperties.networkprofile.networkinterfaces[0].id
        $ep = Get-NatEndpoint -nicName $nicName -groupName $groupName

        $result = New-Object -TypeName psobject
        $result | Add-Member -MemberType NoteProperty -Name VMName -Value $_.Name
        $result | Add-Member -MemberType NoteProperty -Name PrivateIP -Value $ep.privateip
        $result | Add-Member -MemberType NoteProperty -Name SSH-IP -Value $ep.ip
        $result | Add-Member -MemberType NoteProperty -Name SSH-Port -Value $ep.port

        $eps += $result
    }

    return $eps
   
}
function Get-NatEndpoint
{
    <#
    .DESCRIPTION
        Fetches the ssh frontend nat ip and port for the given nic. When SDN integration is configured,
        ssh access to the node is via the inbound nat rule
 
    .PARAMETER nicName
        Networkinterface name
    .PARAMETER groupName
        Group name
    #>

     param (
        [Parameter()]
        [String] $nicName,

        [Parameter()]
        [String] $groupName
    )
   
    $nic = Get-MocNetworkInterface -group $groupName -name $nicName
    $natRuleName = $nic.properties.ipConfigurations[0].properties.loadBalancerInboundNatRules[0].name

    if ($natRuleName -eq $null -or $natRuleName[0] -eq $null)
    {
        return $null
    }

    $ruleMap = Get-NatRuleMap -groupName $groupName
    $ep = $ruleMap[$natRuleName]
    $ep["privateip"] = $nic.properties.ipConfigurations[0].properties.privateIPAddress
    return $ep
}

function Get-AksHciSdnLogs
{

    <#
    .DESCRIPTION
        Collects SDN logs.
    .PARAMETER networkControllerVMName
        NetworkController VM name. If there are more than one VM, use any one of the VM name.
    #>


     param (
        [Parameter()]
        [String] $NetworkControllerVMName,
        [Parameter()]
        [switch] $SkipAksHciLogs,
        [Parameter()]
        [switch] $VirtualMachineLogs,
        [Parameter()]
        [switch] $AgentLogs,
        [Parameter()]
        [switch] $EventLogs,
        [Parameter()]
        [switch] $KvaLogs,
        [Parameter()]
        [switch] $DownloadSdkLogs,
        [Parameter()]
        [switch] $BillingRecords
    )
   
    if ([string]::IsNullOrWhiteSpace($NetworkControllerVMName))
    {
        $NetworkControllerVMName = Get-FirstNcVmName
    }

    if ([string]::IsNullOrWhiteSpace($NetworkControllerVMName))
    {
        Write-Error("Unable to detect NC VM name, please use -networkControllerVMName and provide one of the NC vm name") -ErrorAction Stop
    }

    if(-not $skipAksHciLogs.IsPresent) {
        $zipName = Get-AksHciLogs -EventLogs:$EventLogs.IsPresent -KvaLogs:$KvaLogs.IsPresent -DownloadSdkLogs:$DownloadSdkLogs.IsPresent `
            -BillingRecords:$BillingRecords.IsPresent -VirtualMachineLogs:$VirtualMachineLogs.IsPresent -AgentLogs:$AgentLogs.IsPresent
    }

    $environmentDetails = Get-SdnInfrastructureInfo -NetworkController $NetworkControllerVMName
    Install-SdnDiagnostics -ComputerName ($environmentDetails.FabricNodes | where {$_ -ne $null})
    Start-SdnDataCollection -NetworkController $environmentDetails.NetworkController[0] -Role Gateway,NetworkController,Server,LoadBalancerMux -IncludeLogs -IncludeNetView
    Write-Host "AKS-HCI Log location $zipName" -ForegroundColor Green

}

function Test-IsAksHciInstalled()
{

    try {
        $config = (Get-AksHciConfig -ErrorAction Stop).AksHci
    } catch {
        
        $isNotInstalled = $_.InvocationInfo.Line.Contains("generic_cannot_deploy")

        if ($isNotInstalled) {
            return $false
        }

        if ($_.FullyQualifiedErrorId -eq "CommandNotFoundException") {
            return $false
        }

        throw
    }

    $currentState = $config.installState

    if (-not $currentState)
    {
        return $false
    }

    switch($currentState)
    {
        $([InstallState]::NotInstalled) { return $false }
        $([InstallState]::InstallFailed) { return $false }
    }

    return $true
}

function Write-ConnectionResult([string] $scenario, $result)
{
    $status = $true
    $result | Foreach-Object {
        Write-Msg "Test Connection Result - IP $($_.IP), Port $($_.Port), Status $($_.Status)"
        if ($($_.Status) -eq "Failed") {
            $status = $false
        }
    }  

    if ($status) {
        Write-Msg "Test Connection for $scenario, Status Succeeded"
    }

    return $status
}
function Test-Connectivity([string]$ip, [string]$port)
{
    $ProgressPreference = "SilentlyContinue"
    try {
        $result = Test-NetConnection -ComputerName $ip -Port $port -ErrorAction SilentlyContinue
        Write-Msg -msg "$($MyInvocation.MyCommand) IP [$ip] Port [$port]"
        if ($result -eq $null -or $result.TcpTestSucceeded -ne $true)
        {
            return $false
        }
    }
    finally {
        $ProgressPreference = "Continue"
    }

    return $true
}
function Test-AllLoadbalancingRules([string]$lbName, [string]$group, [bool]$includeSsh)
{
    $testResult = Get-TestResult -testName "LBConnectivity" -testResults @()

    try {
        $lb = Get-MocLoadBalancer -name "$lbName" -group $group
        $feIP = $lb.properties.frontendIPConfigurations[0].properties.ipAddress
        $fePorts = @()
        $lb.properties.loadBalancingRules | foreach-object { 
            $fePorts += $_.Properties.FrontendPort
        }

        if ($includeSsh)
        {
            $fePorts += "22"
        }
        $overallStatus = $true
        $fePorts | foreach-object {            
            $status = Test-Connectivity -ip $feIP -port $_   
            $detailedStatus = "LB [$lbName] feIP [$feIP] fePort [$_] status "
            if ($status) {
                $detailedStatus += "[Success]"
            } else {
                $detailedStatus += "[Failed]"
            }

            $result = Get-Result -scenario "LBConnectivity" -isSuccess $status -detailedStatus $detailedStatus
            $testResult.TestResults += $result
            $overallStatus =  $overallStatus -and $result.Passed
        }
    } catch {
        $testResult.DetailedStatus = Convert-Error $_
        $overallStatus = $false
    }

    $testResult.Passed = $overallStatus
    return $testResult
}

function Test-ManagementControlPlaneConnectivity()
{
    $testResult = Get-TestResult -testName "ManagementControlPlaneConnectivity" -testResults @()
    try {        
        Write-Msg -msg "$($MyInvocation.MyCommand)"
        $kvaName = (Get-KvaConfig).kvaname
        $lbName = "$kvaName-load-balancer"
        $lbTestResult = Test-AllLoadbalancingRules -lbName $lbName -group "clustergroup" -includeSsh $false
        $status = $lbTestResult.Passed
        $testResult.TestResults += $lbTestResult
        $testResult.Passed = $lbTestResult.Passed

        if (-not $status) {
            throw "Management controlplane connectivity test failed"
        }

        Write-Pass -testName "$($MyInvocation.MyCommand)"
    }
    catch {
        Write-Fail -testName "$($MyInvocation.MyCommand)" -err $_
        $testResult.Passed = $false
        $testResult.DetailedStatus = Convert-Error $_
    }

    return $testResult
}


function Test-TargetControlPlaneConnectivity()
{
    param (
        [string]$clusterName
    )

    $testName = "$($MyInvocation.MyCommand)"
    $testResult = Get-TestResult -testName $testName -testResults @() -isSuccess $true
    try {
        Write-Msg -msg "$($MyInvocation.MyCommand)"
        $status = $true
        $clusters = Get-AksHciCluster -ErrorAction Stop -Name $clusterName
        $clusters | ForEach-Object {
            $name = $($_.Name)
            Write-Msg -msg "$($MyInvocation.MyCommand) Cluster [$name]"
            $lbName = "$($_.Name)-load-balancer"
            $lbTestResult = Test-AllLoadbalancingRules -lbName $lbName -group "clustergroup-$name" -includeSsh $false
            $lbTestResult.TestName = "$testName-[$name]"
            $status = $status -and $lbTestResult.Passed
            $testResult.TestResults += $lbTestResult
            $testResult.Passed =  $testResult.Passed -and $lbTestResult.Passed
        }

        if (-not $status) {
            throw "TargetControlPlaneConnectivity test failed"
        }

        Write-Pass -testName "$($MyInvocation.MyCommand)"
    } catch {
        Write-Fail -testName "$($MyInvocation.MyCommand)" -err $_
        $testResult.Passed = $false
        $testResult.DetailedStatus = Convert-Error $_
    }

    return $testResult
}

function Reset-Key
{
     param (
        [string]$sshIP,
        [string]$sshPort
     )

    if ($sshPort -eq "22") {
        $keyGenAddress = $sshIP    
    } else {
        $keyGenAddress = "[$sshIP]`:$sshPort"    
    }

    (& ssh-keygen -R $keyGenAddress) 2>&1>$null
}

function Run-SshCommand
{
       param (
        [string]$sshIP,
        [string]$sshPort,
        [string]$command,
        [bool]$windows = $false
    )
    
    Reset-Key -sshIP $sshIP -sshPort $sshPort 
    $sshKey = (Get-MocConfigCache).sshPrivateKey
    if ($windows) {
        $bytes = [System.Text.Encoding]::Unicode.GetBytes($command)
        $encodedCommand = [Convert]::ToBase64String($bytes)
        $response = (& ssh -oBatchMode=yes -i $sshKey -o LogLevel=ERROR -o StrictHostKeyChecking=no "$script:WindowsUserName@$sshIP" -p $sshPort powershell.exe -NonInteractive -encodedCommand $encodedCommand 2>$null)

    } else {
        $response = (& ssh -oBatchMode=yes -i $sshKey -o LogLevel=ERROR -o StrictHostKeyChecking=no "$script:LinuxUserName@$sshIP" -p $sshPort $command)
    }
    
    return $response
}

function Test-CloudAgentConnectivity
{
    param (
        [string]$sshIP,
        [string]$sshPort,
        [bool]$windows = $false
    )
    
    Reset-Key -sshIP $sshIP -sshPort $sshPort

    $sshKey = (Get-MocConfigCache).sshPrivateKey
    $cloudFqdn = $(Get-MocConfigCache).cloudFqdn
    $cloudAgentAddress = (Resolve-dnsName $cloudFqdn -Type A -DnsOnly -ErrorAction Stop).IPAddress
    $status = $false

    if ($windows) {
        $command ="
        invoke-command -scriptblock {
            try {
                    (Test-NetConnection $cloudAgentAddress -port $script:CloudAgentPort -ErrorAction SilentlyContinue -WarningAction SilentlyContinue).TcpTestSucceeded
                }
                catch {
                }
            }"
       
        $response = Run-SshCommand -sshIP $sshIP -sshPort $sshPort -command $command -windows $windows
        $status = ($response -ieq "True")
    } else {
        $cloudAgentEndpoint = $cloudAgentAddress + ":" + $script:CloudAgentPort
        $command = "curl --connect-timeout 5 $cloudAgentEndpoint --stderr - `| grep curl:"
        $response = Run-SshCommand -sshIP $sshIP -sshPort $sshPort -command $command -windows $windows
        $status = $response -ieq $script:CloudAgentCurlResponse
    }  

    if ($status) {
        $msg = "Cloudagent connectivity to [$cloudAgentAddress] succeeded"
    } else {
        $msg = "Cloudagent connectivity to [$cloudAgentAddress] failed"
    }
    
    $result = Get-Result -scenario "cloudagent" -isSuccess $status -detailedStatus $msg
    return $result
}

function Test-InternetConnectivity
{
    param (
        [string]$sshIP,
        [string]$sshPort,
        [bool]$windows = $false
    )
    Reset-Key -sshIP $sshIP -sshPort $sshPort
    $sshKey = (Get-MocConfigCache).sshPrivateKey

    $status = $false
    $response = $null

    if ($windows) {
        $command ="
        invoke-command -scriptblock {
            try {
                    (Invoke-WebRequest -UseBasicParsing -Uri $($script:InternetUrl) -ErrorAction SilentlyContinue -WarningAction SilentlyContinue).Content
                }
                catch {
                }
            }"

    } else {
        $command = "curl --connect-timeout 5 $($script:InternetUrl) -s -f"
    }

    $response = Run-SshCommand -sshIP $sshIP -sshPort $sshPort -command $command -windows $windows
    $status = $response -ieq $script:NCSIResponse

    if ($status) {
        $msg = "Internet connectivity to [$script:InternetUrl] succeeded"
    } else {
        $msg = "Internet connectivity to [$script:InternetUrl] failed"
    }
    
    $result = Get-Result -scenario "internet" -isSuccess $status -detailedStatus $msg
    return $result 
}

function Test-DnsResolution
{
        param (
        [string]$sshIP,
        [string]$sshPort,
        [bool]$windows = $false
    )

    Write-Verbose "$($MyInvocation.MyCommand) Query [$script:DnsResolutionName]"
    if ($windows) {
        $command ="
        invoke-command -scriptblock {
            try {
                    (Resolve-DnsName -Type A -DnsOnly $script:DnsResolutionName -ErrorAction SilentlyContinue -WarningAction SilentlyContinue).IPAddress
                }
                catch {
                }
            }"

    } else {
        $command = "dig +short A $script:DnsResolutionName | grep -v ';' "
    }

    $response = Run-SshCommand -sshIP $sshIP -sshPort $sshPort -command $command -windows $windows
    Write-Verbose "$($MyInvocation.MyCommand) Response [$response]"
    $status = (-not [String]::IsNullOrWhiteSpace($response)) -and $response.Count -gt 0 

    if ($status) {
        $msg = "DNS resolution to [$script:DnsResolutionName] succeeded"
    } else {
         $msg = "DNS resolution to [$script:DnsResolutionName] failed"
    }
    
    $result = Get-Result -scenario "dns" -isSuccess $status -detailedStatus $msg
    return $result 
}

function Test-ManagementClusterConnectivity
{
    param (
        [ValidateSet("cloudagent", "dns", "internet", "inbound")]
        [string]$scenario
    )

    try {
        
        Write-Msg -msg "$($MyInvocation.MyCommand) Scenario [$scenario]"
        $kvaName = (Get-KvaConfig).kvaname
        $lbName = "$kvaName-load-balancer"
        $lb = Get-MocLoadBalancer -name "$lbName" -group $(GetClusterGroup)
        $feIP = $lb.properties.frontendIPConfigurations[0].properties.ipAddress

        if ($scenario -ne "inbound" -and $script:VmInboundAccessState["managementControlPlane"] -eq $false) {
            Write-Msg -msg "$($MyInvocation.MyCommand) [SKIP NO INBOUND ACCESS] Scenario [$scenario]" -Warning
            throw "Test skipped, no inbound access to the management control plane vm"
        }

        switch ($scenario) {
            "inbound" {
                $status = $(Test-Connectivity  -ip $feIP -port 22)
                $script:VmInboundAccessState["managementControlPlane"] = $result
                if ($status) {
                    $msg = "Inbound connectivity to $feIP succeeded"
                } else {
                    $msg = "Inbound connectivity to $feIP failed"
                }

                $result = Get-Result -scenario "inbound" -isSuccess $status -detailedStatus $msg

                break
            }
            "cloudagent" {
                $result = $(Test-CloudAgentConnectivity -sshIP $feIP  -sshPort "22" -windows $false)
                break
            }
            "dns" {
                $result = $(Test-DnsResolution -sshIP $feIP  -sshPort "22" -windows $false)
                break                
            }
            "internet" {
                $result  = $(Test-InternetConnectivity -sshIP $feIP -sshPort "22" -windows $false)
                break
            }
        }                             
        $status = $result.Passed
        if (-not $status) {
            throw "Test-ManagementClusterConnectivity Scenario [$scenario] failed"
        }

        Write-Pass -testName "$($MyInvocation.MyCommand) Scenario [$scenario]"
        return $result
    } catch {
        Write-Fail -testName "$($MyInvocation.MyCommand) Scenario [$scenario]" -err $_
        return $false
    }
}

function Convert-Error($e)
{
    return $e | Select-Object @{N="Message"; E={$_.Exception.Message}}, @{N="Stack"; E={$_.ScriptStackTrace.Split("`n")[0]}}
}
function Test-VMConnectivity
{
    param (
        [ValidateSet("inbound","cloudagent", "dns", "internet")]
        [string]$scenario,
        [string]$vmName,
        [string]$group
    )

    $result = $false

    try {
        $vmTestResult = Get-VMTestResult -vmName $vmName -isSuccess $result -testName $scenario

        if ($scenario -ne "inbound" -and $script:VmInboundAccessState["$vmName"] -eq $false) {
            $msg = "$($MyInvocation.MyCommand) [SKIP NO INBOUND ACCESS] Scenario [$scenario] Group [$group] VM [$vmName]"
            Write-Msg -msg $msg -Warning
            $vmTestResult.DetailedStatus = $msg
            return $vmTestResult
        }

        Write-Msg -msg "$($MyInvocation.MyCommand) Scenario [$scenario] Group [$group] VM [$vmName]"
        $cloudFqdn = $(Get-MocConfigCache).cloudFqdn
        $cloudAgentAddress = (Resolve-dnsName $cloudFqdn -Type A -DnsOnly -ErrorAction Stop).IPAddress
        $sshKey = $(Get-MocConfigCache).sshPrivateKey
        $lbName = "$name-load-balancer"
        $lb = Get-MocLoadBalancer -name $lbName -group $group
        $feIP = $lb.properties.frontendIPConfigurations[0].properties.ipAddress
        $vm = Get-MocVirtualMachine -group $group -name $vmName
        $nicName = $_.virtualmachineproperties.networkprofile.networkinterfaces[0].id
        $isWindows = $_.virtualmachineproperties.osprofile.osType -eq "Windows"
        $endpoint = Get-NatEndpoint -groupName $group -nicName $nicName
            
        if ($endpoint -eq $null) {
            $msg = "$($MyInvocation.MyCommand) Unable to find ssh endpoint for vm $($_.Name) in cluster $name"
            $vmTestResult.DetailedStatus = $msg
            Write-Fail -msg $msg
            return
        }

        $sshInfo = "[$($endpoint.ip)`:$($endpoint.port)]"
        $vmTestResult.SSHInfo = $sshInfo

        Write-Msg -msg "$($MyInvocation.MyCommand) Scenario [$scenario] Group [$group] VM [$vmName], SSH [$($endpoint.ip)`:$($endpoint.port)]" 
        switch ($scenario) {
            "inbound" {
                $status = $(Test-Connectivity  -ip $($endpoint.ip) -port $($endpoint.port))
                $script:VmInboundAccessState["$vmName"] = $status
                if ($status) {
                    $msg = "Inbound connectivity to $sshInfo succeeded"
                } else {
                    $msg = "Inbound connectivity to $sshInfo failed"
                }
    

                $result = Get-Result -scenario "inbound" -isSuccess $status -detailedStatus $msg
                break
            }
            "cloudagent" {
                $result = $(Test-CloudAgentConnectivity -sshIP $($endpoint.ip) -sshPort $($endpoint.port) -windows $isWindows)
                break
            }
            "dns" {
                $result = $(Test-DnsResolution -sshIP $($endpoint.ip) -sshPort $($endpoint.port) -windows $isWindows)
                break                
            }
            "internet" {
                $result = $(Test-InternetConnectivity -sshIP $($endpoint.ip) -sshPort $($endpoint.port) -windows $isWindows)
                break
            }
        }              

        $vmTestResult.Results = @($result)
        $vmTestResult.Passed = $result.Passed

        if (-not $result.Passed) {
            throw "Test-VMConnectivity Scenario [$scenario] VM [$vmName] failed "
        }

        Write-Pass -testName "$($MyInvocation.MyCommand) Scenario [$scenario] Group [$group] VM [$vmName]"
    } catch {
        Write-Fail -testName "$($MyInvocation.MyCommand) Scenario [$scenario] Group [$group] VM [$vmName]" -err $_
        $vmTestResult.DetailedStatus = Convert-Error $_
        $vmTestResult.Passed = $false
    }

    return $vmTestResult
}

function Test-TargetClusterConnectivity
{
    param (
        [ValidateSet("inbound", "cloudagent", "dns", "internet")]
        [string]$scenario,
        [string]$clusterName
    )

    $testName = "$($MyInvocation.MyCommand)[$scenario]"
    $testResult = Get-TestResult -testName $testName -isSuccess $true -testResults @()
    try {
        Write-Msg -msg "$($MyInvocation.MyCommand) Scenario [$scenario]"
        $cloudFqdn = $(Get-MocConfigCache).cloudFqdn
        $cloudAgentAddress = (Resolve-dnsName $cloudFqdn -Type A -DnsOnly -ErrorAction Stop).IPAddress
        $sshKey = $(Get-MocConfigCache).sshPrivateKey
        $status = $true

        $clusters = Get-AksHciCluster -ErrorAction Stop -Name $clusterName
        $clusters | ForEach-Object {
            $name = $_.Name
            $group = "clustergroup-$name"
            Write-Msg -msg "$($MyInvocation.MyCommand) Cluster [$name]"
            $lbName = "$name-load-balancer"
            $lb = Get-MocLoadBalancer -name $lbName -group $group
            $feIP = $lb.properties.frontendIPConfigurations[0].properties.ipAddress
            $vms = Get-MocVirtualMachine -group $group
            $vms | ForEach-Object {
                $result = Test-VMConnectivity -scenario $scenario -vmName $_.Name -group $group
                $testResult.TestResults += $result
                $status = $status -and $result.Passed
            }
        }

        if (-not $status) {
            throw "Test-TargetClusterConnectivity [$scenario] failed"
        }

        Write-Pass -testName "$($MyInvocation.MyCommand) Scenario [$scenario]"
    }
    catch  {
        $status = $false
        Write-Fail -testName "$($MyInvocation.MyCommand) Scenario [$scenario]" -err $_
        $testResult.DetailedStatus = Convert-Error $_
    } finally  {
        $testResult.Passed = $status
    }
    return $testResult
}

function Test-ManagementClusterNetwork()
{
    param (
        [ValidateSet("cloudagent", "dns", "internet", "inbound", "all")]
        [string]$scenario = "all"
    )

    $status = $true
    $mgmtClusterResult = Get-ClusterTestResult -clusterName "Management" -testResults @()
    try {
        switch($scenario) {
            {$_ -in "inbound", "all"} {
                $result = Test-ManagementClusterConnectivity -scenario "inbound"
                $mgmtClusterResult.TestResults += $result
                $status = $status -and $result.Passed
            }
            {$_ -in "controlplane", "all"} {
                $result = Test-ManagementControlPlaneConnectivity
                $mgmtClusterResult.TestResults += $result
                $status = $status -and $result.Passed
            }

            {$_ -in "cloudagent", "all"} {
                $result = Test-ManagementClusterConnectivity -scenario "cloudagent"
                $mgmtClusterResult.TestResults += $result
                $status = $status -and $result.Passed
            }

            {$_ -in "dns", "all"} {
                $result = Test-ManagementClusterConnectivity -scenario "dns"
                $mgmtClusterResult.TestResults += $result
                $status = $status -and $result.Passed
            }

            {$_ -in "internet", "all"} {
                $result = Test-ManagementClusterConnectivity -scenario "internet"
                $mgmtClusterResult.TestResults += $result
                $status = $status -and $result.Passed
            }
        }
    } catch {
        $status = $false
        $mgmtClusterResult.DetailedStatus = Convert-Error $_
    }

    $mgmtClusterResult.Passed = $status
    return $mgmtClusterResult
}

function Test-TargetClusterNetwork()
{
    param (
        [ValidateSet("controlplane","cloudagent", "dns", "internet", "inbound", "all")]
        [string]$scenario="all",
        [string]$clusterName
    )

    $status = $true
    $testName = "$($MyInvocation.MyCommand)[$scenario]"
    $testResult = Get-TestResult -testName $testName -isSuccess $true -testResults @()

    try {
        $clusters = Get-AksHciCluster -ErrorAction Stop -Name $clusterName
        $clusters | ForEach-Object {
            $name = $_.Name
            $clusterTestResult = Get-ClusterTestResult -clusterName $name -isSuccess $true -testResults @()
            $testResult.TestResults += $clusterTestResult
            switch($scenario) {
                {$_ -in "inbound", "all"} {
                    $result = Test-TargetClusterConnectivity -clusterName $name -scenario "inbound"
                    $clusterTestResult.TestResults += $result
                    $status = $status -and $result.Passed
                }
                {$_ -in "controlplane", "all"} {
                    $result = Test-TargetControlPlaneConnectivity -clusterName $name 
                    $clusterTestResult.TestResults += $result
                    $status = $status -and $result.Passed
                }

                {$_ -in "cloudagent", "all"} {
                    $result = Test-TargetClusterConnectivity -clusterName $name -scenario "cloudagent"
                    $clusterTestResult.TestResults += $result
                    $status = $status -and $result.Passed
                }

                {$_ -in "dns", "all"} {
                    $result = Test-TargetClusterConnectivity -clusterName $name -scenario "dns"
                    $clusterTestResult.TestResults += $result
                    $status = $status -and $result.Passed
                }

                {$_ -in "internet", "all"} {
                    $result = Test-TargetClusterConnectivity -clusterName $name -scenario "internet"
                    $clusterTestResult.TestResults += $result
                    $status = $status -and $result.Passed
                }
            }
        }
    } catch {
        $status = $false
        $testResult.DetailedStatus = Convert-Error $_
    } finally {
        $testResult.Passed = $status
    }

    return $testResult

}

function Test-AksHciSdnNetwork()
{
    <#
    .SYNOPSIS
    Validates various network connectivity when SDN integration with AKS-HCI is present.
    .DESCRIPTION
    Performs the following validations for management and target clusters
    1. DNS resolution for bing.com
    2. Internet connectivity by connecting to "http://www.msftncsi.com/ncsi.txt"
    3. Connectivity from the vm to cloudagent endpoint on the physical host
    4. Inbound ssh connectivity to the VMs from the physical host
    #>


    if (-not $(Test-IsAksHciInstalled)) {
        Write-Warning "$($MyInvocation.MyCommand) AksHci is not installed, skipping validation"
        return
    }

    $status = $true
    $results = @()

    $result = Test-ManagementClusterNetwork
    $results += $result
    $status = $status -and $result.Passed

    $result = Test-TargetClusterNetwork
    $results += $result
    $status = $status -and $result.Passed

    if ($status) {
        Write-Pass -testName "$($MyInvocation.MyCommand)"
    }
    else {
        Write-Fail -testName "$($MyInvocation.MyCommand)" -msg "One or more tests failed"
    }

    Write-TestResults $results

    return $results
}