#Requires -RunAsAdministrator
<#
.SYNOPSIS
    Re-runs GuruScan scanners in clean mode against a previous scan's log folder.
.DESCRIPTION
    Reads the results.json from a prior scan run, then re-launches each scanner
    that completed (or any subset via -Scanners) using clean_args to remove
    detected threats. Writes remediation-results.json to the same log folder.
.PARAMETER LogRoot
    Path to the scan results folder produced by Invoke-GuruScan.ps1.
    This folder must contain a results.json file.
.PARAMETER Scanners
    Run only the named scanners. Names must match the "name" field in
    scanners.json exactly. If omitted, all scanners that previously ran
    successfully are re-run.
.EXAMPLE
    .\Invoke-Remediation.ps1 -LogRoot "C:\ScanLogs\DESKTOP-20260523-143000"
    .\Invoke-Remediation.ps1 -LogRoot "C:\ScanLogs\DESKTOP-20260523-143000" -Scanners AdwCleaner,MSERT
#>
[CmdletBinding()]
param(
    [Parameter(Mandatory = $true)]
    [string]$LogRoot,

    [string[]]$Scanners
)

Set-StrictMode -Version Latest
$ErrorActionPreference = 'Continue'
$ProgressPreference    = 'SilentlyContinue'

# ---------------------------------------------------------------------------
# Helpers (duplicated here so this script is self-contained)
# ---------------------------------------------------------------------------

function Expand-TokenizedArgs {
    param([string[]]$Args, [string]$LogRoot, [string]$ExeDir)
    $out = foreach ($a in $Args) {
        $a = $a -replace '\{LOG_ROOT\}', $LogRoot
        $a = $a -replace '\{EXE_DIR\}',  $ExeDir
        $a
    }
    return $out
}

function Expand-EnvPath {
    param([string]$Path)
    return [System.Environment]::ExpandEnvironmentVariables($Path)
}

function Remove-PathSilent {
    param([string]$Path)
    $expanded = Expand-EnvPath $Path
    if (Test-Path $expanded) {
        Remove-Item -Path $expanded -Recurse -Force -ErrorAction SilentlyContinue
    }
}

function Wait-ProcessWithTimeout {
    param(
        [System.Diagnostics.Process]$Process,
        [int]$TimeoutSeconds
    )
    $deadline = (Get-Date).AddSeconds($TimeoutSeconds)
    while (-not $Process.HasExited) {
        if ((Get-Date) -gt $deadline) {
            try { $Process.Kill() } catch { }
            return $false
        }
        Start-Sleep -Seconds 5
    }
    return $true
}

function Wait-ServicesToStop {
    param([string[]]$ServiceNames, [int]$TimeoutSeconds = 120)
    if (-not $ServiceNames -or $ServiceNames.Count -eq 0) { return }
    $deadline = (Get-Date).AddSeconds($TimeoutSeconds)
    foreach ($svc in $ServiceNames) {
        while ((Get-Date) -lt $deadline) {
            $s = Get-Service -Name $svc -ErrorAction SilentlyContinue
            if (-not $s -or $s.Status -ne 'Running') { break }
            Start-Sleep -Seconds 3
        }
    }
}

function Invoke-HitmanProTrialReset {
    Remove-Item -Path "HKLM:\SOFTWARE\HitmanPro"       -Recurse -Force -ErrorAction SilentlyContinue
    Remove-Item -Path "HKLM:\SOFTWARE\HitmanPro.Alert" -Recurse -Force -ErrorAction SilentlyContinue
    New-Item -Path "HKLM:\SOFTWARE\HitmanPro"       -Force | Out-Null
    New-Item -Path "HKLM:\SOFTWARE\HitmanPro.Alert" -Force | Out-Null
}

function Get-ExitCodeThreats {
    param([string]$ScannerName, [int]$ExitCode)
    switch ($ScannerName) {
        'AdwCleaner' { if ($ExitCode -eq 1) { return 1 } }
        'MSERT'      { if ($ExitCode -ne 0) { return 1 } }
        'TDSSKiller' { if ($ExitCode -eq 1) { return 1 } }
        'HitmanPro'  { if ($ExitCode -eq 1) { return 1 } }
        'Emsisoft'   { if ($ExitCode -ge 1) { return 1 } }
        'Stinger'    { if ($ExitCode -eq 13) { return 1 } }
        'ESET'       { if ($ExitCode -ne 0) { return 1 } }
    }
    return 0
}

