diff --git a/ssh.go b/ssh.go index 818a283..e7fce89 100644 --- a/ssh.go +++ b/ssh.go @@ -3,15 +3,20 @@ 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" ) @@ -28,6 +33,318 @@ type Sshd struct { 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 @@ -158,7 +475,31 @@ func FtpTransferOut(localFilePath, remoteDir string, sftpClient *sftp.Client) er return err } defer srcFile.Close() - var remoteFileName = path.Base(localFilePath) + 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 @@ -166,11 +507,13 @@ func FtpTransferOut(localFilePath, remoteDir string, sftpClient *sftp.Client) er defer dstFile.Close() for { buf := make([]byte, 1024) - n, _ := srcFile.Read(buf) - if n == 0 { + n, err := srcFile.Read(buf) + dstFile.Write(buf[:n]) + num++ + rtefunc(num) + if err == io.EOF { break } - dstFile.Write(buf) } return nil } @@ -181,9 +524,8 @@ func FtpTransferIn(src, dst string, sftpClient *sftp.Client) error { return err } defer srcFile.Close() - - var localFileName = path.Base(src) - dstFile, err := os.Create(path.Join(dst, localFileName)) + var localFileName = filepath.Base(src) + dstFile, err := os.Create(filepath.Join(dst, localFileName)) if err != nil { return err } @@ -196,6 +538,32 @@ func FtpTransferIn(src, dst string, sftpClient *sftp.Client) error { 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 @@ -271,19 +639,16 @@ func (this Sshd) GetResult(maxtime int) (string, string, bool) { } }() } - waittm := 0 for !stop { - time.Sleep(time.Millisecond * 250) - waittm += 1 - if maxtime >= 0 { - if waittm/4 > maxtime { - restr := SedColor(ShellRes) - ShellRes = "" - errstr := SedColor(ShellErr) - ShellErr = "" - return restr, errstr, false - } + 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 @@ -291,7 +656,7 @@ func (this Sshd) GetResult(maxtime int) (string, string, bool) { ShellRes = "" errstr := SedColor(ShellErr) ShellErr = "" - return restr, errstr, true + return strings.TrimSpace(restr), strings.TrimSpace(errstr), true } func (this Sshd) Exec(cmdstr string, maxtime int) (string, string, bool) { @@ -304,6 +669,48 @@ func (this Sshd) WriteCmd(cmdstr string) { 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) + } +} diff --git a/ssh_test.go b/ssh_test.go new file mode 100644 index 0000000..f60e45d --- /dev/null +++ b/ssh_test.go @@ -0,0 +1,23 @@ +package sshd + +import ( + "fmt" + "testing" + "time" +) + +func TestSSH(t *testing.T) { + myssh := new(StarSSH) + err := myssh.Connect("root", "", "9sday.me", "c:\\id_rsa", 22) + if err != nil { + t.Fatalf("%e", err) + } + shell, err := myssh.NewShell() + if err != nil { + t.Fatalf("%e", err) + } + shell.isprint = true + time.Sleep(5 * time.Second) + a, b, err := shell.ShellClear("apt update\n", 20000) + fmt.Println(a, b, err) +}