# Wizard Kit: Windows PE Build Tool

## Init ##
#Requires -Version 3.0
#Requires -RunAsAdministrator
if (Test-Path Env:\DEBUG) {
  Set-PSDebug -Trace 1
}
# TODO REMOVE v
$KitNameShort = "1201"
# TODO REMOVE ^
$Arch = "amd64"
$Host.UI.RawUI.WindowTitle = "Wizard Kit: Windows PE Build Tool"
$WD = $(Split-Path $MyInvocation.MyCommand.Path)
$SetupDir = (Get-Item $WD -Force).Parent.FullName
$Root = (Get-Item $SetupDir -Force).Parent.FullName
$BuildDir = "$SetupDir\BUILD_PE"
$BinDir = "$BuildDir\bin"
$ToolsDir = "$BuildDir\tools"
$OutDir = "$SetupDir\OUT_PE"
$LogDir = "$BuildDir\Logs"
$Temp = "$BuildDir\Temp"
$Date = Get-Date -UFormat "%Y-%m-%d"
$Host.UI.RawUI.BackgroundColor = "Black"
$Host.UI.RawUI.ForegroundColor = "White"
$HostSystem32 = "{0}\System32" -f $Env:SystemRoot
$HostSysWOW64 = "{0}\SysWOW64" -f $Env:SystemRoot
$DISM = "{0}\DISM.exe" -f $Env:DISMRoot
#Enable TLS 1.2
[Net.ServicePointManager]::SecurityProtocol = [Net.SecurityProtocolType]::Tls12


## Functions ##
function Abort {
  Write-Host -ForegroundColor "Red" "`nAborted."
  WKPause "Press Enter to exit..."
  exit
}
function DownloadFile ($Path, $Name, $Url) {
  $OutFile = "{0}\{1}" -f $Path, $Name

  Write-Host ("Downloading: $Name")
  New-Item -Type Directory $Path 2>&1 | Out-Null
  try {
    Invoke-WebRequest -Uri $Url -OutFile $OutFile
  }
  catch {
    Write-Host ("  ERROR: Failed to download file." ) -ForegroundColor "Red"
    $global:DownloadErrors += 1
  }
}
function FindDynamicUrl ($SourcePage, $RegEx) {
  # Get source page
  Invoke-Webrequest -Uri $SourcePage -OutFile "tmp_page"

  # Search for real url
  $Url = Get-Content "tmp_page" | Where-Object {$_ -imatch $RegEx}
  $Url = $Url -ireplace '.*(a |)href="([^"]+)".*', '$2'
  $Url = $Url -ireplace ".*(a |)href='([^']+)'.*", '$2'

  # Remove tmp_page
  Remove-Item "tmp_page"

  $Url | Select-Object -First 1
}
function MakeClean {
  $Folders = @(
    "$BuildDir\mount",
    "$BuildDir\pe_files",
    "$BuildDir\tools")
  # $Folders = @(
    # "$BuildDir\bin",
    # "$BuildDir\mount",
    # "$BuildDir\pe_files",
    # "$BuildDir\tools")
  foreach ($f in $Folders) {
    if (Test-Path $f) {
      Write-Host -ForegroundColor "Yellow" ("Found: {0}" -f $f)
      Remove-Item -Path $f -Recurse -Force
    }
  }
}
function WKPause ($Message = "Press Enter to continue... ") {
  Write-Host $Message -NoNewLine
  Read-Host
}

## Safety Check ##
if ($PSVersionTable.PSVersion.Major -eq 6 -and $PSVersionTable.OS -imatch "Windows 6.1") {
  Write-Host "`nThis script doesn't support PowerShell 6.0 on Windows 7."
	Write-Host "Press Enter to exit... " -NoNewLine
  Abort
}

