Utils.psm1

Import-Module $PSScriptRoot\ClusUtils.psm1
Import-Module $PSScriptRoot\Logger.psm1
Import-Module $PSScriptRoot\PsHelper.psm1

function Copy-FileToRemote {
    param(
        [parameter(Mandatory=$true)][string] $srcPath,
        [parameter(Mandatory=$true)][string] $srcFile,
        [parameter(Mandatory=$true)][string] $destMachine,
        [parameter(Mandatory=$true)][pscredential] $destMachineCred,
        [parameter(Mandatory=$true)][string] $destDirectory
    )

    $s = New-ReliablePsSession -ComputerName $destMachine -Credential $destMachineCred
    try {
        $fileExist = Invoke-Command -Session $s -ScriptBlock {
            param($destDirectory, $destFile)
            if (-not $(Test-Path $destDirectory)) {
                New-Item $destDirectory -ItemType Directory -Force -ErrorAction Stop | Out-Null
                return $false
            }
            $dest = Join-Path $destDirectory $destFile
            Test-Path $dest
        } -ArgumentList ($destDirectory, $srcFile) -ErrorAction Stop

        if ($fileExist) {
            return
        }
        $src = Join-Path $srcPath $srcFile
        if (Test-IsCsvPath -path $destDirectory) {
            Copy-Item -Path $src -Destination $destDirectory
            return
        }
        $fileSharePath = Join-Path "\\$destMachine" $destDirectory.Replace(":","$")
        if (Test-Path $fileSharePath) {
            Copy-Item -Path $src -Destination $fileSharePath
            return
        }
        Copy-Item -Path $src -Destination $destDirectory -ToSession $s

    } finally {
        if ($s -ne $null ) {
            $s | Remove-PSSession
        }
    }
}

function Test-FileExist {

    param(
        [parameter(Mandatory=$true)][string] $path
    )

    if (-not $(Test-Path $path)) {
        throw "File $path not found"
    }
}

function Test-VMSwitchExist {
    param(
        [array] $computerName,
        [pscredential] $credential = [pscredential]::Empty,
        [string] $switchName
    )

    $err = Invoke-Command -ComputerName $computerName -credential $credential -ScriptBlock {
        param($switchName)

        if($(Get-VMSwitch -Name $switchName -ErrorAction SilentlyContinue) -eq $null) {
            return "$(hostName)"
        }
    } -ArgumentList $switchName -ErrorAction Stop

    if ($err -ne $null) {
        throw "Switch $switchName not found on host(s) $err"
    }
}

function Test-InternetConnectivity {

    param (
        [string] $computerName = (hostname),
        [pscredential] $credential = [pscredential]::Empty,
        [string] $srcAddress
    )

    Write-TraceLog "Test-InternetConnectivity: Testing on $computerName from $srcAddress"

    $result = Invoke-ReliableCommand -ComputerName $computerName -Credential $credential -ScriptBlock {
        param($srcAddress)

        $destIP = (Resolve-DnsName bing.com -Type A -ErrorAction SilentlyContinue | Select -First 1).IPAddress

        if ($destIP -eq $null) {
            return $false
        }

        if ([string]::IsNullOrEmpty($srcAddress)) {
            $result = ping $destIP -n 2
        } else {
            $result = ping -S $srcAddress $destIP -n 2
        }

        if (($result | Select-String "unreachable" | Select -First 1) -ne $null) {
            return $false
        }

        return ($result | Select-String "Received = 0") -eq $null
    } -ArgumentList ($srcAddress)

    return $result -eq $true
}

