# Wizard Kit: Windows PE Build Tool

## Init ##
#Requires -Version 3.0
#Requires -RunAsAdministrator
if (Test-Path Env:\DEBUG) {
  Set-PSDebug -Trace 1
}
try {
  Import-Module -Name $Env:DISMRoot -ErrorAction "stop"
}
catch {
  Write-Host -ForegroundColor "Red" "ERROR: Failed to load DISM CmdLet"
  Abort
}
# Dirs
$WorkingDir = $(Split-Path $MyInvocation.MyCommand.Path)
$SetupDir = (Get-Item $WorkingDir -Force).Parent.FullName
$RootDir = (Get-Item $SetupDir -Force).Parent.FullName
$BuildDir = "$SetupDir\BUILD_PE"
$OutDir = "$SetupDir\OUT_PE"
$TempDir = "$BuildDir\temp"
$MountDir = "$BuildDir\mount"
$TargetDir = "$BuildDir\pe_files"
# Misc
$Arch = "amd64"
$Date = Get-Date -UFormat "%Y-%m-%d"
$DISM = "{0}\DISM.exe" -f $Env:DISMRoot
$HostSystem32 = "{0}\System32" -f $Env:SystemRoot
$HostSysWOW64 = "{0}\SysWOW64" -f $Env:SystemRoot
$KitNameFull = (Get-Content "$RootDir\scripts\wk\cfg\main.py" | Where-Object {$_ -imatch '^KIT_NAME_FULL'}) -ireplace '.*=.(.*).$', '$1'
$KitNameShort = (Get-Content "$RootDir\scripts\wk\cfg\main.py" | Where-Object {$_ -imatch '^KIT_NAME_SHORT'}) -ireplace '.*=.(.*).$', '$1'
# Set up UI
$Host.UI.RawUI.WindowTitle = "${KitNameFull}: Windows PE Build Tool"
$Host.UI.RawUI.BackgroundColor = "Black"
$Host.UI.RawUI.ForegroundColor = "White"
# 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 Ask-User ($text = "Kotaero") {
  $text += " [Y/N]"
  while ($true) {
    $answer = read-host $text
    if ($answer -imatch "^(y|yes)$") {
      $answer = $true
      break
    } elseif ($answer -imatch "^(n|no|nope)$") {
      $answer = $false
      break
    }
  }
  $answer
}

function Clean-BuildDir {
  $Folders = @(
    "$BuildDir\additions",
    "$BuildDir\mount",
    "$BuildDir\pe_files",
    "$BuildDir\temp")

  # WIM cleanup
  if (Test-Path "$MountDir") {
    try {
      Dismount-WindowsImage -Path "$MountDir" -Discard
    }
    catch {
      # Ignore
    }
    Start-Process -FilePath $DISM -ArgumentList @("/Cleanup-Mountpoints") -NoNewWindow -Wait
  }

  # Folders
  foreach ($f in $Folders) {
    if (Test-Path $f) {
      Write-Host -ForegroundColor "Yellow" ("Removing: {0}" -f $f)
      Remove-Item -Path $f -Recurse -Force
    }
  }
}

function Download-File ($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"
    Abort
  }
}

function Download-SourceFiles {
  $Sources = Get-Content -Path "$WorkingDir\sources.json" | ConvertFrom-JSON
  foreach ($s in $Sources) {
    $Dest = "$BuildDir\downloads\{0}" -f $s.Name
    if (Test-Path -PathType Leaf -Path "$Dest") {
      if (Test-Hash -File $Dest -Hash $s.Hash) {
        continue
      }
      # Hash didn't match
      Remove-Item $Dest -Force
    }

    # File needs downloaded (again)
    Download-File -Path "$BuildDir\downloads" -Name $s.Name -Url $s.Url

    # Verify download
    if ( -Not (Test-Hash -File $Dest -Hash $s.Hash)) {
      Write-Host ("  ERROR: Download failed hash check." ) -ForegroundColor "Red"
      Abort
    }
  }
}

