From 27f87819f182026879175fcc110f27daebbf36a8 Mon Sep 17 00:00:00 2001
From: 2Shirt <2xShirt@gmail.com>
Date: Sun, 28 Apr 2024 01:27:01 -0700
Subject: [PATCH] Refactor WinPE code for the current Windows ADK

---
 setup/build_pe.cmd                   |  12 +-
 setup/pe/System32/Winpeshl.ini       |   4 +-
 setup/pe/bin/NotepadPlusPlus/npp.cmd |   2 +-
 setup/pe/build_pe.ps1                | 786 ++++++++-------------------
 setup/pe/drivers/README.md           |   3 +
 setup/pe/sources.json                |   7 +
 6 files changed, 243 insertions(+), 571 deletions(-)
 mode change 100644 => 100755 setup/build_pe.cmd
 create mode 100644 setup/pe/drivers/README.md
 create mode 100644 setup/pe/sources.json

diff --git a/setup/build_pe.cmd b/setup/build_pe.cmd
old mode 100644
new mode 100755
index 44903d91..e04718f4
--- a/setup/build_pe.cmd
+++ b/setup/build_pe.cmd
@@ -6,10 +6,6 @@
 setlocal EnableDelayedExpansion
 title WizardKit: Build Tool
 call :CheckFlags %*
-rem TODO: Remove warning
-echo "Windows PE build is currently under development"
-echo "  Proceeding will likely result in errors so be warned"
-pause
 call :CheckElevation || goto Exit
 call :FindKitsRoot || goto ErrorKitNotFound
 
@@ -19,14 +15,8 @@ set "dandi_set_env=%adk_root%\Deployment Tools\DandISetEnv.bat"
 if not exist "%dandi_set_env%" (goto ErrorKitNotFound)
 call "%dandi_set_env%" || goto ErrorUnknown
 
-:EnsureCRLF
-rem Rewrite main.py using PowerShell to have CRLF/`r`n lineendings
-set "script=%~dp0\.bin\Scripts\borrowed\set-eol.ps1"
-set "main=%~dp0\.bin\Scripts\settings\main.py"
-powershell -executionpolicy bypass -noprofile -file %script% -lineEnding win -file %main% || goto ErrorUnknown
-
 :Launch
-set "script=%~dp0\.bin\Scripts\build_pe.ps1"
+set "script=%~dp0\pe\build_pe.ps1"
 powershell -executionpolicy bypass -noprofile -file %script% || goto ErrorUnknown
 goto Exit
 
diff --git a/setup/pe/System32/Winpeshl.ini b/setup/pe/System32/Winpeshl.ini
index 9147e972..5883bef4 100644
--- a/setup/pe/System32/Winpeshl.ini
+++ b/setup/pe/System32/Winpeshl.ini
@@ -2,5 +2,5 @@
 [LaunchApps]
 wpeinit
 wpeutil updatebootinfo
-cd /d "%SystemDrive%\.bin"
-"%SystemDrive%\.bin\ConEmu\ConEmu.exe", /cmd cmd /k cd "%SystemDrive%\.bin" & python "%SystemDrive%\.bin\Scripts\winpe_root_menu.py"
+cd /d "%SystemDrive%"
+"%SystemDrive%\Program Files\ConEmu\ConEmu64.exe", /cmd cmd /k cd "%SystemDrive%"
diff --git a/setup/pe/bin/NotepadPlusPlus/npp.cmd b/setup/pe/bin/NotepadPlusPlus/npp.cmd
index 6996dd3c..d385967f 100644
--- a/setup/pe/bin/NotepadPlusPlus/npp.cmd
+++ b/setup/pe/bin/NotepadPlusPlus/npp.cmd
@@ -1,3 +1,3 @@
 @echo off
 
-start "" %SystemDrive%\.bin\NotepadPlusPlus\notepadplusplus.exe %2 %3 %4 %5 %6 %7 %8 %9
\ No newline at end of file
+start "" "%SystemDrive%\Program Files\NotepadPlusPlus\notepad++.exe" %2 %3 %4 %5 %6 %7 %8 %9
\ No newline at end of file
diff --git a/setup/pe/build_pe.ps1 b/setup/pe/build_pe.ps1
index 2b205093..8bce5112 100644
--- a/setup/pe/build_pe.ps1
+++ b/setup/pe/build_pe.ps1
@@ -6,13 +6,19 @@
 if (Test-Path Env:\DEBUG) {
   Set-PSDebug -Trace 1
 }
