completers/pssuspend_completer/pssuspend_completer.ps1

# pssuspend.exe tab completion for PowerShell
# Native completer for PsSuspend with static switches, remote placeholders, and local process hints.

Set-StrictMode -Version 2.0

if (-not (Get-Variable -Name PsSuspendCompletionCatalog -Scope Script -ErrorAction Ignore)) {
    $script:PsSuspendCompletionCatalog = @{
        SwitchOrder             = @('-r', '-u', '-p', '-nobanner', '-accepteula', '-?', '/?', '--help')
        SwitchInfo              = @{
            '-r'          = 'Resume a suspended process.'
            '-u'          = 'Optional user name for remote login.'
            '-p'          = 'Optional password for the remote login.'
            '-nobanner'   = 'Do not display the startup banner and copyright message.'
            '-accepteula' = 'Suppress the Sysinternals EULA dialog (accepts the license).'
            '-?'          = 'Display pssuspend help.'
            '/?'          = 'Display pssuspend help.'
            '--help'      = 'Display pssuspend help.'
        }
        ProcessEntries          = @()
        ProcessCacheUpdated     = $null
        ProcessCacheTtlSeconds  = 2
    }
}

function New-PsSuspendCompletionResult {
    param(
        [string]$CompletionText,
        [string]$ResultType,
        [string]$ToolTip,
        [string]$ListItemText
    )

    if ([string]::IsNullOrWhiteSpace($ToolTip)) {
        $ToolTip = $CompletionText
    }

    if ([string]::IsNullOrEmpty($ListItemText)) {
        $ListItemText = $CompletionText
    }

    [System.Management.Automation.CompletionResult]::new(
        $CompletionText,
        $ListItemText,
        $ResultType,
        $ToolTip
    )
}

function ConvertFrom-PsSuspendTypedWord {
    # Splits the word typed so far into its value and the opening quote the
    # user typed ('' when bare), undoing that quote style's escapes.
    param([string]$Text)

    $quote = ''
    if ($Text.StartsWith("'") -or $Text.StartsWith('"')) {
        $quote = $Text.Substring(0, 1)
        $Text = $Text.Substring(1)
        if ($Text.EndsWith($quote)) {
            $Text = $Text.Substring(0, $Text.Length - 1)
        }

        $Text = if ($quote -eq "'") { $Text.Replace("''", "'") } else { $Text -replace '`(.)', '$1' }
    }

    [pscustomobject]@{ Value = $Text; Quote = $quote }
}

function ConvertTo-PsSuspendArgument {
    # Renders a value as one PowerShell argument: bare when safe and no quote
    # was typed, otherwise in the typed quote style (single by default).
    param(
        [string]$Value,
        [string]$Quote
    )

    if (-not $Quote) {
        if ($Value -notmatch '[\s{}();,|&<>''"`$]' -and $Value -notmatch '^[@#]') {
            return $Value
        }

        $Quote = "'"
    }

    if ($Quote -eq "'") {
        return "'" + $Value.Replace("'", "''") + "'"
    }

    '"' + ($Value -replace '([`"$])', '`$1') + '"'
}

function Update-PsSuspendProcessCache {
    $lastUpdated = $script:PsSuspendCompletionCatalog.ProcessCacheUpdated
    if ($null -ne $lastUpdated) {
        $cacheAge = (Get-Date) - $lastUpdated
        if ($cacheAge.TotalSeconds -lt $script:PsSuspendCompletionCatalog.ProcessCacheTtlSeconds) {
            return
        }
    }

    $nameSet = [System.Collections.Generic.HashSet[string]]::new([System.StringComparer]::OrdinalIgnoreCase)
    $idSet = [System.Collections.Generic.HashSet[int]]::new()
    $nameEntries = [System.Collections.Generic.List[object]]::new()
    $idEntries = [System.Collections.Generic.List[object]]::new()

    foreach ($process in @(Get-Process -ErrorAction SilentlyContinue)) {
        if ($process.ProcessName -and $nameSet.Add($process.ProcessName)) {
            $nameEntries.Add([pscustomobject]@{
                    CompletionText = $process.ProcessName
                    ResultType     = 'ParameterValue'
                    ToolTip        = "Process name $($process.ProcessName)"
                    SortKey        = $process.ProcessName
                })
        }

        if ($idSet.Add($process.Id)) {
            $idEntries.Add([pscustomobject]@{
                    CompletionText = [string]$process.Id
                    ResultType     = 'ParameterValue'
                    ToolTip        = "Process ID $($process.Id) ($($process.ProcessName))"
                    SortKey        = $process.Id
                })
        }
    }

    # Names first (alphabetical), then PIDs in numeric order, so tab-cycling reaches names before hundreds of numbers.
    $script:PsSuspendCompletionCatalog.ProcessEntries = @(
        @($nameEntries | Sort-Object -Property SortKey) +
        @($idEntries | Sort-Object -Property SortKey)
    )
    $script:PsSuspendCompletionCatalog.ProcessCacheUpdated = Get-Date
}

