Merge pull request #15 from dahlbyk/feature/environment-variable-scope

Allow overriding environment variable scope
dahlbyk · Jul 13, 2018 · 8a97b01 · 8a97b01
2 parents 853e004 + 014fae9
commit 8a97b01
Show file tree

Hide file tree

Showing 3 changed files with 95 additions and 58 deletions.
diff --git a/src/Agent.ps1 b/src/Agent.ps1
@@ -1,58 +1,66 @@
 # Retrieve the current SSH agent PID (or zero). Can be used to determine if there
 # is a running agent.
 function Get-SshAgent() {
-  if ($env:GIT_SSH -imatch 'plink') {
-      $pageantPid = Get-Process | Where-Object { $_.Name -eq 'pageant' } | Select-Object -ExpandProperty Id -First 1
-      if ($null -ne $pageantPid) { return $pageantPid }
-  }
-  else {
-      $agentPid = $Env:SSH_AGENT_PID
-      if ($agentPid) {
-          $sshAgentProcess = Get-Process | Where-Object { ($_.Id -eq $agentPid) -and ($_.Name -eq 'ssh-agent') }
-          if ($null -ne $sshAgentProcess) {
-              return $agentPid
-          }
-          else {
-              setenv 'SSH_AGENT_PID' $null
-              setenv 'SSH_AUTH_SOCK' $null
-          }
-      }
-  }
-
-  return 0
+    if ($env:GIT_SSH -imatch 'plink') {
+        $pageantPid = Get-Process | Where-Object { $_.Name -eq 'pageant' } | Select-Object -ExpandProperty Id -First 1
+        if ($null -ne $pageantPid) { return $pageantPid }
+    }
+    elseif ($native = Get-NativeSshAgent) {
+        return $native
+    }
+    else {
+        $agentPid = $Env:SSH_AGENT_PID
+        if ($agentPid) {
+            $sshAgentProcess = Get-Process | Where-Object { ($_.Id -eq $agentPid) -and ($_.Name -eq 'ssh-agent') }
+            if ($null -ne $sshAgentProcess) {
+                return $agentPid
+            }
+            else {
+                setenv 'SSH_AGENT_PID' $null
+                setenv 'SSH_AUTH_SOCK' $null
+            }
+        }
+    }
+
+    return 0
 }
 
 # Attempt to guess $program's location. For ssh-agent/ssh-add.
 function Find-Ssh($program = 'ssh-agent') {
-  Write-Verbose "$program not in path. Trying to guess location."
-  $gitItem = Get-Command git -CommandType Application -Erroraction SilentlyContinue | Get-Item
-  if ($null -eq $gitItem) {
-      Write-Warning 'git not in path'
-      return
-  }
-
-  $sshLocation = join-path $gitItem.directory.parent.fullname bin/$program
-  if (get-command $sshLocation -Erroraction SilentlyContinue) {
-      return $sshLocation
-  }
-
-  $sshLocation = join-path $gitItem.directory.parent.fullname usr/bin/$program
-  if (get-command $sshLocation -Erroraction SilentlyContinue) {
-      return $sshLocation
-  }
+    Write-Verbose "$program not in path. Trying to guess location."
+    $gitItem = Get-Command git -CommandType Application -Erroraction SilentlyContinue | Get-Item
+    if ($null -eq $gitItem) {
+        Write-Warning 'git not in path'
+        return
+    }
+
+    $sshLocation = join-path $gitItem.directory.parent.fullname bin/$program
+    if (get-command $sshLocation -Erroraction SilentlyContinue) {
+        return $sshLocation
+    }
+
+    $sshLocation = join-path $gitItem.directory.parent.fullname usr/bin/$program
+    if (get-command $sshLocation -Erroraction SilentlyContinue) {
+        return $sshLocation
+    }
 }
 
 # Loosely based on bash script from http://help.github.com/ssh-key-passphrases/
 function Start-SshAgent {
     param(
-       [Parameter(Position = 0)]
-       [ValidateSet("Automatic", "Boot", "Disabled", "Manual", "System")]
-       [string]
-       $StartupType = "Manual",
-
-       [Parameter()]
-       [switch]
-       $Quiet
+        [Parameter(Position = 0)]
+        [ValidateSet("Automatic", "Boot", "Disabled", "Manual", "System")]
+        [string]
+        $StartupType = "Manual",
+
+        [Parameter()]
+        [switch]
+        $Quiet,
+
+        [Parameter()]
+        [ValidateSet("Process", "User")]
+        [string]
+        $Scope = "Process"
     )
 
     # If we're using the win10 native ssh client,
@@ -97,7 +105,7 @@ function Start-SshAgent {
 
         & $sshAgent | ForEach-Object {
             if ($_ -match '(?<key>[^=]+)=(?<value>[^;]+);') {
-                setenv $Matches['key'] $Matches['value']
+                setenv $Matches['key'] $Matches['value'] $Scope
             }
         }
     }
@@ -108,15 +116,20 @@ function Start-SshAgent {
 
 # Stop a running SSH agent
 function Stop-SshAgent() {
-  [int]$agentPid = Get-SshAgent
-  if ($agentPid -gt 0) {
-      # Stop agent process
-      $proc = Get-Process -Id $agentPid -ErrorAction SilentlyContinue
-      if ($null -ne $proc) {
-          Stop-Process $agentPid
-      }
-
-      setenv 'SSH_AGENT_PID' $null
-      setenv 'SSH_AUTH_SOCK' $null
-  }
+    if ($nativeAgent = Get-NativeSshAgent) {
+        Stop-Service $nativeAgent.Name
+        return
+    }
+
+    [int]$agentPid = Get-SshAgent
+    if ($agentPid -gt 0) {
+        # Stop agent process
+        $proc = Get-Process -Id $agentPid -ErrorAction SilentlyContinue
+        if ($null -ne $proc) {
+            Stop-Process $agentPid
+        }
+
+        setenv 'SSH_AGENT_PID' $null
+        setenv 'SSH_AUTH_SOCK' $null
+    }
 }
diff --git a/src/Utils.ps1 b/src/Utils.ps1
@@ -1,7 +1,21 @@
 $ModuleBasePath = "$PSScriptRoot\.."
 
-function setenv($key, $value) {
-    [void][Environment]::SetEnvironmentVariable($key, $value)
+function setenv {
+    param(
+        [Parameter()]
+        [string]
+        $key,
+
+        [Parameter()]
+        [string]
+        $value,
+
+        [Parameter()]
+        [ValidateSet("Process", "User")]
+        [string]
+        $Scope = "Process"
+    )
+    [void][Environment]::SetEnvironmentVariable($key, $value, $Scope)
     Set-TempEnv $key $value
 }
 

diff --git a/test/Ssh.Tests.ps1 b/test/Ssh.Tests.ps1
@@ -63,7 +63,11 @@ Describe 'SSH Function Tests' {
             $result | Should Not Be $null
             $result.Name | Should -Be "ssh-agent"
         }
-
+        It "Finds the service when Get-SshAgent is called" {
+            $result = Get-SshAgent
+            $result | Should Not Be $null
+            $result.Name | Should -Be "ssh-agent"
+        }
         It "Starts the service when stopped and user is admin" {
             Mock Test-Administrator { return $true }
             $result = Start-NativeSshAgent -Quiet
@@ -157,6 +161,12 @@ Describe 'SSH Function Tests' {
                 "$args" -eq 'config --global core.sshCommand "C:/Windows/System32/OpenSSH/ssh.exe"'
             } -Scope It
         }
+        It "Stops the service" {
+            $service.Status = "Running"
+            Mock Stop-Service { $service.Status = "Stopped" } -ParameterFilter { $Name -eq "ssh-agent" }
+            Stop-SshAgent
+            $service.Status | Should -Be "Stopped"
+        }
     }
 
 }