WinBuildUtils.psm1

Import-Module $PSScriptRoot\Logger.psm1
Import-Module $PSScriptRoot\Utils.psm1

$script:WinbuildReleaseShare = "\\winbuilds.ntdev.corp.microsoft.com\release"
function Get-BuildBranch {
    $buildBranch = (Get-ItemProperty "hklm:\SOFTWARE\Microsoft\Windows NT\CurrentVersion").BuildBranch
    return $buildBranch
}

function Get-BuildNumber {
    $buildLabEx = (Get-ItemProperty "hklm:\SOFTWARE\Microsoft\Windows NT\CurrentVersion").BuildLabEx
    $buildLabExGroups = $buildLabEx.Split(".")
    $buildNumber = $buildLabExGroups[0] + "." + $buildLabExGroups[1] + "." + $buildLabExGroups[4]
    return $buildNumber
}

function Get-BuildType {
    $buildLabEx = (Get-ItemProperty "hklm:\SOFTWARE\Microsoft\Windows NT\CurrentVersion").BuildLabEx
    $buildLabExGroups = $buildLabEx.Split(".")
    return $buildLabExGroups[2]
}

function Get-BuildLocation {
    $buildBranch = Get-BuildBranch
    $buildNumber = Get-BuildNumber
    $arch = Get-BuildType

    $buildLocation = Join-Path $script:WinbuildReleaseShare  $BuildBranch
    $buildLocation = Join-Path $buildLocation $buildNumber
    $buildLocation = Join-Path $buildLocation $arch

    return $buildLocation
}

function Get-WttNAOCred {
    if ($Global:CorpCredential -ne $null -and $Global:CorpCredential -ne [pscredential]::Empty) {
        return $Global:CorpCredential
    }

    $user = $env:WTT_USERNAME
    $pass = $env:WTT_PASSWORD

    if (-not [String]::IsNullOrEmpty($user) -and -not [String]::IsNullOrEmpty($pass)) {
        $secPass = ConvertTo-SecureString $pass -AsPlainText -Force
        $cred = New-Object System.Management.Automation.PsCredential($user, $secPass)
        return $cred
    }

    $isWttCmdInstalled = Get-Command wttcmd -ErrorAction SilentlyContinue

    if ($isWttCmdInstalled -eq $null) {
        throw "Unable to detect corp credentials, please run Set-TurnKeySdnCorpCredentials and retry"
    }

    $corpUser = ""
    $corpPassword = ""
    $resStr = wttcmd /getLogicalUser  /LocalName:LLU_NAO
    $resStr | ForEach-Object {
        $line = $_.trim()
        $pair = $line.split(":")
        if ($pair.Count -eq 2) {
            $key = $pair[0].trim()
            $value = $pair[1].trim()
            if ($key -imatch "UserName") {
                $corpUser = $value
            }        

            if ($key -imatch "Password") {
                $corpPassword = ConvertTo-SecureString $value -AsPlainText -Force
            }
        }        
    }

    if ($corpUser -eq "" -or $corpPassword -eq "") {
        throw "Unable to detect corp credentials, please run Set-TurnKeySdnCorpCredentials and retry"
    }
    
    $cred = New-Object System.Management.Automation.PsCredential($corpUser, $corpPassword)

    return $cred
}

function Get-ServerSkuList {

    $thisSku = (Get-ComputerInfo).WindowsEditionId
    $skus = @()
    $skus += $thisSku

    if ($thisSku -ine "serverdatacenter") {
        $skus += "serverdatacenter"
    }

    if ($thisSku -ine "serverdatacentercore") {
        $skus += "serverdatacentercore"
    }

    if ($thisSku -ine "serverazurestackhcicor") {
        $skus += "serverazurestackhcicor"
    }
    
    return $skus
}

function Resolve-BuildVhdxOrVhdFileName {
    param(
        [parameter(Mandatory = $true)][string]$Directory,
        [switch]$IsVhd
    )

    if ($IsVhd.IsPresent) {
        return (Get-ChildItem $Directory -Filter "*.vhd")
    }

    return (Get-ChildItem $Directory -Filter "*.vhdx")
}

