From 44514b2fefc843a91ca04ef47675db7e963aad92 Mon Sep 17 00:00:00 2001
From: Alan Mason <1923621+2Shirt@users.noreply.github.com>
Date: Mon, 27 Nov 2017 16:05:44 -0800
Subject: [PATCH] Initial PowerShell rewrite done.

---
 .bin/Scripts/build_pe.ps1 | 229 +++++++++++++++++++++++++++++++++-----
 .gitignore                |   2 +-
 2 files changed, 201 insertions(+), 30 deletions(-)

diff --git a/.bin/Scripts/build_pe.ps1 b/.bin/Scripts/build_pe.ps1
index de687894..4d1570b8 100644
--- a/.bin/Scripts/build_pe.ps1
+++ b/.bin/Scripts/build_pe.ps1
@@ -2,33 +2,87 @@
 
 ## Init ##
 #Requires -Version 3.0
-#Requires -RunAsAdministrator 
-Clear-Host
+#Requires -RunAsAdministrator
 $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
 $Temp = "{0}\tmp" -f $Bin
-$Errors = 0
-Push-Location "$WD"
-$Host.UI.RawUI.BackgroundColor = "black"
-$Host.UI.RawUI.ForegroundColor = "white"
+$Host.UI.RawUI.BackgroundColor = "Black"
+$Host.UI.RawUI.ForegroundColor = "White"
 $ProgressPreference = 'silentlyContinue'
-# Clear-Host
-foreach ($Var in @('SystemDrive', 'USERPROFILE', 'DISMRoot', 'BCDBootRoot', 'ImagingRoot', 'OSCDImgRoot', 'WdsmcastRoot')) {
-    Write-Host ('{0}: {1}' -f $Var, (Get-ChildItem Env:$Var).Value)
-}
-Write-Host ""
-Write-Host ("wd:   {0}" -f $WD)
-Write-Host ("bin:  {0}" -f $Bin)
-Write-Host ("root: {0}" -f $Root)
-Write-Host ("tmp:  {0}" -f $Temp)
-Read-Host "Bananas?"
-exit
+$SplitWindow = @()
+$DISM = "{0}\DISM.exe" -f $Env:DISMRoot
+$WinPEPackages = @(
+    "WinPE-EnhancedStorage.cab",
+    "WinPE-FMAPI.cab",
+    "WinPE-WMI.cab",
+    "WinPE-EnhancedStorage_en-us.cab",
+    "WinPE-WMI_en-us.cab"
+)
+    # Install WinPE-WMI before you install WinPE-NetFX.
+    # "WinPE-NetFx.cab",
+    # "WinPE-NetFx_en-us.cab",
+
+    # Install WinPE-WMI and WinPE-NetFX before you install WinPE-Scripting.
+    # "WinPE-Scripting.cab",
+    # "WinPE-Scripting_en-us.cab",
+
+    # Install WinPE-WMI, WinPE-NetFX, and WinPE-Scripting before you install WinPE-PowerShell.
+    # "WinPE-PowerShell.cab",
+    # "WinPE-PowerShell_en-us.cab",
+
+    # Install WinPE-WMI, WinPE-NetFX, WinPE-Scripting, and WinPE-PowerShell before you install WinPE-DismCmdlets.
+    # "WinPE-DismCmdlets.cab",
+    # "WinPE-DismCmdlets_en-us.cab",
+
+    # Install WinPE-WMI, WinPE-NetFX, WinPE-Scripting, and WinPE-PowerShell before you install WinPE-SecureBootCmdlets.
+    # "WinPE-SecureBootCmdlets.cab",
+
+    # Install WinPE-WMI, WinPE-NetFX, WinPE-Scripting, and WinPE-PowerShell before you install WinPE-StorageWMI.
+    # "WinPE-StorageWMI.cab",
+    # "WinPE-StorageWMI_en-us.cab",
 
 ## Functions ##
