Merge pull request #2 from nealwon/dev

add file transfer support
nealwon · Jan 30, 2018 · 45ed272 · 45ed272
2 parents d14ba61 + a29e579
commit 45ed272
Show file tree

Hide file tree

Showing 5 changed files with 411 additions and 61 deletions.
diff --git a/README.md b/README.md
@@ -1,47 +1,55 @@
 # optool
-A tool to execute commands on multiple remote hosts
+A tool to execute commands,transfer files on multiple remote hosts
 
 ------------
 ### Usage:
 ```bash
 Usage:
-  -V    print sample configure
+  -V	print sample configure
   -config string
-        set config file path (default "/optool.yml")
+    	set config file path (default "/optool.yml")
   -encrypt
-        encrypt a password/phrase
+    	encrypt a password/phrase
   -g string
-        set default group name for hosts
+    	set default group name for hosts
+  -get string
+    	get a file from remote host
   -gz
-        enable gzip for transfer./usr/bin/gzip must be executable at remote host
+    	enable gzip for transfer./usr/bin/gzip must be executable at remote host
   -host string
-        set run host
+    	set run host
   -key string
-        set private key
+    	set private key
   -nh int
-        (1)1<<0=no header,(2)1<<1=no server ip,3=none
+    	(1)1<<0=no header,(2)1<<1=no server ip,3=none
   -o string
-        set output file (default "-")
+    	set output file (default "-")
+  -override
+    	Override remote file if exists
+  -path string
+    	set path.if get is set this is local path,if put is set this is remote path
   -port int
-        set default ssh port
+    	set default ssh port
+  -put string
+    	put a file to remote host
   -s string
-        read commands from script
+    	read commands from script
   -t string
-        set tagged command
+    	set tagged command
   -ta string
-        append tagged command parameters, overflow params will be dropped, separated by comma(,).
-         to replace in tags use string: _REPLACE_
+    	append tagged command parameters, overflow params will be dropped, separated by comma(,).
+	 to replace in tags use string: _REPLACE_
   -tl
-        list all tags
+    	list all tags
   -tp
-        print tag line
+    	print tag line
   -u string
-        set ssh auth user
-  -v    verbose all configs
+    	set ssh auth user
+  -v	verbose all configs
   -version
-        print version and exit
+    	print version and exit
   -x string
-        execute command directly
+    	execute command directly
 ```
 
 ### Sample configure:
@@ -67,4 +75,5 @@ tags:
   ps: "/bin/ps"
   netstat: "/bin/netstat -lntpu"
   err: "/bin/grep ERROR /var/log/nginx/error.log_REPLACE_"
