star/rmt/remoteCmdServer.go

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: "simple remote shell server",
	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
}