Public/Install-SMBeat.ps1

#Requires -Version 5.1

function Install-SMBeat {
    <#
    .SYNOPSIS
        Copies the SMBeat module into WindowsPowerShell\Modules (local or remote).
    #>

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

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

    begin {
        Assert-SMBeatRemotingExclusive -ComputerName $ComputerName -CimSession $null -Session $Session
        $sourceRoot = $script:SMBeatModuleRoot
        $copyNames = @(Get-SMBeatModuleCopyNames)
    }

    process {
        if ($PSCmdlet.ParameterSetName -eq 'Local') {
            $dest = Copy-SMBeatModuleToDestination -Destination (Get-SMBeatInstallDestination) -SourceRoot $sourceRoot
            Get-Item -LiteralPath $dest
            return
        }

        $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 {
                    $dest = Invoke-Command -Session $s -ScriptBlock {
                        $path = Join-Path $env:ProgramFiles 'WindowsPowerShell\Modules\SMBeat'
                        New-Item -ItemType Directory -Path $path -Force | Out-Null
                        $path
                    }
                    foreach ($name in $copyNames) {
                        $from = Join-Path $sourceRoot $name
                        if (Test-Path -LiteralPath $from) {
                            Copy-Item -LiteralPath $from -Destination $dest -ToSession $s -Recurse -Force
                        }
                    }
                    [PSCustomObject]@{
                        ComputerName = $s.ComputerName
                        Path         = $dest
                        Success      = $true
                    }
                }
                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 Uninstall-SMBeat {
    <#
    .SYNOPSIS
        Removes the SMBeat collector and its data. The module files stay installed.
    .DESCRIPTION
        Stops and unregisters the scheduled task, tears down the ETW session, and
        deletes the data root (default %ProgramData%\SMBeat and any Path from config).
        Does not delete the module under WindowsPowerShell\Modules.
    .PARAMETER KeepData
        Remove the scheduled task only. Sample files stay.
    .PARAMETER Force
        Do not prompt before deleting data.
    #>

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

        [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,
        [switch]$Force,
        [switch]$KeepData
    )

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

    process {
        if ($PSCmdlet.ParameterSetName -eq 'Local') {
            $roots = @(Get-SMBeatDataRootsToRemove -Path $Path)
            $label = $roots -join ', '
            $action = 'Remove SMBeat collector and data (module stays)'
            if ($KeepData) {
                $action = 'Remove SMBeat collector task (samples stay)'
            }
            if (-not $Force -and -not $PSCmdlet.ShouldProcess($label, $action)) {
                return
            }
            $result = Uninstall-SMBeatLocal -Path $Path -KeepData:$KeepData
            if ($KeepData) {
                Write-Host ('SMBeat collector task removed. Samples left under {0}' -f $label)
            }
            else {
                Write-Host ('SMBeat collector and data removed. Module left at {0}' -f $result.ModulePath)
            }
            return $result
        }

        $payload = @{
            Path     = $Path
            Force    = $true
            KeepData = [bool]$KeepData
        }
        $sb = {
            param($P)
            Import-Module SMBeat -ErrorAction Stop
            $splat = @{ Force = $true }
            if ($P.Path) { $splat['Path'] = $P.Path }
            if ($P.KeepData) { $splat['KeepData'] = $true }
            Uninstall-SMBeat @splat
        }

        $invoke = @{
            ScriptBlock   = $sb
            ArgumentList  = @($payload)
            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
    }
}

function Install-SMBeatIfMissing {
    if (Test-SMBeatModuleInstalled) {
        return
    }

    $dest = Get-SMBeatInstallDestination
    if ([string]::IsNullOrWhiteSpace($dest)) {
        throw 'SMBeat install destination is empty.'
    }

    Copy-SMBeatModuleToDestination -Destination $dest -SourceRoot $script:SMBeatModuleRoot | Out-Null
}

function Install-SMBeatToSessionIfMissing {
    param(
        $Session
    )

    $present = Invoke-Command -Session $Session -ScriptBlock {
        Test-Path -LiteralPath (Join-Path $env:ProgramFiles 'WindowsPowerShell\Modules\SMBeat\SMBeat.psd1')
    }
    if ($present) {
        return
    }

    $sourceRoot = $script:SMBeatModuleRoot
    $copyNames = @(Get-SMBeatModuleCopyNames)
    $dest = Invoke-Command -Session $Session -ScriptBlock {
        $path = Join-Path $env:ProgramFiles 'WindowsPowerShell\Modules\SMBeat'
        New-Item -ItemType Directory -Path $path -Force | Out-Null
        $path
    }
    foreach ($name in $copyNames) {
        $from = Join-Path $sourceRoot $name
        if (Test-Path -LiteralPath $from) {
            Copy-Item -LiteralPath $from -Destination $dest -ToSession $Session -Recurse -Force
        }
    }
}