function New-PsSuspendLiteralValueResults {
    param(
        [string]$CurrentValue,
        [string]$Placeholder,
        [string]$ToolTip
    )

    if ([string]::IsNullOrWhiteSpace($CurrentValue)) {
        return @(
            New-PsSuspendCompletionResult -CompletionText $Placeholder -ResultType 'ParameterValue' -ToolTip $ToolTip
        )
    }

    @(
        New-PsSuspendCompletionResult -CompletionText $CurrentValue -ResultType 'ParameterValue' -ToolTip $ToolTip
    )
}

function Get-PsSuspendProcessCompletions {
    param([string]$CurrentWord)

    Update-PsSuspendProcessCache

    $typed = ConvertFrom-PsSuspendTypedWord -Text $CurrentWord
    $typedValue = $typed.Value
    $results = $script:PsSuspendCompletionCatalog.ProcessEntries |
        Where-Object {
            [string]::IsNullOrWhiteSpace($typedValue) -or $_.CompletionText.StartsWith($typedValue, [System.StringComparison]::OrdinalIgnoreCase)
        } |
        ForEach-Object {
            New-PsSuspendCompletionResult -CompletionText (ConvertTo-PsSuspendArgument -Value $_.CompletionText -Quote $typed.Quote) -ListItemText $_.CompletionText -ResultType $_.ResultType -ToolTip $_.ToolTip
        }

    if (@($results).Count -gt 0) {
        return @($results)
    }

    @(New-PsSuspendLiteralValueResults -CurrentValue $CurrentWord -Placeholder '<process-or-pid>' -ToolTip 'Process name or PID.')
}

function Get-PsSuspendCommandState {
    param([object[]]$TokensBeforeCurrent)

    $TokensBeforeCurrent = @($TokensBeforeCurrent)

    $usedSwitchLookup = @{}
    $valueContext = $null
    $remoteTarget = $null
    $remoteUser = $null
    $remotePassword = $null
    $processTarget = $null
    $helpRequested = $false

    for ($index = 0; $index -lt $TokensBeforeCurrent.Count; $index++) {
        $token = [string]$TokensBeforeCurrent[$index]
        if ([string]::IsNullOrWhiteSpace($token)) {
            continue
        }

        $lookup = $token.ToLowerInvariant()

        if ($lookup -in @('-?', '/?', '--help')) {
            $helpRequested = $true
            $usedSwitchLookup[$lookup] = $true
            continue
        }

        if ($lookup -in @('-u', '-p')) {
            $usedSwitchLookup[$lookup] = $true
            if ($index -eq ($TokensBeforeCurrent.Count - 1)) {
                $valueContext = $lookup
                break
            }

            if ($lookup -eq '-u') {
                $remoteUser = [string]$TokensBeforeCurrent[$index + 1]
            } else {
                $remotePassword = [string]$TokensBeforeCurrent[$index + 1]
            }

            $index++
            continue
        }

        if ($lookup.StartsWith('-')) {
            $usedSwitchLookup[$lookup] = $true
            continue
        }

        if (-not $remoteTarget -and $token.StartsWith('\\')) {
            $remoteTarget = $token
            continue
        }

        if (-not $processTarget) {
            $processTarget = $token
        }
    }

    [pscustomobject]@{
        UsedSwitchLookup = $usedSwitchLookup
        ValueContext     = $valueContext
        RemoteTarget     = $remoteTarget
        RemoteUser       = $remoteUser
        RemotePassword   = $remotePassword
        ProcessTarget    = $processTarget
        HelpRequested    = $helpRequested
    }
}

