build-pe-script-launcher/lib/DefenderOptimization.psm1

function Get-DefenderRealTimeProtectionStatus {
    <#
    .SYNOPSIS
    Gets the current Windows Defender real-time protection status.
    #>
    [CmdletBinding()]
    param()

    try {
        $preference = Get-MpPreference -ErrorAction Stop
        return -not $preference.DisableRealtimeMonitoring
    }
    catch {
        Write-Verbose "Could not get Defender status: $_"
        return $null
    }
}

function Disable-DefenderRealTimeProtection {
    <#
    .SYNOPSIS
    Temporarily disables Windows Defender real-time protection and returns original state.

    .OUTPUTS
    Returns the original state (True if enabled, False if already disabled, $null if failed)
    #>
    [CmdletBinding()]
    param()

    try {
        $originalState = Get-DefenderRealTimeProtectionStatus

        if ($null -eq $originalState) {
            Write-Warning "Cannot determine Defender status. Skipping optimization."
            return $null
        }

        if (-not $originalState) {
            Write-Verbose "Real-time protection already disabled"
            return $false
        }

        Set-MpPreference -DisableRealtimeMonitoring $true -ErrorAction Stop
        Write-Verbose "Real-time protection disabled successfully"
        return $true
    }
    catch {
        Write-Warning "Failed to disable Defender real-time protection: $_"
        Write-Warning "Continuing without Defender optimization"
        return $null
    }
}

function Restore-DefenderRealTimeProtection {
    <#
    .SYNOPSIS
    Restores Windows Defender real-time protection to its original state.

    .PARAMETER OriginalState
    The state to restore (from Disable-DefenderRealTimeProtection)
    #>
    [CmdletBinding()]
    param(
        [Parameter(Mandatory = $true)]
        [AllowNull()]
        $OriginalState
    )

    if ($null -eq $OriginalState) {
        Write-Verbose "No original state to restore"
        return
    }

    if ($OriginalState -eq $false) {
        Write-Verbose "Real-time protection was already disabled, leaving as-is"
        return
    }

    try {
        Set-MpPreference -DisableRealtimeMonitoring $false -ErrorAction Stop
        Write-Verbose "Real-time protection restored successfully"
    }
    catch {
        Write-Warning "Failed to restore Defender real-time protection: $_"
        Write-Warning "You may need to manually re-enable it in Windows Security"
    }
}

$script:AddedExclusions = @()

function Add-DefenderExclusions {
    <#
    .SYNOPSIS
    Adds build directories and ADK paths to Windows Defender exclusions.

    .PARAMETER PSScriptRoot
    The root directory of the build script
    #>
    [CmdletBinding()]
    param(
        [Parameter(Mandatory = $true)]
        [string]$PSScriptRoot
    )

    $script:AddedExclusions = @()

    $exclusionPaths = @(
        (Join-Path $PSScriptRoot "work"),
        (Join-Path $PSScriptRoot "downloads"),
        (Join-Path $PSScriptRoot "output"),
        "C:\Program Files (x86)\Windows Kits\10\Assessment and Deployment Kit\Deployment Tools\amd64\DISM",
        "C:\Program Files (x86)\Windows Kits\10\Assessment and Deployment Kit\Windows Preinstallation Environment"
    )

    foreach ($path in $exclusionPaths) {
        try {
            $existingExclusions = (Get-MpPreference).ExclusionPath

            if ($existingExclusions -contains $path) {
                Write-Verbose "Exclusion already exists: $path"
                continue
            }

            Add-MpPreference -ExclusionPath $path -ErrorAction Stop
            $script:AddedExclusions += $path
            Write-Verbose "Added exclusion: $path"
        }
        catch {
            Write-Warning "Failed to add Defender exclusion for $path : $_"
        }
    }

    if ($script:AddedExclusions.Count -gt 0) {
        Write-Verbose "Added $($script:AddedExclusions.Count) Defender exclusions"
    }
    else {
        Write-Warning "No Defender exclusions were added"
    }
}

function Remove-DefenderExclusions {
    <#
    .SYNOPSIS
    Removes Defender exclusions that were added by Add-DefenderExclusions.

    .PARAMETER PSScriptRoot
    The root directory of the build script
    #>
    [CmdletBinding()]
    param(
        [Parameter(Mandatory = $true)]
        [string]$PSScriptRoot
    )

    if ($script:AddedExclusions.Count -eq 0) {
        Write-Verbose "No exclusions to remove"
        return
    }

    $removedCount = 0

    foreach ($path in $script:AddedExclusions) {
        try {
            Remove-MpPreference -ExclusionPath $path -ErrorAction Stop
            Write-Verbose "Removed exclusion: $path"
            $removedCount++
        }
        catch {
            Write-Warning "Failed to remove Defender exclusion for $path : $_"
        }
    }

    Write-Verbose "Removed $removedCount of $($script:AddedExclusions.Count) Defender exclusions"
    $script:AddedExclusions = @()
}

function Test-DefenderAvailable {
    <#
    .SYNOPSIS
    Tests if Windows Defender cmdlets are available on this system.
    #>
    [CmdletBinding()]
    param()

    try {
        $null = Get-MpPreference -ErrorAction Stop
        return $true
    }
    catch {
        return $false
    }
}

Export-ModuleMember -Function @(
    'Get-DefenderRealTimeProtectionStatus',
    'Disable-DefenderRealTimeProtection',
    'Restore-DefenderRealTimeProtection',
    'Add-DefenderExclusions',
    'Remove-DefenderExclusions',
    'Test-DefenderAvailable'
)