Add Get-PromptConnectionInfo as DefaultPromptPrefix

dahlbyk · Jul 9, 2018 · bd3c403 · bd3c403
1 parent 3dd9fdc
commit bd3c403
Show file tree

Hide file tree

Showing 5 changed files with 48 additions and 1 deletion.
diff --git a/src/PoshGitTypes.ps1 b/src/PoshGitTypes.ps1
@@ -273,7 +273,7 @@ class PoshGitPromptSettings {
     [string]$DescribeStyle = ''
     [psobject]$WindowTitle = {param($GitStatus, [bool]$IsAdmin) "$(if ($IsAdmin) {'Admin: '})$(if ($GitStatus) {"$($GitStatus.RepoName) [$($GitStatus.Branch)]"} else {Get-PromptPath}) ~ PowerShell $($PSVersionTable.PSVersion) $([IntPtr]::Size * 8)-bit ($PID)"}
 
-    [PoshGitTextSpan]$DefaultPromptPrefix       = ''
+    [PoshGitTextSpan]$DefaultPromptPrefix       = '$(Get-PromptConnectionInfo)'
     [PoshGitTextSpan]$DefaultPromptPath         = '$(Get-PromptPath)'
     [PoshGitTextSpan]$DefaultPromptBeforeSuffix = ''
     [PoshGitTextSpan]$DefaultPromptDebug        = [PoshGitTextSpan]::new(' [DBG]:', [ConsoleColor]::Magenta)

diff --git a/src/Utils.ps1 b/src/Utils.ps1
@@ -301,6 +301,20 @@ function Get-PromptPath {
     return $currentPath
 }
 
+$sshType = 'System.Management.Automation.Runspaces.SSHConnectionInfo' -as [Type]
+$sshUserName = Get-Runspace | ForEach-Object ConnectionInfo |
+  Where-Object { $sshType -and $_ -is $sshType } | ForEach-Object UserName
+
+function Get-PromptConnectionInfo {
+    if (Test-Path Env:SSH_CONNECTION) {
+        if ($sshUserName -and $sshUserName -ne [System.Environment]::UserName) {
+            "[$sshUserName@$([System.Environment]::MachineName)]: "
+        } else {
+            "[$([System.Environment]::MachineName)]: "
+        }
+    }
+}
+
 function Get-PSModulePath {
     $modulePaths = $Env:PSModulePath -split ';'
     $modulePaths

diff --git a/src/posh-git.psd1 b/src/posh-git.psd1
@@ -29,6 +29,7 @@ FunctionsToExport = @(
     'Get-GitBranchStatusColor',
     'Get-GitStatus',
     'Get-GitDirectory',
+    'Get-PromptConnectionInfo',
     'Get-PromptPath',
     'Update-AllBranches',
     'Write-GitStatus',

diff --git a/src/posh-git.psm1 b/src/posh-git.psm1
@@ -194,6 +194,7 @@ $exportModuleMemberParams = @{
         'Get-GitBranchStatusColor',
         'Get-GitDirectory',
         'Get-GitStatus',
+        'Get-PromptConnectionInfo',
         'Get-PromptPath',
         'Update-AllBranches',
         'Write-GitStatus',

diff --git a/test/Utils.Tests.ps1 b/test/Utils.Tests.ps1
@@ -101,6 +101,37 @@ New-Alias pscore C:\Users\Keith\GitHub\rkeithhill\PowerShell\src\powershell-win-
         }
     }
 
+    Context 'Get-PromptConnectionInfo' {
+        BeforeEach {
+            if (Test-Path Env:SSH_CONNECTION) {
+                [System.Diagnostics.CodeAnalysis.SuppressMessageAttribute('PSUseDeclaredVarsMoreThanAssigments', '')]
+                $ssh_connection = $Env:SSH_CONNECTION
+
+                Remove-Item Env:SSH_CONNECTION
+            }
+        }
+        AfterEach {
+            if ($ssh_connection) {
+                Set-Item Env:SSH_CONNECTION $ssh_connection
+            } elseif (Test-Path Env:SSH_CONNECTION) {
+                Remove-Item Env:SSH_CONNECTION
+            }
+        }
+        It 'Returns null if Env:SSH_CONNECTION is not set' {
+            Get-PromptConnectionInfo | Should BeExactly $null
+        }
+        It 'Returns null if Env:SSH_CONNECTION is empty' {
+            Set-Item Env:SSH_CONNECTION ''
+
+            Get-PromptConnectionInfo | Should BeExactly $null
+        }
+        It 'Returns "[hostname]: " if Env:SSH_CONNECTION is set' {
+            Set-Item Env:SSH_CONNECTION 'test'
+
+            Get-PromptConnectionInfo | Should BeExactly "[$([System.Environment]::MachineName)]: "
+        }
+    }
+
     Context 'Test-PoshGitImportedInScript Tests' {
         BeforeEach {
             [System.Diagnostics.CodeAnalysis.SuppressMessageAttribute('PSUseDeclaredVarsMoreThanAssigments', '')]