# Windows 10 restore. Elevated normal desktop.
# Default run writes Start back and does not enable tasks and does not start WinDefend.
# After that reboot: .\Restore_Defender.PS1 -AfterReboot
# -AfterReboot enables the Defender tasks, then checks the restored state.
# SecurityHealthService is never written.
# A denied or mismatched write stops the pre-reboot run. Later names stay unchanged.
# An empty task list on -AfterReboot is a failure.
# Exit 0 only when that phase prints RESULT: PASS.

function Get-DefenderRestorePlan {
    @(
        @{ Name = 'WdBoot'; Start = 0; Optional = $false }
        @{ Name = 'WdFilter'; Start = 0; Optional = $false }
        @{ Name = 'WinDefend'; Start = 2; Optional = $false }
        @{ Name = 'WdNisDrv'; Start = 3; Optional = $false }
        @{ Name = 'WdNisSvc'; Start = 3; Optional = $false }
        @{ Name = 'Sense'; Start = 3; Optional = $true }
    )
}

function Format-RestoreToken {
    param($Value)
    if ($Value -eq 'NOPROP') { return 'absent' }
    if ($null -eq $Value -or $Value -eq 'UNREADABLE' -or $Value -eq 'NOKEY') {
        return 'unreadable'
    }
    return [string]$Value
}

function Get-DefenderRestoreStart {
    param([Parameter(Mandatory)][string]$Key)
    if (-not (Test-Path -LiteralPath $Key)) { return 'NOKEY' }
    try {
        $item = Get-ItemProperty -LiteralPath $Key -ErrorAction Stop
    } catch {
        return 'UNREADABLE'
    }
    if ($item.PSObject.Properties.Name -notcontains 'Start') { return 'NOPROP' }
    return [int]$item.Start
}