# ---------------------------------------------------------------------------
# Validate inputs
# ---------------------------------------------------------------------------

$ScriptRoot    = $PSScriptRoot
$ConfigPath    = Join-Path $ScriptRoot 'scanners.json'
$PriorResults  = Join-Path $LogRoot 'results.json'

if (-not (Test-Path $LogRoot)) {
    Write-Host "[ERROR] LogRoot not found: $LogRoot" -ForegroundColor Red
    exit 1
}

if (-not (Test-Path $PriorResults)) {
    Write-Host "[ERROR] results.json not found in: $LogRoot" -ForegroundColor Red
    exit 1
}

if (-not (Test-Path $ConfigPath)) {
    Write-Host "[ERROR] scanners.json not found: $ConfigPath" -ForegroundColor Red
    exit 1
}

# ---------------------------------------------------------------------------
# Load config and prior results
# ---------------------------------------------------------------------------

$config       = Get-Content $ConfigPath -Raw | ConvertFrom-Json
$priorData    = Get-Content $PriorResults -Raw | ConvertFrom-Json

$priorScannerNames = $priorData.scanners |
    Where-Object { $_.status -eq 'completed' } |
    ForEach-Object { $_.name }

$scannerList = $config.scanners | Where-Object { $priorScannerNames -contains $_.name }

if ($Scanners -and $Scanners.Count -gt 0) {
    $scannerList = $scannerList | Where-Object { $Scanners -contains $_.name }
}

if (-not $scannerList) {
    Write-Host "[WARNING] No eligible scanners to remediate." -ForegroundColor Yellow
    exit 0
}

Write-Host ""
Write-Host "=== GuruScan Remediation ===" -ForegroundColor Cyan
Write-Host "  Log root : $LogRoot"
Write-Host "  Scanners : $($scannerList.name -join ', ')"
Write-Host ""

# ---------------------------------------------------------------------------
# Run scanners in clean mode
# ---------------------------------------------------------------------------

$results      = [System.Collections.Generic.List[pscustomobject]]::new()
$startedAt    = Get-Date
$totalThreats = 0

