Public/Start-SMBeat.ps1

#Requires -Version 5.1

function Start-SMBeatLocal {
    param(
        $Payload
    )

    Install-SMBeatIfMissing

    $taskPack = Get-SMBeatScheduledTaskInfoSafe
    if (-not $taskPack.Task) {
        Register-SMBeatCollectorLocal -Payload $Payload
        return
    }

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

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

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

function Start-SMBeat {
    <#
    .SYNOPSIS
        Copies the module if missing, registers the collector if needed, otherwise starts it.
    .DESCRIPTION
        Commissioning and resume. If SMBeat.psd1 is missing under Program Files,
        copies the loaded module there (no upgrade if it is already present).
        If the scheduled task is missing, creates and starts it. If the task
        exists, enables and starts it. -CimSession cannot copy files or create
        a new collector; use -ComputerName or -Session.
    .PARAMETER IncludeSmbPerf
        Collects SMB Server / Shares CIM raw counters (kind=smbperf) as a volume check. Default is true.
        Used only when the scheduled task is created for the first time.
    .PARAMETER Wait
        Starts the task immediately and waits until the first heartbeat is fresh (or timeout).
    .PARAMETER MailTo
        Recipient for collector notices via Send-MailMessage (registered, started, stopped, errors).
        Empty keeps the previous value from config.json. Used only when the task is created.
    .PARAMETER SmtpServer
        SMTP host for Send-MailMessage. Required together with MailTo.
    .PARAMETER MailFrom
        Optional sender. Default is MailTo.
    #>

    [CmdletBinding(DefaultParameterSetName = 'Local')]
    param(
        [Parameter(ParameterSetName = 'ComputerName', ValueFromPipeline = $true, ValueFromPipelineByPropertyName = $true)]
        [string[]]$ComputerName,

        [Parameter(ParameterSetName = 'CimSession')]
        [Microsoft.Management.Infrastructure.CimSession[]]$CimSession,

        [Parameter(ParameterSetName = 'Session')]
        [System.Management.Automation.Runspaces.PSSession[]]$Session,

        [pscredential]$Credential,
        [System.Management.Automation.Runspaces.AuthenticationMechanism]$Authentication,
        [switch]$UseSSL,
        [int]$Port,
        [System.Management.Automation.Remoting.PSSessionOption]$SessionOption,
        [int]$ThrottleLimit = 32,

        [string]$Path,
        [int]$IntervalSec = 60,
        [int]$RetentionDays = 90,
        [bool]$IncludeNics = $true,
        [bool]$IncludeSmbPerf = $true,
        [switch]$Wait,
        [int]$WaitTimeoutSec = 90,
        [switch]$PassThru,
        [string]$MailTo,
        [string]$SmtpServer,
        [string]$MailFrom
    )

    begin {
        Assert-SMBeatRemotingExclusive -ComputerName $ComputerName -CimSession $CimSession -Session $Session
        $payload = @{
            Path               = $Path
            IntervalSec        = $IntervalSec
            RetentionDays      = $RetentionDays
            IncludeNics        = $IncludeNics
            IncludeSmbPerf     = $IncludeSmbPerf
            IncludeEtwShares   = $true
            IncludeAdminShares = $false
            Wait               = [bool]$Wait
            WaitTimeoutSec     = $WaitTimeoutSec
            PassThru           = [bool]$PassThru
            MailTo             = $MailTo
            SmtpServer         = $SmtpServer
            MailFrom           = $MailFrom
        }
    }

    process {
        if ($PSCmdlet.ParameterSetName -eq 'Local') {
            Start-SMBeatLocal -Payload $payload
            return
        }

        if ($PSCmdlet.ParameterSetName -eq 'CimSession') {
            foreach ($cs in $CimSession) {
                try {
                    $taskPack = Get-SMBeatScheduledTaskInfoSafe -CimSession $cs
                    if (-not $taskPack.Task) {
                        throw 'SMBeat cannot copy files or register a new collector over a CimSession. Use -ComputerName or -Session.'
                    }
                    Start-SMBeatScheduledTaskInternal -CimSession $cs
                    if ($PassThru -or $Wait) {
                        Get-SMBeatStatusInternal -CimSession $cs
                    }
                }
                catch {
                    $record = New-SMBeatErrorRecord -Message $_.Exception.Message -ComputerName $cs.ComputerName -Exception $_.Exception
                    $PSCmdlet.WriteError($record)
                }
            }
            return
        }

        $remoteSb = {
            param($P)
            Import-Module SMBeat -ErrorAction Stop
            $splat = @{
                IntervalSec    = $P.IntervalSec
                RetentionDays  = $P.RetentionDays
                IncludeNics    = $P.IncludeNics
                IncludeSmbPerf = $P.IncludeSmbPerf
                WaitTimeoutSec = $P.WaitTimeoutSec
            }
            if ($P.Path) { $splat['Path'] = $P.Path }
            if ($P.Wait) { $splat['Wait'] = $true }
            if ($P.PassThru) { $splat['PassThru'] = $true }
            if ($P.MailTo) { $splat['MailTo'] = $P.MailTo }
            if ($P.SmtpServer) { $splat['SmtpServer'] = $P.SmtpServer }
            if ($P.MailFrom) { $splat['MailFrom'] = $P.MailFrom }
            Start-SMBeat @splat
        }

        $sessionsToClose = @()
        $targets = $Session
        if ($PSCmdlet.ParameterSetName -eq 'ComputerName') {
            $newParams = @{ ComputerName = $ComputerName; ErrorAction = 'Stop' }
            if ($Credential) { $newParams['Credential'] = $Credential }
            if ($PSBoundParameters.ContainsKey('Authentication')) { $newParams['Authentication'] = $Authentication }
            if ($UseSSL) { $newParams['UseSSL'] = $true }
            if ($Port) { $newParams['Port'] = $Port }
            if ($SessionOption) { $newParams['SessionOption'] = $SessionOption }
            $targets = @(New-PSSession @newParams)
            $sessionsToClose = $targets
        }

        try {
            foreach ($s in $targets) {
                try {
                    Install-SMBeatToSessionIfMissing -Session $s
                    Invoke-Command -Session $s -ScriptBlock $remoteSb -ArgumentList @($payload)
                }
                catch {
                    $record = New-SMBeatErrorRecord -Message $_.Exception.Message -ComputerName $s.ComputerName -Exception $_.Exception
                    $PSCmdlet.WriteError($record)
                }
            }
        }
        finally {
            if ($sessionsToClose.Count -gt 0) {
                Remove-PSSession -Session $sessionsToClose -ErrorAction SilentlyContinue
            }
        }
    }
}

function Stop-SMBeat {
    <#
    .SYNOPSIS
        Signals the watcher to stop and disables the collector task. Samples are kept.
    #>

    [CmdletBinding(DefaultParameterSetName = 'Local')]
    param(
        [Parameter(ParameterSetName = 'ComputerName', ValueFromPipeline = $true, ValueFromPipelineByPropertyName = $true)]
        [string[]]$ComputerName,
        [Parameter(ParameterSetName = 'CimSession')]
        [Microsoft.Management.Infrastructure.CimSession[]]$CimSession,
        [Parameter(ParameterSetName = 'Session')]
        [System.Management.Automation.Runspaces.PSSession[]]$Session,
        [pscredential]$Credential,
        [System.Management.Automation.Runspaces.AuthenticationMechanism]$Authentication,
        [switch]$UseSSL,
        [int]$Port,
        [System.Management.Automation.Remoting.PSSessionOption]$SessionOption,
        [int]$ThrottleLimit = 32,
        [switch]$PassThru
    )

    begin {
        Assert-SMBeatRemotingExclusive -ComputerName $ComputerName -CimSession $CimSession -Session $Session
    }

    process {
        if ($PSCmdlet.ParameterSetName -eq 'Local') {
            Stop-SMBeatScheduledTaskInternal
            if ($PassThru) { Get-SMBeatStatusInternal }
            return
        }
        if ($PSCmdlet.ParameterSetName -eq 'CimSession') {
            foreach ($cs in $CimSession) {
                Stop-SMBeatScheduledTaskInternal -CimSession $cs
                if ($PassThru) { Get-SMBeatStatusInternal -CimSession $cs }
            }
            return
        }

        $sb = {
            param($Pass)
            Import-Module SMBeat -ErrorAction Stop
            if ($Pass) { Stop-SMBeat -PassThru } else { Stop-SMBeat }
        }
        $invoke = @{ ScriptBlock = $sb; ArgumentList = @([bool]$PassThru); ThrottleLimit = $ThrottleLimit }
        if ($PSCmdlet.ParameterSetName -eq 'Session') { $invoke['Session'] = $Session }
        else {
            $invoke['ComputerName'] = $ComputerName
            if ($Credential) { $invoke['Credential'] = $Credential }
            if ($PSBoundParameters.ContainsKey('Authentication')) { $invoke['Authentication'] = $Authentication }
            if ($UseSSL) { $invoke['UseSSL'] = $true }
            if ($Port) { $invoke['Port'] = $Port }
            if ($SessionOption) { $invoke['SessionOption'] = $SessionOption }
        }
        Invoke-SMBeatRemoteCommand @invoke
    }
}