+# TODO REMOVE v
+$KitNameShort = "WK"
+# TODO REMOVE ^
+$Arch = "amd64"
 $Host.UI.RawUI.WindowTitle = "Wizard Kit: Windows PE Build Tool"
 $WD = $(Split-Path $MyInvocation.MyCommand.Path)
-$Bin = (Get-Item $WD -Force).Parent.FullName
-$Root = (Get-Item $Bin -Force).Parent.FullName
-$Build = "$Root\BUILD_PE"
-$LogDir = "$Build\Logs"
-$Temp = "$Build\Temp"
+$SetupDir = (Get-Item $WD -Force).Parent.FullName
+$Root = (Get-Item $SetupDir -Force).Parent.FullName
+$BuildDir = "$SetupDir\BUILD_PE"
+$BinDir = "$BuildDir\bin"
+$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"
@@ -22,45 +28,13 @@ $DISM = "{0}\DISM.exe" -f $Env:DISMRoot
 #Enable TLS 1.2
 [Net.ServicePointManager]::SecurityProtocol = [Net.SecurityProtocolType]::Tls12
 
+
 ## Functions ##
-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 Abort {
   Write-Host -ForegroundColor "Red" "`nAborted."
-  WKPause "Press Enter to exit... "
+  WKPause "Press Enter to exit..."
   exit
 }