function Resolve-BuildVhdxOrVhdPath {
    param(
        [parameter(Mandatory = $true)][string]$BasePath,
        [switch]$IsVhd
    )

    if ($IsVhd.IsPresent) {
        $prefix = "vhd"
    }
    else {
        $prefix = "vhdx"
    }
    
    $cred = Get-WttNAOCred    
    if ($cred -eq $null) {
        throw "Resolve-BuildVhdxOrVhdPath: Failed to get WTT NAO credential"
    }

    $skuList = Get-ServerSkuList
    foreach ($sku in $skuList) {
        $fullPrefix = $prefix + "_server_" + $sku

        $serverVhdDIr = Join-Path $BasePath  $($fullPrefix + "_en-us_vl")
        if (Test-WinBuildPath -WinbuildPath $serverVhdDIr) {
            return (Resolve-BuildVhdxOrVhdFileName -Directory $serverVhdDIr -IsVhd:($IsVhd.IsPresent))
        }

        $serverVhdDIr = Join-Path $BasePath $($fullPrefix + "_en-us")

        if (Test-WinBuildPath -WinbuildPath $serverVhdDIr) {
            return (Resolve-BuildVhdxOrVhdFileName -Directory $serverVhdDIr -IsVhd:($IsVhd.IsPresent))
        }
    }
}

function Get-BuildVhdxOrVhd {
    param(
        [parameter(Mandatory = $true)][string]$buildLocation
    )

    $vhdxBasePath = Join-Path $buildLocation -ChildPath vhdx
    
    if (Test-WinBuildPath -WinbuildPath $vhdxBasePath) {
        $vhdxFile = Resolve-BuildVhdxOrVhdPath -BasePath $vhdxBasePath

        if ($vhdxFile -ne $null) {
            return $vhdxFile
        }
    }

    $vhdBasePath = Join-Path $buildLocation -ChildPath vhd

    if (-not (Test-WinBuildPath -WinbuildPath $vhdBasePath)) {
        return $null
    }

    $vhdFile = Resolve-BuildVhdxOrVhdPath -BasePath $vhdBasePath
    return $vhdFile
}

function Get-BuildVersionVhd {
    param(
        [string][parameter(Mandatory = $true)]$buildLocation
    )

    Write-TraceLog "Get-BuildVersionVhd: Current build is $buildLocation"

    if (-not $(Test-Path $buildLocation)) {
        throw "Get-BuildVersionVhd: $buildLocation not found, please select a valid build using Get-BestBuild"
        
    }
    
    $vhdFileName = Get-BuildVhdxOrVhd -buildLocation $buildLocation

    if ($vhdFileName -eq $null) {
        Write-TraceLog "Get-BuildVersionVhd: Failed to find vhdx/vhd file in $buildLocation" -Warning
        return $null
    }

    return $vhdFileName.FullName
}
function Test-WinBuildPath {
    param(
        [string][parameter(Mandatory = $true)] $WinbuildPath
    )

    if (Test-Path $WinbuildPath) {
        return $true
    }

    $cred = Get-WttNAOCred    
    if ($cred -eq $null) {
        throw "Test-WinBuildPath: Failed to get WTT NAO credential"
    }

    if ((Connect-SMBShare -RemotePath $WinbuildPath -Credential $cred)) {
        return $true
    }

    return $false
}