function Resolve-InternetSwitch {
    param (
        [string] $computerName = (hostname),
        [pscredential] $credential = [pscredential]::Empty
    )

    $switches = Invoke-ReliableCommand -ComputerName $computerName -Credential $credential -ScriptBlock {
        $switches = (Get-VMSwitch).Name
        return $switches
    }

    foreach ($s in $switches) {
        Write-TraceLog "Resolve-InternetSwitch: Checking switch $s on computer $computerName"

        $ips = Invoke-ReliableCommand -ComputerName $computerName -Credential $credential -ScriptBlock {
            param($s)
            $nics = (Get-VMNetworkAdapter -SwitchName $s -ManagementOS).name
            $allProbableOutBoundIPs = @()
            foreach ($n in $nics) {
                $adapter = Get-NetAdapter | Where-Object {$_.Name -like "*$n*" -and $_.Status -eq "up"};
                if ($adapter -ne $null) {
                    $ip = (Get-NetIPAddress -InterfaceIndex $adapter.ifIndex -AddressFamily IPv4 | Where-Object {-not ($_.IPaddress.StartsWith("169"))}).IPAddress | Select -First 1
                    $allProbableOutBoundIPs += $ip
                }
            }
            return $allProbableOutBoundIPs
        } -ArgumentList($s)

        foreach($ip in $ips) {
            $result = Test-InternetConnectivity -srcAddress $ip -computerName $computerName -credential $credential
            if ($result) {
                Write-TraceLog "Resolve-InternetSwitch: Switch $s on computer $computerName selected"
                return $s
            }
        }
    }

    Write-TraceLog "Resolve-SdnSwitch: No suitable internet switch found on computer $computerName"
}

function New-InternetSwitch {
    param (
        [string] $computerName = (hostname),
        [pscredential] $credential = [pscredential]::Empty,
        [string] $name
    )

    $iswitch = Resolve-InternetSwitch -computerName $computerName -credential $credential

    if ($iswitch -ne $null) {
        return $true
    }

    $ifIndex = Get-InternetAdapterIndex -computerName $computerName -credential $credential

    if ($ifIndex -eq $null) {
        return $false
    }

    $scriptBlock = {
        param($ifIndex, $name)
        $ifdesc = (Get-NetAdapter -InterfaceIndex $ifIndex).InterfaceDescription
        new-vmswitch -Name $name -InterfaceDescription $ifdesc | out-null
    }

    Invoke-ReliableCommand -ComputerName $computerName -credential $credential `
        -ScriptBlock $scriptBlock -ArgumentList @($ifIndex, $name)

    return $true

}
function Resolve-SdnSwitch {
    param (
        [string] $computerName = (hostname),
        [pscredential] $credential = [pscredential]::Empty
    )

    $result = Invoke-ReliableCommand -ComputerName $computerName -Credential $credential -ScriptBlock {

        $allSwitches = Get-VMSwitch -ErrorAction SilentlyContinue
        foreach ($s in $allSwitches) {
            foreach($e in $s.Extensions) {
                if ($e.Name -eq "Microsoft Azure VFP Switch Extension" -and $e.Enabled -eq $true) {
                    $sdnSwitch = $s.Name
                }
            }
        }

        if (-not [String]::IsNullOrEmpty($sdnSwitch)) {
            return $sdnSwitch, "VFP"
        }
        $sdnSwitch = ($allSwitches | `
            Where-Object {$_.Name -ilike "*sdn*"}).Name

        if (-not [String]::IsNullOrEmpty($sdnSwitch)) {
            return $sdnSwitch, "NamePattern"
        }

        $sdnSwitch = ($allSwitches | `
            Where-Object { $_.NetAdapterInterfaceDescriptions -ilike "*Mellanox*" }).Name

        if (-not [String]::IsNullOrEmpty($sdnSwitch)) {
            return $sdnSwitch, "MellanoxNic"
        }

    }

    if ($result -ne $null -and -not [String]::IsNullOrEmpty($result[0])) {
        Write-TraceLog "Resolve-SdnSwitch: Selected $($result[0]) on computer $computerName, Reason $($result[1])"
        return $result[0]
    }

    Write-TraceLog "Resolve-SdnSwitch: No suitable sdn switch found on computer $computerName"
}

function New-SdnSwitch {
    param (
        [string] $computerName = (hostname),
        [pscredential] $credential = [pscredential]::Empty,
        [string] $name
    )

    $sdnswitch = Resolve-SdnSwitch -computerName $computerName -credential $credential

    if ($sdnswitch -ne $null) {
        return $true
    }

    $scriptBlock = {
        param($name)

        $mellanoxNic = (Get-NetAdapter -InterfaceDescription *Mellanox*).InterfaceDescription

        if ($mellanoxNic -eq $null) {
            return $false
        }

        new-vmswitch -Name $name -InterfaceDescription $mellanoxNic | out-null
        return $true
    }

    $result = Invoke-ReliableCommand -ComputerName $computerName -credential $credential `
        -ScriptBlock $scriptBlock -ArgumentList @($name)

    return $result

}

