diff --git a/cmd/root.go b/cmd/root.go index a624975..e834523 100644 --- 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,19 +48,9 @@ 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) } @@ -67,6 +58,17 @@ var rootCmd = &cobra.Command{ }, } +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 index 1efad75..e2a9f1a 100644 --- 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 index 908c5c4..fa30a6d 100644 --- 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 index 872e45b..1bdfae9 100644 --- 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 index f10101a..24600cc 100644 --- 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 index a612356..6274374 100644 --- 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 {