SMBeat

0.3.0

Private/Watcher.ps1

#Requires -Version 5.1

function Get-SMBeatWatcherMutexName {
    'Global\SMBeat-Collector'
}

function Test-SMBeatStopRequested {
    param(
        [Parameter(Mandatory = $true)]
        [string]$DataRoot
    )

    Test-Path -LiteralPath (Get-SMBeatStopRequestPath -DataRoot $DataRoot)
}

function Test-SMBeatFlushRequested {
    param(
        [Parameter(Mandatory = $true)]
        [string]$DataRoot
    )

    Test-Path -LiteralPath (Get-SMBeatFlushRequestPath -DataRoot $DataRoot)
}

function Request-SMBeatWatcherStop {
    param(
        [Parameter(Mandatory = $true)]
        [string]$DataRoot
    )

    $path = Get-SMBeatStopRequestPath -DataRoot $DataRoot
    $dir = Split-Path -Parent $path
    if (-not (Test-Path -LiteralPath $dir)) {
        New-Item -ItemType Directory -Path $dir -Force | Out-Null
    }
    [System.IO.File]::WriteAllText($path, (ConvertTo-SMBeatUtcString -DateTime (Get-SMBeatUtcNow)), $script:SMBeatUtf8NoBom)
}

function Request-SMBeatWatcherFlush {
    param(
        [Parameter(Mandatory = $true)]
        [string]$DataRoot
    )

    $path = Get-SMBeatFlushRequestPath -DataRoot $DataRoot
    $dir = Split-Path -Parent $path
    if (-not (Test-Path -LiteralPath $dir)) {
        New-Item -ItemType Directory -Path $dir -Force | Out-Null
    }
    [System.IO.File]::WriteAllText($path, (ConvertTo-SMBeatUtcString -DateTime (Get-SMBeatUtcNow)), $script:SMBeatUtf8NoBom)
}

function Clear-SMBeatWatcherRequest {
    param(
        [Parameter(Mandatory = $true)]
        [string]$Path
    )

    if (Test-Path -LiteralPath $Path) {
        Remove-Item -LiteralPath $Path -Force -ErrorAction SilentlyContinue
    }
}

function Resolve-SMBeatWatcherDataRoot {
    param(
        [string]$Path
    )

    $root = Get-SMBeatDataRoot -Path $Path
    $configFile = Get-SMBeatConfigPath -DataRoot $root
    $config = Read-SMBeatConfig -DataRoot $root
    if ((Test-Path -LiteralPath $configFile) -and $config.Path -and -not [string]::IsNullOrWhiteSpace([string]$config.Path)) {
        $pointed = [string]$config.Path
        if ($pointed -ne $root) {
            $root = $pointed
            $config = Read-SMBeatConfig -DataRoot $root
        }
    }

    [PSCustomObject]@{
        DataRoot = $root
        Config   = $config
    }
}

function Test-SMBeatWatcherAlive {
    param(
        [string]$Path
    )

    $pack = Resolve-SMBeatWatcherDataRoot -Path $Path
    $hb = Get-SMBeatHeartbeatObject -DataRoot $pack.DataRoot
    if ($null -eq $hb) {
        return $false
    }

    $processId = 0
    if ($hb.Pid) {
        $processId = [int]$hb.Pid
    }
    if ($processId -le 0) {
        return $false
    }

    $proc = Get-Process -Id $processId -ErrorAction SilentlyContinue
    if ($null -eq $proc) {
        return $false
    }

    $interval = [int]$hb.IntervalSec
    if ($interval -le 0) {
        $interval = $script:SMBeatDefaultIntervalSec
    }
    $last = $null
    if ($hb.LastSuccessUtc) {
        $last = ConvertFrom-SMBeatUtcString -Value ([string]$hb.LastSuccessUtc)
    }
    if ($null -eq $last) {
        return $false
    }

    $age = (Get-SMBeatUtcNow) - $last
    return ($age.TotalSeconds -le (2 * $interval))
}

function Invoke-SMBeatWatcherFlush {
    param(
        [Parameter(Mandatory = $true)]
        [string]$DataRoot,
        [Parameter(Mandatory = $true)]
        [hashtable]$Tracker,
        [Parameter(Mandatory = $true)]
        [hashtable]$NicPrev,
        [Parameter(Mandatory = $true)]
        [int]$IntervalSec,
        [Parameter(Mandatory = $true)]
        [int]$RetentionDays,
        [bool]$IncludeNics = $true,
        [switch]$IncludeAdminShares
    )

    $utc = Get-SMBeatUtcNow
    $label = $env:COMPUTERNAME
    $deltas = @(Get-SMBeatTcp445PendingDeltas -Tracker $Tracker -Reset)
    Remove-SMBeatTcp445Stale -Tracker $Tracker
    $records = New-Object System.Collections.Generic.List[object]
    foreach ($row in @(ConvertTo-SMBeatTcp445FlushRecords -Deltas $deltas -Server $label -Utc $utc -IntervalSec $IntervalSec)) {
        $records.Add($row) | Out-Null
    }

    if ($IncludeNics) {
        foreach ($row in @(Get-SMBeatNicDeltaRecords -NicPrev $NicPrev -Server $label -Utc $utc -IntervalSec $IntervalSec)) {
            $records.Add($row) | Out-Null
        }
    }

    Update-SMBeatEtwShareHook
    $shareRows = @(Get-SMBeatEtwSharePendingDeltas -Server $label -Utc $utc -IntervalSec $IntervalSec -IncludeAdminShares:$IncludeAdminShares)
    foreach ($row in $shareRows) {
        $records.Add($row) | Out-Null
    }

    if ($records.Count -gt 0) {
        Add-SMBeatSampleLines -DataRoot $DataRoot -Utc $utc -Records $records.ToArray()
    }
    Invoke-SMBeatRetention -DataRoot $DataRoot -RetentionDays $RetentionDays
    $mode = Get-SMBeatCollectorMode
    Save-SMBeatHeartbeat -DataRoot $DataRoot -Utc $utc -IntervalSec $IntervalSec -ComputerName $label -LastError '' -Success -EmittedCount $records.Count -ProcessId $PID -Mode $mode
    Save-SMBeatWatcherState -DataRoot $DataRoot -Utc $utc -Mode $mode

    [PSCustomObject]@{
        ComputerName = $label
        Utc          = (ConvertTo-SMBeatUtcString -DateTime $utc)
        Path         = $DataRoot
        Emitted      = $records.Count
        Success      = $true
    }
}

