starssh/ssh.go

package starssh

import (
	"bufio"
	"encoding/base64"
	"errors"
	"fmt"
	"golang.org/x/crypto/ssh"
	"io"
	"io/ioutil"
	"net"
	"regexp"
	"strings"
	"sync"
	"time"
)

type StarSSH struct {
	Client       *ssh.Client
	PublicKey    ssh.PublicKey
	PubkeyBase64 string
	Hostname     string
	RemoteAddr   net.Addr
	Banner       string
	LoginInfo    LoginInput
	online       bool
}

type LoginInput struct {
	KeyExchanges    []string
	Ciphers         []string
	MACs            []string
	User            string
	Password        string
	Prikey          string
	PrikeyPwd       string
	Addr            string
	Port            int
	Timeout         time.Duration
	HostKeyCallback func(string, net.Addr, ssh.PublicKey) error
	BannerCallback  func(string) error
}
type StarShell struct {
	Keyword        string
	UseWaitDefault bool
	Session        *ssh.Session
	in             io.Writer
	out            *bufio.Reader
	er             *bufio.Reader
	outbyte        []byte
	errbyte        []byte
	lastout        int64
	errors         error
	isprint        bool
	isfuncs        bool
	iscolor        bool
	isecho         bool
	rw             sync.RWMutex
	funcs          func(string)
}

func Login(info LoginInput) (*StarSSH, error) {
	var (
		auth         []ssh.AuthMethod
		clientConfig *ssh.ClientConfig
		config       ssh.Config
		err          error
	)
	sshInfo := new(StarSSH)
	// get auth method
	auth = make([]ssh.AuthMethod, 0)
	if info.Prikey == "" {
		keyboardInteractiveChallenge := func(
			user,
			instruction string,
			questions []string,
			echos []bool,
		) (answers []string, err error) {
			if len(questions) == 0 {
				return []string{}, nil
			}
			return []string{info.Password}, nil
		}
		auth = append(auth, ssh.Password(info.Password))
		auth = append(auth, ssh.KeyboardInteractive(keyboardInteractiveChallenge))
	} else {
		pemBytes := []byte(info.Prikey)
		var signer ssh.Signer
		if info.PrikeyPwd == "" {
			signer, err = ssh.ParsePrivateKey(pemBytes)
		} else {
			signer, err = ssh.ParsePrivateKeyWithPassphrase(pemBytes, []byte(info.PrikeyPwd))
		}
		if err != nil {
			return nil, err
		}
		auth = append(auth, ssh.PublicKeys(signer))
	}

	if len(info.Ciphers) == 0 {
		config = ssh.Config{
			Ciphers: []string{"aes128-ctr", "aes192-ctr", "aes256-ctr", "aes128-gcm@openssh.com", "arcfour256", "arcfour128", "aes128-cbc", "3des-cbc", "aes192-cbc", "aes256-cbc", "chacha20-poly1305@openssh.com"},
		}
	} else {
		config = ssh.Config{
			Ciphers: info.Ciphers,
		}
	}

	if len(info.MACs) != 0 {
		config.MACs = info.MACs
	}
	if len(info.KeyExchanges) != 0 {
		config.KeyExchanges = info.KeyExchanges
	}

	if info.Timeout == 0 {
		info.Timeout = time.Second * 5
	}
	hostKeycbfunc := func(hostname string, remote net.Addr, key ssh.PublicKey) error {

		sshInfo.PublicKey = key
		sshInfo.RemoteAddr = remote
		sshInfo.Hostname = hostname
		if info.HostKeyCallback != nil {
			return info.HostKeyCallback(hostname, remote, key)
		}
		return nil
	}

	bannercbfunc := func(banner string) error {
		sshInfo.Banner = banner
		if info.BannerCallback != nil {
			return info.BannerCallback(banner)
		}
		return nil
	}

	clientConfig = &ssh.ClientConfig{
		User:            info.User,
		Auth:            auth,
		Timeout:         info.Timeout,
		Config:          config,
		HostKeyCallback: hostKeycbfunc,
		BannerCallback:  bannercbfunc,
	}

	// connet to ssh

	sshInfo.LoginInfo = info
	sshInfo.Client, err = ssh.Dial("tcp", fmt.Sprintf("%s:%d", info.Addr, info.Port), clientConfig)
	if err == nil && sshInfo.PublicKey != nil {
		sshInfo.online = true
		sshInfo.PubkeyBase64 = base64.StdEncoding.EncodeToString(sshInfo.PublicKey.Marshal())
	}
	return sshInfo, err
}