function Disable-LinkLocalNics {
    param (
        [string] $computerName = (hostname),
        [pscredential] $credential = [pscredential]::Empty
    )

    $scriptBlock = {
        $indexes = (Get-NetIPAddress | Where-Object {$_.IPaddress.StartsWith("169") }).InterfaceIndex

        foreach($index in $indexes) {
            $netAdapter = Get-NetAdapter -InterfaceIndex $index -ErrorAction SilentlyContinue

            if ($netAdapter.Name -ilike "*SMB*") {
                continue
            }

            if ($netAdapter -ne $null) {
               $netAdapter | Disable-NetAdapter -ErrorAction SilentlyContinue -Confirm:$false
            }
        }
    }

    Invoke-ReliableCommand -computerName $computerName -credential $credential -scriptblock $scriptBlock
}

function Get-InternetAdapterIndex
{
    param(
        [array]$computerName,
        [pscredential] $credential = [pscredential]::Empty
    )

    $scriptBlock = {
        $retry = 0
        while ($true) {
            try {
                $dnsRes = Resolve-DnsName bing.com -Type A -ErrorAction Stop -DnsOnly | select -First 1
                break
            } catch {
                if ($retry -lt 3) {
                    $retry++
                    Write-Error "Failed to resolve dns, retrying"
                    Start-Sleep 1
                    continue
                }
                Write-Error "Failed to resolve dns halting. Make sure the machine has internet access"
                Write-Error $_ -ErrorAction Stop
            }
        }


        $ip = $dnsRes.IpAddress

        $connect = Test-NetConnection -ComputerName $ip -Port 80
        $corpInterface = $connect.InterfaceAlias

        $adapter = Get-NetAdapter -Name $corpInterface -ErrorAction Stop


        return $adapter.ifIndex
    }

    $index = Invoke-ReliableCommand -computerName $computerName -credential $credential -scriptblock $scriptBlock
    return $index

}

