zerone/_/mainx.go

package main

import (
	"bufio"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"os"
	"os/exec"
	"os/signal"
	"os/user"
	"regexp"
	"strings"
	"syscall"
	"time"

	"git.wecise.com/wecise/common/matrix/logger"
	"git.wecise.com/wecise/common/matrix/util"
	"github.com/creack/pty"
	"golang.org/x/term"
	"gopkg.in/yaml.v2"
)

func init() {
	logger.SetConsole(true)
	logger.SetLevel(logger.TRACE)
	logger.SetFormat("yyyy-MM-dd HH:mm:ss.SSS [level] msg", "\n")
}

type KV struct{ Key, Val string }

func parseArgs(args []string) (kvs []*KV) {
	argk := ""
	argv := ""
	for _, arg := range args {
		if argk != "" {
			argv = arg
		} else if regexp.MustCompile(`^\-\w+$`).MatchString(arg) {
			argk = arg[1:]
			continue
		} else {
			kv := strings.SplitN(arg, "=", 2)
			if len(kv) == 2 {
				argk = kv[0]
				argv = kv[1]
			} else {
				argk = ""
				argv = arg
			}
		}
		kvs = append(kvs, &KV{argk, argv})
		argk, argv = "", ""
	}
	if argk != "" {
		kvs = append(kvs, &KV{argk, argv})
	}
	return
}

func usage() {
	fmt.Println("usage:")
	fmt.Println("  msh p=password|a=passcode [[c=]command [p=password|a=passcode] [x=cmd-end-regexp] [[r=regexp] [o=output|n=outputline]]...]...")
	fmt.Println("  a=passcode should be base64 encoded, or use p=password")
	fmt.Println("  debug info include: p(progress) a(argments) m(match) 1(all)")
}

func main() {
	// get user
	u := func() string {
		user, err := user.Current()
		if err != nil {
			logger.Error(util.ErrorWithSourceLine(err))
			return ""
		}
		return user.Username
	}()
	// get passcode，取默认密码设置
	c := ""
	kvs := parseArgs(os.Args)
	for _, kv := range kvs {
		if kv.Key == "a" {
			c = kv.Val
			if c == "" {
				c = "="
			}
			break
		}
		if kv.Key == "p" {
			c = "=" + kv.Val
			break
		}
	}
	// get password
	p := c
	if p == "" {
		usage()
		return
	} else if p[0:1] == "=" {
		p = p[1:]
	} else {
		x, e := base64.RawStdEncoding.DecodeString(p)
		if e == nil {
			p = string(x)
		}
		// else 不是Base64编码，保持原值
	}
	// explainArgs key, val
	cmds := []*Command{{Cmd: "sh", Password: c, Regexps: []*Matcher{{Regexp: nil, Output: ""}}, Endregx: regxprompt}}
	i := 0
	// 掠过第一个参数，当前执行程序
	kvs = kvs[1:]
	for _, kv := range kvs {
		key, val := kv.Key, kv.Val
		switch key {
		case "", "cmd", "command", "c":
			cmds = append(cmds, &Command{Cmd: val, Password: c, Regexps: []*Matcher{{Regexp: nil, Output: ""}}, Endregx: regxprompt})
		case "ry", "regex_yes_no":
			re, err := regexp.Compile(val)
			if err != nil {
				logger.Error("arg", i, util.ErrorWithSourceLine(err))
				return
			} else {
				regxyesno = &Regexp{re}
			}
		case "rc", "regex_passcode", "regex_password":
			re, err := regexp.Compile(val)
			if err != nil {
				logger.Error("arg", i, util.ErrorWithSourceLine(err))
				return
			} else {
				regxpassword = &Regexp{re}
			}
		case "rp", "regex_prompt":
			re, err := regexp.Compile(val)
			if err != nil {
				logger.Error("arg", i, util.ErrorWithSourceLine(err))
				return
			} else {
				regxprompt = &Regexp{re}
			}
		case "password", "code", "pass", "p":
			cmds[len(cmds)-1].Password = "=" + val
		case "passcode", "b64code", "a":
			cmds[len(cmds)-1].Password = val
		case "re", "r", "regex":
			if val == "" {
				cmds[len(cmds)-1].Regexps = append(cmds[len(cmds)-1].Regexps, &Matcher{Regexp: nil})
			} else {
				re, err := regexp.Compile(val)
				if err != nil {
					logger.Error("arg", i, util.ErrorWithSourceLine(err))
					return
				} else {
					cmds[len(cmds)-1].Regexps = append(cmds[len(cmds)-1].Regexps, &Matcher{Regexp: &Regexp{re}})
				}
			}
		case "x", "end":
			if val == "" {
				cmds[len(cmds)-1].Endregx = regxprompt
			} else {
				re, err := regexp.Compile(val)
				if err != nil {
					logger.Error("arg", i, util.ErrorWithSourceLine(err))
					return
				} else {
					cmds[len(cmds)-1].Endregx = &Regexp{re}
				}
			}
		case "out", "o", "output":
			cmds[len(cmds)-1].Regexps[len(cmds[len(cmds)-1].Regexps)-1].Output += val
		case "outln", "n", "outputline":
			cmds[len(cmds)-1].Regexps[len(cmds[len(cmds)-1].Regexps)-1].Output += val + "\n"
		case "debug", "d":
			cmds[len(cmds)-1].Regexps[len(cmds[len(cmds)-1].Regexps)-1].Debug += val
		}
	}
	if strings.Index(cmds[0].Regexps[0].Debug, "a") >= 0 || strings.Index(cmds[0].Regexps[0].Debug, "1") >= 0 {
		//bs, _ := json.MarshalIndent(cmds, "", "  ")
		bs, _ := yaml.Marshal(cmds)
		logger.Debug("arguments:\n" + string(bs))
	}
	client := &Client{
		User:     u,
		Password: p,
		Commands: cmds,
	}
	client.Run()
}