func LoginSimple(host string, user string, passwd string, prikeyPath string, port int, timeout time.Duration) (*StarSSH, error) {
	var info = LoginInput{
		Addr:    host,
		Port:    port,
		Timeout: timeout,
		User:    user,
	}
	if prikeyPath != "" {
		prikey, err := ioutil.ReadFile(prikeyPath)
		if err != nil {
			return nil, err
		}
		info.Prikey = string(prikey)
		if passwd != "" {
			info.PrikeyPwd = passwd
		}
	} else {
		info.Password = passwd
	}
	return Login(info)
}

func (s *StarShell) ShellWait(cmd string) (string, string, error) {
	var outc, errc string = "  ", "  "
	s.Clear()
	defer s.Clear()
	echo := "echo b7Y85R56TUY6R5UTb612"
	err := s.WriteCommand(cmd)
	if err != nil {
		return "", "", err
	}
	time.Sleep(time.Millisecond * 20)
	err = s.WriteCommand(echo)
	if err != nil {
		return "", "", err
	}
	for {
		time.Sleep(time.Millisecond * 120)
		outs := string(s.outbyte)
		errs := string(s.errbyte)
		outs = strings.TrimSpace(strings.ReplaceAll(outs, "\r\n", "\n"))
		errs = strings.TrimSpace(strings.ReplaceAll(errs, "\r\n", "\n"))
		if len(outs) >= len(cmd+"\n"+echo) && outs[0:len(cmd+"\n"+echo)] == cmd+"\n"+echo {
			outs = outs[len(cmd+"\n"+echo):]
		} else if len(outs) >= len(cmd) && outs[0:len(cmd)] == cmd {
			outs = outs[len(cmd):]
		}
		if len(errs) >= len(cmd) && errs[0:len(cmd)] == cmd {
			errs = errs[len(cmd):]
		}
		if s.UseWaitDefault {
			if strings.Index(string(outs), "b7Y85R56TUY6R5UTb612") >= 0 {
				list := strings.Split(string(outs), "\n")
				for _, v := range list {
					if strings.Index(v, "b7Y85R56TUY6R5UTb612") < 0 {
						outc += v + "\n"
					}
				}
				break
			}
			if strings.Index(string(errs), "b7Y85R56TUY6R5UTb612") >= 0 {
				list := strings.Split(string(errs), "\n")
				for _, v := range list {
					if strings.Index(v, "b7Y85R56TUY6R5UTb612") < 0 {
						errc += v + "\n"
					}
				}
				break
			}
		}
		if s.Keyword != "" {
			if strings.Index(string(outs), s.Keyword) >= 0 {
				list := strings.Split(string(outs), "\n")
				for _, v := range list {
					if strings.Index(v, s.Keyword) < 0 && strings.Index(v, "b7Y85R56TUY6R5UTb612") < 0 {
						outc += v + "\n"
					}
				}
				break
			}
			if strings.Index(string(errs), s.Keyword) >= 0 {
				list := strings.Split(string(errs), "\n")
				for _, v := range list {
					if strings.Index(v, s.Keyword) < 0 && strings.Index(v, "b7Y85R56TUY6R5UTb612") < 0 {
						errc += v + "\n"
					}
				}
				break
			}
		}
	}
	return s.TrimColor(strings.TrimSpace(outc)), s.TrimColor(strings.TrimSpace(errc)), err
}

func (s *StarShell) Close() error {
	return s.Session.Close()
}

func (s *StarShell) SwitchNoColor(is bool) {
	s.iscolor = is
}

func (s *StarShell) SwitchEcho(is bool) {
	s.isecho = is
}

func (s *StarShell) TrimColor(str string) string {
	if s.iscolor {
		return SedColor(str)
	}
	return str
}

/*
本函数控制是否在本地屏幕上打印远程Shell的输出内容[true|false]
*/
func (s *StarShell) SwitchPrint(run bool) {
	s.isprint = run
}

/*
本函数控制是否立即处理远程Shell输出每一行内容[true|false]
*/
func (s *StarShell) SwitchFunc(run bool) {
	s.isfuncs = run
}

func (s *StarShell) SetFunc(funcs func(string)) {
	s.funcs = funcs
}

func (s *StarShell) Clear() {
	defer s.rw.Unlock()
	s.rw.Lock()
	s.outbyte = []byte{}
	s.errbyte = []byte{}
	time.Sleep(time.Millisecond * 15)
}

func (s *StarShell) ShellClear(cmd string, sleep int) (string, string, error) {
	defer s.Clear()
	s.Clear()
	return s.Shell(cmd, sleep)
}