function Get-PsSuspendSwitchCompletions {
    param(
        [string]$CurrentWord,
        [pscustomobject]$State,
        [bool]$NoArgumentsYet
    )

    $prefix = if ([string]::IsNullOrWhiteSpace($CurrentWord)) { '' } else { $CurrentWord }
    $results = [System.Collections.Generic.List[object]]::new()

    foreach ($token in $script:PsSuspendCompletionCatalog.SwitchOrder) {
        if (-not [string]::IsNullOrWhiteSpace($prefix) -and
            -not $token.StartsWith($prefix, [System.StringComparison]::OrdinalIgnoreCase)) {
            continue
        }

        if ($token -in @('-?', '/?', '--help')) {
            if (-not $NoArgumentsYet) {
                continue
            }
        } elseif ($State.UsedSwitchLookup.ContainsKey($token)) {
            continue
        }

        $results.Add((New-PsSuspendCompletionResult -CompletionText $token -ResultType 'ParameterName' -ToolTip $script:PsSuspendCompletionCatalog.SwitchInfo[$token]))
    }

    @($results.ToArray())
}

function Complete-PsSuspend {
    param(
        [string]$wordToComplete,
        [System.Management.Automation.Language.CommandAst]$commandAst,
        [int]$cursorPosition
    )

    # Read the words from the parser: an unterminated quote is one element
    # (and one $wordToComplete), so a quoted multi-word name stays one word.
    $currentWord = $wordToComplete
    $tokensBeforeCurrent = @(
        $commandAst.CommandElements |
            Select-Object -Skip 1 |
            Where-Object { $_.Extent.EndOffset -lt $cursorPosition } |
            ForEach-Object { $_.Extent.Text }
    )

    $state = Get-PsSuspendCommandState -TokensBeforeCurrent $tokensBeforeCurrent

    if ($state.HelpRequested) {
        return @(
            New-PsSuspendCompletionResult -CompletionText '-?' -ResultType 'ParameterName' -ToolTip 'Display pssuspend help.'
        )
    }

    switch ($state.ValueContext) {
        '-u' {
            return @(New-PsSuspendLiteralValueResults -CurrentValue $currentWord -Placeholder '<username>' -ToolTip 'Remote user name.')
        }
        '-p' {
            return @(New-PsSuspendLiteralValueResults -CurrentValue $currentWord -Placeholder '<password>' -ToolTip 'Remote password.')
        }
    }

    if (-not [string]::IsNullOrWhiteSpace($currentWord) -and $currentWord.StartsWith('\\')) {
        return @(
            New-PsSuspendCompletionResult -CompletionText '\\computer' -ResultType 'ParameterValue' -ToolTip 'Remote computer placeholder.'
        )
    }

    if (-not [string]::IsNullOrWhiteSpace($currentWord) -and ($currentWord.StartsWith('-') -or $currentWord.StartsWith('/'))) {
        return @(Get-PsSuspendSwitchCompletions -CurrentWord $currentWord -State $state -NoArgumentsYet:($tokensBeforeCurrent.Count -eq 0))
    }

    if ([string]::IsNullOrWhiteSpace($currentWord)) {
        $results = [System.Collections.Generic.List[object]]::new()
        foreach ($item in @(Get-PsSuspendSwitchCompletions -CurrentWord $currentWord -State $state -NoArgumentsYet:($tokensBeforeCurrent.Count -eq 0))) {
            $results.Add($item)
        }

        if (-not $state.ProcessTarget) {
            if ($state.RemoteTarget) {
                $results.Add((New-PsSuspendCompletionResult -CompletionText '<process-name>' -ResultType 'ParameterValue' -ToolTip 'Remote process name.'))
                $results.Add((New-PsSuspendCompletionResult -CompletionText '<pid>' -ResultType 'ParameterValue' -ToolTip 'Remote process ID.'))
            } else {
                $results.Add((New-PsSuspendCompletionResult -CompletionText '\\computer' -ResultType 'ParameterValue' -ToolTip 'Remote computer placeholder.'))
                foreach ($item in @(Get-PsSuspendProcessCompletions -CurrentWord '')) {
                    $results.Add($item)
                }
            }
        }

        return @($results.ToArray())
    }

    if (-not $state.ProcessTarget) {
        if ($state.RemoteTarget) {
            return @(
                New-PsSuspendLiteralValueResults -CurrentValue $currentWord -Placeholder '<process-or-pid>' -ToolTip 'Remote process name or PID.'
            )
        }

        return @(Get-PsSuspendProcessCompletions -CurrentWord $currentWord)
    }

    @()
}

Register-ArgumentCompleter -Native -CommandName @('pssuspend', 'pssuspend.exe') -ScriptBlock {
    param($wordToComplete, $commandAst, $cursorPosition)

    Complete-PsSuspend -wordToComplete $wordToComplete -commandAst $commandAst -cursorPosition $cursorPosition
}