sshd/ssh.go

package sshd

import (
	"bufio"
	"bytes"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"os"
	"path"
	"path/filepath"
	"regexp"
	"strings"
	"time"

	"b612.me/starainrt"

	"github.com/pkg/sftp"
	"golang.org/x/crypto/ssh"
)

var ShellRes, ShellErr string
var ShellExit bool

type Sshd struct {
	SSHC    *ssh.Session
	infile  io.Writer
	outfile io.Reader
	errfile io.Reader
	thread  bool
	counter int
}

type StarSSH struct {
	Client   *ssh.Client
	user     string
	password string
	host     string
	key      string
	port     int
	online   bool
}

type StarShell struct {
	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
	funcs   func(string)
}

func (this *StarShell) ShellWait(cmd string) (string, string, error) {
	var outc, errc string = "  ", "  "
	this.Clear()
	defer this.Clear()
	echo := "echo b7Y85R56TUY6R5UTb612"
	err := this.WriteCommand(cmd)
	if err != nil {
		return "", "", err
	}
	err = this.WriteCommand(echo)
	if err != nil {
		return "", "", err
	}
	for {
		time.Sleep(time.Millisecond * 100)
		if strings.Index(string(this.outbyte), "b7Y85R56TUY6R5UTb612") >= 0 {
			list := strings.Split(string(this.outbyte), "\n")
			for _, v := range list {
				if strings.Index(v, "b7Y85R56TUY6R5UTb612") < 0 {
					outc += v + "\n"
				}
			}
			break
		}
		if strings.Index(string(this.errbyte), "b7Y85R56TUY6R5UTb612") >= 0 {
			list := strings.Split(string(this.errbyte), "\n")
			for _, v := range list {
				if strings.Index(v, "b7Y85R56TUY6R5UTb612") < 0 {
					errc += v + "\n"
				}
			}
			break
		}
	}
	return this.TrimColor(strings.TrimSpace(outc[0 : len(outc)-1])), this.TrimColor(strings.TrimSpace(errc[0 : len(errc)-1])), err
}

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

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

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

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

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

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

func (this *StarShell) Clear() {
	this.outbyte = []byte{}
	this.errbyte = []byte{}
}

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

func (this *StarShell) Shell(cmd string, sleep int) (string, string, error) {
	if err := this.WriteCommand(cmd); err != nil {
		return "", "", err
	}
	tmp1, tmp2, err := this.GetResult(sleep)
	return this.TrimColor(strings.TrimSpace(string(tmp1))), this.TrimColor(strings.TrimSpace(string(tmp2))), err
}

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

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

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

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

func (this *StarSSH) NewShell() (shell *StarShell, err error) {
	shell = new(StarShell)
	shell.Session, err = this.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.WriteCommand("export PS1= ")
	shell.WriteCommand("export PS2= ")
	go shell.gohub()
	time.Sleep(1 * time.Second)
	shell.Clear()

	return
}

func (this *StarSSH) Connect(user, password, host, key string, port int) error {
	var err error
	if this.online {
		this.online = false
		this.Client.Close()
	}
	this.Client, err = Connect(user, password, host, key, port, []string{})
	if err != nil {
		return err
	}
	this.online = true
	this.host = host
	this.password = password
	this.key = key
	this.port = port
	return nil
}

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

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

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

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

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

func (this *StarSSH) ShellOneShowScreen(cmd string) (string, error) {
	newsess, err := this.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 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:          0,     // disable echoing
		ssh.TTY_OP_ISPEED: 14400, // input speed = 14.4kbaud
		ssh.TTY_OP_OSPEED: 14400, // output speed = 14.4kbaud
	}

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

