You cannot select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
star/rmt/remoteCmdServer.go

167 lines
3.3 KiB
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
}