Set_OSDDiskindex

Begin
# Load Microsoft.SMS.TSEnvironment COM object
{
try {
        $TSEnvironment = New-Object -ComObject Microsoft.SMS.TSEnvironment -ErrorAction Continue
    }
    catch [System.Exception] {
        Write-Warning -Message "Unable to construct Microsoft.SMS.TSEnvironment object"; break
    }
}

Process
{

function Write-CMLogEntry
    {
        param (
            [parameter(Mandatory = $true, HelpMessage = "Value added to the log file.")]
            [ValidateNotNullOrEmpty()]
            [string]$Value,
            [parameter(Mandatory = $true, HelpMessage = "Severity for the log entry. 1 for Informational, 2 for Warning and 3 for Error.")]
            [ValidateNotNullOrEmpty()]
            [ValidateSet("1", "2", "3")]
            [string]$Severity,
            [parameter(Mandatory = $false, HelpMessage = "Name of the log file that the entry will written to.")]
            [ValidateNotNullOrEmpty()]
            [string]$FileName = "Invoke-Set_OSDDiskIndex.log"
        )
        # Determine log file location
        $LogFilePath = Join-Path -Path $Script:TSEnvironment.Value("_SMSTSLogPath") -ChildPath $FileName

        # Construct time stamp for log entry
        $Time = -join @((Get-Date -Format "HH:mm:ss.fff"), "+", (Get-WmiObject -Class Win32_TimeZone | Select-Object -ExpandProperty Bias))

        # Construct date for log entry
        $Date = (Get-Date -Format "MM-dd-yyyy")

        # Construct context for log entry
        $Context = $([System.Security.Principal.WindowsIdentity]::GetCurrent().Name)

        # Construct final log entry
        $LogText = "<![LOG[$($Value)]LOG]!><time=""$($Time)"" date=""$($Date)"" component=""Set_OSDDiskIndex.log"" context=""$($Context)"" type=""$($Severity)"" thread=""$($PID)"" file="""">"

        # Add value to log file
        try
        {
            Out-File -InputObject $LogText -Append -NoClobber -Encoding Default -FilePath $LogFilePath -ErrorAction Stop
        }
        catch [System.Exception] {
            Write-Warning -Message "Unable to append log entry to Invoke-Set_OSDDiskIndex.log file. Error message: $($_.Exception.Message)"
        }
    }


    # Determine OSDDISKINDEX
    # If booted to Windows
    if (($TSEnvironment -ne $null) -and ($TSEnvironment.Value("_SMSTSinWinPE") -ne $true)){

    # Set log file location
    $LogFilePath = Join-Path -Path $TSEnvironment.Value("_SMSTSLogPath") -ChildPath $LogFileName

        try
        {

        $BootInfo = (gwmi Win32_BootConfiguration)
        $BootInfoSplit = ($BootInfo).Caption -split "\\"
        $TargetDisk = $BootInfoSplit[2].Trim('Harddisk')
        $TSEnvironment.Value('OSDDiskIndex') = $TargetDisk

        Write-CMLogEntry -Value "The value of OSDDISKINDEX is being set to $TargetDisk" -Severity 1

        }
        catch [System.Exception]
        {

        Write-CMLogEntry -Value "An error occured while setting the OSDDISKINDEX value. Error message: $($_.Exception.Message)" -Severity 3; exit 1

        }
    }

    #If Booted to WinPE
    if (($TSEnvironment -ne $null) -and ($TSEnvironment.Value("_SMSTSinWinPE") -eq $true)){

    # Set log file location
    $LogFilePath = Join-Path -Path $TSEnvironment.Value("_SMSTSLogPath") -ChildPath $LogFileName


    #Gather disk information. We'll need it further down.
    $DiskInfo = Get-WmiObject Win32_DiskDrive |
        ForEach-Object {
            $disk = $_
            $partitions = "ASSOCIATORS OF " +
            "{Win32_DiskDrive.DeviceID='$($disk.DeviceID)'} " +
            "WHERE AssocClass = Win32_DiskDriveToDiskPartition"
            Get-WmiObject -Query $partitions |
                ForEach-Object {
                $partition = $_
                $drives = "ASSOCIATORS OF " +
                    "{Win32_DiskPartition.DeviceID='$($partition.DeviceID)'} " +
                    "WHERE AssocClass = Win32_LogicalDiskToPartition"
                    Get-WmiObject -Query $drives |
                        ForEach-Object {
                        [PSCustomObject]@{
                        Disk        = $disk.DeviceID
                        DiskSize    = $disk.Size
                        DiskModel   = $disk.Model
                        Partition   = $partition.Name
                        RawSize     = $partition.Size
                        Type        = $partition.Type
                        DriveLetter = $_.DeviceID
                        VolumeName  = $_.VolumeName
                        Size        = $_.Size
                        FreeSpace   = $_.FreeSpace
                        }
                }
            }
        }

    #Check to see if there are multiple installs of Windows
    $BootDrive = (Get-ChildItem Env:SystemDrive | Select-Object -ExpandProperty Value)

    $ArrayDisks = Get-WmiObject -Class Win32_LogicalDisk -Filter "DriveType = 3 AND DeviceID != '$BootDrive'" | Select-Object -ExpandProperty DeviceID

    $Array = foreach ($ArrayDisk in $ArrayDisks)
        {

        (Test-Path "$ArrayDisk\Windows\System32\ntoskrnl.exe")

        }

    If (($Array -match $True).count -gt 1)

        {

        Write-CMLogEntry -Value "Multiple installations of Windows have been detected. Unable to safely set which disk to image from within WinPE." -Severity 3; exit 1

        }


    try
        {

        $OSDDetectedWinDrive = (Get-ChildItem Env:_OSDDetectedWinDrive | Select-Object -ExpandProperty Value)
        $OSDDetectedWinDrive = (($OSDDetectedWinDrive).Trim("\"))
        $TSTargetdisk = ($DiskInfo.Where({$_.Driveletter -eq "$OSDDetectedWinDrive"})).Partition
        $TSTargetdisk = ($TSTargetdisk -split ", ")
        $TSTargetdisk = $TSTargetdisk[0].Trim('Disk #')
        $TSEnvironment.Value('OSDDiskIndex') = $TSTargetdisk

        Write-CMLogEntry -Value "The value of OSDDISKINDEX is being set to $TSTargetDisk" -Severity 1

        }
        catch [System.Exception]
        {
            Write-CMLogEntry -Value "An error occured while setting the OSDDISKINDEX value. Error message: $($_.Exception.Message)" -Severity 3; exit 1
        }

    }

#Test for multiple EFI System Partitons
$EFITest = (gwmi -Query "Select * from Win32_DiskPartition WHERE Type = 'GPT: System'" | Select-Object -ExpandProperty Name)
If ($EFItest.count -gt 1){

Write-CMLogEntry -Value "Multiple EFI boot partitons have been found. This condition will result in the system being unbootable after imaging." -Severity 3; exit 1

}
}