function Extract-SourceFiles {
  $ProgramFilesPE = "$BuildDir\additions\Program Files"
  New-Item -Type Directory $ProgramFilesPE -Force

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

  # ConEmu
  Write-Host "Extracting: ConEmu"
  try {
    $ArgumentList = @(
      "x", "$BuildDir\downloads\ConEmu.7z", "-o`"$ProgramFilesPE`"\ConEmu",
      "-aoa", "-bso0", "-bse0", "-bsp0",
      "ConEmu.exe",
      "ConEmu.map",
      "ConEmu64.exe",
      "ConEmu64.map",
      "ConEmu\CmdInit.cmd",
      "ConEmu\ConEmuC.exe",
      "ConEmu\ConEmuC64.exe",
      "ConEmu\ConEmuCD.dll",
      "ConEmu\ConEmuCD64.dll",
      "ConEmu\ConEmuHk.dll",
      "ConEmu\ConEmuHk64.dll"
      )
    Start-Process -FilePath $SevenZip -ArgumentList $ArgumentList -NoNewWindow -Wait
  }
  catch {
    Write-Host ("  ERROR: Failed to extract files." ) -ForegroundColor "Red"
  }

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

  # NTPWEdit
  Write-Host "Extracting: NTPWEdit"
  try {
    $ArgumentList = @(
      "e", "$BuildDir\downloads\ntpwedit.zip", "-o`"$ProgramFilesPE`"\NTPWEdit",
      "-aoa", "-bso0", "-bse0", "-bsp0",
      "COPYING.txt", "GPL.txt", "ntpwedit64.exe")
    Start-Process -FilePath $SevenZip -ArgumentList $ArgumentList -NoNewWindow -Wait
    Move-Item "$ProgramFilesPE\NTPWEdit\ntpwedit64.exe" "$ProgramFilesPE\NTPWEdit\ntpwedit.exe"
  }
  catch {
    Write-Host ("  ERROR: Failed to extract files." ) -ForegroundColor "Red"
  }

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

function Find-DynamicUrl ($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 Test-Hash ($File, $Hash) {
  Write-Host -BackgroundColor Black -ForegroundColor Cyan "Verifying ${File}..."
  $FileHash = (Get-FileHash -Path $File -Algorithm SHA256 -ErrorAction Stop).Hash
  return ($FileHash.ToLower() -eq $Hash)
}

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 ".") {
  Write-Host "Wizard Kit: Windows PE Build Tool`n`n`n`n`n"
  Push-Location "$WorkingDir"
  New-Item -Type Directory $BuildDir 2>&1 | Out-Null
  Clean-BuildDir
  Copy-Item -Path "$SetupDir\pe\additions" -Destination "$BuildDir\additions" -Recurse -Force

  ## Download Sources ##
  Download-SourceFiles
  Extract-SourceFiles

  ## Build ##

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

  # Remove unwanted items
  foreach ($SubDir in @("media", "media\Boot", "media\EFI\Microsoft\Boot")) {
    foreach ($Item in Get-ChildItem "$TargetDir\$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 $MountDir -ItemType "directory" -Force | Out-Null
  Mount-WindowsImage -Path $MountDir -ImagePath "$TargetDir\media\sources\boot.wim" -Index 1

  # Add drivers
  Write-Host "Adding drivers..."
  Add-WindowsDriver -Path $MountDir -Driver "$SetupDir\pe\drivers" -ForceUnsigned -Recurse

  # Add packages
  Write-Host "Adding packages..."
  $WinPEPackages = @(
    "WinPE-EnhancedStorage",
    "WinPE-FMAPI",
    "WinPE-WMI",
    "WinPE-SecureStartup"
  )
  foreach ($Package in $WinPEPackages) {
    $PackagePath = ("{0}\{1}\WinPE_OCs\{2}.cab" -f $Env:WinPERoot, $Arch, $Package)
    Write-Host "    $Package..."
    Add-WindowsPackage –PackagePath $PackagePath –Path $MountDir
    $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 $MountDir
    }
  }

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

  # Add tools
  Write-Host "Copying tools..."
  Copy-Item -Path "$BuildDir\additions\*" -Destination $MountDir -Recurse -Force
  Copy-Item -Path "$RootDir\Images\WinPE.jpg" -Destination "$MountDir\Program Files\ConEmu\ConEmu.jpg" -Recurse -Force

  # Add System32 items
  $ArgumentList = @("/f", "$MountDir\Windows\System32\winpe.jpg", "/a")
  Start-Process -FilePath "$HostSystem32\takeown.exe" -ArgumentList $ArgumentList -NoNewWindow -Wait
  $ArgumentList = @("$MountDir\Windows\System32\winpe.jpg", "/grant", "Administrators:F")
  Start-Process -FilePath "$HostSystem32\icacls.exe" -ArgumentList $ArgumentList -NoNewWindow -Wait
  Copy-Item -Path "$RootDir\Images\WinPE.jpg" -Destination "$MountDir\Windows\System32\winpe.jpg" -Force

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

  # Configure CMD (command aliases "DOSKEY Macros" and tab completion)
  $RegPath = "HKLM:\WinPE-SW\Microsoft\Command Processor"
  $NewValue = "doskey /macrofile=X:\Windows\System32\custom.doskey"
  New-Item -Path $RegPath -Force | Out-Null
  New-ItemProperty -Path $RegPath -Name "AutoRun" -Value $NewValue -Force | Out-Null
  New-ItemProperty -Path $RegPath -Name "CompletionChar" -Value 9 -Type DWord -Force | Out-Null
  New-ItemProperty -Path $RegPath -Name "PathCompletionChar" -Value 9 -Type DWord -Force | Out-Null

  # 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;%ProgramFiles%\7-Zip;%ProgramFiles%\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 "%ProgramFiles%\NotepadPlusPlus\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 $MountDir -Save

  # Rebuild image
  Write-Host "Recompressing image..."
  Move-Item "$TargetDir\media\sources\boot.wim" "$TargetDir\media\sources\boot-edited.wim"
  Export-WindowsImage -DestinationImagePath "$TargetDir\media\sources\boot.wim" -SourceImagePath "$TargetDir\media\sources\boot-edited.wim" -SourceIndex 1 -CompressionType Max

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

  ## Cleanup ##
  Remove-Item -Path "$MountDir" -Recurse -Force

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