feat: add easyssh and go routine. (#38)

easyssh/easyssh.go

@@ -0,0 +1,193 @@
+// Package easyssh provides a simple implementation of some SSH protocol
+// features in Go. You can simply run a command on a remote server or get a file
+// even simpler than native console SSH client. You don't need to think about
+// Dials, sessions, defers, or public keys... Let easyssh think about it!
+package easyssh
+
+import (
+	"bufio"
+	"fmt"
+	"io"
+	"io/ioutil"
+	"os"
+	"path/filepath"
+	"time"
+
+	"golang.org/x/crypto/ssh"
+)
+
+// MakeConfig Contains main authority information.
+// User field should be a name of user on remote server (ex. john in ssh john@example.com).
+// Server field should be a remote machine address (ex. example.com in ssh john@example.com)
+// Key is a path to private key on your local machine.
+// Port is SSH server port on remote machine.
+// Note: easyssh looking for private key in user's home directory (ex. /home/john + Key).
+// Then ensure your Key begins from '/' (ex. /.ssh/id_rsa)
+type MakeConfig struct {
+	User     string
+	Server   string
+	Key      string
+	KeyPath  string
+	Port     string
+	Password string
+	Timeout  time.Duration
+}
+
+// returns ssh.Signer from user you running app home path + cutted key path.
+// (ex. pubkey,err := getKeyFile("/.ssh/id_rsa") )
+func getKeyFile(keypath string) (ssh.Signer, error) {
+	buf, err := ioutil.ReadFile(keypath)
+	if err != nil {
+		return nil, err
+	}
+
+	pubkey, err := ssh.ParsePrivateKey(buf)
+	if err != nil {
+		return nil, err
+	}
+
+	return pubkey, nil
+}
+
+// connects to remote server using MakeConfig struct and returns *ssh.Session
+func (ssh_conf *MakeConfig) connect() (*ssh.Session, error) {
+	// auths holds the detected ssh auth methods
+	auths := []ssh.AuthMethod{}
+
+	// figure out what auths are requested, what is supported
+	if ssh_conf.Password != "" {
+		auths = append(auths, ssh.Password(ssh_conf.Password))
+	}
+
+	if ssh_conf.KeyPath != "" {
+		if pubkey, err := getKeyFile(ssh_conf.KeyPath); err == nil {
+			auths = append(auths, ssh.PublicKeys(pubkey))
+		}
+	}
+
+	if ssh_conf.Key != "" {
+		signer, _ := ssh.ParsePrivateKey([]byte(ssh_conf.Key))
+		auths = append(auths, ssh.PublicKeys(signer))
+	}
+
+	config := &ssh.ClientConfig{
+		Timeout: ssh_conf.Timeout,
+		User:    ssh_conf.User,
+		Auth:    auths,
+	}
+
+	client, err := ssh.Dial("tcp", ssh_conf.Server+":"+ssh_conf.Port, config)
+	if err != nil {
+		return nil, err
+	}
+
+	session, err := client.NewSession()
+	if err != nil {
+		return nil, err
+	}
+
+	return session, nil
+}
+
+// Stream returns one channel that combines the stdout and stderr of the command
+// as it is run on the remote machine, and another that sends true when the
+// command is done. The sessions and channels will then be closed.
+func (ssh_conf *MakeConfig) Stream(command string) (output chan string, done chan bool, err error) {
+	// connect to remote host
+	session, err := ssh_conf.connect()
+	if err != nil {
+		return output, done, err
+	}
+	// connect to both outputs (they are of type io.Reader)
+	outReader, err := session.StdoutPipe()
+	if err != nil {
+		return output, done, err
+	}
+	errReader, err := session.StderrPipe()
+	if err != nil {
+		return output, done, err
+	}
+	// combine outputs, create a line-by-line scanner
+	outputReader := io.MultiReader(outReader, errReader)
+	err = session.Start(command)
+	scanner := bufio.NewScanner(outputReader)
+	// continuously send the command's output over the channel
+	outputChan := make(chan string)
+	done = make(chan bool)
+	go func(scanner *bufio.Scanner, out chan string, done chan bool) {
+		defer close(outputChan)
+		defer close(done)
+		for scanner.Scan() {
+			outputChan <- scanner.Text()
+		}
+		// close all of our open resources
+		done <- true
+		session.Close()
+	}(scanner, outputChan, done)
+	return outputChan, done, err
+}
+
+// Run command on remote machine and returns its stdout as a string
+func (ssh_conf *MakeConfig) Run(command string) (outStr string, err error) {
+	outChan, doneChan, err := ssh_conf.Stream(command)
+	if err != nil {
+		return outStr, err
+	}
+	// read from the output channel until the done signal is passed
+	stillGoing := true
+	for stillGoing {
+		select {
+		case <-doneChan:
+			stillGoing = false
+		case line := <-outChan:
+			outStr += line + "\n"
+		}
+	}
+	// return the concatenation of all signals from the output channel
+	return outStr, err
+}
+
+// Scp uploads sourceFile to remote machine like native scp console app.
+func (ssh_conf *MakeConfig) Scp(sourceFile string) error {
+	session, err := ssh_conf.connect()
+
+	if err != nil {
+		return err
+	}
+	defer session.Close()
+
+	targetFile := filepath.Base(sourceFile)
+
+	src, srcErr := os.Open(sourceFile)
+
+	if srcErr != nil {
+		return srcErr
+	}
+
+	srcStat, statErr := src.Stat()
+
+	if statErr != nil {
+		return statErr
+	}
+
+	go func() {
+		w, _ := session.StdinPipe()
+
+		fmt.Fprintln(w, "C0644", srcStat.Size(), targetFile)
+
+		if srcStat.Size() > 0 {
+			io.Copy(w, src)
+			fmt.Fprint(w, "\x00")
+			w.Close()
+		} else {
+			fmt.Fprint(w, "\x00")
+			w.Close()
+		}
+	}()
+
+	if err := session.Run(fmt.Sprintf("scp -t %s", targetFile)); err != nil {
+		return err
+	}
+
+	return nil
+}

main.go

@@ -57,11 +57,6 @@ func main() {
 			EnvVar: "PLUGIN_PORT,SSH_PORT",
 			Value:  22,
 		},
-		cli.IntFlag{
-			Name:   "sleep",
-			Usage:  "sleep between hosts",
-			EnvVar: "PLUGIN_SLEEP,SSH_SLEEP",
-		},
 		cli.DurationFlag{
 			Name:   "timeout,t",
 			Usage:  "connection timeout",
@@ -127,7 +122,6 @@ func run(c *cli.Context) error {
 			Password: c.String("password"),
 			Host:     c.StringSlice("host"),
 			Port:     c.Int("port"),
-			Sleep:    c.Int("sleep"),
 			Timeout:  c.Duration("timeout"),
 			Script:   c.StringSlice("script"),
 		},

plugin.go

@@ -2,23 +2,20 @@ package main
 
 import (
 	"fmt"
-	"io/ioutil"
 	"log"
-	"net"
-	"os"
 	"strconv"
 	"strings"
+	"sync"
 	"time"
 
-	"golang.org/x/crypto/ssh"
+	"github.com/appleboy/drone-ssh/easyssh"
 )
 
+var wg sync.WaitGroup
+
 const (
 	missingHostOrUser    = "Error: missing server host or user"
 	missingPasswordOrKey = "Error: can't connect without a private SSH key or password"
-	unableConnectServer  = "Error: Failed to start a SSH session"
-	failParsePrivateKey  = "Error: Failed to parse private key"
-	sshKeyNotFound       = "ssh: no key found"
 )
 
 type (
@@ -30,7 +27,6 @@ type (
 		Password string
 		Host     []string
 		Port     int
-		Sleep    int
 		Timeout  time.Duration
 		Script   []string
 	}
@@ -41,20 +37,8 @@ type (
 	}
 )
 
-// returns ssh.Signer from user you running app home path + cutted key path.
-// (ex. pubkey,err := getKeyFile("/.ssh/id_rsa") )
-func getKeyFile(keypath string) (ssh.Signer, error) {
-	buf, err := ioutil.ReadFile(keypath)
-	if err != nil {
-		return nil, err
-	}
-
-	pubkey, err := ssh.ParsePrivateKey(buf)
-	if err != nil {
-		return nil, err
-	}
-
-	return pubkey, nil
+func (p Plugin) log(host string, message ...interface{}) {
+	log.Printf("%s: %s", host, fmt.Sprintln(message...))
 }
 
 // Exec executes the plugin.
@@ -67,68 +51,49 @@ func (p Plugin) Exec() error {
 		return fmt.Errorf(missingPasswordOrKey)
 	}
 
-	for i, host := range p.Config.Host {
-		addr := net.JoinHostPort(
-			host,
-			strconv.Itoa(p.Config.Port),
-		)
-
-		// auths holds the detected ssh auth methods
-		auths := []ssh.AuthMethod{}
-
-		if p.Config.KeyPath != "" {
-			pubkey, err := getKeyFile(p.Config.KeyPath)
-
-			if err != nil {
-				return err
+	wg.Add(len(p.Config.Host))
+	errChannel := make(chan error, 1)
+	finished := make(chan bool, 1)
+	for _, host := range p.Config.Host {
+		go func(host string) {
+			// Create MakeConfig instance with remote username, server address and path to private key.
+			ssh := &easyssh.MakeConfig{
+				Server:   host,
+				User:     p.Config.User,
+				Password: p.Config.Password,
+				Port:     strconv.Itoa(p.Config.Port),
+				Key:      p.Config.Key,
+				KeyPath:  p.Config.KeyPath,
+				Timeout:  p.Config.Timeout,
 			}
 
-			auths = append(auths, ssh.PublicKeys(pubkey))
-		}
-
-		if p.Config.Key != "" {
-			signer, err := ssh.ParsePrivateKey([]byte(p.Config.Key))
+			p.log(host, "commands: ", strings.Join(p.Config.Script, "\n"))
+			response, err := ssh.Run(strings.Join(p.Config.Script, "\n"))
+			p.log(host, "outputs:", response)
 
 			if err != nil {
-				return fmt.Errorf(failParsePrivateKey)
+				errChannel <- err
 			}
 
-			auths = append(auths, ssh.PublicKeys(signer))
-		}
+			wg.Done()
+		}(host)
+	}
 
-		// figure out what auths are requested, what is supported
-		if p.Config.Password != "" {
-			auths = append(auths, ssh.Password(p.Config.Password))
-		}
-
-		config := &ssh.ClientConfig{
-			Timeout: p.Config.Timeout,
-			User:    p.Config.User,
-			Auth:    auths,
-		}
-
-		log.Printf("+ ssh %s@%s \n", p.Config.User, addr)
-		client, err := ssh.Dial("tcp", addr, config)
+	go func() {
+		wg.Wait()
+		close(finished)
+	}()
 
+	select {
+	case <-finished:
+	case err := <-errChannel:
 		if err != nil {
-			return fmt.Errorf(unableConnectServer)
-		}
-
-		session, _ := client.NewSession()
-		defer session.Close()
-
-		session.Stdout = os.Stdout
-		session.Stderr = os.Stderr
-
-		if err := session.Run(strings.Join(p.Config.Script, "\n")); err != nil {
+			log.Println("drone-ssh error: ", err)
 			return err
 		}
-
-		if p.Config.Sleep != 0 && i != len(p.Config.Host)-1 {
-			log.Printf("+ sleep %d\n", p.Config.Sleep)
-			time.Sleep(time.Duration(p.Config.Sleep) * time.Second)
-		}
 	}
 
+	log.Println("Successfully executed commnads to all host.")
+
 	return nil
 }

plugin_test.go

@@ -42,23 +42,6 @@ func TestIncorrectPassword(t *testing.T) {
 
 	err := plugin.Exec()
 	assert.NotNil(t, err)
-	assert.Equal(t, unableConnectServer, err.Error())
-}
-
-func TestFailParsePrivateKey(t *testing.T) {
-	plugin := Plugin{
-		Config: Config{
-			Host:   []string{"localhost"},
-			User:   "drone-scp",
-			Port:   22,
-			Key:    "123456",
-			Script: []string{"whoami"},
-		},
-	}
-
-	err := plugin.Exec()
-	assert.NotNil(t, err)
-	assert.Equal(t, failParsePrivateKey, err.Error())
 }
 
 func TestSSHScriptFromRawKey(t *testing.T) {
@@ -103,38 +86,6 @@ ib4KbP5ovZlrjL++akMQ7V2fHzuQIFWnCkDA5c2ZAqzlM+ZN+HRG7gWur7Bt4XH1
 	assert.Nil(t, err)
 }
 
-func TestWrongKeyPath(t *testing.T) {
-	plugin := Plugin{
-		Config: Config{
-			Host:    []string{"localhost"},
-			User:    "drone-scp",
-			Port:    22,
-			KeyPath: "/appleboy",
-			Script:  []string{"whoami"},
-		},
-	}
-
-	err := plugin.Exec()
-	assert.NotNil(t, err)
-	assert.Equal(t, "open /appleboy: no such file or directory", err.Error())
-}
-
-func TestWrongKeyFormat(t *testing.T) {
-	plugin := Plugin{
-		Config: Config{
-			Host:    []string{"localhost"},
-			User:    "drone-scp",
-			Port:    22,
-			KeyPath: "./tests/.ssh/id_rsa.pub",
-			Script:  []string{"whoami"},
-		},
-	}
-
-	err := plugin.Exec()
-	assert.NotNil(t, err)
-	assert.Equal(t, sshKeyNotFound, err.Error())
-}
-
 func TestSSHScriptFromKeyFile(t *testing.T) {
 	plugin := Plugin{
 		Config: Config{
@@ -142,26 +93,10 @@ func TestSSHScriptFromKeyFile(t *testing.T) {
 			User:    "drone-scp",
 			Port:    22,
 			KeyPath: "./tests/.ssh/id_rsa",
-			Script:  []string{"whoami"},
-			Sleep:   1,
+			Script:  []string{"whoami", "ls -al"},
 		},
 	}
 
 	err := plugin.Exec()
 	assert.Nil(t, err)
 }
-
-func TestSSHScriptRunError(t *testing.T) {
-	plugin := Plugin{
-		Config: Config{
-			Host:    []string{"localhost"},
-			User:    "drone-scp",
-			Port:    22,
-			KeyPath: "./tests/.ssh/id_rsa",
-			Script:  []string{"whoami", "whoam"},
-		},
-	}
-
-	err := plugin.Exec()
-	assert.NotNil(t, err)
-}