SMBeat

0.5.1

Private/Task.ps1

#Requires -Version 5.1

function Get-SMBeatTaskCimParams {
    param(
        $CimSession
    )

    $params = @{
        TaskName = $script:SMBeatTaskName
        TaskPath = $script:SMBeatTaskPath
    }
    if ($CimSession) {
        $params['CimSession'] = $CimSession
    }
    return $params
}

function Register-SMBeatScheduledTaskInternal {
    param(
        [Parameter(Mandatory = $true)]
        [string]$DataRoot,
        $CimSession
    )

    $exe = Get-SMBeatPowerShellExePath
    $scriptPath = Get-SMBeatCollectorScriptPath
    $arg = '-NoProfile -NonInteractive -ExecutionPolicy Bypass -File "{0}"' -f $scriptPath

    $action = New-ScheduledTaskAction -Execute $exe -Argument $arg
    $startup = New-ScheduledTaskTrigger -AtStartup
    $once = New-ScheduledTaskTrigger -Once -At ((Get-Date).AddSeconds(3))
    $principal = New-ScheduledTaskPrincipal -UserId 'SYSTEM' -LogonType ServiceAccount -RunLevel Highest
    $settings = New-ScheduledTaskSettingsSet `
        -AllowStartIfOnBatteries `
        -DontStopIfGoingOnBatteries `
        -StartWhenAvailable `
        -MultipleInstances IgnoreNew `
        -ExecutionTimeLimit ([TimeSpan]::Zero) `
        -RestartCount 3 `
        -RestartInterval (New-TimeSpan -Minutes 1)
    $desc = 'SMBeat collector {0}' -f (Get-SMBeatModuleVersion)
    $task = New-ScheduledTask -Action $action -Trigger @($startup, $once) -Principal $principal -Settings $settings -Description $desc

    $registerParams = @{
        TaskName    = $script:SMBeatTaskName
        TaskPath    = $script:SMBeatTaskPath
        InputObject = $task
        Force       = $true
    }
    if ($CimSession) {
        $registerParams['CimSession'] = $CimSession
    }

    Register-ScheduledTask @registerParams | Out-Null
}

function Unregister-SMBeatScheduledTaskInternal {
    param(
        $CimSession,
        [string]$Path
    )

    $pack = Resolve-SMBeatWatcherDataRoot -Path $Path
    Request-SMBeatWatcherStop -DataRoot $pack.DataRoot
    $stop = Get-SMBeatTaskCimParams -CimSession $CimSession
    $stop['ErrorAction'] = 'SilentlyContinue'
    Stop-ScheduledTask @stop | Out-Null

    $params = Get-SMBeatTaskCimParams -CimSession $CimSession
    $params['Confirm'] = $false
    $params['ErrorAction'] = 'SilentlyContinue'
    Unregister-ScheduledTask @params | Out-Null
    Stop-SMBeatEtwShareLeftover
}

function Start-SMBeatScheduledTaskInternal {
    param(
        $CimSession,
        [string]$Path
    )

    $pack = Resolve-SMBeatWatcherDataRoot -Path $Path
    Clear-SMBeatWatcherRequest -Path (Get-SMBeatStopRequestPath -DataRoot $pack.DataRoot)
    $enable = Get-SMBeatTaskCimParams -CimSession $CimSession
    Enable-ScheduledTask @enable | Out-Null
    $start = Get-SMBeatTaskCimParams -CimSession $CimSession
    Start-ScheduledTask @start
}

function Stop-SMBeatScheduledTaskInternal {
    param(
        $CimSession,
        [string]$Path
    )

    $pack = Resolve-SMBeatWatcherDataRoot -Path $Path
    Request-SMBeatWatcherStop -DataRoot $pack.DataRoot
    $stop = Get-SMBeatTaskCimParams -CimSession $CimSession
    $stop['ErrorAction'] = 'SilentlyContinue'
    Stop-ScheduledTask @stop
    $disable = Get-SMBeatTaskCimParams -CimSession $CimSession
    Disable-ScheduledTask @disable | Out-Null
}

function Register-SMBeatCollectorLocal {
    param(
        [Parameter(Mandatory = $true)]
        $Payload
    )

    $root = $Payload.Path
    if ([string]::IsNullOrWhiteSpace($root)) {
        $root = Get-SMBeatDefaultDataRoot
    }

    $includeEtw = [bool](Get-SMBeatPropertyValue -Object $Payload -Name 'IncludeEtwShares' -Default $true)
    $includeAdmin = [bool](Get-SMBeatPropertyValue -Object $Payload -Name 'IncludeAdminShares' -Default $false)

    $version = Get-SMBeatModuleVersion
    $config = [PSCustomObject]@{
        IntervalSec        = $Payload.IntervalSec
        RetentionDays      = $Payload.RetentionDays
        Path               = $root
        IncludeNics        = $Payload.IncludeNics
        IncludeSmbPerf     = (Get-SMBeatPropertyValue -Object $Payload -Name 'IncludeSmbPerf' -Default $true)
        IncludeEtwShares   = $includeEtw
        IncludeAdminShares = $includeAdmin
        Mode               = 'tcp445+etw'
        ModuleVersion      = $version
    }

    if (-not (Test-Path -LiteralPath $root)) {
        New-Item -ItemType Directory -Path $root -Force | Out-Null
    }

    Save-SMBeatConfig -DataRoot $root -Config $config
    $defaultRoot = Get-SMBeatDefaultDataRoot
    if ($root -ne $defaultRoot) {
        if (-not (Test-Path -LiteralPath $defaultRoot)) {
            New-Item -ItemType Directory -Path $defaultRoot -Force | Out-Null
        }
        Save-SMBeatConfig -DataRoot $defaultRoot -Config $config
    }

    Clear-SMBeatWatcherRequest -Path (Get-SMBeatStopRequestPath -DataRoot $root)
    Register-SMBeatScheduledTaskInternal -DataRoot $root
    Start-SMBeatScheduledTaskInternal -Path $root

    Write-Host ('SMBeat {0} registered at {1}' -f $version, $root)

    if ($Payload.Wait) {
        return (Wait-SMBeatHeartbeat -Path $root -TimeoutSec $Payload.WaitTimeoutSec -IntervalSec $Payload.IntervalSec)
    }

    if ($Payload.PassThru) {
        return (Get-SMBeatStatusInternal -Path $root)
    }
}

function Test-SMBeatDataRootSafeToRemove {
    param(
        [Parameter(Mandatory = $true)]
        [string]$Path
    )

    if ([string]::IsNullOrWhiteSpace($Path)) {
        return $false
    }

    $full = [System.IO.Path]::GetFullPath($Path)
    $driveRoot = [System.IO.Path]::GetPathRoot($full)
    if ($full.TrimEnd('\', '/') -eq $driveRoot.TrimEnd('\', '/')) {
        return $false
    }

    $blocked = New-Object System.Collections.Generic.List[string]
    if (-not [string]::IsNullOrWhiteSpace($script:SMBeatModuleRoot)) {
        $blocked.Add([System.IO.Path]::GetFullPath($script:SMBeatModuleRoot)) | Out-Null
    }
    $install = Get-SMBeatInstallDestination
    if (-not [string]::IsNullOrWhiteSpace($install)) {
        $blocked.Add([System.IO.Path]::GetFullPath($install)) | Out-Null
    }
    foreach ($item in $blocked) {
        if ($full.StartsWith($item, [System.StringComparison]::OrdinalIgnoreCase)) {
            return $false
        }
        if ($item.StartsWith($full, [System.StringComparison]::OrdinalIgnoreCase)) {
            return $false
        }
    }

    if (Test-Path -LiteralPath (Join-Path $full 'SMBeat.psd1')) {
        return $false
    }
    if (Test-Path -LiteralPath (Join-Path $full 'SMBeat.psm1')) {
        return $false
    }

    $leaf = Split-Path -Path $full -Leaf
    if ($leaf -eq 'SMBeat') {
        return $true
    }
    if ((Test-Path -LiteralPath (Join-Path $full 'config.json')) -or
        (Test-Path -LiteralPath (Join-Path $full 'heartbeat.json')) -or
        (Test-Path -LiteralPath (Join-Path $full 'samples'))) {
        return $true
    }

    return $false
}

function Get-SMBeatDataRootsToRemove {
    param(
        [string]$Path
    )

    $seen = @{}
    $list = New-Object System.Collections.Generic.List[string]
    $candidates = New-Object System.Collections.Generic.List[string]
    $defaultRoot = Get-SMBeatDefaultDataRoot

    try {
        $pack = Resolve-SMBeatWatcherDataRoot -Path $Path
        if ($pack -and $pack.DataRoot) {
            $candidates.Add([string]$pack.DataRoot) | Out-Null
        }
    }
    catch {
    }

    if ([string]::IsNullOrWhiteSpace($Path)) {
        $candidates.Add($defaultRoot) | Out-Null
    }
    else {
        $candidates.Add($Path) | Out-Null
        $defaultConfig = Read-SMBeatConfig -DataRoot $defaultRoot
        $pointed = Get-SMBeatPropertyValue -Object $defaultConfig -Name 'Path' -Default $null
        if ($pointed) {
            $pointedFull = [System.IO.Path]::GetFullPath([string]$pointed)
            $pathFull = [System.IO.Path]::GetFullPath($Path)
            $addDefault = $false
            if ($pointedFull -eq $pathFull) {
                $addDefault = $true
            }
            else {
                foreach ($item in $candidates) {
                    if ([System.IO.Path]::GetFullPath($item) -eq $pointedFull) {
                        $addDefault = $true
                        break
                    }
                }
            }
            if ($addDefault) {
                $candidates.Add($defaultRoot) | Out-Null
            }
        }
    }

    foreach ($item in $candidates) {
        if ([string]::IsNullOrWhiteSpace($item)) {
            continue
        }
        $full = [System.IO.Path]::GetFullPath($item)
        if ($seen.ContainsKey($full)) {
            continue
        }
        $seen[$full] = $true
        $list.Add($full) | Out-Null
    }

    return $list.ToArray()
}

function Remove-SMBeatDataRoot {
    param(
        [Parameter(Mandatory = $true)]
        [string]$Path
    )

    if (-not (Test-SMBeatDataRootSafeToRemove -Path $Path)) {
        return $false
    }
    if (-not (Test-Path -LiteralPath $Path)) {
        return $true
    }

    $attempt = 0
    while ($attempt -lt 8) {
        try {
            Remove-Item -LiteralPath $Path -Recurse -Force -ErrorAction Stop
            return $true
        }
        catch {
            Start-Sleep -Milliseconds 250
            $attempt++
        }
    }

    return $false
}

function Uninstall-SMBeatLocal {
    param(
        [string]$Path
    )

    $roots = @(Get-SMBeatDataRootsToRemove -Path $Path)
    try {
        if (Get-Command Unregister-ScheduledTask -ErrorAction SilentlyContinue) {
            Unregister-SMBeatScheduledTaskInternal -Path $Path
        }
        else {
            $pack = Resolve-SMBeatWatcherDataRoot -Path $Path
            Request-SMBeatWatcherStop -DataRoot $pack.DataRoot
            Stop-SMBeatEtwShareLeftover
        }
    }
    catch {
        try {
            Stop-SMBeatEtwShareLeftover
        }
        catch {
        }
    }

    $deadline = (Get-Date).AddSeconds(8)
    while ((Get-Date) -lt $deadline) {
        if (-not (Test-SMBeatWatcherAlive -Path $Path)) {
            break
        }
        Start-Sleep -Milliseconds 300
    }

    $removed = New-Object System.Collections.Generic.List[string]
    foreach ($root in $roots) {
        if (Remove-SMBeatDataRoot -Path $root) {
            if (-not (Test-Path -LiteralPath $root)) {
                $removed.Add($root) | Out-Null
            }
        }
    }

    [PSCustomObject]@{
        ComputerName = $env:COMPUTERNAME
        Removed      = $removed.ToArray()
        ModulePath   = $script:SMBeatModuleRoot
        Success      = $true
    }
}