Re-added dynamic user lookup if -u flag undefined. Sadly, this means no more cross-compilation. Also added ssh-agent fallback if --key is undefined

Author: Cole Brumley

mssh.go

@@ -6,8 +6,13 @@ import (
 	"github.com/codegangsta/cli"
 	"github.com/fatih/color"
 	"golang.org/x/crypto/ssh"
+	"golang.org/x/crypto/ssh/agent"
 	"io/ioutil"
+	"net"
 	"os"
+	// Use of the os/user package prevents cross-compilation
+	"os/user" // <- https://github.com/golang/go/issues/6376
+	"path/filepath"
 	"strings"
 )
 
@@ -16,7 +21,7 @@ func main() {
 	app.HideVersion = true
 	app.Name = "mssh"
 	app.Usage = "Run SSH commands on multiple machines"
-	app.Version = "0.0.2"
+	app.Version = "0.0.2.1"
 	app.Flags = []cli.Flag{
 		cli.StringFlag{
 			Name:   "user,u",
@@ -43,7 +48,7 @@ func main() {
 		},
 		cli.BoolTFlag{
 			Name:  "color,c",
-			Usage: "Print cmd output in color (true by default)",
+			Usage: "Print cmd output in color (use -c=false to disable)",
 		},
 	}
 	app.Action = defaultAction
@@ -52,41 +57,99 @@ func main() {
 
 func defaultAction(c *cli.Context) {
 	hosts := c.StringSlice("server")
+
+	currentUser, err := user.Current()
+	if err != nil {
+		log.Errorf("Could not get current user: %v", err)
+	}
 	u := c.String("user")
+	// If no flag, get current username
+	if len(u) == 0 {
+		u = currentUser.Username
+		if len(u) == 0 {
+			log.Errorln("No username specified!")
+		}
+	}
 
+	// Getting keys:
+	// --key flag overrides all
+	// If no --key is defined, check for ssh-agent on $SSH_AUTH_SOCK
+	// If no ssh-agent defined, look or key in ~/.ssh/id_rsa
+	// fail if none of the above work
 	key := c.String("key")
+	auths := []ssh.AuthMethod{}
 	if len(key) == 0 {
-		// If no key provided use ~/.ssh/id_rsa
-		key = os.Getenv("HOME") + "/.ssh/id_rsa"
+		// If no key provided see if there's an ssh-agent running
+		if len(os.Getenv("SSH_AUTH_SOCK")) > 0 {
+			log.Println("Attempting to use existing ssh-agent")
+			conn, err := net.Dial("unix", os.Getenv("SSH_AUTH_SOCK"))
+			if err != nil {
+				return
+			}
+			defer conn.Close()
+			ag := agent.NewClient(conn)
+			auths = []ssh.AuthMethod{ssh.PublicKeysCallback(ag.Signers)}
+		} else {
+			// Otherwise use ~/.ssh/id_rsa
+			key = filepath.FromSlash(currentUser.HomeDir + "/.ssh/id_rsa")
+			if !fileExists(key) {
+				log.Errorln("Must specify a key, ~/.ssh/id_rsa does not exist and no ssh-agent is available!")
+				cli.ShowAppHelp(c)
+				os.Exit(2)
+			}
+			pemBytes, err := ioutil.ReadFile(key)
+			if err != nil {
+				log.Errorf("%v", err)
+			}
+			signer, err := ssh.ParsePrivateKey(pemBytes)
+			if err != nil {
+				log.Errorf("%v", err)
+			}
+			auths = []ssh.AuthMethod{ssh.PublicKeys(signer)}
+		}
+	} else {
 		if !fileExists(key) {
-			log.Errorln("Must specify a key, ~/.ssh/id_rsa does not exist!")
-			cli.ShowAppHelp(c)
-			os.Exit(2)
+			log.Errorln("Specified key does not exist!")
+			os.Exit(1)
+		}
+		pemBytes, err := ioutil.ReadFile(key)
+		if err != nil {
+			log.Errorf("%v", err)
+		}
+		signer, err := ssh.ParsePrivateKey(pemBytes)
+		if err != nil {
+			log.Errorf("%v", err)
 		}
+		auths = []ssh.AuthMethod{ssh.PublicKeys(signer)}
 	}
 
-	// user and host(s) required
+	// host(s) is required
 	if len(hosts) == 0 {
-		log.Errorln("At least one host is required")
+		log.Warnln("At least one host is required")
 		cli.ShowAppHelp(c)
 		os.Exit(2)
 	}
+
+	// At least one command is required
 	if len(c.Args().First()) == 0 {
+		log.Warnln("At least one command is required")
 		cli.ShowAppHelp(c)
 		os.Exit(2)
 	}
+
+	// append list of commands to argList
 	argList := append([]string{c.Args().First()}, c.Args().Tail()...)
 	chLen := len(hosts) * len(argList)
 	done := make(chan bool, chLen)
 	for _, h := range hosts {
 		// Hosts are handled asynchronously
-		go func(c *cli.Context, host string, cmds []string) {
+		go func(c *cli.Context, host string, cmds []string, auth []ssh.AuthMethod) {
 			// Set the default port if none specified
 			if len(strings.Split(host, ":")) == 1 {
 				host = host + ":22"
 			}
 			for _, cmd := range cmds {
-				combined, err := runRemoteCmd(u, host, key, cmd)
+				combined, err := runRemoteCmd(u, host, auth, cmd)
 				out := tail(string(combined), c.Int("lines"))
 				if err != nil {
 					pretty := prettyOutput(
@@ -104,7 +167,7 @@ func defaultAction(c *cli.Context) {
 				}
 				done <- true
 			}
-		}(c, h, argList)
+		}(c, h, argList, auths)
 	}
 
 	// Drain chan before exiting
@@ -150,18 +213,10 @@ func tail(s string, n int) string {
 	return strings.Join(lines[len(lines)-(n):], "\n")
 }
 
-func runRemoteCmd(user string, addr string, key string, cmd string) (out []byte, err error) {
-	pemBytes, err := ioutil.ReadFile(key)
-	if err != nil {
-		return
-	}
-	signer, err := ssh.ParsePrivateKey(pemBytes)
-	if err != nil {
-		return
-	}
+func runRemoteCmd(user string, addr string, auth []ssh.AuthMethod, cmd string) (out []byte, err error) {
 	client, err := ssh.Dial("tcp", addr, &ssh.ClientConfig{
 		User: user,
-		Auth: []ssh.AuthMethod{ssh.PublicKeys(signer)},
+		Auth: auth,
 	})
 	if err != nil {
 		return
@@ -176,7 +231,7 @@ func runRemoteCmd(user string, addr string, key string, cmd string) (out []byte,
 	return
 }
 
-// FileExists returns a bool if os.Stat returns an IsNotExist error
+// fileExists returns a bool if os.Stat returns an IsNotExist error
 func fileExists(name string) bool {
 	if _, err := os.Stat(name); err != nil {
 		if os.IsNotExist(err) {

shippable.yml

@@ -1,18 +1,10 @@
 language: go
-
 go:
   - release
-
-env:
-  - CGO_ENABLED=0
-
 before_install:
   - export GOPATH=$SHIPPABLE_GOPATH
-
 install:
-  - go get -t -d -v ./... && go build -x -o mssh -a -installsuffix cgo ./...
-
+  - go get -t -d -v ./... && go build -x -o mssh ./...
 script:
-  - cp mssh shippable/
-
-archive: true
+  - cp mssh shippable/buildoutput/
+archive: true