FFU/FFUDevelopment/USBImagingToolCreator.ps1

[CmdletBinding()]
param(
    [Parameter(Mandatory = $True, Position = 0)]
    [io.fileinfo] $DeployISOPath,
    [Switch]$DisableAutoPlay
)

function WriteLog($LogText) { 
    $LogFileName = '\Script.log'
    $LogFile = $DevelopmentPath + $LogFilename
    Add-Content -path $LogFile -value "$((Get-Date).ToString()) $LogText" -Force -ErrorAction SilentlyContinue
    Write-Verbose $LogText
}

function Get-USBDrive {
    $USBDrives = (Get-WmiObject -Class Win32_DiskDrive -Filter "MediaType='Removable Media'")
    If ($USBDrives -and ($null -eq $USBDrives.count)) {
        $USBDrivesCount = 1
    } else {
        $USBDrivesCount = $USBDrives.Count
    }
    WriteLog "Found $USBDrivesCount USB drives"

    if ($null -eq $USBDrives) {
        WriteLog "No removable USB drive found. Exiting"
        Write-Error "No removable USB drive found. Exiting"
        exit 1
    }
    return $USBDrives, $USBDrivesCount
}

function Build-DeploymentUSB {
    param(
        [Array]$Drives       
    )
    writelog "Checking if ffu files are present in the ffu folder"
    $Images = Get-ChildItem -Path $FFUPath -Filter "*.ffu" -File -Recurse
    writelog "Checking if drivers are present in the drivers folder"
    $Drivers = Get-ChildItem -Path $DriversPath -Recurse
    $DrivesCount = $Drives.Count
    Writelog "Creating partitions..."
    foreach ($USBDrive in $Drives) {
        $DriveNumber = $USBDrive.DeviceID.Replace("\\.\PHYSICALDRIVE", "")
        $Model = $USBDrive.model
        $ScriptBlock = {
            param($DriveNumber)
            Clear-Disk -Number $DriveNumber -RemoveData -RemoveOEM -Confirm:$false
            $Disk = Get-Disk -Number $DriveNumber
            $PartitionStyle = $Disk.PartitionStyle
            if ($PartitionStyle -ne 'MBR') {
                $Disk | Set-Disk -PartitionStyle MBR
            }
            $BootPartition = New-Partition -DiskNumber $DriveNumber -Size 2GB -IsActive -AssignDriveLetter
            $DeployPartition = New-Partition -DiskNumber $DriveNumber -UseMaximumSize -AssignDriveLetter
            Format-Volume -Partition $BootPartition -FileSystem FAT32 -NewFileSystemLabel "Boot" -Confirm:$false
            Format-Volume -Partition $DeployPartition -FileSystem NTFS -NewFileSystemLabel "Deploy" -Confirm:$false
        }
        WriteLog "Start job to create BOOT and Deploy partitions on drive number $DriveNumber"
        Start-Job -ScriptBlock $ScriptBlock -ArgumentList $DriveNumber | Out-Null
    }
    writelog "Wait for partitioning jobs to complete"
    Get-Job | Wait-Job | Out-Null
    if ($DrivesCount -gt 1) {
        writelog "Get file system information for all drives"
        $Partitions = Get-Partition | Get-Volume
    } else {
        writelog "Get file system information for drive number $DiskNumber"
        $Partitions = Get-Partition -DiskNumber $DriveNumber | Get-Volume
    }
    writelog "Get drive letter for all volumes labeled:BOOT"
    $BootDrives = ($Partitions | Where-Object { $_.FileSystemLabel -eq "BOOT" }).DriveLetter
    writelog "Get drive letter for all volumes labeled:Deploy"
    $DeployDrives = ($Partitions | Where-Object { $_.FileSystemLabel -eq "Deploy" }).DriveLetter
    writelog "Mount Deployment .iso image"
    $ISOMountPoint = (Mount-DiskImage -ImagePath "$DeployISOPath" -PassThru | Get-Volume).DriveLetter + ":\"
    writelog "Copying boot files to all drives labeled BOOT concurrently"
    foreach ($Drive in $BootDrives) {
        $Destination = $Drive + ":\"
        $jobScriptBlock = {
            param (
                [string]$SFolder,
                [string]$DFolder
            )
            Robocopy $SFolder $DFolder /E /COPYALL /R:5 /W:5 /J
        }
        WriteLog "Start job to copy all boot files to $Destination"
        Start-Job -ScriptBlock $jobScriptBlock -ArgumentList $ISOMountPoint, $Destination | Out-Null
    }
    if ($Images) {
        writelog "Copying FFU image files to all drives labeled deploy concurrently"
        foreach ($Drive in $DeployDrives) {
            $Destination = $Drive + ":\"
            $jobScriptBlock = {
                param (
                    [string]$SFolder,
                    [string]$DFolder
                )
                Robocopy $SFolder $DFolder /E /COPYALL /R:5 /W:5 /J
            }

            WriteLog "Start job to copy all FFU files to $Destination"
            Start-Job -ScriptBlock $jobScriptBlock -ArgumentList $FFUPath, $Destination | Out-Null
        }
    }
    if ($Drivers) {
        writelog "Copying driver files to all drives labeled deploy concurrently"
        foreach ($Drive in $DeployDrives) {
            $Destination = $Drive + ":\Drivers"
            $jobScriptBlock = {
                param (
                    [string]$SFolder,
                    [string]$DFolder
                )
                New-Item -Path $DFolder -ItemType Directory -Force -Confirm: $false | Out-Null
                Robocopy $SFolder $DFolder /E /COPYALL /R:5 /W:5 /J
            }
            WriteLog "Start job to copy all drivers to $Destination"
            Start-Job -ScriptBlock $jobScriptBlock -ArgumentList $DriversPath, $Destination | Out-Null
        }
    }
    if (!($Drivers)) {
        foreach ($Drive in $DeployDrives) {
            WriteLog "Create drivers directory"
            $drivepath = $Drive + ":\"
            New-Item -Path "$drivepath" -Name Drivers -ItemType Directory -Force -Confirm: $false | Out-Null
        }
    }
    if ($DrivesCount -gt 1) {
        Writelog "Building $DrivesCount drives concurrently...Please be patient..."
    } else {
        Writelog "Building the imaging tool on $model...Please be patient..."
    }
    Get-Job | Wait-Job | Out-Null

    Dismount-DiskImage -ImagePath $DeployISOPath | Out-Null
    Writelog "Drive creation jobs completed..."
}

