fix(v1.0): support 实时输出

cuisongliu · Mar 25, 2020 · e71b344 · e71b344
1 parent 1615e61
commit e71b344
Show file tree

Hide file tree

Showing 6 changed files with 68 additions and 32 deletions.
diff --git a/cmd/root.go b/cmd/root.go
@@ -20,6 +20,7 @@ import (
 	"github.com/wonderivan/logger"
 	"golang.org/x/crypto/ssh"
 	"os"
+	"strings"
 	"sync"
 
 	"github.com/spf13/cobra"
@@ -47,26 +48,27 @@ var rootCmd = &cobra.Command{
 			wg.Add(1)
 			go func(host string) {
 				defer wg.Done()
-				switch mode {
-				case "ssh":
-					sshType.Cmd(host, command)
-				case "scp":
-					sshType.CopyForMD5(host, localFilePath, remoteFilePath, "")
-				case "ssh|scp":
-					sshType.Cmd(host, command)
-					sshType.CopyForMD5(host, localFilePath, remoteFilePath, "")
-				case "scp|ssh":
-					sshType.CopyForMD5(host, localFilePath, remoteFilePath, "")
-					sshType.Cmd(host, command)
-				default:
-					sshType.Cmd(host, command)
+				modes := strings.Split(mode, "|")
+				for i, _ := range modes {
+					exec(sshType, modes[i], host)
 				}
 			}(node)
 		}
 		wg.Wait()
 	},
 }
 
+func exec(ssh *sshutil.SSH, mode, host string) {
+	switch mode {
+	case "ssh":
+		ssh.Cmd(host, command)
+	case "scp":
+		ssh.CopyForMD5(host, localFilePath, remoteFilePath, "")
+	case "sshAsync":
+		_ = ssh.CmdAsync(host, command)
+	}
+}
+
 //validate host is connect
 func validate(tSSH *sshutil.SSH) {
 	if len(host) == 0 {
@@ -119,5 +121,5 @@ func init() {
 	rootCmd.Flags().StringVar(&command, "cmd", "", "exec shell")
 	rootCmd.Flags().StringVar(&localFilePath, "local-path", "", "local path , ex /etc/local.txt")
 	rootCmd.Flags().StringVar(&remoteFilePath, "remote-path", "", "local path , ex /etc/local.txt")
-	rootCmd.Flags().StringVar(&mode, "mode", "ssh", "mode type ,use | spilt . ex ssh scp ssh|scp scp|ssh")
+	rootCmd.Flags().StringVar(&mode, "mode", "ssh", "mode type ,use | spilt . ex ssh sshAsync scp ssh|scp scp|ssh")
 }
diff --git a/pkg/cmd/cmd.go b/pkg/cmd/cmd.go
@@ -16,7 +16,7 @@ func Cmd(name string, arg ...string) {
 	cmd.Stdout = os.Stdout
 	err := cmd.Run()
 	if err != nil {
-		logger.Error("os call error.", err)
+		logger.Error("[os]os call error.", err)
 	}
 }
 
@@ -30,7 +30,7 @@ func CmdToString(name string, arg ...string) string {
 	cmd.Stderr = &b
 	err := cmd.Run()
 	if err != nil {
-		logger.Error("os call error.", err)
+		logger.Error("[os]os call error.", err)
 		return ""
 	}
 	return b.String()

diff --git a/pkg/cmd/watch.go b/pkg/cmd/watch.go
@@ -27,7 +27,7 @@ func LoggerFileSize(filename string, size int) {
 			}
 			lengthFloat := float64(lengthByte)
 			value, _ := strconv.ParseFloat(fmt.Sprintf("%.2f", lengthFloat/oneMBByte), 64)
-			logger.Alert("[%s]transfer total size is: %.2f%s", filename, value, "MB")
+			logger.Alert("[os][%s]transfer total size is: %.2f%s", filename, value, "MB")
 		}
 	}
 }