foreach ($s in $scannerList) {
    Write-Host "------------------------------------------------------------" -ForegroundColor DarkGray
    Write-Host "  Scanner : $($s.name) [CLEAN MODE]" -ForegroundColor Cyan

    # ------------------------------------------------------------------
    # Resolve EXE paths
    # ------------------------------------------------------------------
    $mainExePath      = $null
    $installerExePath = $null

    if ($s.installer_exe) {
        $installerExePath = if ([System.IO.Path]::IsPathRooted($s.installer_exe)) {
            $s.installer_exe
        } else {
            Join-Path $ScriptRoot $s.installer_exe
        }
    }

    $rawExe = $s.exe
    if ([System.IO.Path]::IsPathRooted($rawExe)) {
        $mainExePath = $rawExe
    } else {
        $mainExePath = Join-Path $ScriptRoot $rawExe
    }

    $exeToCheck = if ($installerExePath) { $installerExePath } else { $mainExePath }

    if (-not (Test-Path $exeToCheck)) {
        Write-Host "  [WARNING] EXE not found — skipping: $exeToCheck" -ForegroundColor Yellow
        $results.Add([pscustomobject]@{
            Scanner      = $s.name
            Status       = 'SKIPPED (missing)'
            ExitCode     = $null
            ThreatsFound = 0
            Duration     = '0 min'
            LogPath      = ''
        })
        continue
    }

    # ------------------------------------------------------------------
    # Resolve args (always clean_args)
    # ------------------------------------------------------------------
    $exeDir      = Split-Path $mainExePath -Parent
    $expandedArgs = Expand-TokenizedArgs -Args @($s.clean_args) -LogRoot $LogRoot -ExeDir $exeDir

    # ------------------------------------------------------------------
    # HitmanPro trial reset
    # ------------------------------------------------------------------
    if ($s.hitmanpro_trial_reset -eq $true) {
        Write-Host "  [INFO] Resetting HitmanPro trial registry..." -ForegroundColor Cyan
        Invoke-HitmanProTrialReset
    }

    # ------------------------------------------------------------------
    # Pre-clean
    # ------------------------------------------------------------------
    foreach ($p in $s.pre_clean_paths) {
        Remove-PathSilent $p
    }

    # ------------------------------------------------------------------
    # Timeout
    # ------------------------------------------------------------------
    $timeoutSec = $s.timeout_min * 60

    # ------------------------------------------------------------------
    # Two-step install (Emsisoft)
    # ------------------------------------------------------------------
    if ($installerExePath) {
        Write-Host "  [INFO] Running installer: $installerExePath" -ForegroundColor Cyan
        try {
            $instProc = Start-Process -FilePath $installerExePath -ArgumentList '/S' -PassThru -NoNewWindow
            $instCompleted = Wait-ProcessWithTimeout -Process $instProc -TimeoutSeconds 300
            if (-not $instCompleted) {
                Write-Host "  [WARNING] Installer timed out — skipping" -ForegroundColor Yellow
                $results.Add([pscustomobject]@{
                    Scanner      = $s.name
                    Status       = 'FAILED (installer timeout)'
                    ExitCode     = $null
                    ThreatsFound = 0
                    Duration     = '0 min'
                    LogPath      = ''
                })
                continue
            }
        }
        catch {
            Write-Host "  [ERROR] Installer failed: $_" -ForegroundColor Red
            $results.Add([pscustomobject]@{
                Scanner      = $s.name
                Status       = "FAILED (installer error)"
                ExitCode     = $null
                ThreatsFound = 0
                Duration     = '0 min'
                LogPath      = ''
            })
            continue
        }

        if (-not (Test-Path $mainExePath)) {
            Write-Host "  [ERROR] Main EXE missing after install: $mainExePath" -ForegroundColor Red
            $results.Add([pscustomobject]@{
                Scanner      = $s.name
                Status       = 'FAILED (post-install EXE missing)'
                ExitCode     = $null
                ThreatsFound = 0
                Duration     = '0 min'
                LogPath      = ''
            })
            continue
        }

        Write-Host "  [INFO] Updating Emsisoft definitions..." -ForegroundColor Cyan
        try {
            $updateProc = Start-Process -FilePath $mainExePath -ArgumentList '/update' -PassThru -NoNewWindow
            Wait-ProcessWithTimeout -Process $updateProc -TimeoutSeconds 300 | Out-Null
        }
        catch {
            Write-Host "  [WARNING] Update step failed: $_ — continuing with existing definitions" -ForegroundColor Yellow
        }
    }

    # ------------------------------------------------------------------
    # Randomize EXE (TDSSKiller)
    # ------------------------------------------------------------------
    $randomizedExePath = $null
    $launchExe = $mainExePath

    if ($s.randomize_exe -eq $true) {
        $tempName = [System.IO.Path]::GetRandomFileName() -replace '\..*$', ''
        $randomizedExePath = Join-Path $env:TEMP "$tempName.exe"
        Copy-Item -Path $mainExePath -Destination $randomizedExePath -Force
        $launchExe = $randomizedExePath
        Write-Host "  [INFO] TDSSKiller randomized to: $randomizedExePath" -ForegroundColor Cyan
    }

    # ------------------------------------------------------------------
    # Launch
    # ------------------------------------------------------------------
    Write-Host "  [....] Launching $($s.name) in clean mode..." -ForegroundColor Cyan
    $scanStart = Get-Date
    $proc      = $null
    $status    = 'completed'
    $exitCode  = $null

    try {
        $startParams = @{
            FilePath    = $launchExe
            PassThru    = $true
            NoNewWindow = $true
        }
        if ($expandedArgs -and $expandedArgs.Count -gt 0) {
            $startParams['ArgumentList'] = $expandedArgs
        }

        $proc = Start-Process @startParams

        $completed = Wait-ProcessWithTimeout -Process $proc -TimeoutSeconds $timeoutSec
        if (-not $completed) {
            $status = 'TIMED OUT'
            Write-Host "  [WARNING] $($s.name) timed out after $($s.timeout_min) min" -ForegroundColor Yellow
        } else {
            $exitCode = $proc.ExitCode
        }
    }
    catch {
        $status = 'FAILED'
        Write-Host "  [ERROR] $($s.name) failed: $_" -ForegroundColor Red
    }

    # ------------------------------------------------------------------
    # Wait for services
    # ------------------------------------------------------------------
    if ($s.service_names -and $s.service_names.Count -gt 0) {
        Wait-ServicesToStop -ServiceNames $s.service_names -TimeoutSeconds 180
    }

    # ------------------------------------------------------------------
    # Duration
    # ------------------------------------------------------------------
    $durMin = [math]::Round(((Get-Date) - $scanStart).TotalMinutes, 2)

    # ------------------------------------------------------------------
    # Collect logs
    # ------------------------------------------------------------------
    $logDestDir = Join-Path $LogRoot "$($s.name)_Remediation_Logs"

    if ($s.log_src) {
        $expandedLogSrc = Expand-TokenizedArgs -Args @($s.log_src) -LogRoot $LogRoot -ExeDir $exeDir
        $expandedLogSrc = Expand-EnvPath $expandedLogSrc[0]
        if (Test-Path $expandedLogSrc) {
            New-Item -ItemType Directory -Path $logDestDir -Force | Out-Null
            Copy-Item -Path $expandedLogSrc -Destination $logDestDir -Recurse -Force -ErrorAction SilentlyContinue
        }
    }

    # ------------------------------------------------------------------
    # Post-clean
    # ------------------------------------------------------------------
    foreach ($p in $s.post_clean_paths) {
        Remove-PathSilent $p
    }

    # ------------------------------------------------------------------
    # Randomized EXE cleanup
    # ------------------------------------------------------------------
    if ($randomizedExePath -and (Test-Path $randomizedExePath)) {
        Remove-Item $randomizedExePath -Force -ErrorAction SilentlyContinue
    }

    # ------------------------------------------------------------------
    # Threats
    # ------------------------------------------------------------------
    $threatsFound = 0
    if ($null -ne $exitCode) {
        $threatsFound = Get-ExitCodeThreats -ScannerName $s.name -ExitCode $exitCode
    }
    $totalThreats += $threatsFound

    $color = if ($status -eq 'completed') { 'Green' } elseif ($status -eq 'TIMED OUT') { 'Yellow' } else { 'Red' }
    Write-Host "  [$status] ExitCode=$exitCode Threats=$threatsFound Duration=${durMin}min" -ForegroundColor $color

    $results.Add([pscustomobject]@{
        Scanner      = $s.name
        Status       = $status
        ExitCode     = $exitCode
        ThreatsFound = $threatsFound
        Duration     = "$durMin min"
        LogPath      = $logDestDir
    })
}

