ing

ssh/ssh.go

@@ -5,6 +5,8 @@ import (
 	"fmt"
 	"golang.org/x/crypto/ssh"
 	"strconv"
+	"io/ioutil"
+	"time"
 )
 
 type SSHConfig struct {
@@ -19,13 +21,23 @@ type SSHClient struct {
 	conf    *SSHConfig
 }
 
-func New(ip, port, user, pw string) (*SSHClient, error) {
+func New(ip, port, user, pw, keyFilePath string) (*SSHClient, error) {
 	p, _ := strconv.Atoi(port)
+
+	auth := make([]ssh.AuthMethod, 0)
+
+	if keyFilePath != "" {
+		key, err := parsePrivateKey(keyFilePath, pw)
+		if err != nil {
+			return nil, err
+		}
+		auth = append(auth, ssh.PublicKeys(key))
+	}
+	auth = append(auth, ssh.Password(pw))
+
 	conf := &SSHConfig{
 		User: user,
-		Auth: []ssh.AuthMethod{
-			ssh.Password(pw),
-		},
+		Auth: auth,
 		Host: ip,
 		Port: p,
 	}
@@ -38,6 +50,7 @@ func (cli *SSHClient) Session() error {
 		User:            cli.conf.User,
 		Auth:            cli.conf.Auth,
 		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
+		Timeout: time.Second * 10,
 	}
 
 	connection, err := ssh.Dial("tcp", fmt.Sprintf("%s:%d", cli.conf.Host, cli.conf.Port), sshConfig)
@@ -83,3 +96,15 @@ func (cli *SSHClient) RunCommand(command string) ([]byte, error) {
 
 	return b.Bytes(), nil
 }
+
+
+func parsePrivateKey(keyPath, pw string) (ssh.Signer, error) {
+	buff, err := ioutil.ReadFile(keyPath)
+	if err != nil {
+		return nil, err
+	}
+	if pw == "" {
+		return ssh.ParsePrivateKey(buff)
+	}
+	return ssh.ParsePrivateKeyWithPassphrase(buff, []byte(pw))
+}