diff --git a/Update-BootWIM.ps1 b/Update-BootWIM.ps1
index 7300819..9fc87de 100644
--- a/Update-BootWIM.ps1
+++ b/Update-BootWIM.ps1
@@ -1,33 +1,42 @@
-<#
-    Stripped down script from https://learn.microsoft.com/en-us/windows/deployment/update/media-dynamic-update
+<#
+    Stripped down version of https://learn.microsoft.com/en-us/windows/deployment/update/media-dynamic-update
     to update WinPE/boot image/boot.wim.
     You need to download the approproiate CU for the boot image version you are using.
-    Copy the CU to $LCU_PATH. Copy your boot.wim file to $BOOT_WIM.
 
-    Version 0.1 - 2023-05-10
+    Version 0.2 - 2023-05-14
     Sassan Fanai
-
 #>
 
 #Requires -RunAsAdministrator
 
+param (
+    $UpdateFile = "C:\mediaRefresh\packages\LCU.msu",
+    $WimFile = "C:\Boot Image Backup\boot.wim",
+    $StagingPath = "C:\mediaRefresh\Temp"
+)
+
 function Get-TS { return "{0:HH:mm:ss}" -f [DateTime]::Now }
 
-Write-Output "$(Get-TS): Starting media refresh"
+if (!(Test-Path $UpdateFile)) {
+    Write-Output "$(Get-TS): No update file found @ [$UpdateFile]. Exiting..."
+    break
+}
+
+if (!(Test-Path $WimFile)) {
+    Write-Output "$(Get-TS): No WIM file found @ [$WimFile]. Exiting..."
+    break
+}
+
+Write-Output "$(Get-TS): Starting media refresh of [$WimFile]"
 
 # Declare Dynamic Update packages
-$LCU_PATH        = "C:\mediaRefresh\packages\LCU.msu"
-$BOOT_WIM        = "C:\Boot Image Backup\boot.wim"
+$LCU_PATH        = $UpdateFile
+$BOOT_WIM        = $WimFile
 
 # Declare folders for mounted images and temp files
-$WORKING_PATH    = "C:\mediaRefresh\temp"
-$WINPE_MOUNT     = "C:\mediaRefresh\temp\WinPEMount"
+$WORKING_PATH    = "$StagingPath"
+$WINPE_MOUNT     = "$StagingPath\WinPEMount"
 
-# Check for LCU MSU
-if (!(Test-Path $LCU_PATH)) {
-    Write-Output "$(Get-TS): No LCU.msu found. You nned to download and copy the CU for your boot image version to $LCU_PATH and rename the file to LCU.msu"
-    break
-}
 
 # Create folders for mounting images and storing temporary files
 if (!(Test-Path $WORKING_PATH)) {
@@ -47,12 +56,12 @@ $WINPE_IMAGES = Get-WindowsImage -ImagePath $BOOT_WIM
 Foreach ($IMAGE in $WINPE_IMAGES) {
 
     # update WinPE
-    Write-Output "$(Get-TS): Mounting WinPE, image index $($IMAGE.ImageIndex)"
+    Write-Output "$(Get-TS): Mounting WinPE, image index [$($IMAGE.ImageIndex)]"
     Mount-WindowsImage -ImagePath $BOOT_WIM  -Index $IMAGE.ImageIndex -Path $WINPE_MOUNT -ErrorAction stop | Out-Null
 
     try
     {
-        Write-Output "$(Get-TS): Adding package $LCU_PATH"
+        Write-Output "$(Get-TS): Adding package [$LCU_PATH]"
         Add-WindowsPackage -Path $WINPE_MOUNT -PackagePath $LCU_PATH | Out-Null
     }
     Catch
@@ -76,9 +85,9 @@ Foreach ($IMAGE in $WINPE_IMAGES) {
     Dismount-WindowsImage -Path $WINPE_MOUNT -Save -ErrorAction stop | Out-Null
 
     #Export WinPE
-    Write-Output "$(Get-TS): Exporting image to $WORKING_PATH\boot2.wim"
-    Export-WindowsImage -SourceImagePath $BOOT_WIM  -SourceIndex $IMAGE.ImageIndex -DestinationImagePath $WORKING_PATH"\boot2.wim" -ErrorAction stop | Out-Null
-
+    Write-Output "$(Get-TS): Exporting image to [$WORKING_PATH\boot2.wim]"
+    Export-WindowsImage -SourceImagePath $BOOT_WIM -SourceIndex $IMAGE.ImageIndex -DestinationImagePath $WORKING_PATH"\boot2.wim" -ErrorAction stop | Out-Null
 }
 
-Move-Item -Path $WORKING_PATH"\boot2.wim" -Destination $BOOT_WIM  -Force -ErrorAction stop | Out-Null
\ No newline at end of file
+Write-Output "$(Get-TS): Moving updated and exported boot2.wim from staging area [$StagingPath] to [$BOOT_WIM]"
+Move-Item -Path $WORKING_PATH"\boot2.wim" -Destination $BOOT_WIM  -Force -ErrorAction stop | Out-Null