From 829e6de049e3b1c482aa9c1d9d325dd845a67952 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C5=A0t=C4=9Bp=C3=A1n=20Pila=C5=99?= <ja@mingan.name>
Date: Wed, 7 Mar 2018 17:37:28 +0100
Subject: [PATCH] Use parsed identity file

https://github.com/pressly/sup/issues/128
---
 ssh.go | 57 ++++++++++++++++++++++++++++++++-------------------------
 sup.go |  7 ++++---
 2 files changed, 36 insertions(+), 28 deletions(-)

diff --git a/ssh.go b/ssh.go
index eb3cefb..43c4546 100644
--- a/ssh.go
+++ b/ssh.go
@@ -21,6 +21,7 @@ type SSHClient struct {
 	sess         *ssh.Session
 	user         string
 	host         string
+	identityFile string
 	remoteStdin  io.WriteCloser
 	remoteStdout io.Reader
 	remoteStderr io.Reader
@@ -80,34 +81,40 @@ var initAuthMethodOnce sync.Once
 var authMethod ssh.AuthMethod
 
 // initAuthMethod initiates SSH authentication method.
-func initAuthMethod() {
-	var signers []ssh.Signer
-
-	// If there's a running SSH Agent, try to use its Private keys.
-	sock, err := net.Dial("unix", os.Getenv("SSH_AUTH_SOCK"))
-	if err == nil {
-		agent := agent.NewClient(sock)
-		signers, _ = agent.Signers()
-	}
-
-	// Try to read user's SSH private keys form the standard paths.
-	files, _ := filepath.Glob(os.Getenv("HOME") + "/.ssh/id_*")
-	for _, file := range files {
-		if strings.HasSuffix(file, ".pub") {
-			continue // Skip public keys.
-		}
-		data, err := ioutil.ReadFile(file)
-		if err != nil {
-			continue
+func initAuthMethod(identityFilePath string) func() {
+	return func() {
+		var signers []ssh.Signer
+
+		// If there's a running SSH Agent, try to use its Private keys.
+		sock, err := net.Dial("unix", os.Getenv("SSH_AUTH_SOCK"))
+		if err == nil {
+			agent := agent.NewClient(sock)
+			signers, _ = agent.Signers()
 		}
-		signer, err := ssh.ParsePrivateKey(data)
-		if err != nil {
-			continue
+
+		// Try to read user's SSH private keys form the standard paths.
+		files, _ := filepath.Glob(os.Getenv("HOME") + "/.ssh/id_*")
+		// Add nonstandard path
+		if identityFilePath != "" {
+			files = append(files, identityFilePath)
 		}
-		signers = append(signers, signer)
+		for _, file := range files {
+			if strings.HasSuffix(file, ".pub") {
+				continue // Skip public keys.
+			}
+			data, err := ioutil.ReadFile(file)
+			if err != nil {
+				continue
+			}
+			signer, err := ssh.ParsePrivateKey(data)
+			if err != nil {
+				continue
+			}
+			signers = append(signers, signer)
 
+		}
+		authMethod = ssh.PublicKeys(signers...)
 	}
-	authMethod = ssh.PublicKeys(signers...)
 }
 
 // SSHDialFunc can dial an ssh server and return a client
@@ -127,7 +134,7 @@ func (c *SSHClient) ConnectWith(host string, dialer SSHDialFunc) error {
 		return fmt.Errorf("Already connected")
 	}
 
-	initAuthMethodOnce.Do(initAuthMethod)
+	initAuthMethodOnce.Do(initAuthMethod(c.identityFile))
 
 	err := c.parseHost(host)
 	if err != nil {
diff --git a/sup.go b/sup.go
index d815068..2f6db1f 100644
--- a/sup.go
+++ b/sup.go
@@ -70,9 +70,10 @@ func (sup *Stackup) Run(network *Network, envVars EnvList, commands ...*Command)
 
 			// SSH client.
 			remote := &SSHClient{
-				env:   env + `export SUP_HOST="` + host + `";`,
-				user:  network.User,
-				color: Colors[i%len(Colors)],
+				env:          env + `export SUP_HOST="` + host + `";`,
+				user:         network.User,
+				color:        Colors[i%len(Colors)],
+				identityFile: network.IdentityFile,
 			}
 
 			if bastion != nil {