Simple experiment to parallel ssh-ing into remote machines with Go

parassh.go

// Simple experiment to parallel ssh-ing into remote machines with Go.
//
// (c) 2017 by Sascha L. Teichmann
// The is Free Software covered by the terms of the MIT License.
// See https://opensource.org/licenses/MIT for details.
//
package main

import (
	"bufio"
	"bytes"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net"
	"os"
	"strings"
	"sync"

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

func parseHosts(r io.Reader) ([]string, error) {
	var hosts []string
	s := bufio.NewScanner(r)
	for s.Scan() {
		line := strings.TrimSpace(s.Text())
		if line == "" || strings.HasPrefix(line, "#") {
			continue
		}
		hosts = append(hosts, line)
	}
	return hosts, s.Err()
}

func parseHostsFromFile(fname string) ([]string, error) {
	f, err := os.Open(fname)
	if err != nil {
		return nil, err
	}
	defer f.Close()
	return parseHosts(f)
}

func parseKeyFile(fname string) (ssh.Signer, error) {
	pem, err := ioutil.ReadFile(fname)
	if err != nil {
		return nil, err
	}
	return ssh.ParsePrivateKey(pem)
}

func doSession(client *ssh.Client) error {
	session, err := client.NewSession()
	if err != nil {
		return err
	}
	defer session.Close()
	var buf bytes.Buffer
	session.Stdout = &buf
	if err = session.Run("apt-cache pkgnames"); err != nil {
		return err
	}
	fmt.Printf("%s: %d\n", client.RemoteAddr(), buf.Len())
	return nil
}

func doSSH(hostCh <-chan string, wg *sync.WaitGroup, config *ssh.ClientConfig) {
	defer wg.Done()
	for h := range hostCh {
		fmt.Printf("connect to host %s\n", h)
		client, err := ssh.Dial("tcp", h+":22", config)
		if err != nil {
			log.Printf("Failed to dial: %v", err)
			continue
		}
		if err = doSession(client); err != nil {
			log.Printf("Session failed: %v\n", err)
		}
	}
}

func main() {
	var (
		user           string
		keyFile        string
		hostsFile      string
		knownHostsFile string
		maxConnections int
		useAgent       bool
	)
	flag.StringVar(&user, "user", "root", "Name of the remote user.")
	flag.StringVar(&keyFile, "key", "key", "Name of the certificate file")
	flag.StringVar(&hostsFile, "hosts", "hosts.txt", "File with hosts to connect to.")
	flag.StringVar(&knownHostsFile, "known", "known_hosts", "File with known hosts.")
	flag.IntVar(&maxConnections, "conns", 0, "Number of max open SSH connections.")
	flag.BoolVar(&useAgent, "agent", false, "Use SSH agent.")

	flag.Parse()

	hosts, err := parseHostsFromFile(hostsFile)
	if err != nil {
		log.Fatalf("error: %v\n", err)
	}

	knownHosts, err := knownhosts.New(knownHostsFile)
	if err != nil {
		log.Fatalf("error: %v\n", err)
	}

	var auth ssh.AuthMethod

	if useAgent {
		agentEnv, ok := os.LookupEnv("SSH_AUTH_SOCK")
		if !ok {
			log.Fatalln("Variable SSH_AUTH_SOCK not set.")
		}

		agentSocket, err := net.Dial("unix", agentEnv)
		if err != nil {
			log.Fatalf("error: %v\n", err)
		}
		defer agentSocket.Close()
		ag := agent.NewClient(agentSocket)
		auth = ssh.PublicKeysCallback(ag.Signers)

	} else {
		signer, err := parseKeyFile(keyFile)
		if err != nil {
			log.Fatalf("error: %v\n", err)
		}
		auth = ssh.PublicKeys(signer)
	}

	config := &ssh.ClientConfig{
		User:            user,
		Auth:            []ssh.AuthMethod{auth},
		HostKeyCallback: knownHosts,
	}

	if maxConnections <= 0 || maxConnections > len(hosts) {
		maxConnections = len(hosts)
	}

	hostCh := make(chan string)

	var wg sync.WaitGroup
	for i := 0; i < maxConnections; i++ {
		wg.Add(1)
		go doSSH(hostCh, &wg, config)
	}
	for _, h := range hosts {
		hostCh <- h
	}
	close(hostCh)
	wg.Wait()
}