type Regexp struct {
	*regexp.Regexp
}

func (r *Regexp) MarshalJSON() ([]byte, error) {
	if r == nil || r.Regexp == nil {
		return json.Marshal(nil)
	}
	return json.Marshal(r.String())
}

func (r *Regexp) MarshalYAML() (interface{}, error) {
	if r == nil || r.Regexp == nil {
		return nil, nil
	}
	return r.String(), nil
}

type Matcher struct {
	Regexp *Regexp
	Output string `yaml:"output"`
	Debug  string `yaml:"debug"`
}

type Command struct {
	Cmd      string `yaml:"cmd"`
	Password string `yaml:"password"`
	Regexps  []*Matcher
	Endregx  *Regexp
}

var regxyesno = &Regexp{regexp.MustCompile(`.*(\(yes\/no\)\?)\s*$`)}
var regxpassword = &Regexp{regexp.MustCompile(`.*([Pp]assword:)\s*$`)}
var regxprompt = &Regexp{regexp.MustCompile(`.*([%#\$\>])\s*$`)}

type Client struct {
	User     string
	Password string
	Commands []*Command
}

func Pipe(reader io.Reader, writer io.Writer, pfs ...func(lastbs []byte, sin string) (sout string)) {
	btr := NewBTReader("P:", reader, 30*time.Millisecond, 1024)
	lastbs := []byte{'\n'}
	for {
		bs, err := btr.Read()
		if err != nil {
			if err == io.EOF {
				return
			}
			logger.Error(util.ErrorWithSourceLine(err))
			return
		}
		if len(bs) > 0 {
			xbs := bs
			for _, pf := range pfs {
				if pf != nil {
					s := pf(lastbs, string(xbs))
					xbs = []byte(s)
				}
			}
			if writer != nil {
				writer.Write(xbs)
			}
			lastbs = bs
		}
	}
}

type TerminalIO struct {
	termin  io.Reader
	termout io.Writer
}

func (tio *TerminalIO) Read(bs []byte) (n int, e error) {
	return tio.termin.Read(bs)
}

func (tio *TerminalIO) Write(bs []byte) (n int, e error) {
	return tio.termout.Write(bs)
}

