package rmt import ( "b612.me/notify" "b612.me/notify/starnotify" "b612.me/starlog" "b612.me/staros" "fmt" "github.com/spf13/cobra" "os" "os/signal" "runtime" "strings" "time" ) type RmtCmd struct { Env []string WorkDir string Cmd string Daemon bool Stream bool } type RmtCmdBack struct { RetCode int OutPut string } var ( rmtListenPort string rmtPol string ) func init() { notify.RegisterName("remoteback", RmtCmdBack{}) notify.RegisterName("remotecmd", RmtCmd{}) Cmds.Flags().StringVarP(&rmtListenPort, "port", "p", "5780", "Listen Port") Cmds.Flags().StringVarP(&rmtPol, "protocol", "o", "tcp", "Remote protocol") } var Cmds = &cobra.Command{ Use: "rmts", Short: "远程命令执行服务端", Run: func(cmd *cobra.Command, args []string) { if rmtListenPort == "" { starlog.Errorln("Please Enter Port") os.Exit(1) } err := starnotify.NewServer("s").Listen(rmtPol, "0.0.0.0:"+rmtListenPort) if err != nil { starlog.Errorln("Create Listener Failed", err) os.Exit(2) } starnotify.S("s").SetLink("cmd", ListenAndServe) starlog.Infoln("Service Running") stopSig := make(chan os.Signal) signal.Notify(stopSig, os.Kill, os.Interrupt) <-stopSig starlog.Noticeln("Recv Stop Sig") starnotify.S("s").Stop() }, } func ListenAndServe(msg *notify.Message) { data, err := msg.Value.ToInterface() if err != nil { return } cmd, ok := data.(RmtCmd) if !ok { return } var sysName, sysArg string = "bash", "-c" if runtime.GOOS == "windows" { sysName = "cmd.exe" sysArg = "/c" } if cmd.Cmd == "" { msg.ReplyObj(RmtCmdBack{ RetCode: 255, OutPut: "Please enter the command", }) return } starlog.Noticef("Recv Command:%+v\n", cmd) var myCmd *staros.StarCmd myCmd, err = staros.Command(sysName, sysArg, cmd.Cmd) if err != nil { msg.ReplyObj(RmtCmdBack{ RetCode: 255, OutPut: err.Error(), }) return } if cmd.WorkDir != "" { myCmd.CMD.Dir = cmd.WorkDir } if len(cmd.Env) > 0 { myCmd.CMD.Env = cmd.Env } if cmd.Daemon { err := myCmd.Release() if err != nil { msg.ReplyObj(RmtCmdBack{ RetCode: 254, OutPut: err.Error(), }) return } pid := myCmd.CMD.Process.Pid msg.ReplyObj(RmtCmdBack{ RetCode: 0, OutPut: fmt.Sprintf("Runned,PID is %d", pid), }) return } err = myCmd.Start() if err != nil { msg.ReplyObj(RmtCmdBack{ RetCode: 254, OutPut: err.Error(), }) return } for myCmd.IsRunning() { time.Sleep(time.Millisecond * 10) if cmd.Stream { std, err := myCmd.NowAllOutput() if err != nil { std += "\n" + err.Error() } std = strings.TrimSpace(std) if len(std) == 0 { continue } msg.ClientConn.Server().SendObj(msg.ClientConn, "stream", std) } } time.Sleep(time.Millisecond * 50) if !cmd.Stream { std := myCmd.AllStdOut() if myCmd.AllStdErr() != nil { std += "\n" + myCmd.AllStdErr().Error() } msg.ReplyObj(RmtCmdBack{ RetCode: myCmd.ExitCode(), OutPut: std, }) } else { std, err := myCmd.NowAllOutput() if err != nil { std += "\n" + err.Error() } msg.ClientConn.Server().SendObj(msg.ClientConn, "stream", std) time.Sleep(time.Millisecond * 1000) err = msg.ReplyObj(RmtCmdBack{ RetCode: myCmd.ExitCode(), OutPut: "", }) if err != nil { starlog.Warningln("Reply failed:", err) } } return }