#//—————————————————————————-
#//  This script by Diagg/OSDC.
#//   http://osd-couture.com/
#// 
#//  Version : 2.4
#//  Release Date : 13/06/2014
#//  Latest Update : 25/07/2014
#//  Usage : 
#//  Warning :
#//  History :
#//		- v1.0 : Initial release
#//		- v2.0 : Added Local WDS support
#//		- v2.1 : Fixed an issue with Network share !
#//		- v2.2 : Changed sharing method to stay compatible with posh 2
#//     - V2.3 : CustomSetting.ini variables renamed
#//		- V2.4 : Fixed an issue when a network card is set to public
#//            : Fixed an issue with Get-Psdrive 
#//            : Added more controls and error handling
#//
#//—————————————————————————


	# Test values for debug
	#Set-Item env:DEPLOYROOT -Value "C:\DeploymentShare"
	#Set-Item env:STAGE -Value "ISO"
    #Set-Item env:PLATFORM -Value "x86"
	#Set-Item env:CONTENT -Value "C:\DeploymentShare\Boot\LiteTouchPE_x86.wim"

#region Usefull Functions
Function Write-Log
	{
		Param([string]$Mess,[Bool]$WriteToLog,[Bool]$WriteToHost)
		# $Mess = Message to be Written on the log file
		# $WriteToLog = If true the Write to logfile
		# $WriteToHost =If true the Write to Console
		
		If ( $WriteToLog -eq $true)
			{
				$Date =Get-Date
				Write-output $("$Date -     $Mess")|Out-File -width 500 -FilePath $logFile -Append
				#Write-Verbose $("$Date -     $Mess")
			}
		If ( $WriteToHost -eq $true)
			{
				Write-host $("$Date -     $Mess")
			}	
	}

Function Exit-Script	
	{
		#Clean up

		Write-log -Mess "---------------------------------------------------" -WriteToLog 1 -WriteToHost 1
        Write-log -Mess "####                 Clean-up                  ####" -WriteToLog 1 -WriteToHost 1
		Write-log -Mess "---------------------------------------------------" -WriteToLog 1 -WriteToHost 1
		
        If (-not $WDS_Local)
            { 
                If ($CopyOnWDS)
                    {
				        Write-log -Mess ("Removing file from remote WDS : " + $WDS_LocalDrive + ":\LiteTouchPE_" + $Platform +".wim") -WriteToLog 1 -WriteToHost 1
				        Remove-Item ($WDS_LocalDrive + ":\LiteTouchPE_" + $Platform +".wim") -verbose
						If (!(Test-Path($WDS_LocalDrive + ":\LiteTouchPE_" + $Platform +".wim")))
							{
								 Write-log -Mess ( $WDS_LocalDrive + ":\LiteTouchPE_" + $Platform +".wim, removed successfully !!") -WriteToLog 1 -WriteToHost 1	
							}
						Else
							{
								Write-log -Mess ("Error, Unable to remove " + $WDS_LocalDrive + ":\LiteTouchPE_" + $Platform +".wim") -WriteToLog 1 -WriteToHost 1	
							}
                    }
				
                # Remove Mapped Drive
                If ($Drv)
                    {
                        Write-log -Mess ("Removing mapped Drive " + $Drv) -WriteToLog 1 -WriteToHost 1
                        $Net.RemoveNetworkDrive($Drv)
                    }

				# Close the Remote Connection
				Write-log -Mess ("Removing Ps Session " + $RemoteSession) -WriteToLog 1 -WriteToHost 1
				Remove-PSSession -Session $RemoteSession -verbose
				If ($RemoteSession.State -eq "Closed")
					{
						Write-log -Mess "Session closed Successfully !!! "  -WriteToLog 1 -WriteToHost 1
					}
				Else
					{
						Write-log -Mess "Error Closing session" -WriteToLog 1 -WriteToHost 1
					}
		    }
 
		Write-log -Mess "---------------------------------------------------" -WriteToLog 1 -WriteToHost 1
		Write-log -Mess "#                                                 #" -WriteToLog 1 -WriteToHost 1
		Write-log -Mess "---------------------------------------------------" -WriteToLog 1 -WriteToHost 1
		Exit
	}
	
