SMBeat

0.4.3

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

    $params = Get-SMBeatTaskCimParams -CimSession $CimSession
    $params['Confirm'] = $false
    $params['ErrorAction'] = 'SilentlyContinue'
    Unregister-ScheduledTask @params
    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
        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)
    }
}