func (c *Client) Run() {
	// Create arbitrary command.
	cmd := exec.Command("sh")

	// Start the command with a pty.
	ptmx, err := pty.Start(cmd)
	if err != nil {
		logger.Error(err)
		return
	}
	// Make sure to close the pty at the end.
	defer func() { _ = ptmx.Close() }() // Best effort.

	// Handle pty size.
	ch := make(chan os.Signal, 1)
	signal.Notify(ch, syscall.SIGWINCH)
	go func() {
		for range ch {
			if err := pty.InheritSize(os.Stdin, ptmx); err != nil {
				logger.Error("error resizing pty:", err)
			}
		}
	}()
	ch <- syscall.SIGWINCH                        // Initial resize.
	defer func() { signal.Stop(ch); close(ch) }() // Cleanup signals when done.

	// Set stdin in raw mode.
	oldState, err := term.MakeRaw(int(os.Stdin.Fd()))
	if err != nil {
		logger.Error(err)
	}
	defer func() { _ = term.Restore(int(os.Stdin.Fd()), oldState) }() // Best effort.

	go func() { _, _ = io.Copy(ptmx, os.Stdin) }()
	//_, _ = io.Copy(os.Stdout, ptmx)
	Pipe(ptmx, os.Stdout)

	return
}

func (c *Client) RunX() {
	logger.Info("tsh ready")
	exit := -1
	cmdsctrlreader, cmdsctrlwriter := io.Pipe()
	terminreader, terminwriter := io.Pipe()
	termoutreader, termoutwriter := io.Pipe()
	termerrreader, termerrwriter := io.Pipe()
	go Pipe(os.Stdin, terminwriter)
	go Pipe(cmdsctrlreader, terminwriter, func(lastbs []byte, sin string) (sout string) {
		sout = sin
		termoutwriter.Write([]byte(sout))
		return
	})
	// autocmds := []string{"pwd", "echo `ls`", "sh -i", "exit"}
	// go Pipe(termoutreader, os.Stdout, func(lastbs []byte, sin string) (sout string) {
	// 	sout = sin
	// 	return
	// })
	// go Pipe(termerrreader, os.Stderr)
	eoc, eos := c.RunShell(terminreader, termoutwriter, termerrwriter)
	c.IOProc(termoutreader, termerrreader, os.Stdout, os.Stderr, cmdsctrlwriter)
	for exit < 0 {
		select {
		case <-eoc:
			// if len(autocmds) > 0 {
			// 	cmdsctrlwriter.Write([]byte(autocmds[0] + "\n"))
			// 	autocmds = autocmds[1:]
			// }
		case exit = <-eos:
		}
	}
	logger.Info("tsh exit")
}

func (c *Client) RunShell(stdin io.Reader, cmdout, cmderr io.Writer) (eoc <-chan int, exit <-chan int) {
	endofcmd := make(chan int)
	endofshell := make(chan int)
	prompt := util.Hostname() + ":> "
	// cmd := exec.Command("bash", "-c", "export LANG=en_US.utf8; export LC_ALL=en_US.utf8; export PS1='\\h"+prompt+"'; bash")
	cmd := exec.Command("bash")
	// cmd.Stdin = stdin
	cmdinwriter, err := cmd.StdinPipe()
	if err != nil {
		logger.Error(err)
	}
	newcmd := true
	eocflag := "<end-of-command-in-zerone-t-shell>"
	cmdeoc := "echo '" + eocflag + "'\n"
	go Pipe(stdin, cmdinwriter, func(lastbs []byte, sin string) (sout string) {
		sout = sin
		if newcmd {
			newcmd = false
			sout = sout[:len(sout)-1] + ";" + cmdeoc
		}
		return
	})
	cmdoutreader, err := cmd.StdoutPipe()
	if err != nil {
		logger.Error(err)
	}
	go Pipe(cmdoutreader, cmdout, func(lastbs []byte, sin string) (sout string) {
		sout = sin
		// sout = strings.ReplaceAll(strings.ReplaceAll(sin, "\r", "\rS:"), "\n", "\nO:")
		// if strings.HasSuffix(sout, "\nO:") {
		// 	sout = sout[:len(sout)-3]
		// }
		// if lastbs[len(lastbs)-1] == '\n' {
		// 	sout = "\nO:" + sout
		// }
		if strings.HasSuffix(sout, eocflag+"\n") {
			sout = sout[:len(sout)-len(eocflag)-1]
			sout += prompt
			newcmd = true
			endofcmd <- 0
		}
		return
	})
	cmderrreader, err := cmd.StderrPipe()
	go Pipe(cmderrreader, cmderr, func(lastbs []byte, sin string) (sout string) {
		sout = sin
		// sout = strings.ReplaceAll(strings.ReplaceAll(sin, "\r", "\rQ:"), "\n", "\nE:")
		// if strings.HasSuffix(sout, "\nE:") {
		// 	sout = sout[:len(sout)-3]
		// }
		// if lastbs[len(lastbs)-1] == '\n' {
		// 	sout = "\nE:" + sout
		// }
		// sout = strings.ReplaceAll(sou, "no job control in this shell\n", "ready\n")
		return
	})
	go func() {
		err := cmd.Run()
		if err != nil {
			logger.Error(err)
		}
		endofshell <- cmd.ProcessState.ExitCode()
	}()
	cmdinwriter.Write([]byte(cmdeoc))
	return endofcmd, endofshell
}