## PowerShell equivalent of Python's "if __name__ == '__main__'"
# Code based on StackOverflow comments
# Question:     https://stackoverflow.com/q/4693947
# Using answer: https://stackoverflow.com/a/5582692
# Asked by:     https://stackoverflow.com/users/65164/mark-mascolino
# Answer by:    https://stackoverflow.com/users/696808/bacon-bits
if ($MyInvocation.InvocationName -ne ".") {
  Clear-Host
  Write-Host "WizardKit: Build Tool`n`n`n`n`n"

  ## Prep ##
  try {
    Import-Module -Name $Env:DISMRoot -ErrorAction "stop"
  }
  catch {
    Write-Host -ForegroundColor "Red" "ERROR: Failed to load DISM CmdLet"
    Abort
  }
  Push-Location "$WD"
  MakeClean
  New-Item -Type Directory $BuildDir 2>&1 | Out-Null
  New-Item -Type Directory $LogDir 2>&1 | Out-Null

  ## Sources ##
  $Sources = Get-Content -Path "$WD\sources.json" | ConvertFrom-JSON

  ## Download ##
  $DownloadErrors = 0
  DownloadFile -Path $Temp -Name "7z-installer.msi" -Url $Sources.'7-Zip'
  DownloadFile -Path $Temp -Name "ConEmuPack.7z" -Url $Sources.'ConEmu'
  DownloadFile -Path $Temp -Name "notepadplusplus.7z" -Url $Sources.'Notepad++'
  DownloadFile -Path $Temp -Name "ntpwedit.zip" -Url $Sources.'NTPWEdit'
  DownloadFile -Path $Temp -Name "wimlib.zip" -Url $Sources.'wimlib'

  ## Bail ##
  # If errors were encountered during downloads
  if ($DownloadErrors -gt 0) {
    Abort
  }

  ## Extract ##
  Copy-Item -Path "$SetupDir\pe\bin" -Destination "$BinDir" -Recurse -Force
  Copy-Item -Path "$SetupDir\pe\tools" -Destination "$ToolsDir" -Recurse -Force

  # 7-Zip
  Write-Host "Extracting: 7-Zip"
  try {
    $ArgumentList = @("/a", "$Temp\7z-installer.msi", "TARGETDIR=$Temp\7zi", "/qn")
    Start-Process -FilePath "$HostSystem32\msiexec.exe" -ArgumentList $ArgumentList -Wait
    New-Item -Type Directory "$BinDir\7-Zip" 2>&1 | Out-Null
    Move-Item "$Temp\7zi\Files\7-Zip\7z.dll" "$BinDir\7-Zip\7z.dll"
    Move-Item "$Temp\7zi\Files\7-Zip\7z.exe" "$BinDir\7-Zip\7z.exe"
    Move-Item "$Temp\7zi\Files\7-Zip\License.txt" "$BinDir\7-Zip\License.txt"
    # Remove-Item "$Temp\7z*" -Recurse
    $SevenZip = "$BinDir\7-Zip\7z.exe"
  }
  catch {
    Write-Host ("  ERROR: Failed to extract files." ) -ForegroundColor "Red"
  }

  # ConEmu
  Write-Host "Extracting: ConEmu"
  try {
    $ArgumentList = @(
      "x", "$Temp\ConEmuPack.7z", "-o$BinDir\ConEmu",
      "-aoa", "-bso0", "-bse0", "-bsp0")
    Start-Process -FilePath $SevenZip -ArgumentList $ArgumentList -NoNewWindow -Wait
    # Remove-Item "$Temp\ConEmuPack.7z"
  }
  catch {
    Write-Host ("  ERROR: Failed to extract files." ) -ForegroundColor "Red"
  }

  # Notepad++
  Write-Host "Extracting: Notepad++"
  try {
    $ArgumentList = @(
      "x", "$Temp\notepadplusplus.7z", "-o$BinDir\NotepadPlusPlus",
      "-aoa", "-bso0", "-bse0", "-bsp0")
    Start-Process -FilePath $SevenZip -ArgumentList $ArgumentList -NoNewWindow -Wait
    # Remove-Item "$Temp\notepadplusplus.7z"
  }
  catch {
    Write-Host ("  ERROR: Failed to extract files." ) -ForegroundColor "Red"
  }

  # NTPWEdit
  Write-Host "Extracting: NTPWEdit"
  try {
    $ArgumentList = @(
      "x", "$Temp\ntpwedit.zip", "-o$BinDir\NTPWEdit",
      "-aoa", "-bso0", "-bse0", "-bsp0")
    Start-Process -FilePath $SevenZip -ArgumentList $ArgumentList -NoNewWindow -Wait
    # Remove-Item "$Temp\ntpwedit.zip"
    Remove-Item "$BinDir\NTPWEdit\ntpwedit.exe"
    Move-Item "$BinDir\NTPWEdit\ntpwedit64.exe" "$BinDir\NTPWEdit\ntpwedit.exe"
  }
  catch {
    Write-Host ("  ERROR: Failed to extract files." ) -ForegroundColor "Red"
  }

  # wimlib
  Write-Host "Extracting: wimlib"
  try {
    $ArgumentList = @(
      "x", "$Temp\wimlib.zip", "-o$BinDir\wimlib",
      "-aoa", "-bso0", "-bse0", "-bsp0")
    Start-Process -FilePath $SevenZip -ArgumentList $ArgumentList -NoNewWindow -Wait
    # Remove-Item "$Temp\wimlib.zip"
  }
  catch {
    Write-Host ("  ERROR: Failed to extract files." ) -ForegroundColor "Red"
  }

  ## Build ##
  $Drivers = "$SetupDir\pe\drivers"
  $Mount = "$BuildDir\mount"
  $PEFiles = "$BuildDir\pe_files"

  # Copy WinPE files
  Write-Host "Copying files..."
  $Cmd = ("{0}\copype.cmd" -f $Env:WinPERoot)
  Start-Process -FilePath $Cmd -ArgumentList @($Arch, $PEFiles) -NoNewWindow -Wait

  # Remove unwanted items
  foreach ($SubDir in @("media", "media\Boot", "media\EFI\Microsoft\Boot")) {
    foreach ($Item in Get-ChildItem "$PEFiles\$SubDir") {
      if ($Item.Name -inotmatch "^(boot|efi|en-us|sources|fonts|resources|bcd|memtest)") {
        Remove-Item -Path $Item.FullName -Recurse -Force
      }
    }
  }

  # Mount image
  Write-Host "Mounting image..."
  New-Item -Path $Mount -ItemType "directory" -Force | Out-Null
  Mount-WindowsImage -Path $Mount -ImagePath "$PEFiles\media\sources\boot.wim" -Index 1 -LogPath "$LogDir\DISM.log"

  # Add drivers
  Add-WindowsDriver -Path $Mount -Driver $Drivers -Recurse -LogPath "$LogDir\DISM.log"

  # Add packages
  Write-Host "Adding packages:"
  $WinPEPackages = @(
    "WinPE-EnhancedStorage",
    "WinPE-FMAPI",
    "WinPE-WMI",
    "WinPE-SecureStartup"
  )
  # $WinPEPackages = @(
    # "WinPE-EnhancedStorage",
    # "WinPE-FMAPI",
    # "WinPE-WMI",
    # "WinPE-SecureStartup",
    # "WinPE-NetFX",
    # "WinPE-Scripting",
    # "WinPE-PowerShell",
    # "WinPE-StorageWMI"
  # )
  foreach ($Package in $WinPEPackages) {
    $PackagePath = ("{0}\{1}\WinPE_OCs\{2}.cab" -f $Env:WinPERoot, $Arch, $Package)
    Write-Host "    $Package..."
    Add-WindowsPackage –PackagePath $PackagePath –Path $Mount -LogPath "$LogDir\DISM.log"
    $LangPackagePath = ("{0}\{1}\WinPE_OCs\en-us\{2}_en-us.cab" -f $Env:WinPERoot, $Arch, $Package)
    if (Test-Path $LangPackagePath) {
      Add-WindowsPackage –PackagePath $LangPackagePath –Path $Mount -LogPath "$LogDir\DISM.log"
    }
  }

  # Set RamDisk size
  $ArgumentList = @(
    ('/Image:"{0}"' -f $Mount),
    "/Set-ScratchSpace:512",
    ('/LogPath:"{0}\DISM.log"' -f $LogDir)
  )
  Start-Process -FilePath $DISM -ArgumentList $ArgumentList -NoNewWindow -Wait

  # Add tools
  Write-Host "Copying tools..."
  Copy-Item -Path "$BinDir\*" -Destination "$Mount\Program Files" -Recurse -Force
  New-Item -Path "$Mount\tools" -ItemType "directory" -Force | Out-Null
  Copy-Item -Path "$ToolsDir\*" -Destination "$Mount\tools" -Recurse -Force
  Copy-Item -Path "$Root\Images\WinPE.jpg" -Destination "$Mount\Program Files\ConEmu\ConEmu.jpg" -Recurse -Force

  # Add System32 items
  $HostSystem32 = "{0}\System32" -f $Env:SystemRoot
  Copy-Item -Path "$SetupDir\pe\System32\*" -Destination "$Mount\Windows\System32" -Recurse -Force
  $ArgumentList = @("/f", "$Mount\Windows\System32\winpe.jpg", "/a")
  Start-Process -FilePath "$HostSystem32\takeown.exe" -ArgumentList $ArgumentList -NoNewWindow -Wait
  $ArgumentList = @("$Mount\Windows\System32\winpe.jpg", "/grant", "Administrators:F")
  Start-Process -FilePath "$HostSystem32\icacls.exe" -ArgumentList $ArgumentList -NoNewWindow -Wait
  Copy-Item -Path "$Root\Images\WinPE.jpg" -Destination "$Mount\Windows\System32\winpe.jpg" -Recurse -Force

  # Load registry hives
  Write-Host "Updating Registry..."
  $Reg = "$HostSystem32\reg.exe"
  $ArgumentList = @("load", "HKLM\WinPE-SW", "$Mount\Windows\System32\config\SOFTWARE")
  Start-Process -FilePath $Reg -ArgumentList $ArgumentList -NoNewWindow -Wait
  $ArgumentList = @("load", "HKLM\WinPE-SYS", "$Mount\Windows\System32\config\SYSTEM")
  Start-Process -FilePath $Reg -ArgumentList $ArgumentList -NoNewWindow -Wait

  # Add tools to path
  ## .NET code to properly handle REG_EXPAND_SZ values
  ## Credit: https://www.sepago.com/blog/2013/08/22/reading-and-writing-regexpandsz-data-with-powershell
  ## By: Marius Gawenda
  $Hive = [Microsoft.Win32.Registry]::LocalMachine
  $RegPath = "WinPE-SYS\ControlSet001\Control\Session Manager\Environment"
  $RegKey = $Hive.OpenSubKey($RegPath)
  $CurValue = $RegKey.GetValue(
    "Path", $false, [Microsoft.Win32.RegistryValueOptions]::DoNotExpandEnvironmentNames)
  $NewValue = "$CurValue;%SystemDrive%\Program Files\7-Zip;%SystemDrive%\Program Files\wimlib"
  Set-ItemProperty -Path "HKLM:\$RegPath" -Name "Path" -Value $NewValue -Force | Out-Null
  $Hive.close()
  $RegKey.close()

  # Replace Notepad
  $RegPath = "HKLM:\WinPE-SW\Microsoft\Windows NT\CurrentVersion\Image File Execution Options\notepad.exe"
  $NewValue = 'cmd /c "%SystemDrive%\tools\npp.cmd"'
  New-Item -Path $RegPath -Force | Out-Null
  New-ItemProperty -Path $RegPath -Name "Debugger" -Value $NewValue -Force | Out-Null

  # Run garbage collection to release potential stale handles
  ## Credit: https://jrich523.wordpress.com/2012/03/06/powershell-loading-and-unloading-registry-hives/
  Start-Sleep -Seconds 2
  [gc]::collect()

  # Unload registry hives
  Start-Sleep -Seconds 2
  Start-Process -FilePath $Reg -ArgumentList @("unload", "HKLM\WinPE-SW") -NoNewWindow -Wait
  Start-Process -FilePath $Reg -ArgumentList @("unload", "HKLM\WinPE-SYS") -NoNewWindow -Wait

  # Unmount image
  Write-Host "Dismounting image..."
  Dismount-WindowsImage -Path $Mount -Save -LogPath "$LogDir\DISM.log"

  # Create ISO
  New-Item -Type Directory "$SetupDir\OUT_PE" 2>&1 | Out-Null
  $ArgumentList = @("/iso", $PEFiles, "$SetupDir\OUT_PE\$KitNameShort-WinPE-$Date-$Arch.iso")
  $Cmd = "{0}\MakeWinPEMedia.cmd" -f $Env:WinPERoot
  Start-Process -FilePath $Cmd -ArgumentList $ArgumentList -NoNewWindow -Wait

  ## Cleanup ##
  Remove-Item -Path "$BuildDir\mount" -Recurse -Force
  # Remove-Item -Path "$BuildDir\pe_files" -Recurse -Force

  ## Done ##
  Pop-Location
  Write-Host "`nDone."
  WKPause "Press Enter to exit... "
}