zerone/_/main.go

package main

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

	"git.wecise.com/wecise/common/matrix/logger"
	"git.wecise.com/wecise/common/matrix/util"
	"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{}
	i := 0
	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)) {
	lastbs := []byte{'\n'}
	for {
		bs := make([]byte, 1024)
		n, err := reader.Read(bs)
		if err != nil {
			if err == io.EOF {
				return
			}
			logger.Error(util.ErrorWithSourceLine(err))
			return
		}
		if n > 0 {
			xbs := bs[:n]
			for _, pf := range pfs {
				if pf != nil {
					s := pf(lastbs, string(xbs))
					xbs = []byte(s)
				}
			}
			if writer != nil {
				writer.Write(xbs)
			}
			lastbs = bs[:n]
		}
	}
}

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() {
	logger.Info("msh ready")

	fdstdin := int(os.Stdin.Fd())

	var err error
	oldTermState, err := term.MakeRaw(fdstdin)
	if err != nil {
		fmt.Println(err)
		return
	}

	// cmdinreader, cmdinwriter := io.Pipe()

	terminreader, terminwriter := io.Pipe()
	go Pipe(os.Stdin, terminwriter)

	termoutreader, termoutwriter := io.Pipe()
	go Pipe(termoutreader, os.Stdout)

	tio := &TerminalIO{termin: terminreader, termout: termoutwriter}
	tm := term.NewTerminal(tio, util.Hostname()+":> ")
	rawState, err := term.GetState(fdstdin)
	if err != nil {
		fmt.Println(err)
		return
	}
	exit := -1
	// term.Restore(fdstdin, oldTermState)
	// eoc, eos := c.RunShell(cmdinreader, termoutwriter, termoutwriter)
	// select {
	// case <-eoc:
	// case exit = <-eos:
	// }
	// term.Restore(fdstdin, rawState)

	for exit < 0 {
		cmdline, err := tm.ReadLine()
		if err != nil {
			break
		}
		cmdline = strings.TrimSpace(cmdline)
		if cmdline == "" {
			continue
		}
		term.Restore(fdstdin, oldTermState)
		if cmdline == "exit" {
			break
		}
		// cmdinwriter.Write([]byte(cmdline + "\n"))
		// select {
		// case <-eoc:
		// case exit = <-eos:
		// }
		cmd := exec.Command("sh", "-c", cmdline)
		cmd.Stdin = os.Stdin
		cmd.Stdout = os.Stdout
		cmd.Stderr = os.Stderr
		cmd.Run()
		term.Restore(fdstdin, rawState)
	}

	term.Restore(fdstdin, oldTermState)

	logger.Info("msh 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 := ":> "
	cmd := exec.Command("sh", "-c", "export LANG=en_US.utf8; export LC_ALL=en_US.utf8; export PS1='\\h"+prompt+"'; sh -i")
	// cmd := exec.Command("sh")
	// cmd.Stdin = stdin
	cmdinwriter, err := cmd.StdinPipe()
	if err != nil {
		logger.Error(err)
	}
	go Pipe(stdin, cmdinwriter)
	cmdoutreader, err := cmd.StdoutPipe()
	if err != nil {
		logger.Error(err)
	}
	go Pipe(cmdoutreader, cmdout, func(lastbs []byte, sin string) (sout string) {
		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
		}
		return
	})
	cmderrreader, err := cmd.StderrPipe()
	go Pipe(cmderrreader, cmderr, func(lastbs []byte, sin string) (sout string) {
		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
		}
		if strings.HasSuffix(sout, prompt) {
			endofcmd <- 0
		}
		return
	})
	go func() {
		err := cmd.Run()
		if err != nil {
			logger.Error(err)
		}
		endofshell <- cmd.ProcessState.ExitCode()
	}()
	return endofcmd, endofshell
}

func (c *Client) IOProc(cmdout_pipe_reader, cmderr_pipe_reader io.Reader, inputEcho io.Writer, inputPipe io.Writer) {
	cmdidx := 0
	outputproc := func(stdoutbtr *BTReader, inputEcho io.Writer, inputPipe io.Writer) {
		as := ""
		change := false
		for {
			bs, err := stdoutbtr.ReadTimeout(10 * time.Millisecond)
			if err != nil {
				if err == io.EOF {
					return
				}
				logger.Error(util.ErrorWithSourceLine(err))
				return
			}
			if len(bs) > 0 {
				change = true
				s := string(bs)
				inputEcho.Write([]byte(s))
				as += s
			} else if change {
				change = false
				if cmdidx < len(c.Commands) {
					if msi := c.Commands[cmdidx].Endregx.FindStringSubmatchIndex(as); msi != nil {
						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) {
							continue
						}
						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)
						}
						inputEcho.Write([]byte(c.Commands[cmdidx].Cmd + "\n"))
						inputPipe.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 != "" {
								inputEcho.Write([]byte(regex.Output))
								inputPipe.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 != "" {
								inputEcho.Write([]byte(regex.Output))
								inputPipe.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+"'")
					}
					os.Stdout.Write([]byte("yes\n"))
					inputPipe.Write([]byte("yes\n"))
				}
				if msi := regxpassword.FindStringSubmatchIndex(as); msi != nil {
					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 != "" {
							os.Stdout.Write([]byte(p + "\n"))
						}
						inputPipe.Write([]byte(p + "\n"))
					}
				}
				if len(as) > 1024 {
					as = as[len(as)-1024:]
				}
			}
		}
	}
	go outputproc(&BTReader{Reader: bufio.NewReader(cmdout_pipe_reader)}, os.Stdout, inputPipe)
	go outputproc(&BTReader{Reader: bufio.NewReader(cmderr_pipe_reader)}, os.Stderr, inputPipe)
}

type BTReader struct {
	*bufio.Reader
	bufop int32
	chbs  chan []byte
}

func (me *BTReader) ReadTimeout(d time.Duration) (rbs []byte, err error) {
	if me.chbs == nil {
		me.chbs = make(chan []byte)
		go func() {
			n := 0
			bs := make([]byte, me.Size())
			for {
				_, err = me.ReadByte()
				if err != nil {
					return
				}
				err = me.UnreadByte()
				if err != nil {
					return
				}
				n, err = me.Read(bs[0:me.Buffered()])
				if err != nil {
					return
				}
				me.chbs <- bs[0:n]
			}
		}()
	}
	t := time.NewTimer(d)
	select {
	case rbs = <-me.chbs:
		return
	case <-t.C:
		return
	}
}