function New-DeploymentUSB {
    param(
        [Array]$Drives,
        [int]$Count,
        [String]$FFUPath = "$DevelopmentPath\FFU",
        [String]$DriversPath = "$DevelopmentPath\Drivers"
    
    )

    $Drivelist = @()
    writelog "Creating a USB drive selection list"
    for ($i = 0; $i -le $Count - 1; $i++) {
        $DriveModel = $Drives[$i].Model
        $DriveSize = [math]::round($Drives[$i].size / 1GB, 2)
        $DiskNumber = $Drives[$i].DeviceID.Replace("\\.\PHYSICALDRIVE", "")
        $Properties = [ordered]@{Number = $i + 1  ; DriveNumber = $DiskNumber ; DriveModel = $driveModel ; 'Size (GB)' = $DriveSize } 

        $Drivelist += New-Object PSObject -Property $Properties
    }
    if ($Count -gt 1) {
        $Last = $Count + 1
        $Drivelist += New-Object -TypeName PSObject -Property @{ Number = "$last"; DriveModel = "Select this option to use all ($count) inserted USB Drives" }       
    }
    $Drivelist | Format-Table -AutoSize -Property Number, DriveModel , 'Size (GB)'
    do {
        try {
            $var = $true
            $DriveSelected = Read-Host 'Enter the drive number to apply the .iso to'
            $DriveSelected = ($DriveSelected -as [int]) - 1
            if ($Last) {
                writelog "All drives selected"
            }
            else {
                writelog "Drive $DriveSelected selected"
            }
        }
        catch {
            Write-Host 'Input was not in correct format. Please enter a valid FFU number'
            $var = $false
        }
    } until (($DriveSelected -le $Count - 1 -or $last) -and $var)

    $DisableAutoPlayCurrentSetting = (Get-ItemProperty -Path "HKCU:\SOFTWARE\Microsoft\Windows\CurrentVersion\Explorer\AutoplayHandlers" -Name DisableAutoplay).DisableAutoplay
    if ($DisableAutoPlay -and $DisableAutoPlayCurrentSetting -ne 1) {
        writelog "Disable autoPlay current setting is $DisableAutoPlayCurrentSetting"
        WriteLog "Setting the registry key to disable autoplay for all drives"
        Set-ItemProperty -Path "HKCU:\SOFTWARE\Microsoft\Windows\CurrentVersion\Explorer\AutoplayHandlers" -Name "DisableAutoplay" -Value 1 -Type DWORD
    }
    WriteLog "Closing all MMC windows to prevent drive lock errors"
    Stop-Process -Name mmc -ErrorAction SilentlyContinue
    WriteLog "Closing all Diskpart windows to prevent drive lock errors"
    Stop-Process -Name diskpart -ErrorAction SilentlyContinue
    $Selection = $Drivelist[$DriveSelected].Number

    if ($Selection -eq $last) {
        Read-Host -Prompt "ALL DRIVES SELECTED! WILL ERASE ALL CURRENTLY CONNECTED USB DRIVES!! Press ENTER to continue"
        Build-DeploymentUSB -Drives $Drives
    } else {
        Read-Host -Prompt "Drive number $Selection was selected. Press ENTER to continue"
        Build-DeploymentUSB -Drives $Drives[$DriveSelected]
    }
    WriteLog "Setting the registry key to re-enable autoplay for all drives"
    if ($DisableAutoPlay) {
        Writelog "Setting disable autoplay setting back to $DisableAutoPlayCurrentSetting"
        Set-ItemProperty -Path "HKCU:\SOFTWARE\Microsoft\Windows\CurrentVersion\Explorer\AutoplayHandlers" -Name "DisableAutoplay" -Value $DisableAutoPlayCurrentSetting -Type DWORD
    }
    Writelog "Completed!"
}