function Update-CorpNicDnsServerAddress
{
    param(
        [array]$hosts,
        [pscredential] $credential = [pscredential]::Empty,
        [array] $dnsServers
    )

    if ($dnsServers -eq $null -or $dnsServers.Count -eq 0) {
        return
    }
    
    $updateDns = {
        param(
            $index,
            [array]$dnsServers
        )
        Set-DnsClientServerAddress -ServerAddresses $dnsServers -InterfaceIndex $index -ErrorAction Stop -Verbose
    } 

    $getIfIndex = {
        [array]$ifIndexes = (Get-DnsClientServerAddress | `
            Where-Object  { $_.AddressFamily -eq 2 -and $_.ServerAddresses.Count -gt 0}).InterfaceIndex | `
            Select-Object -Unique
        return $ifIndexes
    }

    $hosts | ForEach-Object {
        $ifIndexes = Invoke-ReliableCommand -computerName $_ -credential $credential -scriptblock $getIfIndex
        foreach($ifIndex in $ifIndexes) {       
            Write-TraceLog "Update-HostDnsServerAddress: Host ($_) adapter $ifIndex with dns $dnsServers"
            Invoke-ReliableCommand -computerName $_ -credential $credential -scriptblock $updateDns -ArgumentList @($ifIndex, $dnsServers)
        }
    }
}


function Get-NextMacAddress
{
    param(
        [String] $MacAddress
    )

    return ("{0:X12}" -f ([convert]::ToInt64($MacAddress.ToUpper().Replace(":", "").Replace("-", ""), 16) + 1)).Insert(2, "-").Insert(5, "-").Insert(8, "-").Insert(11, "-").Insert(14, "-")
}

function Restart-LanmanWorkstationService {

    Write-TraceLog "Restart-LanmanWorkstationService: Restarting LanmanWorkstation service"
    Get-SmbMapping | Remove-SmbMapping -Force | Out-Null
    Restart-Service LanmanWorkstation -Force | Out-Null

    $notStarted = $true
    $retryCount = 10
    while($notStarted) {
        $retryCount--
        Write-TraceLog "Restart-LanmanWorkstationService: Waiting for LanmanWorkstation service to start"
        Start-Sleep 5
        $notStarted = (get-service LanmanWorkstation).Status -ne "Running"

        if ($retryCount -lt 0) {
            throw "Restart-LanmanWorkstationService: Failed waiting for LanmanWorkstation service to start"
        }
    }
}

function Connect-SMBShare {
    param (
        [Parameter(Mandatory=$true)] [pscredential] $Credential,
        [Parameter(Mandatory=$true)] [string] $RemotePath,
        [Parameter(Mandatory=$false)] [switch] $RestartLanManServer
    )

    Write-TraceLog "Connect-SMBShare: Connecting to $RemotePath, User: $($Credential.UserName)"
    # net use sometimes hung with multiple connections etc, restart the service to fix it

    if ($RestartLanManServer.IsPresent) {
      Restart-LanmanWorkstationService
    }

    $retry = $true
    while($true) {
        try {
            New-SmbMapping -RemotePath $RemotePath -UserName $($Credential.UserName) -Password $($Credential.GetNetworkCredential().Password) -Persistent $true | Out-Null
            break
        } catch {

            Write-TraceLog "Connect-SMBShare: Failed to connect to $RemotePath"
            Write-TraceLog "Connect-SMBShare: Error $_"
            if ($_.Exception.Message -like "*Multiple connections to a server*" -and $retry) {
                $retry = $false
                Write-TraceLog "Connect-SMBShare: restarting lanmanworkstation service and retry"
                Restart-LanmanWorkstationService
                continue
            }

            return $false
        }
    }

    if($(Test-Path  $RemotePath)) {
        Write-TraceLog "Connect-SMBShare: Connected to $RemotePath"
    } else {
        Write-TraceLog "Connect-SMBShare: Failed to connect to $RemotePath"
        return $false
    }

    return $true
}

function Get-VMHostName {
    param (
       [Parameter(Mandatory=$true)] [array] $Computers,
       [Parameter(Mandatory=$true)] [pscredential] $Credential,
       [Parameter(Mandatory=$true)] [string] $VMName
       )

   foreach($c in $Computers) {
       $scriptBlock = {
           param($VMName)
           $vm = Get-VM -Name $VMName -ErrorAction SilentlyContinue
           return ($vm -ne $null)
       }

       $result = Invoke-ReliableCommand -ComputerName $c -Credential $Credential -ScriptBlock $scriptBlock -ArgumentList $VMName

       if ($result) {
           return $c
       }
   }

   return $null
}

function Install-SshKey {
    param(
        [Parameter(Mandatory=$true)] [array] $Computers,
        [Parameter(Mandatory=$true)] [pscredential] $Credential,
        [Parameter(Mandatory=$true)] [string] $SshKey,
        [Parameter(Mandatory=$true)] [string] $KeyPath
    )

    $scriptBlock = {
        param($SshKey, $KeyPath, $KeyUser)
        $startLine = "-----BEGIN OPENSSH PRIVATE KEY-----`n"
        $endLine = "-----END OPENSSH PRIVATE KEY-----`n"

        $content = $startLine

        $SshKey.Replace(" ", "`n") | ForEach-Object {
            $content+= $_
        }

        $content += "`n"
        $content += $endLine

        Remove-Item $KeyPath -ErrorAction SilentlyContinue
        $Utf8NoBomEncoding = New-Object System.Text.UTF8Encoding $False
        [System.IO.File]::WriteAllLines($KeyPath, $content, $Utf8NoBomEncoding)
        $acl = Get-Acl $KeyPath
        $acl.SetAccessRuleProtection($true,$false)
        $accessRule = New-Object System.Security.AccessControl.FileSystemAccessRule($KeyUser, "FullControl", "Allow")
        $acl.SetAccessRule($accessRule)
        $acl | Set-Acl $KeyPath
    }

    foreach($c in $Computers) {
        Invoke-ReliableCommand -ComputerName $c -Credential $Credential -ScriptBlock $scriptBlock -ArgumentList @($SshKey, $KeyPath, $Credential.UserName) | Out-Null
    }
}

function Enable-DefaultFirewallRules {
    param(
        [Parameter(Mandatory=$true)] [array] $ComputerName,
        [Parameter(Mandatory=$false)] [pscredential] $Credential = [pscredential]::Empty
    )

    foreach($c in $ComputerName) {
        Invoke-ReliableCommand -ComputerName $c -Credential $Credential -ScriptBlock {
            Import-Module NetSecurity
            Set-NetFirewallRule -Name "FPS-*" -enabled True
        }
    }
}

function Remove-NCResources {
    param(
        [Parameter(Mandatory=$true)] [string] $NCUri,
        [Parameter(Mandatory=$false)] [pscredential] $Credential = [pscredential]::Empty
    )

    Write-TraceLog "Remove-NCResources: Removing NC resources"
    $nics = Get-NetworkControllerNetworkInterface -ConnectionUri $NCUri -ErrorAction SilentlyContinue -PassInnerException

    foreach($nic in $nics)
    {
        if ($nic.ResourceId -ilike "*gw*backend" -or  $nic.ResourceId -ilike "*gw*frontend") {
            continue
        }
        
        Write-TraceLog "Remove-NCResources: Removing nic $($nic.ResourceId)"
        Remove-NetworkControllerNetworkInterface -ConnectionUri $NCUri -ResourceId $nic.ResourceId -Force -ErrorAction SilentlyContinue -PassInnerException
    }


    $vnets = Get-NetworkControllerVirtualNetwork -ConnectionUri $NCUri -ErrorAction SilentlyContinue -PassInnerException
    foreach($vnet in $vnets)
    {
        Write-TraceLog "Remove-NCResources: Removing vnet $($vnet.ResourceId)"
        Remove-NetworkControllerVirtualNetwork -ConnectionUri $NCUri -ResourceId $vnet.ResourceId -Force -ErrorAction SilentlyContinue -PassInnerException
    }

    $lbs = Get-NetworkControllerLoadBalancer -ConnectionUri $NCUri -ErrorAction SilentlyContinue -PassInnerException
    foreach($lb in $lbs)
    {
        Write-TraceLog "Remove-NCResources: Removing lb $($lb.ResourceId)"
        Remove-NetworkControllerLoadBalancer -ConnectionUri $NCUri -ResourceId $lb.ResourceId -Force -ErrorAction SilentlyContinue -PassInnerException
    }

    $vgs = Get-NetworkControllerVirtualGateway -ConnectionUri $NCUri -ErrorAction SilentlyContinue -PassInnerException

    foreach($vg in $vgs)
    {
        Write-TraceLog "Remove-NCResources: Removing vg $($vg.ResourceId)"
        Remove-NetworkControllerVirtualGateway -ConnectionUri $NCUri -ResourceId $vg.ResourceId -Force -ErrorAction SilentlyContinue -PassInnerException
    }

    $firewalls = Get-NetworkControllerAccessControlList -ConnectionUri $NCUri -ErrorAction SilentlyContinue -PassInnerException

    foreach($rule in $firewalls) {
        Write-TraceLog "Remove-NCResources: Removing firewall $($rule.ResourceId)"
        Remove-NetworkControllerAccessControlList -ConnectionUri $NCUri -ResourceId $rule.ResourceId -Force -ErrorAction SilentlyContinue -PassInnerException
    }

    $lnets = Get-NetworkControllerLogicalNetwork -ConnectionUri $NCUri -ErrorAction SilentlyContinue -PassInnerException

    foreach($lnet in $lnets) {
        if ($lnet.ResourceId -ieq "HNVPA" -or $lnet.ResourceId -ieq "PublicVip" -or $lnet.ResourceId -eq "PrivateVip") {
            continue
        }
        Write-TraceLog "Remove-NCResources: Removing lnet $($lnet.ResourceId)"
        Remove-NetworkControllerLogicalNetwork -ConnectionUri $NCUri -ResourceId $lnet.ResourceId -Force -ErrorAction SilentlyContinue -PassInnerException
    }

    $publicIps = Get-NetworkControllerPublicIPAddress -ConnectionUri $NCUri -ErrorAction SilentlyContinue -PassInnerException

    foreach($publicIp in $publicIps) {
        Write-TraceLog "Remove-NCResources: Removing public ip $($publicIp.ResourceId)"
        Remove-NetworkControllerPublicIPAddress -ConnectionUri $NCUri -ResourceId $publicIp.ResourceId -Force -ErrorAction SilentlyContinue -PassInnerException
    }

}

function Enable-HostStorageNetwork {

    <#
    .SYNOPSIS
    Enables RDMA on the host and creates vnics for RDMA traffic
    #>


    param(
        [Parameter(Mandatory=$true)] [Array] $ComputerName,
        [Parameter(Mandatory=$false)] [pscredential] $Credential = [pscredential]::Empty,
        [Parameter(Mandatory=$true)] [string] $switchName
    )

    $scriptBlock ={

        param($switchName)

        $team = Get-VMSwitch -name $switchName | Get-VMSwitchTeam -ErrorAction SilentlyContinue

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

        $rdmaNics = Get-NetAdapter -ErrorAction SilentlyContinue | `
            Where-Object { $_.ComponentID -eq "vms_mp" } | `
            Get-NetAdapterRdma -ErrorAction SilentlyContinue | Where-Object {$_.Enabled -eq $True }

        if ($rdmaNics -ne $null -and $rdmaNics.Count -gt 0) {
            return
        }

        $state = (Get-WindowsFeature -Name "Data-Center-Bridging" ).InstallState

        if ($state -ne "Installed") {
            Install-WindowsFeature -Name "Data-Center-Bridging"
        }

        $nicDesc = $team.NetAdapterInterfaceDescription
        $nics = (Get-NetAdapter | Where-Object { $nicDesc -contains $_.InterfaceDescription }).Name

        $policy = Get-NetQosPolicy -Name "SMB" -ErrorAction SilentlyContinue
        if ($policy -eq $null) {
            New-NetQosPolicy "SMB" -NetDirectPortMatchCondition 445 -PriorityValue8021Action 3
            Enable-NetQosFlowControl -priority 3
        }
        $trafficClass = Get-NetQosTrafficClass "SMB" -ErrorAction SilentlyContinue
        if ($trafficClass -eq $null) {
            New-NetQosTrafficClass "SMB" -priority 3 -bandwidthpercentage 50 -algorithm ETS
        }

        Enable-NetAdapterQos -InterfaceAlias $nics
        Set-NetQosDcbxSetting -InterfaceAlias $_ -Willing $false -Confirm:$false

        $defaultPolicy = Get-NetQosPolicy -Name "DEFAULT" -ErrorAction SilentlyContinue
        if ($defaultPolicy -eq $null) {
            New-NetQosPolicy "DEFAULT" -Default -PriorityValue8021Action 0
            Disable-NetQosFlowControl -priority 0,1,2,4,5,6,7
        }

        Enable-NetAdapterRdma $nics

        $smbNic1 = Get-VMNetworkAdapter -Name "SMB1" -ManagementOS -ErrorAction SilentlyContinue
        if ($smbNic1 -eq $null) {
            Add-VMNetworkAdapter -SwitchName $switchName -Name "SMB1" -managementOS
        }

        $smbNic2 = Get-VMNetworkAdapter -Name "SMB2" -ManagementOS -ErrorAction SilentlyContinue
        if ($smbNic2 -eq $null) {
            Add-VMNetworkAdapter -SwitchName $switchName -Name "SMB2" -managementOS
        }

        Disable-NetAdapterBinding -Name "vEthernet (SMB1)","vEthernet (SMB2)" -DisplayName "Internet Protocol Version 4 (TCP/IPv4)"

        Enable-NetAdapterRDMA -Name "vEthernet (SMB1)","vEthernet (SMB2)"
    }

    foreach($c in $ComputerName) {
        Write-TraceLog "Enable-HostStorageNetwork: Enabling host storage network on $c, switch $switchName"
        $s = New-ReliablePsSession -ComputerName $c -Credential $Credential
        Invoke-Command -Session $s -ScriptBlock $scriptBlock -ArgumentList @($switchName) | Out-Null
    }
}

function Invoke-PowershellCommandOnVm {
    param(

        [ValidateNotNullOrEmpty()]
        [string] $vmName,

        [ValidateNotNullOrEmpty()]
        [string] $hostName,

        [ValidateNotNullOrEmpty()]
        [string] $cmd,

        [pscredential] $hostCred,
        [pscredential] $vmCred
    )


    Invoke-PowershellCommandOnVmInternal -vmName $vmName `
         -hostName $hostName `
         -hostCred $hostCred `
         -vmCred $vmCred `
         -scriptBlock $scriptBlock `
         -cmd $cmd
}