function Watch-SMBeatCollector {
    [CmdletBinding()]
    param(
        [string]$Path,
        [int]$PollSec
    )

    if ($PollSec -le 0) {
        $PollSec = $script:SMBeatPollSec
    }

    $pack = Resolve-SMBeatWatcherDataRoot -Path $Path
    $dataRoot = $pack.DataRoot
    $config = $pack.Config
    $interval = [int]$config.IntervalSec
    if ($interval -le 0) {
        $interval = $script:SMBeatDefaultIntervalSec
    }
    $retention = [int]$config.RetentionDays
    if ($retention -le 0) {
        $retention = $script:SMBeatDefaultRetentionDays
    }
    $includeNics = $true
    if ($null -ne $config.IncludeNics) {
        $includeNics = [bool]$config.IncludeNics
    }
    $includeEtw = $true
    if ($null -ne $config.IncludeEtwShares) {
        $includeEtw = [bool]$config.IncludeEtwShares
    }
    $includeAdminShares = $false
    if ($null -ne $config.IncludeAdminShares) {
        $includeAdminShares = [bool]$config.IncludeAdminShares
    }

    $mutex = $null
    try {
        $mutex = New-Object System.Threading.Mutex($false, (Get-SMBeatWatcherMutexName))
        if (-not $mutex.WaitOne(0)) {
            Write-SMBeatLog -DataRoot $dataRoot -Message 'Watcher already running (mutex).'
            $mutex.Dispose()
            return
        }
    }
    catch {
        Write-SMBeatLog -DataRoot $dataRoot -Message ('Mutex failed: {0}' -f $_.Exception.Message)
        throw
    }

    try {
        Clear-SMBeatWatcherRequest -Path (Get-SMBeatStopRequestPath -DataRoot $dataRoot)
        Clear-SMBeatWatcherRequest -Path (Get-SMBeatFlushRequestPath -DataRoot $dataRoot)

        $utc = Get-SMBeatUtcNow
        if ($includeEtw) {
            Start-SMBeatEtwShare -DataRoot $dataRoot | Out-Null
        }
        $mode = Get-SMBeatCollectorMode
        Save-SMBeatHeartbeat -DataRoot $dataRoot -Utc $utc -IntervalSec $interval -ComputerName $env:COMPUTERNAME -LastError '' -Success -EmittedCount 0 -ProcessId $PID -Mode $mode

        $tracker = New-SMBeatTcp445Tracker
        $nicPrev = @{}
        $lastFlush = $utc

        while (-not (Test-SMBeatStopRequested -DataRoot $dataRoot)) {
            try {
                $live = @(Get-SMBeatTcp445LiveConnections)
                Update-SMBeatTcp445Tracker -Tracker $tracker -Live $live
                Update-SMBeatEtwShareHook

                $now = Get-SMBeatUtcNow
                $due = ($now - $lastFlush).TotalSeconds -ge $interval
                $flushNow = Test-SMBeatFlushRequested -DataRoot $dataRoot
                if ($due -or $flushNow) {
                    Invoke-SMBeatWatcherFlush -DataRoot $dataRoot -Tracker $tracker -NicPrev $nicPrev -IntervalSec $interval -RetentionDays $retention -IncludeNics $includeNics -IncludeAdminShares:$includeAdminShares | Out-Null
                    $lastFlush = $now
                    if ($flushNow) {
                        Clear-SMBeatWatcherRequest -Path (Get-SMBeatFlushRequestPath -DataRoot $dataRoot)
                    }
                }
            }
            catch {
                $msg = $_.Exception.Message
                Write-SMBeatLog -DataRoot $dataRoot -Message $msg
                Save-SMBeatHeartbeat -DataRoot $dataRoot -Utc (Get-SMBeatUtcNow) -IntervalSec $interval -ComputerName $env:COMPUTERNAME -LastError $msg -EmittedCount 0 -ProcessId $PID -Mode (Get-SMBeatCollectorMode)
            }

            Start-Sleep -Seconds $PollSec
        }

        try {
            Invoke-SMBeatWatcherFlush -DataRoot $dataRoot -Tracker $tracker -NicPrev $nicPrev -IntervalSec $interval -RetentionDays $retention -IncludeNics $includeNics -IncludeAdminShares:$includeAdminShares | Out-Null
        }
        catch {
            Write-SMBeatLog -DataRoot $dataRoot -Message $_.Exception.Message
        }
        Clear-SMBeatWatcherRequest -Path (Get-SMBeatStopRequestPath -DataRoot $dataRoot)
    }
    finally {
        try { Stop-SMBeatEtwShare } catch { }
        try { $mutex.ReleaseMutex() } catch { }
        $mutex.Dispose()
    }
}