+# transfer_max_size: 1099511627776 #100MB
 ```
diff --git a/common/command.go b/common/command.go
@@ -1,12 +1,12 @@
 package common
 
 import (
+	"bytes"
 	"compress/gzip"
 	"fmt"
 	"io"
 	"io/ioutil"
 	"log"
-	"os"
 	"strconv"
 	"strings"
 	"sync"
@@ -54,48 +54,16 @@ func NewRemoteCommand(hosts []string, cmd string) *RemoteCommand {
 }
 
 // Start run remote command
-func (rc *RemoteCommand) Start() error {
+func (rc *RemoteCommand) Start() (err error) {
 	cfg := &ssh.ClientConfig{
 		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
 		Timeout:         time.Second * 10,
 	}
-	password := C.Auth.Password
-	if !C.Auth.PlainPassword {
-		password = string(Decrypt(C.Auth.Password))
-	}
 	if C.Auth.User != "" {
 		cfg.User = C.Auth.User
-		if C.Auth.PrivateKey != "" {
-			if _, err := os.Stat(C.Auth.PrivateKey); err != nil {
-				return err
-			}
-			key, err := ioutil.ReadFile(C.Auth.PrivateKey)
-			if err != nil {
-				return err
-			}
-			var signer ssh.Signer
-			if C.Auth.PrivateKeyPhrase == "" {
-				signer, err = ssh.ParsePrivateKey(key)
-			} else {
-				passphrase := []byte(C.Auth.PrivateKeyPhrase)
-				if !C.Auth.PlainPassword {
-					passphrase = Decrypt(C.Auth.PrivateKeyPhrase)
-				}
-				signer, err = ssh.ParsePrivateKeyWithPassphrase(key, passphrase)
-			}
-			if err != nil {
-				return err
-			}
-			cfg.Auth = []ssh.AuthMethod{
-				ssh.PublicKeys(signer),
-			}
-			if password != "" {
-				cfg.Auth = append(cfg.Auth, ssh.Password(password))
-			}
-		} else {
-			cfg.Auth = []ssh.AuthMethod{
-				ssh.Password(password),
-			}
+		cfg.Auth, err = GetAuth()
+		if err != nil {
+			return err
 		}
 	}
 	for _, host := range rc.Hosts {
@@ -172,7 +140,12 @@ func (rc *RemoteCommand) PrettyPrint(wo io.Writer, we io.Writer, noHeader bool,
 			we.Write([]byte("================================= ERROR =================================\n"))
 		}
 		for h, e := range rc.Error {
-			fmt.Fprintln(we, h, ":\n", e)
+			e = strings.TrimRight(e, "\n")
+			if strings.Contains(e, "\n") {
+				fmt.Fprintln(we, h, ":\n", e)
+			} else {
+				fmt.Fprintln(we, h, ":", e)
+			}
 		}
 	}
 	if len(rc.Output) > 0 {
@@ -191,15 +164,23 @@ func (rc *RemoteCommand) PrettyPrint(wo io.Writer, we io.Writer, noHeader bool,
 				if err != nil {
 					log.Println(err)
 				}
+				data = bytes.TrimRight(data, "\n")
 				if !noHost {
-					wo.Write([]byte(h + ": \n"))
+					fmt.Fprintf(wo, "%15s: ", h)
+					if bytes.Contains(data, []byte("\n")) {
+						wo.Write([]byte("\n"))
+					}
 				}
 				wo.Write(data)
 				wo.Write([]byte("\n"))
 				continue
 			}
+			o = strings.TrimRight(o, "\n")
 			if !noHost {
-				wo.Write([]byte(h + ": \n"))
+				fmt.Fprintf(wo, "%15s: ", h)
+				if strings.Contains(o, "\n") {
+					wo.Write([]byte("\n"))
+				}
 			}
 			wo.Write([]byte(o))
 			wo.Write([]byte("\n"))

diff --git a/common/config.go b/common/config.go
@@ -2,6 +2,9 @@ package common
 
 import (
 	"io/ioutil"
+	"os"
+
+	"golang.org/x/crypto/ssh"
 
 	"github.com/go-yaml/yaml"
 )
@@ -24,7 +27,10 @@ type Configure struct {
 	Tags map[string]string `yaml:"tags"` // shortcut for frequently used commands
 	Gzip bool              `yaml:"-"`    // enable gzip transfer
 	//DefaultGroup string              `yaml:"default_group"` // set default host group
+	TransferMaxSize int64 `yaml:"transfer_max_size"`
 }
+
+// Server server groups and default port/group config
 type Server struct {
 	DefaultGroup string              `yaml:"default_group"`
 	DefaultPort  int                 `yaml:"default_port"`
@@ -50,3 +56,44 @@ func ParseConfig(f string) error {
 	}
 	return nil
 }
+
+// GetAuth get auth method list from configs
+func GetAuth() (auth []ssh.AuthMethod, err error) {
+	password := C.Auth.Password
+	if !C.Auth.PlainPassword {
+		password = string(Decrypt(C.Auth.Password))
+	}
+	if C.Auth.PrivateKey != "" {
+		if _, err := os.Stat(C.Auth.PrivateKey); err != nil {
+			return nil, err
+		}
+		key, err := ioutil.ReadFile(C.Auth.PrivateKey)
+		if err != nil {
+			return nil, err
+		}
+		var signer ssh.Signer
+		if C.Auth.PrivateKeyPhrase == "" {
+			signer, err = ssh.ParsePrivateKey(key)
+		} else {
+			passphrase := []byte(C.Auth.PrivateKeyPhrase)
+			if !C.Auth.PlainPassword {
+				passphrase = Decrypt(C.Auth.PrivateKeyPhrase)
+			}
+			signer, err = ssh.ParsePrivateKeyWithPassphrase(key, passphrase)
+		}
+		if err != nil {
+			return nil, err
+		}
+		auth = []ssh.AuthMethod{
+			ssh.PublicKeys(signer),
+		}
+		if password != "" {
+			auth = append(auth, ssh.Password(password))
+		}
+	} else {
+		auth = []ssh.AuthMethod{
+			ssh.Password(password),
+		}
+	}
+	return
+}