zerone/mssh/main.go

package main

import (
	"encoding/base64"
	"fmt"
	"os"
	"os/user"
	"regexp"
	"strings"

	"git.wecise.com/wecise/common/matrix/logger"
	"git.wecise.com/wecise/common/matrix/util"
	"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")
}

func uchi() (u, c, h string, i int) {
	if len(os.Args) < 2 {
		return
	}
	if msi := regexp.MustCompile(`^([^:]+):(.*)@([^@]*)$`).FindStringSubmatchIndex(os.Args[1]); msi != nil {
		u = os.Args[1][msi[2]:msi[3]]
		c = "=" + os.Args[1][msi[4]:msi[5]]
		h = os.Args[1][msi[6]:msi[7]]
		i = 2
		return
	}
	if msi := regexp.MustCompile(`^([^:]+)@([^@]*)$`).FindStringSubmatchIndex(os.Args[1]); msi != nil {
		u = os.Args[1][msi[2]:msi[3]]
		h = os.Args[1][msi[6]:msi[7]]
		i = 2
	} else if os.Args[1][0] != '-' && strings.Index(os.Args[1], "=") < 0 {
		h = os.Args[1]
		i = 2
	} else {
		i = 1
	}
	if u == "" {
		user, err := user.Current()
		if err != nil {
			logger.Error(util.ErrorWithSourceLine(err))
			return "", "", "", 0
		}
		u = user.Username
	}
	if h == "" {
		h = "127.0.0.1"
	}
	return
}

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 main() {
	u, c, h, i := uchi()
	if i == 0 {
		fmt.Println("usage:")
		fmt.Println("  msh user:password@host|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)")
		return
	}
	kvs := parseArgs(os.Args[i:])
	if c == "" {
		for _, kv := range kvs {
			if kv.Key == "a" {
				c = kv.Val
				if c == "" {
					c = "="
				}
				break
			}
			if kv.Key == "p" {
				c = "=" + kv.Val
				break
			}
		}
	}
	p := c
	if len(p) > 0 && 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: "ssh " + u + ":" + p + "@" + h, Password: c, Regexps: []*Matcher{{Regexp: nil, Output: ""}}, Endregx: regxprompt}}
	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":
			re, err := regexp.Compile(val)
			if err != nil {
				logger.Error("arg", i, util.ErrorWithSourceLine(err))
				return
			} else {
				regxyesno = &Regexp{re}
			}
		case "rc":
			re, err := regexp.Compile(val)
			if err != nil {
				logger.Error("arg", i, util.ErrorWithSourceLine(err))
				return
			} else {
				regxpassword = &Regexp{re}
			}
		case "rp":
			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))
	}
	node := &Node{
		Name:       "x",
		Alias:      "x",
		Host:       h,
		User:       u,
		Port:       22,
		KeyPath:    "",
		Passphrase: "",
		Password:   p,
		Commands:   cmds,
		Children:   []*Node{},
		Jump:       []*Node{},
	}
	sshc := NewClient(node)
	sshc.Login()
}