gistify723013.go

/* -.-.-.-.-.-.-.-.-.-.-.-.-.-.-.-.-.-.-.-.

* File Name : ssh.go

* Purpose :

* Creation Date : 11-18-2013

* Last Modified : Thu Dec  5 23:12:09 2013

* Created By : Kiyor

_._._._._._._._._._._._._._._._._._._._._.*/

package main

import (
	"bytes"

	"code.google.com/p/go.crypto/ssh"

	"fmt"

	"github.com/wsxiaoys/terminal/color"

	"io"
	"io/ioutil"
	"runtime"
	"strings"
	"sync"
	"time"
)

func strip(v string) string {
	return strings.TrimSpace(strings.Trim(v, "\n"))
}

type keychain struct {
	keys []ssh.Signer
}

func (k *keychain) Key(i int) (ssh.PublicKey, error) {
	if i < 0 || i >= len(k.keys) {
		return nil, nil
	}
	return k.keys[i].PublicKey(), nil
}

func (k *keychain) Sign(i int, rand io.Reader, data []byte) (sig []byte, err error) {
	return k.keys[i].Sign(rand, data)
}

func (k *keychain) add(key ssh.Signer) {
	k.keys = append(k.keys, key)
}

func (k *keychain) loadPEM(file string) error {
	buf, err := ioutil.ReadFile(file)
	if err != nil {
		return err
	}
	key, err := ssh.ParsePrivateKey(buf)
	if err != nil {
		return err
	}
	k.add(key)
	return nil
}

func Sshcmd(hosts []string, command string, background bool, debug bool) {
	k := new(keychain)
	// Add path to id_rsa file
	err := k.loadPEM("/home/user/.ssh/id_rsa")

	if err != nil {
		panic("Cannot load key: " + err.Error())
	}

	// config := ssh.ClientAuth

	// auth := make(ssh.Client) ssh.ClientAuthKeyring(k)

	// Switch out username
	config := &ssh.ClientConfig{
		// Change to your username
		User: "user",
		Auth: []ssh.ClientAuth{
			ssh.ClientAuthKeyring(k),
		},
	}

	command = fmt.Sprintf("/usr/bin/sudo bash <<CMD\nexport PATH=/usr/local/sbin:/usr/local/bin:/sbin:/bin:/usr/sbin:/usr/bin:/root/bin\n%s\nCMD", command)

	if background {
		command = fmt.Sprintf("/usr/bin/nohup bash -c \\\n\"%s\" `</dev/null` >nohup.out 2>&1 &", command)
	}

	if debug {
		color.Printf("@{b}%s\n", command)
	}

	var wg sync.WaitGroup
	queue := make(chan Result)
	count := new(int)
	// Change if you need
	workers := 24
	var results []Result

	// Filter hosts if you need
	// Should be like ["hostname1:22","hostname2:22"]
	conns := FilterHosts(hosts)

	for _, conn := range conns {
		wg.Add(1)
		*count++
		if debug {
			color.Printf("@{y}%s\t\tcounter %3d\n", conn, *count)
		}
		for *count >= workers {
			time.Sleep(10 * time.Millisecond)
		}
		go func(h string) {
			defer wg.Done()
			var r Result

			r.Host = h
			client, err := ssh.Dial("tcp", h, config)
			if err != nil {
				color.Printf("@{!r}%s: Failed to connect: %s\n", h, err.Error())
				*count--
				if debug {
					color.Printf("@{y}%s\t\tcounter %3d\n", conn, *count)
				}
				return
			}

			session, err := client.NewSession()
			if err != nil {
				color.Printf("@{!r}%s: Failed to create session: %s\n", h, err.Error())
				*count--
				if debug {
					color.Printf("@{y}%s\t\tcounter %3d\n", conn, *count)
				}
				return
			}
			defer session.Close()

			// This not working, need debug
			//session.Setenv("PATH", "/usr/local/sbin:/usr/local/bin:/sbin:/bin:/usr/sbin:/usr/bin:/root/bin")

			var b bytes.Buffer
			var e bytes.Buffer
			session.Stdout = &b
			session.Stderr = &e
			if err := session.Run(command); err != nil {
				color.Printf("@{!r}%s: Failed to run: %s\n", h, err.Error())
				color.Printf("@{!r}%s\n", strip(e.String()))
				*count--
				if debug {
					color.Printf("@{y}%s\t\tcounter %3d\n", conn, *count)
				}
				return
			}

			if !background {
				r.Res = strip(b.String())
			} else {
				r.Res = "command success and out put in remote server's /home/user/nohup.out"
			}
			color.Printf("@{!g}%s\n", r.Host)
			fmt.Println(r.Res)

			*count--
			if debug {
				color.Printf("@{y}%s\t\tcounter %3d\n", conn, *count)
			}

			runtime.Gosched()
			queue <- r
		}(conn)
	}
	go func() {
		defer wg.Done()
		for r := range queue {
			results = append(results, r)
		}
	}()
	wg.Wait()
	// Print res if you need
	//fmt.Println(results)
}