优化原生ssh连接代码

pkg/term/next_terminal.go

@@ -0,0 +1,99 @@
+package term
+
+import (
+	"github.com/pkg/sftp"
+	"golang.org/x/crypto/ssh"
+	"io"
+)
+
+type NextTerminal struct {
+	SshClient  *ssh.Client
+	SshSession *ssh.Session
+	StdinPipe  io.WriteCloser
+	SftpClient *sftp.Client
+	Recorder   *Recorder
+	NextWriter *NextWriter
+}
+
+func NewNextTerminal(ip string, port int, username, password, privateKey, passphrase string, rows, cols int, recording string) (*NextTerminal, error) {
+
+	sshClient, err := NewSshClient(ip, port, username, password, privateKey, passphrase)
+	if err != nil {
+		return nil, err
+	}
+
+	sshSession, err := sshClient.NewSession()
+	if err != nil {
+		return nil, err
+	}
+	//defer sshSession.Close()
+
+	modes := ssh.TerminalModes{
+		ssh.ECHO:          1,
+		ssh.TTY_OP_ISPEED: 14400,
+		ssh.TTY_OP_OSPEED: 14400,
+	}
+
+	if err := sshSession.RequestPty("xterm-256color", rows, cols, modes); err != nil {
+		return nil, err
+	}
+
+	var nextWriter NextWriter
+	sshSession.Stdout = &nextWriter
+	sshSession.Stderr = &nextWriter
+
+	stdinPipe, err := sshSession.StdinPipe()
+	if err != nil {
+		return nil, err
+	}
+
+	if err := sshSession.Shell(); err != nil {
+		return nil, err
+	}
+
+	var recorder *Recorder
+	if recording != "" {
+		recorder, err = CreateRecording(recording, rows, cols)
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	terminal := NextTerminal{
+		SshClient:  sshClient,
+		SshSession: sshSession,
+		Recorder:   recorder,
+		StdinPipe:  stdinPipe,
+		NextWriter: &nextWriter,
+	}
+
+	return &terminal, nil
+}
+
+func (ret *NextTerminal) Write(p []byte) (int, error) {
+	return ret.StdinPipe.Write(p)
+}
+
+func (ret *NextTerminal) Read() ([]byte, int, error) {
+	return ret.NextWriter.Read()
+}
+
+func (ret *NextTerminal) Close() error {
+	if ret.SshSession != nil {
+		return ret.SshSession.Close()
+	}
+
+	if ret.SshClient != nil {
+		return ret.SshClient.Close()
+	}
+
+	if ret.Recorder != nil {
+		return ret.Close()
+	}
+
+	return nil
+}
+
+func (ret *NextTerminal) WindowChange(h int, w int) error {
+	return ret.SshSession.WindowChange(h, w)
+}

pkg/term/next_writer.go

@@ -0,0 +1,30 @@
+package term
+
+import (
+	"bytes"
+	"sync"
+)
+
+type NextWriter struct {
+	b  bytes.Buffer
+	mu sync.Mutex
+}
+
+func (w *NextWriter) Write(p []byte) (int, error) {
+	w.mu.Lock()
+	defer w.mu.Unlock()
+	return w.b.Write(p)
+}
+
+func (w *NextWriter) Read() ([]byte, int, error) {
+	w.mu.Lock()
+	defer w.mu.Unlock()
+	p := w.b.Bytes()
+	buf := make([]byte, len(p))
+	read, err := w.b.Read(buf)
+	w.b.Reset()
+	if err != nil {
+		return nil, 0, err
+	}
+	return buf, read, err
+}

pkg/term/recording.go

@@ -0,0 +1,122 @@
+package term
+
+import (
+	"encoding/json"
+	"next-terminal/pkg/utils"
+	"os"
+	"time"
+)
+
+type Env struct {
+	Shell string `json:"SHELL"`
+	Term  string `json:"TERM"`
+}
+
+type Header struct {
+	Title     string `json:"title"`
+	Version   int    `json:"version"`
+	Height    int    `json:"height"`
+	Width     int    `json:"width"`
+	Env       Env    `json:"env"`
+	Timestamp int    `json:"Timestamp"`
+}
+
+type Recorder struct {
+	File      *os.File
+	Timestamp int
+}
+
+func NewRecorder(recoding string) (recorder *Recorder, err error) {
+	recorder = &Recorder{}
+
+	parentDirectory := utils.GetParentDirectory(recoding)
+
+	if utils.FileExists(parentDirectory) {
+		if err := os.RemoveAll(parentDirectory); err != nil {
+			return nil, err
+		}
+	}
+
+	if err = os.MkdirAll(parentDirectory, 0777); err != nil {
+		return
+	}
+
+	var file *os.File
+	file, err = os.Create(recoding)
+	if err != nil {
+		return nil, err
+	}
+
+	recorder.File = file
+	return recorder, nil
+}
+
+func (recorder *Recorder) Close() {
+	if recorder.File != nil {
+		recorder.File.Close()
+	}
+}
+
+func (recorder *Recorder) WriteHeader(header *Header) (err error) {
+	var p []byte
+
+	if p, err = json.Marshal(header); err != nil {
+		return
+	}
+
+	if _, err := recorder.File.Write(p); err != nil {
+		return err
+	}
+	if _, err := recorder.File.Write([]byte("\n")); err != nil {
+		return err
+	}
+
+	recorder.Timestamp = header.Timestamp
+
+	return
+}
+
+func (recorder *Recorder) WriteData(data string) (err error) {
+	now := int(time.Now().UnixNano())
+
+	delta := float64(now-recorder.Timestamp*1000*1000*1000) / 1000 / 1000 / 1000
+
+	row := make([]interface{}, 0)
+	row = append(row, delta)
+	row = append(row, "o")
+	row = append(row, data)
+
+	var s []byte
+	if s, err = json.Marshal(row); err != nil {
+		return
+	}
+	if _, err := recorder.File.Write(s); err != nil {
+		return err
+	}
+	if _, err := recorder.File.Write([]byte("\n")); err != nil {
+		return err
+	}
+	return
+}
+
+func CreateRecording(recordingPath string, h int, w int) (*Recorder, error) {
+	recorder, err := NewRecorder(recordingPath)
+	if err != nil {
+		return nil, err
+	}
+
+	header := &Header{
+		Title:     "",
+		Version:   2,
+		Height:    h,
+		Width:     w,
+		Env:       Env{Shell: "/bin/bash", Term: "xterm-256color"},
+		Timestamp: int(time.Now().Unix()),
+	}
+
+	if err := recorder.WriteHeader(header); err != nil {
+		return nil, err
+	}
+
+	return recorder, nil
+}

pkg/term/ssh.go

@@ -0,0 +1,53 @@
+package term
+
+import (
+	"fmt"
+	"golang.org/x/crypto/ssh"
+	"time"
+)
+
+func NewSshClient(ip string, port int, username, password, privateKey, passphrase string) (*ssh.Client, error) {
+	var authMethod ssh.AuthMethod
+	if username == "-" || username == "" {
+		username = "root"
+	}
+	if password == "-" {
+		password = ""
+	}
+	if privateKey == "-" {
+		privateKey = ""
+	}
+	if passphrase == "-" {
+		passphrase = ""
+	}
+
+	var err error
+	if privateKey != "" {
+		var key ssh.Signer
+		if len(passphrase) > 0 {
+			key, err = ssh.ParsePrivateKeyWithPassphrase([]byte(privateKey), []byte(passphrase))
+			if err != nil {
+				return nil, err
+			}
+		} else {
+			key, err = ssh.ParsePrivateKey([]byte(privateKey))
+			if err != nil {
+				return nil, err
+			}
+		}
+		authMethod = ssh.PublicKeys(key)
+	} else {
+		authMethod = ssh.Password(password)
+	}
+
+	config := &ssh.ClientConfig{
+		Timeout:         1 * time.Second,
+		User:            username,
+		Auth:            []ssh.AuthMethod{authMethod},
+		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
+	}
+
+	addr := fmt.Sprintf("%s:%d", ip, port)
+
+	return ssh.Dial("tcp", addr, config)
+}

pkg/term/test/test_ssh.go

@@ -0,0 +1,174 @@
+package main
+
+import (
+	"fmt"
+	"io"
+	"os"
+	"time"
+
+	"golang.org/x/crypto/ssh"
+	"golang.org/x/crypto/ssh/terminal"
+)
+
+type SSHTerminal struct {
+	Session *ssh.Session
+	exitMsg string
+	stdout  io.Reader
+	stdin   io.Writer
+	stderr  io.Reader
+}
+
+func main() {
+	sshConfig := &ssh.ClientConfig{
+		User: "root",
+		Auth: []ssh.AuthMethod{
+			ssh.Password("root"),
+		},
+		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
+	}
+
+	client, err := ssh.Dial("tcp", "172.16.101.32:22", sshConfig)
+	if err != nil {
+		fmt.Println(err)
+	}
+	defer client.Close()
+
+	err = New(client)
+	if err != nil {
+		fmt.Println(err)
+	}
+}
+
+func (t *SSHTerminal) updateTerminalSize() {
+
+	go func() {
+		// SIGWINCH is sent to the process when the window size of the terminal has
+		// changed.
+		sigwinchCh := make(chan os.Signal, 1)
+		//signal.Notify(sigwinchCh, syscall.SIN)
+
+		fd := int(os.Stdin.Fd())
+		termWidth, termHeight, err := terminal.GetSize(fd)
+		if err != nil {
+			fmt.Println(err)
+		}
+
+		for {
+			select {
+			// The client updated the size of the local PTY. This change needs to occur
+			// on the server side PTY as well.
+			case sigwinch := <-sigwinchCh:
+				if sigwinch == nil {
+					return
+				}
+				currTermWidth, currTermHeight, err := terminal.GetSize(fd)
+
+				// Terminal size has not changed, don't do anything.
+				if currTermHeight == termHeight && currTermWidth == termWidth {
+					continue
+				}
+
+				t.Session.WindowChange(currTermHeight, currTermWidth)
+				if err != nil {
+					fmt.Printf("Unable to send window-change reqest: %s.", err)
+					continue
+				}
+
+				termWidth, termHeight = currTermWidth, currTermHeight
+
+			}
+		}
+	}()
+
+}
+
+func (t *SSHTerminal) interactiveSession() error {
+
+	defer func() {
+		if t.exitMsg == "" {
+			fmt.Fprintln(os.Stdout, "the connection was closed on the remote side on ", time.Now().Format(time.RFC822))
+		} else {
+			fmt.Fprintln(os.Stdout, t.exitMsg)
+		}
+	}()
+
+	fd := int(os.Stdin.Fd())
+	state, err := terminal.MakeRaw(fd)
+	if err != nil {
+		return err
+	}
+	defer terminal.Restore(fd, state)
+
+	termWidth, termHeight, err := terminal.GetSize(fd)
+	if err != nil {
+		return err
+	}
+
+	termType := os.Getenv("TERM")
+	if termType == "" {
+		termType = "xterm-256color"
+	}
+
+	err = t.Session.RequestPty(termType, termHeight, termWidth, ssh.TerminalModes{})
+	if err != nil {
+		return err
+	}
+
+	t.updateTerminalSize()
+
+	t.stdin, err = t.Session.StdinPipe()
+	if err != nil {
+		return err
+	}
+	t.stdout, err = t.Session.StdoutPipe()
+	if err != nil {
+		return err
+	}
+	t.stderr, err = t.Session.StderrPipe()
+
+	go io.Copy(os.Stderr, t.stderr)
+	go io.Copy(os.Stdout, t.stdout)
+	go func() {
+		buf := make([]byte, 128)
+		for {
+			n, err := os.Stdin.Read(buf)
+			if err != nil {
+				fmt.Println(err)
+				return
+			}
+			if n > 0 {
+				_, err = t.stdin.Write(buf[:n])
+				if err != nil {
+					fmt.Println(err)
+					t.exitMsg = err.Error()
+					return
+				}
+			}
+		}
+	}()
+
+	err = t.Session.Shell()
+	if err != nil {
+		return err
+	}
+	err = t.Session.Wait()
+	if err != nil {
+		return err
+	}
+	return nil
+}
+
+func New(client *ssh.Client) error {
+
+	session, err := client.NewSession()
+	if err != nil {
+		return err
+	}
+	defer session.Close()
+
+	s := SSHTerminal{
+		Session: session,
+	}
+
+	return s.interactiveSession()
+}