完善命令判断

ssh.go

@@ -3,6 +3,7 @@ package sshd
 import (
 	"bufio"
 	"bytes"
+	"encoding/base64"
 	"errors"
 	"fmt"
 	"io"
@@ -11,6 +12,7 @@ import (
 	"os"
 	"regexp"
 	"strings"
+	"sync"
 	"time"
 
 	"b612.me/starainrt"
@@ -33,6 +35,8 @@ type Sshd struct {
 
 type StarSSH struct {
 	Client       *ssh.Client
+	PublicKey    ssh.PublicKey
+	PubkeyBase64 string
 	user         string
 	password     string
 	host         string
@@ -42,6 +46,8 @@ type StarSSH struct {
 }
 
 type StarShell struct {
+	Keyword        string
+	UseWaitDefault bool
 	Session        *ssh.Session
 	in             io.Writer
 	out            *bufio.Reader
@@ -54,6 +60,7 @@ type StarShell struct {
 	isfuncs        bool
 	iscolor        bool
 	isecho         bool
+	rw             sync.RWMutex
 	funcs          func(string)
 }
 
@@ -66,14 +73,28 @@ func (this *StarShell) ShellWait(cmd string) (string, string, error) {
 	if err != nil {
 		return "", "", err
 	}
+	time.Sleep(time.Millisecond * 20)
 	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")
+		time.Sleep(time.Millisecond * 120)
+		outs := string(this.outbyte)
+		errs := string(this.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 this.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"
@@ -81,8 +102,8 @@ func (this *StarShell) ShellWait(cmd string) (string, string, error) {
 				}
 				break
 			}
-		if strings.Index(string(this.errbyte), "b7Y85R56TUY6R5UTb612") >= 0 {
-			list := strings.Split(string(this.errbyte), "\n")
+			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"
@@ -91,7 +112,28 @@ func (this *StarShell) ShellWait(cmd string) (string, string, error) {
 				break
 			}
 		}
-	return this.TrimColor(strings.TrimSpace(outc[0 : len(outc)-1])), this.TrimColor(strings.TrimSpace(errc[0 : len(errc)-1])), err
+		if this.Keyword != "" {
+			if strings.Index(string(outs), this.Keyword) >= 0 {
+				list := strings.Split(string(outs), "\n")
+				for _, v := range list {
+					if strings.Index(v, this.Keyword) < 0 && strings.Index(v, "b7Y85R56TUY6R5UTb612") < 0 {
+						outc += v + "\n"
+					}
+				}
+				break
+			}
+			if strings.Index(string(errs), this.Keyword) >= 0 {
+				list := strings.Split(string(errs), "\n")
+				for _, v := range list {
+					if strings.Index(v, this.Keyword) < 0 && strings.Index(v, "b7Y85R56TUY6R5UTb612") < 0 {
+						errc += v + "\n"
+					}
+				}
+				break
+			}
+		}
+	}
+	return this.TrimColor(strings.TrimSpace(outc)), this.TrimColor(strings.TrimSpace(errc)), err
 }
 
 func (this *StarShell) Close() error {
@@ -132,9 +174,11 @@ func (this *StarShell) SetFunc(funcs func(string)) {
 }
 
 func (this *StarShell) Clear() {
+	defer this.rw.Unlock()
+	this.rw.Lock()
 	this.outbyte = []byte{}
 	this.errbyte = []byte{}
-	time.Sleep(time.Millisecond * 5)
+	time.Sleep(time.Millisecond * 15)
 }
 
 func (this *StarShell) ShellClear(cmd string, sleep int) (string, string, error) {
@@ -232,8 +276,10 @@ func (this *StarShell) gohub() {
 			this.errors = err
 			return
 		}
+		this.rw.Lock()
 		this.outbyte = append(this.outbyte, read)
 		cache = append(cache, read)
+		this.rw.Unlock()
 		if read == '\n' {
 			if this.isfuncs {
 				go this.funcs(strings.TrimSpace(string(cache)))
@@ -246,6 +292,23 @@ func (this *StarShell) gohub() {
 	}
 }
 
+func (this *StarShell) GetUid() string {
+	res, _, _ := this.ShellWait(`id | grep -oP "(?<=uid\=)\d+"`)
+	return strings.TrimSpace(res)
+}
+func (this *StarShell) GetGid() string {
+	res, _, _ := this.ShellWait(`id | grep -oP "(?<=gid\=)\d+"`)
+	return strings.TrimSpace(res)
+}
+func (this *StarShell) GetUser() string {
+	res, _, _ := this.ShellWait(`id | grep -oP "(?<=\().*?(?=\))" | head -n 1`)
+	return strings.TrimSpace(res)
+}
+func (this *StarShell) GetGroup() string {
+	res, _, _ := this.ShellWait(`id | grep -oP "(?<=\().*?(?=\))" | head -n 2 | tail -n 1`)
+	return strings.TrimSpace(res)
+}
+
 func (this *StarSSH) NewShell() (shell *StarShell, err error) {
 	shell = new(StarShell)
 	shell.Session, err = this.NewSession()
@@ -260,11 +323,13 @@ func (this *StarSSH) NewShell() (shell *StarShell, err error) {
 	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(1 * time.Second)
+	time.Sleep(500 * time.Millisecond)
 	shell.Clear()
 	shell.Clear()
 	return
@@ -276,10 +341,11 @@ func (this *StarSSH) Connect(user, password, host, key string, port int) error {
 		this.online = false
 		this.Client.Close()
 	}
-	this.Client, err = Connect(user, password, host, key, port, []string{})
+	this.Client, this.PublicKey, err = Connect(user, password, host, key, port, []string{})
 	if err != nil {
 		return err
 	}
+	this.PubkeyBase64 = base64.StdEncoding.EncodeToString(this.PublicKey.Marshal())
 	this.online = true
 	this.host = host
 	this.password = password
@@ -317,6 +383,24 @@ func (this *StarSSH) Exists(filepath string) bool {
 		return false
 	}
 }
+
+func (this *StarSSH) GetUid() string {
+	res, _ := this.ShellOne(`id | grep -oP "(?<=uid\=)\d+"`)
+	return strings.TrimSpace(res)
+}
+func (this *StarSSH) GetGid() string {
+	res, _ := this.ShellOne(`id | grep -oP "(?<=gid\=)\d+"`)
+	return strings.TrimSpace(res)
+}
+func (this *StarSSH) GetUser() string {
+	res, _ := this.ShellOne(`id | grep -oP "(?<=\().*?(?=\))" | head -n 1`)
+	return strings.TrimSpace(res)
+}
+func (this *StarSSH) GetGroup() string {
+	res, _ := this.ShellOne(`id | grep -oP "(?<=\().*?(?=\))" | head -n 2 | tail -n 1`)
+	return strings.TrimSpace(res)
+}
+
 func (this *StarSSH) IsFile(filepath string) bool {
 	res, _ := this.ShellOne(`echo 1 && [ ! -f "` + filepath + `" ] && echo 2`)
 	if res == "1" {
@@ -380,6 +464,51 @@ func (this *StarSSH) ShellOneShowScreen(cmd string) (string, error) {
 	return strings.TrimSpace(string(bytes)), err
 }
 
+func (this *StarSSH) ShellOneToFunc(cmd string, callback func(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
+			}
+			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
@@ -398,13 +527,13 @@ func NewSession(client *ssh.Client) (*ssh.Session, error) {
 		ssh.TTY_OP_OSPEED: 14400, // output speed = 14.4kbaud
 	}
 
-	if err := session.RequestPty("vt100", 80, 40, modes); err != nil {
+	if err := session.RequestPty("xterm", 500, 250, modes); err != nil {
 		return nil, err
 	}
 	return session, nil
 }
 
-func Connect(user, password, host, key string, port int, cipherList []string) (*ssh.Client, error) {
+func Connect(user, password, host, key string, port int, cipherList []string) (*ssh.Client, ssh.PublicKey, error) {
 	var (
 		auth         []ssh.AuthMethod
 		addr         string
@@ -412,6 +541,7 @@ func Connect(user, password, host, key string, port int, cipherList []string) (*
 		client       *ssh.Client
 		config       ssh.Config
 		err          error
+		pubkey       ssh.PublicKey
 	)
 	// get auth method
 	auth = make([]ssh.AuthMethod, 0)
@@ -432,7 +562,7 @@ func Connect(user, password, host, key string, port int, cipherList []string) (*
 	} else {
 		pemBytes, err := ioutil.ReadFile(key)
 		if err != nil {
-			return nil, err
+			return nil, nil, err
 		}
 
 		var signer ssh.Signer
@@ -442,7 +572,7 @@ func Connect(user, password, host, key string, port int, cipherList []string) (*
 			signer, err = ssh.ParsePrivateKeyWithPassphrase(pemBytes, []byte(password))
 		}
 		if err != nil {
-			return nil, err
+			return nil, nil, err
 		}
 		auth = append(auth, ssh.PublicKeys(signer))
 	}
@@ -463,6 +593,7 @@ func Connect(user, password, host, key string, port int, cipherList []string) (*
 		Timeout: 10 * time.Second,
 		Config:  config,
 		HostKeyCallback: func(hostname string, remote net.Addr, key ssh.PublicKey) error {
+			pubkey = key
 			return nil
 		},
 	}
@@ -471,10 +602,9 @@ func Connect(user, password, host, key string, port int, cipherList []string) (*
 	addr = fmt.Sprintf("%s:%d", host, port)
 
 	if client, err = ssh.Dial("tcp", addr, clientConfig); err != nil {
-		return nil, err
+		return nil, nil, err
 	}
-
-	return client, nil
+	return client, pubkey, nil
 }
 
 func ScpTransfer(src, dst string, session *ssh.Session) error {
@@ -527,6 +657,16 @@ func FtpTransferOut(localFilePath, remotePath string, sftpClient *sftp.Client) e
 	return nil
 }
 
+func FtpTransferOutByte(localData []byte, remotePath string, sftpClient *sftp.Client) error {
+	dstFile, err := sftpClient.Create(remotePath)
+	if err != nil {
+		return err
+	}
+	defer dstFile.Close()
+	_, err = dstFile.Write(localData)
+	return err
+}
+
 func FtpTransferOutFunc(localFilePath, remotePath string, bufcap int, rtefunc func(float64), sftpClient *sftp.Client) error {
 	num := 0
 	srcFile, err := os.Open(localFilePath)
@@ -558,6 +698,17 @@ func FtpTransferOutFunc(localFilePath, remotePath string, bufcap int, rtefunc fu
 	return nil
 }
 
+func FtpTransferInByte(remotePath string, sftpClient *sftp.Client) ([]byte, error) {
+	dstFile, err := sftpClient.Open(remotePath)
+	if err != nil {
+		return []byte{}, err
+	}
+	defer dstFile.Close()
+	buf := new(bytes.Buffer)
+	_, err = dstFile.WriteTo(buf)
+	return buf.Bytes(), err
+}
+
 func FtpTransferIn(src, dst string, sftpClient *sftp.Client) error {
 	srcFile, err := sftpClient.Open(src)
 	if err != nil {

ssh_test.go

@@ -7,7 +7,7 @@ import (
 
 func TestSSH(t *testing.T) {
 	myssh := new(StarSSH)
-	err := myssh.Connect("root", "sakua", "9sday.me", "C:\\id_rsa", 22)
+	err := myssh.Connect("root", "128je980", "nasf.b612.me", "", 22000)
 	if err != nil {
 		t.Fatalf("%e", err)
 	}