func (s *StarShell) Shell(cmd string, sleep int) (string, string, error) {
	if err := s.WriteCommand(cmd); err != nil {
		return "", "", err
	}
	tmp1, tmp2, err := s.GetResult(sleep)
	tmps := s.TrimColor(strings.TrimSpace(string(tmp1)))
	if s.isecho {
		n := len(strings.Split(cmd, "\n"))
		if n == 1 {
			list := strings.SplitN(tmps, "\n", 2)
			if len(list) == 2 {
				tmps = list[1]
			}
		} else {
			list := strings.Split(tmps, "\n")
			cmds := strings.Split(cmd, "\n")
			for _, v := range cmds {
				for k, v2 := range list {
					if strings.TrimSpace(v2) == strings.TrimSpace(v) {
						list[k] = ""
						break
					}
				}
			}
			tmps = ""
			for _, v := range list {
				if v != "" {
					tmps += v + "\n"
				}
			}
			tmps = tmps[0 : len(tmps)-1]
		}
	}
	return tmps, s.TrimColor(strings.TrimSpace(string(tmp2))), err
}

func (s *StarShell) GetResult(sleep int) ([]byte, []byte, error) {
	if s.errors != nil {
		s.Session.Close()
		return s.outbyte, s.errbyte, s.errors
	}
	if sleep > 0 {
		time.Sleep(time.Millisecond * time.Duration(sleep))
	}
	return s.outbyte, s.errbyte, nil
}

func (s *StarShell) WriteCommand(cmd string) error {
	return s.Write([]byte(cmd + "\n"))
}

func (s *StarShell) Write(bstr []byte) error {
	if s.errors != nil {
		s.Session.Close()
		return s.errors
	}
	_, err := s.in.Write(bstr)
	return err
}

func (s *StarShell) gohub() {
	go func() {
		var cache []byte
		for {
			read, err := s.er.ReadByte()
			if err != nil {
				s.errors = err
				return
			}
			s.errbyte = append(s.errbyte, read)
			if s.isprint {
				fmt.Print(string([]byte{read}))
			}
			cache = append(cache, read)
			if read == '\n' {
				if s.isfuncs {
					go s.funcs(s.TrimColor(strings.TrimSpace(string(cache))))
					cache = []byte{}
				}
			}
		}
	}()
	var cache []byte
	for {
		read, err := s.out.ReadByte()
		if err != nil {
			s.errors = err
			return
		}
		s.rw.Lock()
		s.outbyte = append(s.outbyte, read)
		cache = append(cache, read)
		s.rw.Unlock()
		if read == '\n' {
			if s.isfuncs {
				go s.funcs(strings.TrimSpace(string(cache)))
				cache = []byte{}
			}
		}
		if s.isprint {
			fmt.Print(string([]byte{read}))
		}
	}
}

func (s *StarShell) GetUid() string {
	res, _, _ := s.ShellWait(`id | grep -oP "(?<=uid\=)\d+"`)
	return strings.TrimSpace(res)
}
func (s *StarShell) GetGid() string {
	res, _, _ := s.ShellWait(`id | grep -oP "(?<=gid\=)\d+"`)
	return strings.TrimSpace(res)
}
func (s *StarShell) GetUser() string {
	res, _, _ := s.ShellWait(`id | grep -oP "(?<=\().*?(?=\))" | head -n 1`)
	return strings.TrimSpace(res)
}
func (s *StarShell) GetGroup() string {
	res, _, _ := s.ShellWait(`id | grep -oP "(?<=\().*?(?=\))" | head -n 2 | tail -n 1`)
	return strings.TrimSpace(res)
}

func (s *StarSSH) NewShell() (shell *StarShell, err error) {
	shell = new(StarShell)
	shell.Session, err = s.NewSession()
	if err != nil {
		return
	}
	shell.in, _ = shell.Session.StdinPipe()
	tmp, _ := shell.Session.StdoutPipe()
	shell.out = bufio.NewReader(tmp)
	tmp, _ = shell.Session.StderrPipe()
	shell.er = bufio.NewReader(tmp)
	err = shell.Session.Shell()
	shell.isecho = true
	go shell.Session.Wait()
	shell.UseWaitDefault = true
	shell.WriteCommand("bash")
	time.Sleep(500 * time.Millisecond)
	shell.WriteCommand("export PS1= ")
	shell.WriteCommand("export PS2= ")
	go shell.gohub()
	time.Sleep(500 * time.Millisecond)
	shell.Clear()
	shell.Clear()
	return
}

func (s *StarSSH) Close() error {
	if s.online {
		return s.Client.Close()
	}
	return nil
}

func (s *StarSSH) NewSession() (*ssh.Session, error) {
	return NewSession(s.Client)
}

func (s *StarSSH) ShellOne(cmd string) (string, error) {
	newsess, err := s.NewSession()
	if err != nil {
		return "", err
	}
	data, err := newsess.CombinedOutput(cmd)
	newsess.Close()
	return strings.TrimSpace(string(data)), err
}

func (s *StarSSH) Exists(filepath string) bool {
	res, _ := s.ShellOne(`echo 1 && [ ! -e "` + filepath + `" ] && echo 2`)
	if res == "1" {
		return true
	} else {
		return false
	}
}

