Add Driver to Offline Image Functions

function Test-IsElevated {

    [CmdletBinding()]
    [OutputType([bool])]
    param ()

    $currentIdentity = [Security.Principal.WindowsPrincipal] [Security.Principal.WindowsIdentity]::GetCurrent()
    $currentIdentity.IsInRole([Security.Principal.WindowsBuiltInRole]::Administrator)
}


function Get-OfflineImageDriverFolder {

    [CmdletBinding()]
    [OutputType([IO.DirectoryInfo])]
    param (
        [Parameter(Mandatory)]
        [ValidateScript({ $_ -notin [IO.Path]::GetInvalidFileNameChars() })]
        [string] $FolderName,

        [string] $Computer
    )

    $cimParams = @{
        Class       = 'Win32_Volume'
        Computer    = $Computer
        Filter      = 'DriveType = "2"'
        ErrorAction = 'Stop'
    }

    $usbDrives = Get-CimInstance @cimParams

    foreach ($usb in $usbDrives) {
        [IO.DirectoryInfo] $folder = [IO.Path]::Combine($usb.Name, $FolderName)

        if ($folder.Exists) {
            Write-Debug ('Found driver folder: {0}' -f $folder.FullName)
            $folder
        }
    }
}


function Add-OfflineImageDriver {

    [CmdletBinding(SupportsShouldProcess, ConfirmImpact = 'Medium')]
    [OutputType([void])]
    param (
        [Parameter(Mandatory, ValueFromPipelineByPropertyName)]
        [Alias('FullName')]
        [ValidateScript({ [IO.Directory]::Exists($_) })]
        [string[]] $Path
    )

    begin {
        if (!(Test-IsElevated)) {
            $errorMsg = 'The requested operation requires elevation.'
            $errorRecord = [Management.Automation.ErrorRecord]::new($errorMsg, $null, [Management.Automation.ErrorCategory]::SecurityError, $null)
            $PSCmdlet.ThrowTerminatingError($errorRecord)
        }

        try {
            $tsEnv = New-Object -ComObject Microsoft.SMS.TSEnvironment
        } catch  {
            $PSCmdlet.ThrowTerminatingError($_)
        }

        $dismPath = [IO.Path]::Combine([Environment]::GetFolderPath([Environment+SpecialFolder]::System), 'dism.exe')
    }

    process {
        foreach ($inputPath in $Path) {
            $dismArgs = @(
                '/Image:"{0}\"' -f $tsEnv.Value('OSDTargetSystemDrive')
                '/LogPath:"{0}"' -f $tsEnv.Value('_SMSTSLogPath')
                '/Add-Driver'
                '/Driver:"{0}"' -f $inputPath
                '/Recurse'
            ) -join ' '

            if ($PSCmdlet.ShouldProcess($dismPath, $dismArgs)) {
                $output = & $dismPath $dismArgs

                if ($LASTEXITCODE -ne 0) {
                    $errorMsg = $output -join "`n"
                    $errorRecord = [Management.Automation.ErrorRecord]::new($errorMsg, $null, [Management.Automation.ErrorCategory]::InvalidResult, $null)
                    $PSCmdlet.WriteError($errorRecord)
                    continue
                }

                Write-Debug ('Added drivers from {0} to {1}.' -f $inputPath, $tsEnv.Value('OSDTargetSystemDrive'))
                $tsEnv.Value('OSDCustomDriversApplied') = 'True'
            }
        }
    }
}