diff --git a/ssh.go b/ssh.go index 5deaa3b..320c7e4 100644 --- a/ssh.go +++ b/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" @@ -32,29 +34,34 @@ type Sshd struct { } type StarSSH struct { - Client *ssh.Client - user string - password string - host string - key string - port int - online bool + Client *ssh.Client + PublicKey ssh.PublicKey + PubkeyBase64 string + 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 - isecho bool - funcs func(string) + 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 (this *StarShell) ShellWait(cmd string) (string, string, error) { @@ -66,32 +73,67 @@ 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") - for _, v := range list { - if strings.Index(v, "b7Y85R56TUY6R5UTb612") < 0 { - outc += v + "\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" + } } + 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 } - 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" + 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 } - break } } - return this.TrimColor(strings.TrimSpace(outc[0 : len(outc)-1])), this.TrimColor(strings.TrimSpace(errc[0 : len(errc)-1])), err + 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 { diff --git a/ssh_test.go b/ssh_test.go index 457fcec..c22eb05 100644 --- a/ssh_test.go +++ b/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) }