func (s *StarSSH) GetUid() string {
	res, _ := s.ShellOne(`id | grep -oP "(?<=uid\=)\d+"`)
	return strings.TrimSpace(res)
}
func (s *StarSSH) GetGid() string {
	res, _ := s.ShellOne(`id | grep -oP "(?<=gid\=)\d+"`)
	return strings.TrimSpace(res)
}
func (s *StarSSH) GetUser() string {
	res, _ := s.ShellOne(`id | grep -oP "(?<=\().*?(?=\))" | head -n 1`)
	return strings.TrimSpace(res)
}
func (s *StarSSH) GetGroup() string {
	res, _ := s.ShellOne(`id | grep -oP "(?<=\().*?(?=\))" | head -n 2 | tail -n 1`)
	return strings.TrimSpace(res)
}

func (s *StarSSH) IsFile(filepath string) bool {
	res, _ := s.ShellOne(`echo 1 && [ ! -f "` + filepath + `" ] && echo 2`)
	if res == "1" {
		return true
	} else {
		return false
	}
}

func (s *StarSSH) IsFolder(filepath string) bool {
	res, _ := s.ShellOne(`echo 1 && [ ! -d "` + filepath + `" ] && echo 2`)
	if res == "1" {
		return true
	} else {
		return false
	}
}

func (s *StarSSH) ShellOneShowScreen(cmd string) (string, error) {
	newsess, err := s.NewSession()
	if err != nil {
		return "", err
	}
	var bytes, errbytes []byte
	tmp, _ := newsess.StdoutPipe()
	reader := bufio.NewReader(tmp)
	tmp, _ = newsess.StderrPipe()
	errder := bufio.NewReader(tmp)
	err = newsess.Start(cmd)
	if err != nil {
		return "", err
	}
	c := make(chan int, 1)
	go newsess.Wait()
	go func() {
		for {
			byt, err := reader.ReadByte()
			if err != nil {
				break
			}
			fmt.Print(string([]byte{byt}))
			bytes = append(bytes, byt)
		}
		c <- 1
	}()
	for {
		byt, err := errder.ReadByte()
		if err != nil {
			break
		}
		fmt.Print(string([]byte{byt}))
		errbytes = append(errbytes, byt)
	}
	_ = <-c
	newsess.Close()
	if len(errbytes) != 0 {
		err = errors.New(strings.TrimSpace(string(errbytes)))
	} else {
		err = nil
	}
	return strings.TrimSpace(string(bytes)), err
}

func (s *StarSSH) ShellOneToFunc(cmd string, callback func(string)) (string, error) {
	newsess, err := s.NewSession()
	if err != nil {
		return "", err
	}
	var bytes, errbytes []byte
	tmp, _ := newsess.StdoutPipe()
	reader := bufio.NewReader(tmp)
	tmp, _ = newsess.StderrPipe()
	errder := bufio.NewReader(tmp)
	err = newsess.Start(cmd)
	if err != nil {
		return "", err
	}
	c := make(chan int, 1)
	go newsess.Wait()
	go func() {
		for {
			byt, err := reader.ReadByte()
			if err != nil {
				break
			}
			callback(string([]byte{byt}))
			bytes = append(bytes, byt)
		}
		c <- 1
	}()
	for {
		byt, err := errder.ReadByte()
		if err != nil {
			break
		}
		callback(string([]byte{byt}))
		errbytes = append(errbytes, byt)
	}
	_ = <-c
	newsess.Close()
	if len(errbytes) != 0 {
		err = errors.New(strings.TrimSpace(string(errbytes)))
	} else {
		err = nil
	}
	return strings.TrimSpace(string(bytes)), err
}

func NewTransferSession(client *ssh.Client) (*ssh.Session, error) {
	session, err := client.NewSession()
	return session, err
}

func NewSession(client *ssh.Client) (*ssh.Session, error) {
	var session *ssh.Session
	var err error
	// create session
	if session, err = client.NewSession(); err != nil {
		return nil, err
	}
	modes := ssh.TerminalModes{
		ssh.ECHO: 1, // 还是要强制开启
		//ssh.IGNCR:         0,
		ssh.TTY_OP_ISPEED: 14400, // input speed = 14.4kbaud
		ssh.TTY_OP_OSPEED: 14400, // output speed = 14.4kbaud
	}

	if err := session.RequestPty("xterm", 500, 250, modes); err != nil {
		return nil, err
	}
	return session, nil
}

func SedColor(str string) string {
	reg := regexp.MustCompile(`\x1B\[([0-9]{1,2}(;[0-9]{1,2})?)?[m|K]`)
	//fmt.Println("regexp:", reg.Match([]byte(str)))
	return string(reg.ReplaceAll([]byte(str), []byte("")))
}