From 4fbfc9f3c89aa85cb9dbc72ea9e3c8b62079a7fa Mon Sep 17 00:00:00 2001
From: Sergey Kvetko <s_kvetko@wargaming.net>
Date: Thu, 23 Mar 2023 08:56:29 +0200
Subject: [PATCH] Added support of ssh_agent and passphrase

---
 helper.go    | 36 ++++++++++++++++++++++++++++++++++--
 vssh_test.go | 21 ++++++++++++++++++++-
 2 files changed, 54 insertions(+), 3 deletions(-)

diff --git a/helper.go b/helper.go
index 39993a0..4fa8587 100644
--- a/helper.go
+++ b/helper.go
@@ -5,9 +5,10 @@ package vssh
 
 import (
 	"fmt"
-	"io/ioutil"
-
 	"golang.org/x/crypto/ssh"
+	"golang.org/x/crypto/ssh/agent"
+	"io/ioutil"
+	"net"
 )
 
 // GetConfigPEM returns SSH configuration that uses the given private key.
@@ -26,6 +27,37 @@ func GetConfigPEM(user, keyFile string) (*ssh.ClientConfig, error) {
 	return getConfig(user, ssh.PublicKeys(signer)), nil
 }
 
+// GetConfigPEMWithPassphrase returns SSH configuration that uses the given private key and passphrase.
+func GetConfigPEMWithPassphrase(user, keyFile string, passphrase string) (*ssh.ClientConfig, error) {
+	key, err := ioutil.ReadFile(keyFile)
+	if err != nil {
+		return nil, fmt.Errorf("unable to read private key: %v", err)
+	}
+
+	signer, err := ssh.ParsePrivateKeyWithPassphrase(key, []byte(passphrase))
+	if err != nil {
+		return nil, fmt.Errorf("unable to parse private key: %v", err)
+	}
+
+	return getConfig(user, ssh.PublicKeys(signer)), nil
+}
+
+// GetConfigSSHAgent returns SSH configuration from SSH Agent.
+// default socket can get from env: os.Getenv("SSH_AUTH_SOCK")
+func GetConfigSSHAgent(user, socket string) (*ssh.ClientConfig, error) {
+	conn, err := net.Dial("unix", socket)
+	if err != nil {
+		return nil, fmt.Errorf("unable to connect to ssh agent: %v", err)
+	}
+	agentconn := agent.NewClient(conn)
+	signers, err := agentconn.Signers()
+	if err != nil {
+		return nil, fmt.Errorf("unable to get signers: %v", err)
+	}
+
+	return getConfig(user, ssh.PublicKeys(signers...)), nil
+}
+
 // GetConfigUserPass returns SSH configuration that uses the given
 // username and password.
 func GetConfigUserPass(user, password string) *ssh.ClientConfig {
diff --git a/vssh_test.go b/vssh_test.go
index 2ffb516..5580b8a 100644
--- a/vssh_test.go
+++ b/vssh_test.go
@@ -16,7 +16,7 @@ func TestForceReConn(t *testing.T) {
 		t.Error(err)
 	}
 
-	vs.ForceReConn("127.0.0.1:22")
+	_ = vs.ForceReConn("127.0.0.1:22")
 	if len(vs.actionQ) != 2 {
 		t.Fatal("expect to have two tasks but got", len(vs.actionQ))
 	}
@@ -140,3 +140,22 @@ func TestGetConfigPEM(t *testing.T) {
 		t.Error("expect error but nil")
 	}
 }
+
+func TestGetConfigPEMWithPassphrase(t *testing.T) {
+	_, err := GetConfigPEMWithPassphrase("vssh", "notexitfile", "pass1")
+	if err == nil {
+		t.Error("expect error but nil")
+	}
+
+	_, err = GetConfigPEMWithPassphrase("vssh", "vssh.go", "pass2")
+	if err == nil {
+		t.Error("expect error but nil")
+	}
+}
+
+func TestGetConfigSSHAgent(t *testing.T) {
+	_, err := GetConfigSSHAgent("root", "/path/to/socket/file")
+	if err == nil {
+		t.Error("expect error but nil")
+	}
+}