#EndRegion


Cls
# Log file path
$logFile = "c:\Windows\logs\WDS-BootImageUpdate.log"


# "Update Deployment Share" current phase
$DeployRoot = (Get-ChildItem env:DEPLOYROOT).value
$Platform = (Get-ChildItem env:PLATFORM).value
$Content = (Get-ChildItem env:CONTENT).value
$Stage = (Get-ChildItem env:STAGE).value


#Clean Up log File
If ((Test-Path -LiteralPath $($logFile)) -and ($Stage -eq "WIM") -and ($Platform -eq "x86"))
	{
		Remove-Item $logFile
	}

#Start-Transcript -path $logFile -Append
Write-log -Mess "---------------------------------------------------" -WriteToLog 1 -WriteToHost 1
Write-log -Mess "####               Updating $Platform                 ####" -WriteToLog 1 -WriteToHost 1
Write-log -Mess "---------------------------------------------------" -WriteToLog 1 -WriteToHost 1
Write-log -Mess "Starting WDS Update" -WriteToLog 1 -WriteToHost 1
Write-log -Mess ("DEPLOYROOT: " + $DeployRoot) -WriteToLog 1 -WriteToHost 1
Write-log -Mess ("PLATFORM: " + $Platform) -WriteToLog 1 -WriteToHost 1
Write-log -Mess ("CONTENT: " + $Content) -WriteToLog 1 -WriteToHost 1
Write-log -Mess ("STAGE: " + $Stage) -WriteToLog 1 -WriteToHost 1


