update shell

1 files changed, 46 insertions, 53 deletions

diff --git a/shell.go b/shell.go
index 3492b33..d750136 100644
--- a/shell.go
+++ b/shell.go
@@ -2,69 +2,62 @@ package main
 
 import (
 	"fmt"
-	"net"
+	"io"
 	"os"
 	"os/exec"
 
-	"golang.org/x/crypto/ssh"
+	"github.com/creack/pty"
+	"github.com/gliderlabs/ssh"
 )
 
-func shell(conn net.Conn, config *ssh.ServerConfig) {
-	sshConn, chans, reqs, err := ssh.NewServerConn(conn, config)
-	if err != nil {
-		fmt.Println("不能创建连接:", err)
+func shell(s ssh.Session) {
+	ptyReq, winCh, isPty := s.Pty()
+
+	if !isPty {
+		fmt.Fprintln(s, "Must be PTY")
+		s.Exit(1)
 		return
 	}
-	defer sshConn.Close()
-
-	fmt.Println("New connection from", sshConn.RemoteAddr(), "with client version", sshConn.ClientVersion())
-
-	go ssh.DiscardRequests(reqs)
-
-	for newChannel := range chans {
-		if newChannel.ChannelType() != "session" {
-			newChannel.Reject(ssh.UnknownChannelType, "unknown channel type")
-			continue
-		}
-
-		channel, requests, err := newChannel.Accept()
-		if err != nil {
-			fmt.Println("Can not accept channel:", err)
-			continue
-		}
-		defer channel.Close()
 
-		shell := os.Getenv("SHELL")
-		if shell == "" {
-			shell = "cmd.exe"
-		}
+	shell := os.Getenv("SHELL")
+	if shell == "" {
+		shell = "/bin/sh"
+	}
 
-		command := exec.Command(shell)
-		command.Stdin = channel
-		command.Stdout = channel
-		command.Stderr = channel
+	cmd := exec.Command(shell)
+	cmd.Env = append(cmd.Env, fmt.Sprintf("TERM=%s", ptyReq.Term))
 
-		if err := command.Start(); err != nil {
-			fmt.Println("Failed to start shell:", err)
-			return
+	ptmx, err := pty.Start(cmd)
+	if err != nil {
+		fmt.Fprintln(s, "Can not start shell")
+		fmt.Println(err)
+		s.Exit(1)
+		return
+	}
+	defer func() { _ = ptmx.Close() }()
+
+	go func() {
+		for win := range winCh {
+			pty.Setsize(ptmx, &pty.Winsize{
+				Rows: uint16(win.Height),
+				Cols: uint16(win.Width),
+			})
 		}
-
-		go func() {
-			if err := command.Wait(); err != nil {
-				fmt.Println("Run shell failed:", err)
-			}
-			channel.Close()
-		}()
-
-		go func() {
-			for req := range requests {
-				switch req.Type {
-				case "shell":
-					req.Reply(true, nil)
-				default:
-					req.Reply(false, nil)
-				}
-			}
-		}()
+	}()
+
+	go func() {
+		io.Copy(s, ptmx)
+		s.Close()
+	}()
+
+	go func() {
+		io.Copy(ptmx, s)
+		ptmx.Close()
+	}()
+
+	if err := cmd.Wait(); err != nil {
+		s.Exit(1)
+	} else {
+		s.Exit(0)
 	}
 }