SMBeat

0.1.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,
        [int]$IntervalSec = 60,
        $CimSession
    )

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

    $action = New-ScheduledTaskAction -Execute $exe -Argument $arg
    $start = (Get-Date).AddMinutes(-1)
    $repeat = New-TimeSpan -Seconds $IntervalSec
    $duration = New-TimeSpan -Days 9999
    $trigger = New-ScheduledTaskTrigger -Once -At $start -RepetitionInterval $repeat -RepetitionDuration $duration
    $principal = New-ScheduledTaskPrincipal -UserId 'SYSTEM' -LogonType ServiceAccount -RunLevel Highest
    $settings = New-ScheduledTaskSettingsSet -AllowStartIfOnBatteries -DontStopIfGoingOnBatteries -StartWhenAvailable -MultipleInstances IgnoreNew
    $task = New-ScheduledTask -Action $action -Trigger $trigger -Principal $principal -Settings $settings

    $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
    )

    $params = Get-SMBeatTaskCimParams -CimSession $CimSession
    $params['Confirm'] = $false
    $params['ErrorAction'] = 'SilentlyContinue'
    Unregister-ScheduledTask @params
}

function Start-SMBeatScheduledTaskInternal {
    param(
        $CimSession
    )

    $enable = Get-SMBeatTaskCimParams -CimSession $CimSession
    Enable-ScheduledTask @enable | Out-Null
    $start = Get-SMBeatTaskCimParams -CimSession $CimSession
    Start-ScheduledTask @start
}

function Stop-SMBeatScheduledTaskInternal {
    param(
        $CimSession
    )

    $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
    }

    $config = [PSCustomObject]@{
        IntervalSec        = $Payload.IntervalSec
        RetentionDays      = $Payload.RetentionDays
        Path               = $root
        IncludeSessions    = $Payload.IncludeSessions
        IncludeAdminShares = $Payload.IncludeAdminShares
        IncludeNics        = $Payload.IncludeNics
    }

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

    Save-SMBeatConfig -DataRoot $root -Config $config
    Register-SMBeatScheduledTaskInternal -DataRoot $root -IntervalSec $Payload.IntervalSec
    Start-SMBeatScheduledTaskInternal
    Invoke-SMBeatSample -Path $root -IntervalSec $Payload.IntervalSec -RetentionDays $Payload.RetentionDays -IncludeSessions $Payload.IncludeSessions -IncludeAdminShares $Payload.IncludeAdminShares -IncludeNics $Payload.IncludeNics | Out-Null

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

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