$Host.UI.RawUI.WindowTitle = 'USB Imaging Tool Creator'

# Check if path is relative
if (-not [System.IO.Path]::IsPathRooted($DeployISOPath)) {
    # Path is relative, Build full path
    [io.fileinfo] $DeployISOPath = (Resolve-Path $DeployISOPath).Path
    Write-Verbose "Path to ISO was relative. Attempting to resolve at: $($DeployISOPath.FullName)"
}

if ($DeployISOPath.Exists) {
    $DevelopmentPath = $DeployISOPath.Directory.FullName

    #Get USB Drive and create log file
    if (Test-Path "$DevelopmentPath\Script.log") {
        Remove-Item -Path "$DevelopmentPath\Script.log" -Force -Confirm:$false
        New-item -Path $DevelopmentPath -Name 'Script.log' -ItemType "file" -Force | Out-Null
    }

    WriteLog 'Begin Logging'
    WriteLog 'Getting USB drive information and usb drive count'

    $USBDrives, $USBDrivesCount = Get-USBDrive

    New-DeploymentUSB -Drives $USBDrives -Count $USBDrivesCount
    
    Read-Host -Prompt "USB drive creation complete. Press ENTER to exit"
    Exit
} else {
    Write-Error "Unable to locate ISO file at: $($DeployISOPath.FullName)"
}