@@ -44,7 +44,7 @@ func IsFilExist(filepath string) bool {
 	count, err := strconv.Atoi(strings.TrimSpace(data))
 	defer func() {
 		if r := recover(); r != nil {
-			logger.Error("[%s]RemoteFilExist:%s", filepath, err)
+			logger.Error("[os][%s]RemoteFilExist:%s", filepath, err)
 		}
 	}()
 	if err != nil {

diff --git a/pkg/sshutil/scp.go b/pkg/sshutil/scp.go
@@ -19,17 +19,17 @@ func (ss *SSH) CopyForMD5(host, localFilePath, remoteFilePath, md5 string) bool
 	if md5 == "" {
 		md5 = md5sum.FromLocal(localFilePath)
 	}
-	logger.Debug("source file md5 value is %s", md5)
+	logger.Debug("[ssh]source file md5 value is %s", md5)
 	ss.Copy(host, localFilePath, remoteFilePath)
 	remoteMD5 := ss.Md5Sum(host, remoteFilePath)
-	logger.Debug("host: %s , remote md5: %s", host, remoteMD5)
+	logger.Debug("[ssh]host: %s , remote md5: %s", host, remoteMD5)
 	remoteMD5 = strings.TrimSpace(remoteMD5)
 	md5 = strings.TrimSpace(md5)
 	if remoteMD5 == md5 {
-		logger.Info("md5 validate true")
+		logger.Info("[ssh]md5 validate true")
 		return true
 	}
-	logger.Error("md5 validate false")
+	logger.Error("[ssh]md5 validate false")
 	return false
 }
 func (ss *SSH) Md5Sum(host, remoteFilePath string) string {
@@ -43,7 +43,7 @@ func (ss *SSH) Copy(host, localFilePath, remoteFilePath string) {
 	sftpClient, err := ss.sftpConnect(host)
 	defer func() {
 		if r := recover(); r != nil {
-			logger.Error("[%s]scpCopy: %s", host, err)
+			logger.Error("[ssh][%s]scpCopy: %s", host, err)
 		}
 	}()
 	if err != nil {
@@ -53,7 +53,7 @@ func (ss *SSH) Copy(host, localFilePath, remoteFilePath string) {
 	srcFile, err := os.Open(localFilePath)
 	defer func() {
 		if r := recover(); r != nil {
-			logger.Error("[%s]scpCopy: %s", host, err)
+			logger.Error("[ssh][%s]scpCopy: %s", host, err)
 		}
 	}()
 	if err != nil {
@@ -64,7 +64,7 @@ func (ss *SSH) Copy(host, localFilePath, remoteFilePath string) {
 	dstFile, err := sftpClient.Create(remoteFilePath)
 	defer func() {
 		if r := recover(); r != nil {
-			logger.Error("[%s]scpCopy: %s", host, err)
+			logger.Error("[ssh][%s]scpCopy: %s", host, err)
 		}
 	}()
 	if err != nil {
@@ -80,7 +80,7 @@ func (ss *SSH) Copy(host, localFilePath, remoteFilePath string) {
 		}
 		length, _ := dstFile.Write(buf[0:n])
 		totalMB += length / oneMBByte
-		logger.Alert("[%s]transfer total size is: %d%s", host, totalMB, "MB")
+		logger.Alert("[ssh][%s]transfer total size is: %d%s", host, totalMB, "MB")
 	}
 }
 

diff --git a/pkg/sshutil/ssh.go b/pkg/sshutil/ssh.go
@@ -1,17 +1,18 @@
 package sshutil
 
 import (
+	"bufio"
 	"github.com/wonderivan/logger"
 	"strings"
 )
 
 //Cmd is in host exec cmd
 func (ss *SSH) Cmd(host string, cmd string) []byte {
-	logger.Info("[%s]exec cmd is : %s", host, cmd)
+	logger.Info("[ssh][%s]exec cmd is : %s", host, cmd)
 	session, err := ss.Connect(host)
 	defer func() {
 		if r := recover(); r != nil {
-			logger.Error("[%s]Error create ssh session failed,%s", host, err)
+			logger.Error("[ssh][%s]Error create ssh session failed,%s", host, err)
 		}
 	}()
 	if err != nil {
@@ -20,10 +21,10 @@ func (ss *SSH) Cmd(host string, cmd string) []byte {
 	defer session.Close()
 
 	b, err := session.CombinedOutput(cmd)
-	logger.Debug("[%s]command result is: %s", host, string(b))
+	logger.Debug("[ssh][%s]command result is: %s", host, string(b))
 	defer func() {
 		if r := recover(); r != nil {
-			logger.Error("[%s]Error exec command failed: %s", host, err)
+			logger.Error("[ssh][%s]Error exec command failed: %s", host, err)
 		}
 	}()
 	if err != nil {
@@ -32,6 +33,39 @@ func (ss *SSH) Cmd(host string, cmd string) []byte {
 	return b
 }
 
+func (ss *SSH) CmdAsync(host string, cmd string) error {
+	logger.Info("[ssh][%s]exec cmd is : %s", host, cmd)
+	session, err := ss.Connect(host)
+	if err != nil {
+		logger.Error("[ssh][%s]Error create ssh session failed,%s", host, err)
+		return err
+	}
+	defer session.Close()
+	stdout, err := session.StdoutPipe()
+	if err != nil {
+		logger.Error("[ssh][%s]Unable to request StdoutPipe(): %s", host, err)
+		return err
+	}
+	if err := session.Start(cmd); err != nil {
+		logger.Error("[ssh][%s]Unable to execute command: %s", host, err)
+		return err
+	}
+	done := make(chan bool, 1)
+	go func() {
+		for {
+			r := bufio.NewReader(stdout)
+			line, _, err := r.ReadLine()
+			if line == nil || err != nil {
+				done <- true
+			} else {
+				logger.Info("[ssh][%s]: %s", host, line)
+			}
+		}
+	}()
+	<-done
+	return nil
+}
+
 //CmdToString is in host exec cmd and replace to spilt str
 func (ss *SSH) CmdToString(host, cmd, spilt string) string {
 	data := ss.Cmd(host, cmd)

diff --git a/pkg/sshutil/watch.go b/pkg/sshutil/watch.go
@@ -28,7 +28,7 @@ func (ss *SSH) LoggerFileSize(host, filename string, size int) {
 			}
 			lengthFloat := float64(lengthByte)
 			value, _ := strconv.ParseFloat(fmt.Sprintf("%.2f", lengthFloat/oneMBByte), 64)
-			logger.Alert("[%s]transfer total size is: %.2f%s", host, value, "MB")
+			logger.Alert("[ssh][%s]transfer total size is: %.2f%s", host, value, "MB")
 		}
 	}
 }
@@ -46,7 +46,7 @@ func (ss *SSH) IsFilExist(host, remoteFilePath string) bool {
 	count, err := strconv.Atoi(string(data))
 	defer func() {
 		if r := recover(); r != nil {
-			logger.Error("[%s]RemoteFilExist:%s", host, err)
+			logger.Error("[ssh][%s]RemoteFilExist:%s", host, err)
 		}
 	}()
 	if err != nil {