function Invoke-DefenderRestorePreReboot {
    param(
        [string]$ServicesRoot = 'HKLM:\SYSTEM\CurrentControlSet\Services',
        [scriptblock]$WriteStart
    )
    if ($env:DEFENDER_PROCEDURE_UNDER_TEST -eq '1' -and $ServicesRoot -like 'HKLM:*') {
        throw 'Refusing to write HKLM during tests.'
    }
    $root = $ServicesRoot.TrimEnd('\')
    $plan = @(Get-DefenderRestorePlan)
    $lines = New-Object System.Collections.Generic.List[string]
    $failed = $false
    $stoppedEarly = $false
    for ($index = 0; $index -lt $plan.Count; $index++) {
        $entry = $plan[$index]
        $name = [string]$entry.Name
        $target = [int]$entry.Start
        $key = "$root\$name"
        $current = Get-DefenderRestoreStart -Key $key
        if ($current -eq 'NOKEY') {
            [void]$lines.Add("missing: $name")
            if ($entry.Optional) { continue }
            $failed = $true
            $stoppedEarly = ($index -lt ($plan.Count - 1))
            break
        }
        if ($current -eq 'UNREADABLE') {
            [void]$lines.Add("${name}: unreadable -> denied")
            $failed = $true
            $stoppedEarly = ($index -lt ($plan.Count - 1))
            break
        }
        try {
            if ($null -ne $WriteStart) {
                & $WriteStart $key $target
            } else {
                Set-ItemProperty -LiteralPath $key -Name Start -Value $target -Type DWord -ErrorAction Stop
            }
        } catch {
            $still = Get-DefenderRestoreStart -Key $key
            $shown = Format-RestoreToken $still
            [void]$lines.Add("${name}: denied (still $shown)")
            $failed = $true
            $stoppedEarly = ($index -lt ($plan.Count - 1))
            break
        }
        $after = Get-DefenderRestoreStart -Key $key
        if ($after -ne $target) {
            $shown = Format-RestoreToken $after
            [void]$lines.Add("${name}: readback $shown")
            $failed = $true
            $stoppedEarly = ($index -lt ($plan.Count - 1))
            break
        }
        [void]$lines.Add("$name = $after")
    }
    if ($stoppedEarly) { [void]$lines.Add('Later names were not written.') }
    [void]$lines.Add('Tasks were not changed.')
    if ($failed) {
        [void]$lines.Add('Do not reboot for a restore while RESULT is FAIL.')
        [void]$lines.Add('RESULT: FAIL')
        $code = 1
    } else {
        [void]$lines.Add('Reboot, then run Restore_Defender.PS1 -AfterReboot')
        [void]$lines.Add('RESULT: PASS')
        $code = 0
    }
    [pscustomobject]@{
        Lines        = $lines.ToArray()
        Passed       = (-not $failed)
        Code         = $code
        TasksEnabled = $false
    }
}

function Get-RestoreScStartType {
    param([string]$Text)
    if ([string]::IsNullOrWhiteSpace($Text)) { return $null }
    $match = [regex]::Match($Text, '(?m)START_TYPE\s*:\s*(\d+)')
    if (-not $match.Success) { return $null }
    return [int]$match.Groups[1].Value
}

function Get-RestoreScState {
    param([string]$Text)
    if ([string]::IsNullOrWhiteSpace($Text)) { return $null }
    $match = [regex]::Match($Text, '(?m)STATE\s*:\s*(\d+)')
    if (-not $match.Success) { return $null }
    return [int]$match.Groups[1].Value
}

function Test-RestoreWdFilterListed {
    param([string]$FltmcText)
    if ([string]::IsNullOrWhiteSpace($FltmcText)) { return $null }
    foreach ($line in ($FltmcText -split '\r?\n')) {
        $trim = $line.Trim()
        if ($trim -match '^(?i)WdFilter(\s|$)') { return $true }
    }
    return $false
}

function Test-DefenderRestored {
    param($Snapshot)
    $failures = New-Object System.Collections.Generic.List[string]
    if ($Snapshot.InSafeMode -eq $true) {
        [void]$failures.Add('Safe Mode is still set. Verify only after a normal boot.')
    } elseif ($Snapshot.InSafeMode -ne $false) {
        [void]$failures.Add('Safe Mode state was not read.')
    }
    $winDefend = $Snapshot.Services['WinDefend']
    if ($null -eq $winDefend) {
        [void]$failures.Add('WinDefend was not found.')
    } else {
        $status = [string]$winDefend.Status
        $startType = [string]$winDefend.StartType
        $auto = $startType -eq 'Automatic' -or $startType -eq 'AutomaticDelayedStart'
        if ($status -ne 'Running' -or -not $auto) {
            [void]$failures.Add("WinDefend is $status, StartType $startType. Required: Running, Automatic.")
        }
    }
    $nis = $Snapshot.Services['WdNisSvc']
    if ($null -eq $nis) {
        [void]$failures.Add('WdNisSvc was not found.')
    } else {
        $status = [string]$nis.Status
        $startType = [string]$nis.StartType
        $alive = $status -eq 'Stopped' -or $status -eq 'Running'
        if ($startType -ne 'Manual' -or -not $alive) {
            [void]$failures.Add(
                "WdNisSvc is $status, StartType $startType. " +
                'Required: Manual, and Stopped or Running.')
        }
    }
    $sense = $Snapshot.Services['Sense']
    if ($null -ne $sense) {
        $status = [string]$sense.Status
        $startType = [string]$sense.StartType
        $alive = $status -eq 'Stopped' -or $status -eq 'Running'
        if ($startType -ne 'Manual' -or -not $alive) {
            [void]$failures.Add(
                "Sense is $status, StartType $startType. " +
                'Required: Manual, and Stopped or Running, or missing.')
        }
    }
    if ($Snapshot.MsMpEngRunning -eq $false) {
        [void]$failures.Add('MsMpEng.exe is not running.')
    } elseif ($Snapshot.MsMpEngRunning -ne $true) {
        [void]$failures.Add('MsMpEng.exe state was not read.')
    }
    $expected = @{ WdFilter = 0; WdBoot = 0; WdNisDrv = 3 }
    foreach ($name in @('WdFilter', 'WdBoot', 'WdNisDrv')) {
        $start = Get-RestoreScStartType -Text $Snapshot.QcText[$name]
        if ($start -ne $expected[$name]) {
            $shown = if ($null -eq $start) { 'unreadable' } else { [string]$start }
            [void]$failures.Add("$name START_TYPE is $shown. Required: $($expected[$name]).")
        }
    }
    $filterState = Get-RestoreScState -Text $Snapshot.QueryText['WdFilter']
    if ($filterState -ne 4) {
        $shown = if ($null -eq $filterState) { 'unreadable' } else { [string]$filterState }
        [void]$failures.Add("WdFilter STATE is $shown. Required: 4 RUNNING.")
    }
    $listed = Test-RestoreWdFilterListed -FltmcText $Snapshot.FltmcText
    if ($listed -ne $true) {
        [void]$failures.Add('fltmc does not list WdFilter.')
    }
    if ($null -eq $Snapshot.Tasks) {
        [void]$failures.Add('No Defender tasks were returned. An empty task list fails the restore.')
    } else {
        $real = @($Snapshot.Tasks | Where-Object { $null -ne $_ })
        if ($real.Count -eq 0) {
            [void]$failures.Add('No Defender tasks were returned. An empty task list fails the restore.')
        }
        foreach ($task in $real) {
            $state = [string]$task.State
            $taskName = [string]$task.TaskName
            if ([string]::IsNullOrWhiteSpace($taskName)) { $taskName = '(unnamed task)' }
            $enabled = $state -eq 'Ready' -or $state -eq 'Running' -or $state -eq 'Queued'
            if (-not $enabled) {
                $shown = if ([string]::IsNullOrWhiteSpace($state)) { 'unknown' } else { $state }
                [void]$failures.Add("$taskName = $shown. Required: Ready, Running, or Queued.")
            }
        }
    }
    [pscustomobject]@{
        Passed   = ($failures.Count -eq 0)
        Failures = $failures.ToArray()
    }
}

function Format-DefenderRestoredReport {
    param(
        $Snapshot,
        [string[]]$ExtraFailures
    )
    $verdict = Test-DefenderRestored -Snapshot $Snapshot
    $failures = New-Object System.Collections.Generic.List[string]
    foreach ($failure in @($verdict.Failures)) { [void]$failures.Add([string]$failure) }
    foreach ($extra in @($ExtraFailures)) {
        if (-not [string]::IsNullOrWhiteSpace([string]$extra)) { [void]$failures.Add([string]$extra) }
    }
    $lines = New-Object System.Collections.Generic.List[string]
    if ($Snapshot.InSafeMode -eq $true) { [void]$lines.Add('Safe Mode: yes') }
    elseif ($Snapshot.InSafeMode -eq $false) { [void]$lines.Add('Safe Mode: no') }
    else { [void]$lines.Add('Safe Mode: unread') }
    foreach ($name in @('WinDefend', 'WdNisSvc')) {
        $svc = $Snapshot.Services[$name]
        if ($null -eq $svc) { [void]$lines.Add("${name}: not found") }
        else { [void]$lines.Add("${name}: $($svc.Status), $($svc.StartType)") }
    }
    $sense = $Snapshot.Services['Sense']
    if ($null -eq $sense) { [void]$lines.Add('Sense: missing') }
    else { [void]$lines.Add("Sense: $($sense.Status), $($sense.StartType)") }
    if ($Snapshot.MsMpEngRunning -eq $true) { [void]$lines.Add('MsMpEng.exe: running') }
    elseif ($Snapshot.MsMpEngRunning -eq $false) { [void]$lines.Add('MsMpEng.exe: not running') }
    else { [void]$lines.Add('MsMpEng.exe: unread') }
    foreach ($name in @('WdFilter', 'WdBoot', 'WdNisDrv')) {
        $start = Get-RestoreScStartType -Text $Snapshot.QcText[$name]
        $shown = if ($null -eq $start) { 'unreadable' } else { [string]$start }
        [void]$lines.Add("${name} START_TYPE: $shown")
    }
    $filterState = Get-RestoreScState -Text $Snapshot.QueryText['WdFilter']
    $shownState = if ($null -eq $filterState) { 'unreadable' } else { [string]$filterState }
    [void]$lines.Add("WdFilter STATE: $shownState")
    [void]$lines.Add('WdBoot runtime is informational. START_TYPE 0 is the pass test.')
    $listed = Test-RestoreWdFilterListed -FltmcText $Snapshot.FltmcText
    if ($listed -eq $true) { [void]$lines.Add('fltmc WdFilter: listed') }
    elseif ($listed -eq $false) { [void]$lines.Add('fltmc WdFilter: absent') }
    else { [void]$lines.Add('fltmc WdFilter: unread') }
    [void]$lines.Add('Tasks:')
    if ($null -eq $Snapshot.Tasks) {
        [void]$lines.Add('No tasks returned')
    } else {
        $real = @($Snapshot.Tasks | Where-Object { $null -ne $_ })
        if ($real.Count -eq 0) { [void]$lines.Add('No tasks returned') }
        foreach ($task in $real) { [void]$lines.Add("$($task.TaskName) = $($task.State)") }
    }
    [void]$lines.Add('Get-MpComputerStatus (non-authoritative):')
    if ([string]::IsNullOrWhiteSpace([string]$Snapshot.MpText)) {
        [void]$lines.Add('(no status)')
    } else {
        foreach ($statusLine in ([string]$Snapshot.MpText -split '\r?\n')) {
            if (-not [string]::IsNullOrWhiteSpace($statusLine)) { [void]$lines.Add($statusLine.TrimEnd()) }
        }
    }
    foreach ($failure in $failures) { [void]$lines.Add("FAIL: $failure") }
    $passed = ($failures.Count -eq 0)
    if ($passed) {
        [void]$lines.Add('RESULT: PASS')
        [void]$lines.Add('Tamper Protection and real-time protection are still manual steps.')
        $code = 0
    } else {
        [void]$lines.Add('RESULT: FAIL')
        $code = 1
    }
    [pscustomobject]@{
        Lines  = $lines.ToArray()
        Passed = $passed
        Code   = $code
    }
}

function Enable-DefenderScheduledTasks {
    param(
        [scriptblock]$GetTasks,
        [scriptblock]$EnableOne
    )
    if (-not $GetTasks) {
        if ($env:DEFENDER_PROCEDURE_UNDER_TEST -eq '1') {
            throw 'Refusing to query scheduled tasks during tests.'
        }
        $GetTasks = {
            Get-ScheduledTask -TaskPath '\Microsoft\Windows\Windows Defender\' -ErrorAction Stop
        }
    }
    if (-not $EnableOne) {
        $EnableOne = {
            param($Task)
            Enable-ScheduledTask -InputObject $Task -ErrorAction Stop
        }
    }
    $lines = New-Object System.Collections.Generic.List[string]
    $failed = $false
    $tasks = @()
    try {
        $tasks = @(& $GetTasks)
    } catch {
        [void]$lines.Add('No Defender tasks were returned. An empty task list fails the restore.')
        return [pscustomobject]@{ Lines = $lines.ToArray(); Failed = $true }
    }
    $real = @($tasks | Where-Object { $null -ne $_ })
    if ($real.Count -eq 0) {
        [void]$lines.Add('No Defender tasks were returned. An empty task list fails the restore.')
        return [pscustomobject]@{ Lines = $lines.ToArray(); Failed = $true }
    }
    foreach ($task in $real) {
        $name = [string]$task.TaskName
        if ([string]::IsNullOrWhiteSpace($name)) { $name = '(unnamed task)' }
        $updated = $null
        try {
            $updated = & $EnableOne $task
        } catch {
            [void]$lines.Add("$name = denied")
            $failed = $true
            continue
        }
        $state = $null
        if ($null -ne $updated -and $updated.PSObject.Properties.Name -contains 'State') {
            $state = [string]$updated.State
        } elseif ($task.PSObject.Properties.Name -contains 'State') {
            $state = [string]$task.State
        }
        $enabled = $state -eq 'Ready' -or $state -eq 'Running' -or $state -eq 'Queued'
        if (-not $enabled) {
            $shown = if ([string]::IsNullOrWhiteSpace($state)) { 'unknown' } else { $state }
            [void]$lines.Add("$name = $shown")
            $failed = $true
            continue
        }
        [void]$lines.Add("$name = $state")
    }
    [pscustomobject]@{
        Lines  = $lines.ToArray()
        Failed = [bool]$failed
    }
}

function Invoke-DefenderRestoreAfterReboot {
    param(
        [scriptblock]$GetTasks,
        [scriptblock]$EnableOne,
        [scriptblock]$GetSnapshot
    )
    $enableParams = @{}
    if ($PSBoundParameters.ContainsKey('GetTasks')) { $enableParams.GetTasks = $GetTasks }
    if ($PSBoundParameters.ContainsKey('EnableOne')) { $enableParams.EnableOne = $EnableOne }
    $enabled = Enable-DefenderScheduledTasks @enableParams
    if ($PSBoundParameters.ContainsKey('GetSnapshot')) {
        $snapshot = & $GetSnapshot
    } else {
        $snapshot = Get-DefenderRestoredSnapshot
    }
    $extra = @()
    if ($enabled.Failed) { $extra += 'Task enable did not succeed.' }
    $report = Format-DefenderRestoredReport -Snapshot $snapshot -ExtraFailures $extra
    $lines = New-Object System.Collections.Generic.List[string]
    foreach ($line in @($enabled.Lines)) { [void]$lines.Add([string]$line) }
    foreach ($line in @($report.Lines)) { [void]$lines.Add([string]$line) }
    [pscustomobject]@{
        Lines  = $lines.ToArray()
        Passed = [bool]$report.Passed
        Code   = [int]$report.Code
    }
}

function Get-RestoreNativeText {
    param([Parameter(Mandatory)][scriptblock]$Command)
    try {
        $output = & $Command 2>&1 | Out-String
        if ($LASTEXITCODE -ne 0) { return $null }
        return $output
    } catch {
        return $null
    }
}

function Get-RestoreServiceState {
    param([Parameter(Mandatory)][string]$Name)
    try {
        $svc = Get-Service -Name $Name -ErrorAction Stop
        return @{ Status = [string]$svc.Status; StartType = [string]$svc.StartType }
    } catch {
        return $null
    }
}

function Get-DefenderRestoredSnapshot {
    if ($env:DEFENDER_PROCEDURE_UNDER_TEST -eq '1') {
        throw 'Refusing to query the live machine during tests.'
    }
    $snapshot = @{
        InSafeMode     = $false
        Services       = @{}
        MsMpEngRunning = $false
        QcText         = @{}
        QueryText      = @{}
        FltmcText      = $null
        MpText         = $null
        Tasks          = $null
    }
    try {
        $safeBootKey = 'HKLM:\SYSTEM\CurrentControlSet\Control\SafeBoot\Option'
        $opt = Get-ItemProperty -LiteralPath $safeBootKey -ErrorAction Stop
        if ($opt.PSObject.Properties.Name -contains 'OptionValue' -and $null -ne $opt.OptionValue) {
            $snapshot.InSafeMode = $true
        }
    } catch {
        $snapshot.InSafeMode = $false
    }
    foreach ($name in @('WinDefend', 'WdNisSvc', 'Sense')) {
        $snapshot.Services[$name] = Get-RestoreServiceState -Name $name
    }
    try {
        $processes = @(Get-Process -Name MsMpEng -ErrorAction SilentlyContinue | Where-Object { $null -ne $_ })
        $snapshot.MsMpEngRunning = ($processes.Count -gt 0)
    } catch {
        $snapshot.MsMpEngRunning = $null
    }
    foreach ($name in @('WdFilter', 'WdBoot', 'WdNisDrv')) {
        $snapshot.QcText[$name] = Get-RestoreNativeText -Command ([scriptblock]::Create("sc.exe qc $name"))
    }
    $snapshot.QueryText['WdFilter'] = Get-RestoreNativeText -Command { sc.exe query WdFilter }
    $snapshot.FltmcText = Get-RestoreNativeText -Command { fltmc }
    try {
        $mp = Get-MpComputerStatus -ErrorAction Stop |
            Select-Object AMServiceEnabled, AntivirusEnabled, RealTimeProtectionEnabled, IsTamperProtected |
            Format-List | Out-String
        $snapshot.MpText = $mp.Trim()
    } catch {
        $snapshot.MpText = "Get-MpComputerStatus threw: $($_.Exception.Message)"
    }
    try {
        $taskPath = '\Microsoft\Windows\Windows Defender\'
        $snapshot.Tasks = @(Get-ScheduledTask -TaskPath $taskPath -ErrorAction Stop)
    } catch {
        $snapshot.Tasks = $null
    }
    return $snapshot
}

# Script-scope entry. Dot-sourcing loads the functions and does not change the machine.
# DEFENDER_PROCEDURE_UNDER_TEST is set only by the offline test runner.
if ($MyInvocation.InvocationName -ne '.') {
    if ($env:DEFENDER_PROCEDURE_UNDER_TEST -eq '1') {
        throw 'Refusing to run the live procedure during tests.'
    }
    if ($args -contains '-AfterReboot') {
        $liveResult = Invoke-DefenderRestoreAfterReboot
    } else {
        $liveResult = Invoke-DefenderRestorePreReboot
    }
    foreach ($liveLine in @($liveResult.Lines)) { Write-Output $liveLine }
    if ($PSCommandPath) { exit $liveResult.Code }
}