func Connect(user, password, host, key string, port int, cipherList []string) (*ssh.Client, error) {
	var (
		auth         []ssh.AuthMethod
		addr         string
		clientConfig *ssh.ClientConfig
		client       *ssh.Client
		config       ssh.Config
		err          error
	)
	// get auth method
	auth = make([]ssh.AuthMethod, 0)
	if key == "" {
		keyboardInteractiveChallenge := func(
			user,
			instruction string,
			questions []string,
			echos []bool,
		) (answers []string, err error) {
			if len(questions) == 0 {
				return []string{}, nil
			}
			return []string{password}, nil
		}
		auth = append(auth, ssh.Password(password))
		auth = append(auth, ssh.KeyboardInteractive(keyboardInteractiveChallenge))
	} else {
		pemBytes, err := ioutil.ReadFile(key)
		if err != nil {
			return nil, err
		}

		var signer ssh.Signer
		if password == "" {
			signer, err = ssh.ParsePrivateKey(pemBytes)
		} else {
			signer, err = ssh.ParsePrivateKeyWithPassphrase(pemBytes, []byte(password))
		}
		if err != nil {
			return nil, err
		}
		auth = append(auth, ssh.PublicKeys(signer))
	}

	if len(cipherList) == 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"},
		}
	} else {
		config = ssh.Config{
			Ciphers: cipherList,
		}
	}

	clientConfig = &ssh.ClientConfig{
		User:    user,
		Auth:    auth,
		Timeout: 10 * time.Second,
		Config:  config,
		HostKeyCallback: func(hostname string, remote net.Addr, key ssh.PublicKey) error {
			return nil
		},
	}

	// connet to ssh
	addr = fmt.Sprintf("%s:%d", host, port)

	if client, err = ssh.Dial("tcp", addr, clientConfig); err != nil {
		return nil, err
	}

	return client, nil
}

func ScpTransfer(src, dst string, session *ssh.Session) error {

	go func() {
		//Buf := make([]byte, 1024)
		w, _ := session.StdinPipe()
		defer w.Close()
		File, err := os.Open(src)
		if err != nil {
			panic(err)
		}
		info, _ := File.Stat()
		fmt.Fprintln(w, "C0777", info.Size(), info.Name())
		io.Copy(w, File)
		fmt.Fprintln(w, "\x00")
	}()
	if err := session.Run("/usr/bin/scp -qrt " + dst); err != nil {
		return err
	}
	return nil
}

func CreateSftp(client *ssh.Client) (*sftp.Client, error) {

	sftpClient, err := sftp.NewClient(client)
	return sftpClient, err
}

func FtpTransferOut(localFilePath, remoteDir string, sftpClient *sftp.Client) error {
	srcFile, err := os.Open(localFilePath)
	if err != nil {
		return err
	}
	defer srcFile.Close()
	var remoteFileName = filepath.Base(localFilePath)
	dstFile, err := sftpClient.Create(path.Join(remoteDir, remoteFileName))
	if err != nil {
		return err
	}
	defer dstFile.Close()
	for {
		buf := make([]byte, 1024)
		n, err := srcFile.Read(buf)
		dstFile.Write(buf[:n])
		if err == io.EOF {
			break
		}
	}
	return nil
}

func FtpTransferOutFunc(localFilePath, remoteDir string, bufcap int, rtefunc func(int), sftpClient *sftp.Client) error {
	num := 0
	srcFile, err := os.Open(localFilePath)
	if err != nil {
		return err
	}
	defer srcFile.Close()
	var remoteFileName = filepath.Base(localFilePath)
	dstFile, err := sftpClient.Create(path.Join(remoteDir, remoteFileName))
	if err != nil {
		return err
	}
	defer dstFile.Close()
	for {
		buf := make([]byte, 1024)
		n, err := srcFile.Read(buf)
		dstFile.Write(buf[:n])
		num++
		rtefunc(num)
		if err == io.EOF {
			break
		}
	}
	return nil
}

func FtpTransferIn(src, dst string, sftpClient *sftp.Client) error {
	srcFile, err := sftpClient.Open(src)
	if err != nil {
		return err
	}
	defer srcFile.Close()
	var localFileName = filepath.Base(src)
	dstFile, err := os.Create(filepath.Join(dst, localFileName))
	if err != nil {
		return err
	}
	defer dstFile.Close()

	if _, err = srcFile.WriteTo(dstFile); err != nil {
		return err
	}

	return nil
}