func (c *Client) IOProc(cmdout_pipe_reader, cmderr_pipe_reader io.Reader, stdout io.Writer, stderr io.Writer, cmdinwriter io.Writer) {
	cmdidx := 0
	outputproc := func(flag string, cmdout io.Reader, stdout io.Writer, cmdinwriter io.Writer) {
		as := ""
		change := false
		stdoutbtr := NewBTReader(flag, cmdout, 50*time.Millisecond, 1024)
		for {
			bs, err := stdoutbtr.Read()
			if err != nil {
				if err == io.EOF {
					return
				}
				logger.Error(util.ErrorWithSourceLine(err))
				return
			}
			if len(bs) > 0 {
				change = true
				s := string(bs)
				stdout.Write([]byte(s))
				as += s
				logger.Warn(s)
			} else if change {
				change = false
				if cmdidx < len(c.Commands) {
					if msi := c.Commands[cmdidx].Endregx.FindStringSubmatchIndex(as); msi != nil {
						logger.Error("EndregexMatch", as)
						match := as[msi[0]:msi[1]]
						if strings.Index(c.Commands[cmdidx].Regexps[0].Debug, "m") >= 0 || strings.Index(c.Commands[cmdidx].Regexps[0].Debug, "1") >= 0 {
							logger.Trace("match end:", "'"+match+"'")
						}
						as = "" // 全清，开始新的命令
						cmdidx++
						if cmdidx >= len(c.Commands) {
							return
						}
						if strings.Index(c.Commands[cmdidx].Regexps[0].Debug, "p") >= 0 || strings.Index(c.Commands[cmdidx].Regexps[0].Debug, "1") >= 0 {
							logger.Trace("command:", c.Commands[cmdidx].Cmd)
						}
						// stdout.Write([]byte(c.Commands[cmdidx].Cmd + "\n"))
						cmdinwriter.Write([]byte(c.Commands[cmdidx].Cmd + "\n"))
						continue
					}
					for _, regex := range c.Commands[cmdidx].Regexps {
						if regex == nil {
							continue
						}
						if regex.Regexp == nil {
							// like match all
							if regex.Output != "" {
								// stdout.Write([]byte(regex.Output))
								cmdinwriter.Write([]byte(regex.Output))
							}
						} else if msi := regex.Regexp.FindStringSubmatchIndex(as); msi != nil {
							match := as[msi[0]:msi[1]]
							if len(msi) >= 4 {
								match = as[msi[2]:msi[3]]
								as = as[msi[3]:] // 清除已处理完的内容
							} else {
								as = as[msi[1]:] // 清除已处理完的内容
							}
							if strings.Index(regex.Debug, "m") >= 0 || strings.Index(regex.Debug, "1") >= 0 {
								logger.Trace("match regex:", "'"+match+"'")
							}
							if regex.Output != "" {
								// stdout.Write([]byte(regex.Output))
								cmdinwriter.Write([]byte(regex.Output))
							}
						}
					}
				}
				if msi := regxyesno.FindStringSubmatchIndex(as); msi != nil {
					match := as[msi[0]:msi[1]]
					if len(msi) >= 4 {
						as = as[msi[3]:] // 清除已处理完的内容
					} else {
						as = as[msi[1]:] // 清除已处理完的内容
					}
					if strings.Index(c.Commands[0].Regexps[0].Debug, "m") >= 0 || strings.Index(c.Commands[0].Regexps[0].Debug, "1") >= 0 ||
						cmdidx < len(c.Commands) &&
							(strings.Index(c.Commands[cmdidx].Regexps[0].Debug, "m") >= 0 || strings.Index(c.Commands[cmdidx].Regexps[0].Debug, "1") >= 0) {
						logger.Trace("match yesno:", "'"+match+"'")
					}
					// stdout.Write([]byte("yes\n"))
					cmdinwriter.Write([]byte("yes\n"))
				}
				if msi := regxpassword.FindStringSubmatchIndex(as); msi != nil {
					logger.Error(as)
					match := as[msi[0]:msi[1]]
					if len(msi) >= 4 {
						as = as[msi[3]:] // 清除已处理完的内容
					} else {
						as = as[msi[1]:] // 清除已处理完的内容
					}
					p := c.Commands[0].Password
					if cmdidx < len(c.Commands) {
						p = c.Commands[cmdidx].Password
					}
					if strings.Index(c.Commands[0].Regexps[0].Debug, "m") >= 0 || strings.Index(c.Commands[0].Regexps[0].Debug, "1") >= 0 ||
						cmdidx < len(c.Commands) &&
							(strings.Index(c.Commands[cmdidx].Regexps[0].Debug, "m") >= 0 || strings.Index(c.Commands[cmdidx].Regexps[0].Debug, "1") >= 0) {
						logger.Trace("match password:", "'"+match+"'")
					}
					if p != "" {
						if p[0:1] == "=" {
							p = p[1:]
						} else {
							x, e := base64.RawStdEncoding.DecodeString(p)
							if e == nil {
								p = string(x)
							}
							// else 不是Base64编码，保持原值
						}
						// don't echo password
						if c.Commands[0].Regexps[0].Debug != "" || cmdidx < len(c.Commands) && c.Commands[cmdidx].Regexps[0].Debug != "" {
							stdout.Write([]byte(p + "\n"))
						}
						cmdinwriter.Write([]byte(p + "\n"))
					}
				}
				if len(as) > 1024 {
					as = as[len(as)-1024:]
				}
			}
		}
	}
	go outputproc("O:", cmdout_pipe_reader, stdout, cmdinwriter)
	go outputproc("E:", cmderr_pipe_reader, stderr, cmdinwriter)
}