function Connect-Winbuild {
    try {

        Write-TraceLog "Connect-Winbuild: Test-Path $($script:WinbuildReleaseShare) ..."
        if ((Test-Path $script:WinbuildReleaseShare)) {
            Write-TraceLog "Connect-Winbuild: Connection to $($script:WinbuildReleaseShare) exists"
            return
        }
   
        $cred = Get-WttNAOCred    
        if ($cred -eq $null) {
            throw "Connect-Winbuild: Failed to get WTT NAO credential"
        }

        $connected = Connect-SMBShare -RemotePath $script:WinbuildReleaseShare -Credential $cred -RestartLanManServer

        if (-not $connected) {
            throw "Connect-Winbuild: Failed to connect to winbuilds, please check the credential and try again"
        }

    }
    catch {
        Write-TraceLog "Connect-Winbuild: Failed to connect to winbuild, error $_"
        Write-TraceLog "Connect-Winbuild: Please set a valid corp credentials using [Set-TurnKeySdnCorpCredentials] and retry"
        throw
    }
}

# Test if the build has bin and vhd folder
function Test-Build {
    param(
        [string] $build
    )

    Write-TraceLog "Test-Build: $build"

    if (-not $(Test-WinBuildPath -WinbuildPath $build)) {
        Write-TraceLog "Test-Build: $build path not found"
        return $false
    }

    $binLoc = Join-Path $build "bin"
    if (-not $(Test-WinBuildPath $binLoc)) {
        Write-TraceLog "Test-Build: $binLoc path not found"
        return $false
    }

    $binenus = Join-Path $build "en-us\bin"
    
    if (-not $(Test-WinBuildPath $binenus)) {
        Write-TraceLog "Test-Build: $binenus path not found"
        return $false
    }

    $vhd = Get-BuildVersionVhd -buildLocation $build

    if ($vhd -eq $null) {
        Write-TraceLog "Test-Build: No suitable vhd found in build"
        return $false
    }

    return $true

}

function Get-LastBuilds {
    param(
        [int] $count,
        [string] $branch
    )

    if ($count -eq $null -or $count -eq 0 ) {
        $count = 5
    }
    
    if ([String]::IsNullOrEmpty($branch)) {
        $branch = Get-BuildBranch
    }

    $buildBranchLocation = Join-Path $script:WinbuildReleaseShare $branch
    if (-not $(Test-WinBuildPath -WinbuildPath $buildBranchLocation)) {
        Write-TraceLog "Get-LastBuilds: $buildBranchLocation not found"
        return $null
    }

    $builds = Get-ChildItem $buildBranchLocation -Directory | Sort-Object -Property LastWriteTime -Descending | Select-Object -First $count
    return $builds
    
}
function Get-BestBuild {
    param(
        [string] $SpecificBuild,
        [string] $WindowsBuildBranch
    )

    Write-TraceLog "Get-BestBuild SpecificBuild <$SpecificBuild> WindowsBuildBranch <$WindowsBuildBranch>"
    Connect-Winbuild

    if (-not [String]::IsNullOrEmpty($SpecificBuild)) {
        $buildElements = $SpecificBuild.Split((@("\")), [StringSplitOptions]::RemoveEmptyEntries).Trim()
        if ($buildElements.Count -lt 4) {
            throw "Invalid build location $SpecificBuild"
        }

        $SpecificBuild = $SpecificBuild.Trim()
        $SpecificBuild = $SpecificBuild.TrimEnd("\")
        
        if ($buildElements[0] -ieq "winbuilds") {
            $SpecificBuild = $SpecificBuild -replace "winbuilds", "winbuilds.ntdev.corp.microsoft.com"
        }

        if (-not ($SpecificBuild.EndsWith("amd64fre"))) {
            $SpecificBuild = Join-Path $SpecificBuild amd64fre
        }
        
        if ($(Test-Build $SpecificBuild)) {
            return $SpecificBuild
        }

        return $null
    } 

    if ([String]::IsNullOrEmpty($WindowsBuildBranch)) {
        $branch = Get-BuildBranch
    }
    else {
        $branch = $WindowsBuildBranch
    }


    $builds = Get-LastBuilds -count 10 -branch $branch
    $arch = Get-BuildType

    foreach ($build in $builds) {
        $buildLocation = Join-Path $build.FullName $arch
        if ($(Test-Build $buildLocation)) {
            return $buildLocation
        }
    }
}