func FtpTransferInFunc(src, dst string, bufcap int, rtefunc func(int), sftpClient *sftp.Client) error {
	num := 0
	srcFile, err := sftpClient.Open(src)
	if err != nil {
		return err
	}
	defer srcFile.Close()
	var localFileName = filepath.Base(src)
	dstFile, err := os.Create(filepath.Join(dst, localFileName))
	if err != nil {
		return err
	}
	defer dstFile.Close()
	for {
		buf := make([]byte, 1024)
		n, err := srcFile.Read(buf)
		dstFile.Write(buf[:n])
		num++
		rtefunc(num)
		if err == io.EOF {
			break
		}
	}
	return nil
}

func Command(session *ssh.Session, cmdstr string) (string, error) {
	var res bytes.Buffer
	session.Stdout = &res
	err := session.Run(cmdstr)
	if err != nil {
		return res.String(), err
	}
	return res.String(), nil
}

func SSHPipeShell(session *ssh.Session, cmdstr string) (*Sshd, error) {
	var err error
	lovessh := Sshd{}
	lovessh.SSHC = session
	lovessh.infile, err = lovessh.SSHC.StdinPipe()
	if err != nil {
		return &lovessh, err
	}
	lovessh.outfile, err = lovessh.SSHC.StdoutPipe()
	if err != nil {
		return &lovessh, err
	}
	lovessh.errfile, err = lovessh.SSHC.StderrPipe()
	if err != nil {
		return &lovessh, err
	}
	if err := lovessh.SSHC.Start(cmdstr); err != nil {
		return &lovessh, err
	}
	go func() {
		lovessh.SSHC.Wait()
	}()
	ShellExit = false
	lovessh.thread = false
	return &lovessh, 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("")))
}

func (this Sshd) GetResult(maxtime int) (string, string, bool) {
	var stop bool
	reader := bufio.NewReader(this.outfile)
	erreader := bufio.NewReader(this.errfile)
	if !this.thread {
		this.thread = true
		go func() {
			var line2 string
			var stack bool = false
			stop = false
			for {
				if !stack {
					go func() {
						stack = true
						if erreader.Size() > 0 {
							line2, _ = erreader.ReadString('\n')
							ShellErr += line2
							line2 = ""
						}
						stack = false
					}()
				}
				line, err2 := reader.ReadString('\n')
				if err2 != nil || io.EOF == err2 {
					stop = true
					break
				}
				this.counter++
				ShellRes += line
			}
		}()
	}
	for !stop {
		if maxtime > 0 {
			time.Sleep(time.Millisecond * time.Duration(maxtime))
		}
		restr := SedColor(ShellRes)
		ShellRes = ""
		errstr := SedColor(ShellErr)
		ShellErr = ""
		return strings.TrimSpace(restr), strings.TrimSpace(errstr), false

	}
	ShellExit = true
	this.thread = false
	restr := SedColor(ShellRes)
	ShellRes = ""
	errstr := SedColor(ShellErr)
	ShellErr = ""
	return strings.TrimSpace(restr), strings.TrimSpace(errstr), true
}

func (this Sshd) Exec(cmdstr string, maxtime int) (string, string, bool) {
	this.infile.Write([]byte(cmdstr + "\n"))
	return this.GetResult(maxtime)
}

func (this Sshd) WriteCmd(cmdstr string) {
	this.infile.Write([]byte(cmdstr + "\n"))
	return
}

func (this Sshd) NeverTimeOut() {
	this.infile.Write([]byte("export TMOUT=0 \n"))
}

func (this Sshd) SetBash() {
	this.infile.Write([]byte("export PS1= \n"))
}

func (this Sshd) Clear() {
	this.GetResult(50)
}

func (this Sshd) IsExit() bool {
	return ShellExit
}

func (this Sshd) Client(sakura func(string)) {
	var stop bool
	reader := bufio.NewReader(this.outfile)
	if !this.thread {
		this.thread = true
		stop = false
		go func() {
			for {
				line, err2 := reader.ReadByte()
				if err2 != nil {
					stop = true
					break
				}
				sakura(string(line))
			}
			this.counter++
		}()
	}
	for !stop {
		cmd := starainrt.MessageBox("", "")
		if cmd == "ctrlc" {
			cmd = string('\x03')
		}
		if cmd == `\ctrlc` {
			cmd = "ctrlc"
		}
		this.WriteCmd(cmd)
	}
}