-function MakeClean {
-  $Folders = @(
-    "$Build\Mount",
-    "$Build\PEFiles")
-  $Clean = $false
-  foreach ($f in $Folders) {
-    if (Test-Path $f) {
-      Write-Host -ForegroundColor "Yellow" ("Found: {0}" -f $f)
-      $Clean = $true
-    }
-  }
-  if (($Clean) -and (Ask-User "Delete the above folder(s)?")) {
-    foreach ($f in $Folders) {
-      if (Test-Path $f) {
-        Remove-Item -Path $f -Recurse -Force
-      }
-    }
-  }
-}
 function DownloadFile ($Path, $Name, $Url) {
   $OutFile = "{0}\{1}" -f $Path, $Name
 
@@ -88,11 +62,30 @@ function FindDynamicUrl ($SourcePage, $RegEx) {
 
   $Url | Select-Object -First 1
 }
+function MakeClean {
+  $Folders = @(
+    "$BuildDir\bin",
+    "$BuildDir\mount",
+    "$BuildDir\pe_files")
+  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
@@ -101,7 +94,7 @@ function WKPause ($Message = "Press Enter to continue... ") {
 # Answer by:    https://stackoverflow.com/users/696808/bacon-bits
 if ($MyInvocation.InvocationName -ne ".") {
   Clear-Host
-  Write-Host "Wizard Kit: Windows PE Build Tool`n`n`n`n`n"
+  Write-Host "WizardKit: Build Tool`n`n`n`n`n"
 
   ## Prep ##
   try {
@@ -113,542 +106,221 @@ if ($MyInvocation.InvocationName -ne ".") {
   }
   Push-Location "$WD"
   MakeClean
-  New-Item -Type Directory $Build 2>&1 | Out-Null
+  New-Item -Type Directory $BuildDir 2>&1 | Out-Null
   New-Item -Type Directory $LogDir 2>&1 | Out-Null
 
-  ## main.py ##
-  if (!(Test-Path "$Build\main.py") -or (Ask-User "Replace existing main.py?")) {
-    Copy-Item -Path "$Bin\Scripts\settings\main.py" -Destination "$Build\main.py" -Force
+  ## 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.zip" -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
   }
-  WKPause "Press Enter to open settings..."
-  Start-Process "$HostSystem32\notepad.exe" -ArgumentList @("$Build\main.py") -Wait
-  $KitNameFull = (Get-Content "$Build\main.py" | Where-Object {$_ -match 'FULL'}) -replace ".*'(.*)'$", '$1'
-  $KitNameShort = (Get-Content "$Build\main.py" | Where-Object {$_ -match 'SHORT'}) -replace ".*'(.*)'$", '$1'
 
-  if (Ask-User "Update Tools?") {
-    $DownloadErrors = 0
+  ## Extract ##
+  Copy-Item -Path "$SetupDir\pe\bin" -Destination "$BinDir" -Recurse -Force
 
-    ## Download Tools ##
-    $ToolSources = @(
-      # 7-Zip
-      @("7z-installer.msi", "https://www.7-zip.org/a/7z1900.msi"),
-      @("7z-extra.7z", "https://www.7-zip.org/a/7z1900-extra.7z"),
-      # Blue Screen View
-      @("bluescreenview32.zip", "http://www.nirsoft.net/utils/bluescreenview.zip"),
-      @("bluescreenview64.zip", "http://www.nirsoft.net/utils/bluescreenview-x64.zip"),
-      # ConEmu
-      @("ConEmuPack.7z", "https://github.com/Maximus5/ConEmu/releases/download/v19.03.10/ConEmuPack.190310.7z"),
-      # Fast Copy
-      @("fastcopy.zip", "http://ftp.vector.co.jp/71/31/2323/FastCopy363_installer.exe"),
-      # HWiNFO
-      @("hwinfo.zip", "http://files2.majorgeeks.com/377527622c5325acc1cb937fb149d0de922320c0/systeminfo/hwi_602.zip"),
-      # Killer Network Drivers
-      @(
-        "killerinf.zip",
-        ("https://www.killernetworking.com"+(FindDynamicUrl "https://www.killernetworking.com/killersupport/category/other-downloads" "Download Killer-Ethernet").replace('&amp;', '&'))
-      ),
-      # Notepad++
-      @("npp_x86.7z", "https://notepad-plus-plus.org/repository/7.x/7.6.4/npp.7.6.4.bin.minimalist.7z"),
-      @("npp_amd64.7z", "https://notepad-plus-plus.org/repository/7.x/7.6.4/npp.7.6.4.bin.minimalist.x64.7z"),
-      # NT Password Editor
-      @("ntpwed.zip", "http://cdslow.org.ru/files/ntpwedit/ntpwed07.zip"),
-      # Prime95
-      @("prime95_32.zip", "http://www.mersenne.org/ftp_root/gimps/p95v294b7.win32.zip"),
-      @("prime95_64.zip", "http://www.mersenne.org/ftp_root/gimps/p95v294b8.win64.zip"),
-      # ProduKey
-      @("produkey32.zip", "http://www.nirsoft.net/utils/produkey.zip"),
-      @("produkey64.zip", "http://www.nirsoft.net/utils/produkey-x64.zip"),
-      # Python
-      @("python32.zip", "https://www.python.org/ftp/python/3.7.2/python-3.7.2.post1-embed-win32.zip"),
-      @("python64.zip", "https://www.python.org/ftp/python/3.7.2/python-3.7.2.post1-embed-amd64.zip"),
-      # Python: psutil
-      @(
-        "psutil64.whl",
-        (FindDynamicUrl "https://pypi.org/project/psutil/" "href=.*-cp37-cp37m-win_amd64.whl")
-      ),
-      @(
-        "psutil32.whl",
-        (FindDynamicUrl "https://pypi.org/project/psutil/" "href=.*-cp37-cp37m-win32.whl")
-      ),
-      # Q-Dir
-      @("qdir32.zip", "https://www.softwareok.com/Download/Q-Dir_Portable.zip"),
-      @("qdir64.zip", "https://www.softwareok.com/Download/Q-Dir_Portable_x64.zip"),
-      # TestDisk / PhotoRec
-      @("testdisk32.zip", "https://www.cgsecurity.org/testdisk-7.1-WIP.win.zip"),
-      @("testdisk64.zip", "https://www.cgsecurity.org/testdisk-7.1-WIP.win64.zip"),
-      # VirtIO drivers
-      @("virtio-win.iso", "https://fedorapeople.org/groups/virt/virtio-win/direct-downloads/latest-virtio/virtio-win.iso"),
-      # Visual C++ Runtimes
-      @("vcredist_x86.exe", "https://aka.ms/vs/15/release/vc_redist.x86.exe"),
-      @("vcredist_x64.exe", "https://aka.ms/vs/15/release/vc_redist.x64.exe"),
-      # wimlib-imagex
-      @("wimlib32.zip", "https://wimlib.net/downloads/wimlib-1.13.0-windows-i686-bin.zip"),
-      @("wimlib64.zip", "https://wimlib.net/downloads/wimlib-1.13.0-windows-x86_64-bin.zip")
-    )
-    foreach ($Tool in $ToolSources) {
-      DownloadFile -Path $Temp -Name $Tool[0] -Url $Tool[1]
-    }
+  # 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"
+  }
 
-    ## Bail ##
-    # If errors were encountered during downloads
-    if ($DownloadErrors -gt 0) {
-      Abort
-    }
+  # 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"
+  }
 
-    ## Install ##
-    # Visual C++ Runtimes
-    Write-Host "Installing: Visual C++ Runtimes"
-    $ArgumentList = @("/install", "/passive", "/norestart")
-    Start-Process -FilePath "$Temp\vcredist_x86.exe" -ArgumentList $ArgumentList -Wait
-    Start-Process -FilePath "$Temp\vcredist_x64.exe" -ArgumentList $ArgumentList -Wait
+  # 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"
+  }
 
-    ## Extract ##
-    # 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
-      $SevenZip = "$Temp\7zi\Files\7-Zip\7z.exe"
-      $ArgumentList = @(
-        "e", "$Temp\7z-extra.7z", "-o$Build\bin\amd64\7-Zip",
-        "-aoa", "-bso0", "-bse0", "-bsp0",
-        "x64\7za.exe", "*.txt")
-      Start-Process -FilePath $SevenZip -ArgumentList $ArgumentList -NoNewWindow -Wait
-      $ArgumentList = @(
-        "e", "$Temp\7z-extra.7z", "-o$Build\bin\x86\7-Zip",
-        "-aoa", "-bso0", "-bse0", "-bsp0",
-        "7za.exe", "*.txt")
-      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 = @(
+      "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"
+  }
 
-    # Blue Screen View
-    Write-Host "Extracting: BlueScreenView"
-    try {
-      $ArgumentList = @(
-        "x", "$Temp\bluescreenview64.zip", "-o$Build\bin\amd64\BlueScreenView",
-        "-aoa", "-bso0", "-bse0", "-bsp0")
-      Start-Process -FilePath $SevenZip -ArgumentList $ArgumentList -NoNewWindow -Wait
-      $ArgumentList = @(
-        "x", "$Temp\bluescreenview32.zip", "-o$Build\bin\x86\BlueScreenView",
-        "-aoa", "-bso0", "-bse0", "-bsp0")
-      Start-Process -FilePath $SevenZip -ArgumentList $ArgumentList -NoNewWindow -Wait
-    }
-    catch {
-      Write-Host ("  ERROR: Failed to extract files." ) -ForegroundColor "Red"
-    }
-
-    # ConEmu
-    Write-Host "Extracting: ConEmu"
-    try {
-      $ArgumentList = @(
-        "x", "$Temp\ConEmuPack.7z", "-o$Build\bin\amd64\ConEmu",
-        "-aoa", "-bso0", "-bse0", "-bsp0")
-      Start-Process -FilePath $SevenZip -ArgumentList $ArgumentList -NoNewWindow -Wait
-      Remove-Item "$Build\bin\amd64\ConEmu\ConEmu.exe"
-      Remove-Item "$Build\bin\amd64\ConEmu\ConEmu.map"
-      Move-Item "$Build\bin\amd64\ConEmu\ConEmu64.exe" "$Build\bin\amd64\ConEmu\ConEmu.exe" -Force
-      Move-Item "$Build\bin\amd64\ConEmu\ConEmu64.map" "$Build\bin\amd64\ConEmu\ConEmu.map" -Force
-      $ArgumentList = @(
-        "x", "$Temp\ConEmuPack.7z", "-o$Build\bin\x86\ConEmu",
-        "-aoa", "-bso0", "-bse0", "-bsp0")
-      Start-Process -FilePath $SevenZip -ArgumentList $ArgumentList -NoNewWindow -Wait
-      Remove-Item "$Build\bin\x86\ConEmu\ConEmu64.exe"
-      Remove-Item "$Build\bin\x86\ConEmu\ConEmu64.map"
-    }
-    catch {
-      Write-Host ("  ERROR: Failed to extract files." ) -ForegroundColor "Red"
-    }
-
-    # Fast Copy
-    Write-Host "Extracting: FastCopy"
-    try {
-      # Extract Installer
-      $ArgumentList = @(
-        "e", "$Temp\fastcopy.zip", "-o$Temp",
-        "-aoa", "-bso0", "-bse0", "-bsp0")
-      Start-Process -FilePath $SevenZip -ArgumentList $ArgumentList -NoNewWindow -Wait
-
-      # Extract 64-bit
-      $ArgumentList = @(
-        "/NOSUBDIR", "/DIR=$Build\bin\amd64\FastCopy",
-        "/EXTRACT64")
-      Start-Process -FilePath "$TEMP\FastCopy354_installer.exe" -ArgumentList $ArgumentList -NoNewWindow -Wait
-      Remove-Item "$Build\bin\amd64\FastCopy\setup.exe" -Force
-
-      # Extract 32-bit
-      $ArgumentList = @(
-        "/NOSUBDIR", "/DIR=$Build\bin\x86\FastCopy",
-        "/EXTRACT32")
-      Start-Process -FilePath "$TEMP\FastCopy354_installer.exe" -ArgumentList $ArgumentList -NoNewWindow -Wait
-      Remove-Item "$Build\bin\x86\FastCopy\setup.exe" -Force
-    }
-    catch {
-      Write-Host ("  ERROR: Failed to extract files." ) -ForegroundColor "Red"
-    }
-
-
-    # Killer Network Driver
-    Write-Host "Extracting: Killer Network Driver"
-    try {
-      $ArgumentList = @(
-        "e", "$Temp\killerinf.zip", "-o$Build\Drivers\amd64\Killer",
-        "-aoa", "-bso0", "-bse0", "-bsp0",
-        "Production\Windows10-x64\Eth\*")
-      Start-Process -FilePath $SevenZip -ArgumentList $ArgumentList -NoNewWindow -Wait
-      $ArgumentList = @(
-        "e", "$Temp\killerinf.zip", "-o$Build\Drivers\x86\Killer",
-        "-aoa", "-bso0", "-bse0", "-bsp0",
-        "Production\Windows10-x86\Eth\*")
-      Start-Process -FilePath $SevenZip -ArgumentList $ArgumentList -NoNewWindow -Wait
-    }
-    catch {
-      Write-Host ("  ERROR: Failed to extract files." ) -ForegroundColor "Red"
-    }
-
-    # HWiNFO
-    Write-Host "Extracting: HWiNFO"
-    try {
-      $ArgumentList = @(
-        "e", "$Temp\hwinfo.zip", "-o$Build\bin\amd64\HWiNFO",
-        "-aoa", "-bso0", "-bse0", "-bsp0", "HWiNFO64.exe")
-      Start-Process -FilePath $SevenZip -ArgumentList $ArgumentList -NoNewWindow -Wait
-      $ArgumentList = @(
-        "e", "$Temp\hwinfo.zip", "-o$Build\bin\x86\HWiNFO",
-        "-aoa", "-bso0", "-bse0", "-bsp0", "HWiNFO32.exe")
-      Start-Process -FilePath $SevenZip -ArgumentList $ArgumentList -NoNewWindow -Wait
-      Move-Item "$Build\bin\amd64\HWiNFO\HWiNFO64.exe" "$Build\bin\amd64\HWiNFO\HWiNFO.exe" -Force
-      Move-Item "$Build\bin\x86\HWiNFO\HWiNFO32.exe" "$Build\bin\x86\HWiNFO\HWiNFO.exe" -Force
-    }
-    catch {
-      Write-Host ("  ERROR: Failed to extract files." ) -ForegroundColor "Red"
-    }
-
-    # Notepad++
-    Write-Host "Extracting: Notepad++"
-    try {
-      $ArgumentList = @(
-        "x", "$Temp\npp_amd64.7z", "-o$Build\bin\amd64\NotepadPlusPlus",
-        "-aoa", "-bso0", "-bse0", "-bsp0")
-      Start-Process -FilePath $SevenZip -ArgumentList $ArgumentList -NoNewWindow -Wait
-      $ArgumentList = @(
-        "x", "$Temp\npp_x86.7z", "-o$Build\bin\x86\NotepadPlusPlus",
-        "-aoa", "-bso0", "-bse0", "-bsp0")
-      Start-Process -FilePath $SevenZip -ArgumentList $ArgumentList -NoNewWindow -Wait
-      Move-Item "$Build\bin\amd64\NotepadPlusPlus\notepad++.exe" "$Build\bin\amd64\NotepadPlusPlus\notepadplusplus.exe" -Force
-      Move-Item "$Build\bin\x86\NotepadPlusPlus\notepad++.exe" "$Build\bin\x86\NotepadPlusPlus\notepadplusplus.exe" -Force
-    }
-    catch {
-      Write-Host ("  ERROR: Failed to extract files." ) -ForegroundColor "Red"
-    }
-
-    # NT Password Editor
-    Write-Host "Extracting: NT Password Editor"
-    try {
-      $ArgumentList = @(
-        "e", "$Temp\ntpwed.zip", ('-o"{0}\bin\amd64\NT Password Editor"' -f $Build),
-        "-aoa", "-bso0", "-bse0", "-bsp0",
-        "ntpwedit64.exe", "*.txt")
-      Start-Process -FilePath $SevenZip -ArgumentList $ArgumentList -NoNewWindow -Wait
-      Move-Item "$Build\bin\amd64\NT Password Editor\ntpwedit64.exe" "$Build\bin\amd64\NT Password Editor\ntpwedit.exe" -Force
-      $ArgumentList = @(
-        "e", "$Temp\ntpwed.zip", ('-o"{0}\bin\x86\NT Password Editor"' -f $Build),
-        "-aoa", "-bso0", "-bse0", "-bsp0",
-        "ntpwedit.exe", "*.txt")
-      Start-Process -FilePath $SevenZip -ArgumentList $ArgumentList -NoNewWindow -Wait
-    }
-    catch {
-      Write-Host ("  ERROR: Failed to extract files." ) -ForegroundColor "Red"
-    }
-
-    # PhotoRec / TestDisk
-    Write-Host "Extracting: PhotoRec / TestDisk"
-    try {
-      $ArgumentList = @(
-        "x", "$Temp\testdisk64.zip", "-o$Build\bin\amd64\TestDisk",
-        "-aoa", "-bso0", "-bse0", "-bsp0")
-      # Remove destination since Move-Item -Force can't handle this recursive merge
-      Remove-Item "$Build\bin\amd64\TestDisk" -Recurse -Force 2>&1 | Out-Null
-      Start-Process -FilePath $SevenZip -ArgumentList $ArgumentList -NoNewWindow -Wait
-      Move-Item "$Build\bin\amd64\TestDisk\testdisk-7.1-WIP\*" "$Build\bin\amd64\TestDisk" -Force
-      Remove-Item "$Build\bin\amd64\TestDisk\testdisk-7.1-WIP" -Recurse -Force
-      $ArgumentList = @(
-        "x", "$Temp\testdisk32.zip", "-o$Build\bin\x86\TestDisk",
-        "-aoa", "-bso0", "-bse0", "-bsp0")
-      # Remove destination since Move-Item -Force can't handle this recursive merge
-      Remove-Item "$Build\bin\x86\TestDisk" -Recurse -Force 2>&1 | Out-Null
-      Start-Process -FilePath $SevenZip -ArgumentList $ArgumentList -NoNewWindow -Wait
-      Move-Item "$Build\bin\x86\TestDisk\testdisk-7.1-WIP\*" "$Build\bin\x86\TestDisk" -Force
-      Remove-Item "$Build\bin\x86\TestDisk\testdisk-7.1-WIP" -Recurse -Force
-    }
-    catch {
-      Write-Host ("  ERROR: Failed to extract files." ) -ForegroundColor "Red"
-    }
-
-    # Prime95
-    Write-Host "Extracting: Prime95"
-    try {
-      $ArgumentList = @(
-        "x", "$Temp\prime95_64.zip", "-o$Build\bin\amd64\Prime95",
-        "-aoa", "-bso0", "-bse0", "-bsp0")
-      Start-Process -FilePath $SevenZip -ArgumentList $ArgumentList -NoNewWindow -Wait
-      $ArgumentList = @(
-        "x", "$Temp\prime95_32.zip", "-o$Build\bin\x86\Prime95",
-        "-aoa", "-bso0", "-bse0", "-bsp0")
-      Start-Process -FilePath $SevenZip -ArgumentList $ArgumentList -NoNewWindow -Wait
-    }
-    catch {
-      Write-Host ("  ERROR: Failed to extract files." ) -ForegroundColor "Red"
-    }
-
-    # ProduKey
-    try {
-      $ArgumentList = @(
-        "x", "$Temp\produkey64.zip", "-o$Build\bin\amd64\ProduKey",
-        "-aoa", "-bso0", "-bse0", "-bsp0")
-      Start-Process -FilePath $SevenZip -ArgumentList $ArgumentList -NoNewWindow -Wait
-      $ArgumentList = @(
-        "x", "$Temp\produkey32.zip", "-o$Build\bin\x86\ProduKey",
-        "-aoa", "-bso0", "-bse0", "-bsp0")
-      Start-Process -FilePath $SevenZip -ArgumentList $ArgumentList -NoNewWindow -Wait
-    }
-    catch {
-      Write-Host ("  ERROR: Failed to extract files." ) -ForegroundColor "Red"
-    }
-
-    # Python (x64)
-    Write-Host "Extracting: Python (x64)"
-    try {
-      $ArgumentList = @(
-        "x", "$Temp\python64.zip", "-o$Build\bin\amd64\python",
-        "-aoa", "-bso0", "-bse0", "-bsp0")
-      Start-Process -FilePath $SevenZip -ArgumentList $ArgumentList -NoNewWindow -Wait
-      $ArgumentList = @(
-        "x", "$Temp\psutil64.whl", "-o$Build\bin\amd64\python",
-        "-aoa", "-bso0", "-bse0", "-bsp0")
-      Start-Process -FilePath $SevenZip -ArgumentList $ArgumentList -NoNewWindow -Wait
-
-    }
-    catch {
-      Write-Host ("  ERROR: Failed to extract files." ) -ForegroundColor "Red"
-    }
-    try {
-      Copy-Item -Path "$HostSystem32\vcruntime140.dll" -Destination "$Build\bin\amd64\python\vcruntime140.dll" -Force
-    }
-    catch {
-      Write-Host ("  ERROR: Failed to copy Visual C++ Runtime DLL." ) -ForegroundColor "Red"
-    }
-
-    # Python (x32)
-    Write-Host "Extracting: Python (x32)"
-    try {
-      $ArgumentList = @(
-        "x", "$Temp\python32.zip", "-o$Build\bin\x86\python",
-        "-aoa", "-bso0", "-bse0", "-bsp0")
-      Start-Process -FilePath $SevenZip -ArgumentList $ArgumentList -NoNewWindow -Wait
-      $ArgumentList = @(
-        "x", "$Temp\psutil32.whl", "-o$Build\bin\x86\python",
-        "-aoa", "-bso0", "-bse0", "-bsp0")
-      Start-Process -FilePath $SevenZip -ArgumentList $ArgumentList -NoNewWindow -Wait
-
-    }
-    catch {
-      Write-Host ("  ERROR: Failed to extract files." ) -ForegroundColor "Red"
-    }
-    try {
-      Copy-Item -Path "$HostSysWOW64\vcruntime140.dll" -Destination "$Build\bin\x86\python\vcruntime140.dll" -Force
-    }
-    catch {
-      Write-Host ("  ERROR: Failed to copy Visual C++ Runtime DLL." ) -ForegroundColor "Red"
-    }
-
-    # Q-Dir
-    Write-Host "Extracting: Q-Dir"
-    try {
-      $ArgumentList = @(
-        "x", "$Temp\qdir64.zip", "-o$Build\bin\amd64",
-        "-aoa", "-bso0", "-bse0", "-bsp0")
-      Start-Process -FilePath $SevenZip -ArgumentList $ArgumentList -NoNewWindow -Wait
-      Move-Item "$Build\bin\amd64\Q-Dir\Q-Dir_x64.exe" "$Build\bin\amd64\Q-Dir\Q-Dir.exe" -Force
-      $ArgumentList = @(
-        "x", "$Temp\qdir32.zip", "-o$Build\bin\x86",
-        "-aoa", "-bso0", "-bse0", "-bsp0")
-      Start-Process -FilePath $SevenZip -ArgumentList $ArgumentList -NoNewWindow -Wait
-    }
-    catch {
-      Write-Host ("  ERROR: Failed to extract files." ) -ForegroundColor "Red"
-    }
-
-    # VirtIO Drivers
-    Write-Host "Extracting: VirtIO Drivers"
-    try {
-      $ArgumentList = @(
-        "e", "$Temp\virtio-win.iso", "-o$Build\Drivers\amd64\VirtIO",
-        "-aoa", "-bso0", "-bse0", "-bsp0",
-        "*\w10\amd64\*")
-      Start-Process -FilePath $SevenZip -ArgumentList $ArgumentList -NoNewWindow -Wait
-      $ArgumentList = @(
-        "e", "$Temp\virtio-win.iso", "-o$Build\Drivers\x86\VirtIO",
-        "-aoa", "-bso0", "-bse0", "-bsp0",
-        "*\w10\x86\*")
-      Start-Process -FilePath $SevenZip -ArgumentList $ArgumentList -NoNewWindow -Wait
-    }
-    catch {
-      Write-Host ("  ERROR: Failed to extract files." ) -ForegroundColor "Red"
-    }
-
-    # wimlib-imagex
-    try {
-      $ArgumentList = @(
-        "x", "$Temp\wimlib64.zip", "-o$Build\bin\amd64\wimlib",
-        "-aoa", "-bso0", "-bse0", "-bsp0")
-      Start-Process -FilePath $SevenZip -ArgumentList $ArgumentList -NoNewWindow -Wait
-      $ArgumentList = @(
-        "x", "$Temp\wimlib32.zip", "-o$Build\bin\x86\wimlib",
-        "-aoa", "-bso0", "-bse0", "-bsp0")
-      Start-Process -FilePath $SevenZip -ArgumentList $ArgumentList -NoNewWindow -Wait
-    }
-    catch {
-      Write-Host ("  ERROR: Failed to extract files." ) -ForegroundColor "Red"
-    }
-
-    ## Cleanup ##
-    if (Ask-User "Delete temp files?") {
-      Remove-Item "$Temp" -Recurse
-    }
+  # 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 ##
-  foreach ($Arch in @("amd64", "x86")) {
-    $Drivers = "$Build\Drivers\$Arch"
-    $Mount = "$Build\Mount"
-    $PEFiles = "$Build\PEFiles\$Arch"
+  $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
+  # 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
-        }
+  # 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"
-    )
-    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 "$Build\bin\$Arch" -Destination "$Mount\.bin" -Recurse -Force
-    Copy-Item -Path "$Root\.pe_items\_include\*" -Destination "$Mount\.bin" -Recurse -Force
-    if ($Arch -eq "amd64") {
-      $DestIni = "$Mount\.bin\HWiNFO\HWiNFO64.INI"
-    } else {
-      $DestIni = "$Mount\.bin\HWiNFO\HWiNFO32.INI"
-    }
-    Move-Item -Path "$Mount\.bin\HWiNFO\HWiNFO.INI" -Destination $DestIni -Force
-    Copy-Item -Path "$Root\Images\WinPE.jpg" -Destination "$Mount\.bin\ConEmu\ConEmu.jpg" -Recurse -Force
-    Copy-Item -Path "$Bin\Scripts" -Destination "$Mount\.bin\Scripts" -Recurse -Force
-    Copy-Item -Path "$Build\main.py" -Destination "$Mount\.bin\Scripts\settings\main.py" -Force
-
-    # Add System32 items
-    $HostSystem32 = "{0}\System32" -f $Env:SystemRoot
-    Copy-Item -Path "$Root\.pe_items\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%\.bin\7-Zip;%SystemDrive%\.bin\python;%SystemDrive%\.bin\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%\.bin\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 $Mount -Save -LogPath "$LogDir\DISM.log"
-
-    # Create ISO
-    New-Item -Type Directory "$Root\OUT_PE" 2>&1 | Out-Null
-    $ArgumentList = @("/iso", $PEFiles, "$Root\OUT_PE\$KitNameShort-WinPE-$Date-$Arch.iso")
-    $Cmd = "{0}\MakeWinPEMedia.cmd" -f $Env:WinPERoot
-    Start-Process -FilePath $Cmd -ArgumentList $ArgumentList -NoNewWindow -Wait
   }
 
+  # 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"
+  )
+  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
+  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()
+
+  # Hasleo Disk Clone
+  # $ArgumentList = @("import", "$SetupDir\pe\Hasleo.reg")
+  # Start-Process -FilePath $Reg -ArgumentList $ArgumentList -NoNewWindow -Wait
+
+  # Replace Notepad
+  $RegPath = "HKLM:\WinPE-SW\Microsoft\Windows NT\CurrentVersion\Image File Execution Options\notepad.exe"
+  $NewValue = 'cmd /c "%SystemDrive%\.bin\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 $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 "$Build\Mount" -Recurse -Force
-  Remove-Item -Path "$Build\PEFiles" -Recurse -Force
+  Remove-Item -Path "$BuildDir\mount" -Recurse -Force
+  # Remove-Item -Path "$BuildDir\pe_files" -Recurse -Force
 
   ## Done ##
   Pop-Location
diff --git a/setup/pe/drivers/README.md b/setup/pe/drivers/README.md
new file mode 100644
index 00000000..2e85bbb6
--- /dev/null
+++ b/setup/pe/drivers/README.md
@@ -0,0 +1,3 @@
+# WizardKit: WinPE - Drivers #
+
+All drivers added to this folder will be injected into the WinPE image (recursively)
diff --git a/setup/pe/sources.json b/setup/pe/sources.json
new file mode 100644
index 00000000..044979d3
--- /dev/null
+++ b/setup/pe/sources.json
@@ -0,0 +1,7 @@
+{
+  "7-Zip": "https://7-zip.org/a/7z2301-x64.msi",
+  "ConEmu": "https://github.com/Maximus5/ConEmu/releases/download/v23.07.24/ConEmuPack.230724.7z",
+  "Notepad++": "https://github.com/notepad-plus-plus/notepad-plus-plus/releases/download/v8.6.2/npp.8.6.2.portable.x64.7z",
+  "NTPWEdit": "http://cdslow.org.ru/files/ntpwedit/ntpwed07.zip",
+  "wimlib": "https://wimlib.net/downloads/wimlib-1.14.3-windows-x86_64-bin.zip"
+}
\ No newline at end of file