pctl

main.go

package main

import (
	"bufio"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net"
	"os"
	"sync"
	"time"

	ssh "golang.org/x/crypto/sssh"
	agent "golang.org/x/crypto/ssh/agent"
)

var (
	cmd, _   = ioutil.ReadAll(os.Stdin)
	hosts    = os.Args[1:]
	results  = make(chan string, 1000)
	timeout  = time.After(1000 * time.Second)
	con, err = net.Dial("unix", os.Getenv("SSH_AUTH_SOCK"))
	user     = "miked"
	wg       sync.WaitGroup
)

func main() {
	check(err, "problem opening SSH_AUTH_SOCK")
	wg.Add(len(hosts))
	defer con.Close()
	executeAll(makeConfig())
}

func executeAll(config *ssh.ClientConfig) {
	for _, hostname := range hosts {
		go func(hostname string) {
			executeCmd(string(cmd), hostname, config, results)
		}(hostname)
	}
	go func() {
		for {
			select {
			case res := <-results:
				fmt.Println(res)
			case <-timeout:
				fmt.Println("Timed out!")
				return
			}
		}
	}()
	wg.Wait()
}

func executeCmd(cmd, hostname string, config *ssh.ClientConfig, results chan string) {
	conn, err := ssh.Dial("tcp", hostname+":22", config)
	check(err, hostname)

	session, err := conn.NewSession()
	check(err)
	defer session.Close()

	r, w := io.Pipe()
	defer w.Close()

	go func() {
		defer wg.Done()
		s := bufio.NewScanner(r)
		defer r.Close()
		for s.Scan() {
			results <- hostname + ": " + string(s.Text())
		}
	}()

	session.Stdout = w
	session.Run(cmd)
}

func makeConfig() *ssh.ClientConfig {
	config := ssh.ClientConfig{
		User: user,
		Auth: []ssh.AuthMethod{makeKeyring()},
	}
	config.SetDefaults()
	return &config
}

func makeKeyring() ssh.AuthMethod {
	a := agent.NewClient(con)
	signers, err := a.Signers()
	check(err)

	return ssh.PublicKeys(signers...)
}

func check(err error, args ...interface{}) {
	if err != nil {
		// log.Panicln(err, args)
		log.Fatalln(err, args)
	}
}