If ($Stage -eq "ISO")
	{

		# Initialisation
		$Customsettings = ($DeployRoot+"\Control\CustomSettings.ini")
		$TempContent = ($Content + "\Sources\Boot.wim")
		# Following line is for debug puropse only, don't forget to comment back when using in MDT.
		#$TempContent = $Content
		
		
		# WDSServerURI
		$WDS = (Select-String -path $Customsettings -Pattern WDS-Update|ForEach {$_.Line -replace "WDS-Update=",""})
		
		If ($WDS -like "\\*\*")
		 	{
                #Path is UNC
                $WDS_Path = $WDS
			    $WDS_RemoteFolder = $WDS_path.Split("\")[($WDS_path.Split("\").Length)-1]
			    $WDS = $WDS.split("\")[2]
                $WDS_Local = $false
                Write-log -Mess ("WDS is remote with path: " + $WDS_Path) -WriteToLog 1 -WriteToHost 1
			    
			}
		ElseIf ($WDS -like "*:\*")
			{
				#Path is Local
                $WDS_Path = $WDS
                $WDS_Local = $true
                Write-log -Mess ("WDS is local with path: " + $WDS_Path) -WriteToLog 1 -WriteToHost 1
			}
        ElseIf ($WDS -like "*\*")
        	{
				#Path is UNC but "\\" are missing
                $WDS_Path = "\\"+$WDS
				$WDS_RemoteFolder = $WDS_path.Split("\")[($WDS_path.Split("\").Length)-1]
				$WDS = $WDS.split("\")[0]
                $WDS_Local = $false
                Write-log -Mess ("WDS is remote with path: " + $WDS_Path) -WriteToLog 1 -WriteToHost 1
			}
        Else
            {
                #Path is not valid !!!
                Write-log -Mess "Unable to locate WDS path in Customsettings.ini, Aborting !" -WriteToLog 1 -WriteToHost 1
                Exit-Script
            }
            
		# Login to WDS Server if not local
        If (-not $WDS_Local)
            {
		        $user = (Select-String -path $Customsettings -Pattern WDS-User|ForEach {$_.Line -replace "WDS-User=",""})
		        # Password to WDS Server
		        $Pword = (Select-String -path $Customsettings -Pattern WDS-Password|ForEach {$_.Line -replace "WDS-Password=",""})
		        $Pword = [System.Text.Encoding]::UNICODE.GetString([System.Convert]::FromBase64String($Pword))
                $Pword2 = $Pword
		        $Pword = ConvertTo-SecureString -String $Pword -AsPlainText -Force
		        # Cypher credential
		        $Credential = New-Object -TypeName System.Management.Automation.PSCredential -ArgumentList $user, $Pword
		    }

        # Image Name to update/create 
	 	$WDS_Image = ((Select-String -path $Customsettings -Pattern WDS-Image|ForEach {$_.Line -replace "WDS-Image=",""}) + "-" + $Platform )
        If (!$WDS_Image)
            {
                #Image Name is not valid !!!
                Write-log -Mess "WDS Image name not specified or invalid, Aborting !" -WriteToLog 1 -WriteToHost 1
                Exit-Script
            }

	

		# Check for local PS Remoting
		If ((Get-Service WinRm).Status -ne "Running")
			{
				try
					{
					    Enable-PSRemoting –force -verbose
					}
				catch
					{
					   If ($_.Exception.ToString() -like '*this machine is set to Public*')
					   	{
							Get-NetConnectionProfile -NetworkCategory public|Set-NetConnectionProfile -NetworkCategory Private
							Enable-PSRemoting –force -verbose
						}
					}
			}
				
		# Check if any remote computer can be member of the trusted list
        If (-not $WDS_Local)		
	        {
                If ( -not (Get-ChildItem wsman:localhost\client\TrustedHosts).value -like "*")
				    {
						Set-WSManInstance -ResourceURI winrm/config/client -ValueSet @{TrustedHosts = "*"}
				    }
            }

		#Establishing a persistent remote session
        If(!$WDS_Local)
            { 
		        $RemoteSession = New-PSSession -ComputerName $WDS -Credential $Credential -Verbose
				$toto = $RemoteSession.state
		
		
		        # Map Network drive to copy the new image

		        If ((Get-PSDrive).Displayroot -like $WDS_Path -or (Get-PSDrive).root -like $WDS_Path )
			        {
				        $WDS_LocalDrive = Get-PSDrive|where{$_.Displayroot -eq $WDS_Path -or $_.Root -eq $WDS_Path}
				        $WDS_LocalDrive = ($WDS_LocalDrive).Name
			        }
		        Else
			        {
				       
                       # Using an alternative method to stay comaptible with powershell 2
                        If ((host).version.Major -gt 3)
                            {
                                #choose a random drive letter (By Shay Lavy)
                                $Drv = ls function:[d-z]: -n|?{!(test-path $_)}|random
                                
                                $net = new-object -ComObject WScript.Network
                                $net.MapNetworkDrive($Drv, $WDS_Path, $false, $user, $Pword2)
                                $WDS_LocalDrive = ($net.EnumNetworkDrives()|where {$_ -like $Drv}).trim(":")
                                if (!$WDS_LocalDrive -eq $Drv)
                                    {
                                        Write-log -Mess ("Unable to map Network Drive, Aborting..." ) -WriteToLog 1 -WriteToHost 1
                                        Exit-Script
                                    }
                            }
                        Else
                            {
                                $WDS_LocalDrive = New-PSDrive -Name "WDS001" -PSProvider FileSystem -Root $WDS_Path -Credential $Credential -verbose
				                $WDS_LocalDrive = ($WDS_LocalDrive).Name
                                If (!$WDS_LocalDrive)
                                    {
                                        Write-log -Mess ("Unable to map Network Drive, Aborting..." ) -WriteToLog 1 -WriteToHost 1
                                        Exit-Script
                                    }
                            }
			        }
		        Write-log -Mess ("Remote WDS mapped with Letter : " + $WDS_LocalDrive) -WriteToLog 1 -WriteToHost 1	
			
			
			
		        #Copy File to The WDS
		        Write-log -Mess ("Copying file " + $TempContent + " to " +$WDS_LocalDrive + ":\LiteTouchPE_" + $Platform +".wim") -WriteToLog 1 -WriteToHost 1
		        Copy-Item $TempContent -Destination ($WDS_LocalDrive + ":\LiteTouchPE_" + $Platform +".wim") -verbose
				If (Test-Path($WDS_LocalDrive + ":\LiteTouchPE_" + $Platform +".wim"))
					{
                		Write-log -Mess ("File " + $WDS_LocalDrive + ":\LiteTouchPE_" + $Platform +".wim, copied successfully !!!") -WriteToLog 1 -WriteToHost 1
						$CopyOnWDS = $true
					}
				Else
					{
						Write-log -Mess ("Error, unable to copy " + $WDS_LocalDrive + ":\LiteTouchPE_" + $Platform +".wim, aborting !!!") -WriteToLog 1 -WriteToHost 1
					}
		    }			
		
		# Check if the boot image already exist on the remote WDS 
		If(!$WDS_Local)
            { 
                $Iretval = Invoke-Command -Session $RemoteSession -ArgumentList $WDS_RemoteFolder -ScriptBlock `
			        {
				        Param($WDS_RemoteFolder)
				        $BootImages = (Get-WdsBootImage).name
				        # Set a variable remotely with the WDS folder path
				        $WDS_RemotePath = ((gwmi -class win32_share)|where { $_.Name -like $WDS_RemoteFolder}).path
				        $BootImages
			        }
			        Write-log -Mess ("list of Remote images found : " + $Iretval ) -WriteToLog 1 -WriteToHost 1
            }
        Else 
            {
                $Iretval = (Get-WdsBootImage).name
            }

		# If Image already exists, Update the image (Remove it)
		If ($Iretval -like $WDS_Image)
			{
				Write-log -Mess ("The image " + $WDS_Image + " was found on the remote WDS and will be updated" ) -WriteToLog 1 -WriteToHost 1
				# Remove the Old Image
				If (!$WDS_Local)
                    {
                        $Iretval = Invoke-Command -Session $RemoteSession -ArgumentList $WDS_Image, $Platform -ScriptBlock `
					        {
						        Param($WDS_Image, $Platform)
						        Remove-WdsBootImage -imageName $WDS_Image -architecture $Platform -verbose
						        $BootImages = (Get-WdsBootImage).name
						        $BootImages
					        }
                    }
                 Else
                    {
                        Remove-WdsBootImage -imageName $WDS_Image -architecture $Platform -verbose
						$Iretval = (Get-WdsBootImage).name                        
                    }              
				
				# verify if the image was correctly removed
				If ($Iretval -like $WDS_Image)
					{
						Write-log -Mess ("Unable to remove old Image" + $WDS_Image + "...Exit Program !!!") -WriteToLog 1 -WriteToHost 1
						Exit-Script
					}
				Else 	
					{
						Write-log -Mess ("Old Image " + $WDS_Image + " was removed Successfully !") -WriteToLog 1 -WriteToHost 1
					}
				
            }
	
		# Import the new image
		$Date = Get-Date
        If (!$WDS_Local)
            {
				$Iretval = Invoke-Command -Session $RemoteSession -ArgumentList $WDS_Image, $Platform, $Date -ScriptBlock `
					{
						Param($WDS_Image, $Platform, $Date)
						Import-WdsBootImage -Path ($WDS_RemotePath + "\LiteTouchPE_" + $Platform + ".wim") -NewImageName $WDS_Image -NewDescription ("Updated By MDT: " + $Date) –SkipVerify -verbose
						$BootImages = (Get-WdsBootImage).name
						$BootImages
					}
            }
        Else
            {
				Import-WdsBootImage -Path ($TempContent) -NewImageName $WDS_Image -NewDescription ("Updated By MDT: " + $Date) –SkipVerify -verbose
				$Iretval = (Get-WdsBootImage).name                
            }
					
		# Verify if the new image was correctly uploaded
		If ($Iretval -like $WDS_Image)
			{
				Write-log -Mess ("Image " + $WDS_Image + " uploaded Successfully on the remote WDS!") -WriteToLog 1 -WriteToHost 1
				Write-log -Mess ("Image updated with the following comment => Updated By MDT: " + $Date) -WriteToLog 1 -WriteToHost 1
				Write-log -Mess ("List of remote Images : " + $Iretval)  -WriteToLog 1 -WriteToHost 1					
			}
		Else 	
			{
				Write-log -Mess ("Unable to upload new Image" + $WDS_Image + "...Exit Program !!!") -WriteToLog 1 -WriteToHost 1
				Exit-Script
			}	
								
		Exit-Script		
	}
Else
	{
		Write-log -Mess "Awaiting next stage to proceed" -WriteToLog 1 -WriteToHost 1
	}