type BTReader struct {
	*bufio.Reader
	flag    string
	chbs    chan []byte
	cher    chan error
	timeout time.Duration
	sizeout int
	err     error
}

func NewBTReader(flag string, reader io.Reader, timeout time.Duration, sizeout int) (me *BTReader) {
	me = &BTReader{Reader: bufio.NewReader(reader),
		flag:    flag,
		chbs:    make(chan []byte),
		cher:    make(chan error),
		timeout: timeout,
		sizeout: sizeout,
	}
	go func() {
		bs := make([]byte, me.Size())
		for {
			n, err := me.Reader.Read(bs[:1])
			if err != nil {
				me.cher <- err
				return
			}
			x, err := me.Reader.Read(bs[1 : me.Reader.Buffered()+1])
			if err != nil {
				me.cher <- err
				return
			}
			n += x
			abs := make([]byte, n)
			copy(abs, bs[:n])
			me.chbs <- abs
		}
	}()
	return
}

// 指定时间内没有新数据进入，且有积累数据，或积累数据超过指定数量，即返回
func (me *BTReader) Read() (rbs []byte, err error) {
	if me.err != nil {
		return nil, me.err
	}
	for {
		t := time.NewTimer(me.timeout)
		select {
		case me.err = <-me.cher:
			if len(rbs) > 0 {
				// 返回最后的数据，下次读时才返回错误
				return rbs, nil
			}
			return nil, me.err
		case abs := <-me.chbs:
			rbs = append(rbs, abs...)
			if len(rbs) > me.sizeout {
				return
			}
			t.Stop()
			t.Reset(me.timeout)
		case <-t.C:
			if len(rbs) == 0 {
				t.Stop()
				t.Reset(me.timeout)
			}
			return
		}
	}
}