# ---------------------------------------------------------------------------
# Write remediation-results.json
# ---------------------------------------------------------------------------

$completedAt = Get-Date

$scannerResults = foreach ($r in $results) {
    $durValue = 0
    if ($r.Duration -match '^([\d.]+)') {
        $durValue = [double]$Matches[1]
    }
    [ordered]@{
        name          = $r.Scanner
        status        = $r.Status
        exit_code     = $r.ExitCode
        threats_found = $r.ThreatsFound
        duration_min  = $durValue
        log_path      = $r.LogPath
    }
}

$remObj = [ordered]@{
    scan_id         = $priorData.scan_id
    machine         = $env:COMPUTERNAME
    started_at      = $startedAt.ToUniversalTime().ToString('o')
    completed_at    = $completedAt.ToUniversalTime().ToString('o')
    total_threats   = $totalThreats
    reboot_required = $false
    scan_mode       = 'clean'
    scanners        = @($scannerResults)
}

$remResultsPath = Join-Path $LogRoot 'remediation-results.json'
$remObj | ConvertTo-Json -Depth 10 | Set-Content -Path $remResultsPath -Encoding UTF8

# ---------------------------------------------------------------------------
# Summary
# ---------------------------------------------------------------------------

Write-Host ""
Write-Host "============================================================" -ForegroundColor Cyan
Write-Host "  REMEDIATION COMPLETE" -ForegroundColor Cyan
Write-Host "============================================================" -ForegroundColor Cyan
Write-Host ""
$results | Format-Table -AutoSize
Write-Host ""

if ($totalThreats -gt 0) {
    Write-Host "  [WARNING] $totalThreats threat indicator(s) still detected after clean pass." -ForegroundColor Yellow
    Write-Host "            Review logs carefully — some threats may require a reboot before" -ForegroundColor Yellow
    Write-Host "            they can be fully removed. Reboot and re-scan." -ForegroundColor Yellow
} else {
    Write-Host "  [OK] Clean pass complete — no threats detected in output." -ForegroundColor Green
}

Write-Host ""
Write-Host "  Remediation results: $remResultsPath" -ForegroundColor Gray
Write-Host ""

exit 0