-function Download-File {
-    param ([String]$Path, [String]$Name, [String]$Url)
+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... "
+    exit
+}
+function MakeClean {
+    $Folders = @(
+        "$Root\Mount",
+        "$Root\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
 
     Write-Host ("Downloading: {0}" -f $Name)
@@ -38,11 +92,9 @@ function Download-File {
     }
     catch {
         Write-Host ("  ERROR: Failed to download file." ) -ForegroundColor "Red"
-        $Errors += 1
     }
 }
-function Find-DynamicUrl {
-    param ([String]$SourcePage, [String]$RegEx)
+function FindDynamicUrl ($SourcePage, $RegEx) {
     $Url = ""
 
     # Get source page
@@ -58,14 +110,133 @@ function Find-DynamicUrl {
 
     return $Url
 }
-function WK-Pause {
-    param([string]$Message = "Press Enter to continue... ")
-    Write-Host $Message
+function WKPause ($Message = "Press Enter to continue... ") {
+    Write-Host $Message -NoNewLine
     Read-Host
 }
+function WKRun ($Cmd, $ArgumentList, [switch]$SplitWindow=$false) {
+    if (Test-Path Env:\ConEmuBuild) {
+        $ArgumentList += "-new_console:ns33V"
+    }
+    Start-Process $Cmd -ArgumentList $ArgumentList -NoNewWindow -Wait
+}
 
-## Build ##
-# TODO #
+## 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 "Wizard Kit: Windows PE Build Tool`n"
+    
+    ## Prep ##
+    Push-Location "$WD"
+    $Date = Get-Date -UFormat "%Y-%m-%d"
+    MakeClean
+    
+    ## Build ##
+    foreach ($Arch in @("amd64", "x86")) {
+        $Drivers = "$Root\Drivers\%arch"
+        $Mount = "$Root\Mount"
+        $PEFiles = "$Root\PEFiles\$arch"
+        
+        # Copy WinPE files
+        Write-Host "Copying files..."
+        $Cmd = ("{0}\copype.cmd" -f $Env:WinPERoot)
+        WKRun -Cmd $Cmd -ArgumentList @($Arch, $PEFiles) -SplitWindow
+        
+        # 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..."
+        $ArgumentList = @(
+            "/Mount-Image",
+            "/ImageFile:'$PEFiles\media\sources\boot.wim'",
+            "/Index:1",
+            "/MountDir:'$Mount'"
+        )
+        WKRun -Cmd $DISM -ArgumentList $ArgumentList -SplitWindow
+        
+        # Add packages
+        Write-Host "Adding packages..."
+        foreach ($Package in $WinPEPackages) {
+            $PackagePath = "{0}\{1}\WinPE_OCs" -f $Env:WinPERoot, $Arch
+            $ArgumentList = @(
+                "/Add-Package",
+                "/Image:'$Mount'",
+                "/PackagePath:'$PackagePath'"
+            )
+            WKRun -Cmd $DISM -ArgumentList $ArgumentList -SplitWindow
+        }
+        
+        # Set RamDisk size
+        $ArgumentList = @("/Image:'$Mount'", "/Set-ScratchSpace:512")
+        WKRun -Cmd $DISM -ArgumentList $ArgumentList
+        
+        # Add WK tools
+        Write-Host "Copying tools..."
+        Copy-Item -Path "$Root\WK\$Arch" -Destination "$Mount\WK" -Recurse -Force
+        Copy-Item -Path "$Root\WK\_include\*" -Destination "$Mount\WK" -Recurse -Force
+        if ($Arch -eq "amd64") {
+            $DestIni = "$Mount\WK\HWiNFO\HWiNFO64.INI"
+        } else {
+            $DestIni = "$Mount\WK\HWiNFO\HWiNFO32.INI"
+        }
+        Move-Item -Path "$Root\WK\HWiNFO\HWiNFO.INI" -Destination $DestIni -Force
+        Copy-Item -Path "$Root\WinPE.jpg" -Destination "$Mount\WK\ConEmu\ConEmu.jpg" -Recurse -Force
+        Copy-Item -Path "$Root\Scripts" -Destination "$Mount\WK\Scripts" -Recurse -Force
+        
+        # Add System32 items
+        Copy-Item -Path "$Root\System32" -Destination "$Mount\Windows\System32" -Recurse -Force
+        $ArgumentList = @("/f", "$Mount\Windows\System32\winpe.jpg", "/a")
+        WKRun -Cmd "C:\Windows\System32\takeown.exe" -ArgumentList $ArgumentList
+        $ArgumentList = @("$Mount\Windows\System32\winpe.jpg", "/grant", "Administrators:F")
+        WKRun -Cmd "C:\Windows\System32\icacls.exe" -ArgumentList $ArgumentList
+        Copy-Item -Path "$Root\WinPE.jpg" -Destination "$Mount\Windows\System32\winpe.jpg" -Recurse -Force
+        
+        # Update registry
+        Write-Host "Updating Registry..."
+        $Reg = "C:\Windows\System32\reg.exe"
+        WKRun -Cmd $Reg -ArgumentList @("load", "HKLM\WinPE-SW", "$Mount\Windows\System32\config\SOFTWARE")
+        WKRun -Cmd $Reg -ArgumentList @("load", "HKLM\WinPE-SYS", "$Mount\Windows\System32\config\SYSTEM")
+        
+            # Add 7-Zip and Python to path
+            $RegPath = "HKLM:\WinPE-SYS\ControlSet001\Control\Session Manager\Environment"
+            $RegKey = Get-ItemProperty -Path $RegPath
+            $NewValue = "{0};%SystemDrive%\WK\7-Zip;%SystemDrive%\WK\python;%SystemDrive%\WK\wimlib" -f $RegKey.Path
+            Set-ItemProperty -Path $RegPath -Name "Path" -Value $NewValue -Force
+            
+            # Replace Notepad
+            $RegPath = "HKLM:\WinPE-SW\Microsoft\Windows NT\CurrentVersion\Image File Execution Options\notepad.exe"
+            $NewValue = 'wscript "X:\WK\NotepadPlusPlus\npp.vbs"'
+            New-Item -Path $RegPath -Force
+            New-ItemProperty -Path $RegPath -Name "Debugger" -Value $NewValue -Force
 
-## Done ##
-Pop-Location
+        # Unload registry hives
+        Start-Sleep -Seconds 1
+        WKRun -Cmd $Reg -ArgumentList @("unload", "HKLM\WinPE-SW")
+        WKRun -Cmd $Reg -ArgumentList @("unload", "HKLM\WinPE-SYS")
+        
+        # Unmount image
+        Write-Host "Dismounting image..."
+        $ArgumentList = @("/Unmount-Image", "/MountDir:'$Mount'", "/Commit")
+        WKRun -Cmd $DISM -ArgumentList $ArgumentList -SplitWindow
+        
+        # Create ISO
+        $ArgumentList = @("/iso", $PEFiles, "wk-winpe-$Date-$Arch.iso")
+        $Cmd = "{0}\MakeWinPEMedia.cmd" -f $Env:WinPERoot
+        WKRun -Cmd $Cmd -ArgumentList $ArgumentList -SplitWindow
+    }
+
+    ## Done ##
+    Pop-Location
+}
diff --git a/.gitignore b/.gitignore
index 28435552..c9141162 100644
--- a/.gitignore
+++ b/.gitignore
@@ -4,6 +4,6 @@
 .bin/tmp
 Drivers
 Logs
+PEFiles/
 Scripts/__pycache__
 mount
-pe_files
\ No newline at end of file