From 1a1fe7692fab66e194aa2a02606a56871998feb5 Mon Sep 17 00:00:00 2001
From: Hatter Jiang <jht5945@gmail.com>
Date: Sat, 17 Jul 2021 21:14:20 +0800
Subject: [PATCH] feat: pty ok

---
 go.mod  |  5 ++++-
 go.sum  |  2 ++
 main.go | 33 +++++++++++++++++++++++++++++++++
 3 files changed, 39 insertions(+), 1 deletion(-)

diff --git a/go.mod b/go.mod
index 85a1093..bdf4969 100644
--- a/go.mod
+++ b/go.mod
@@ -2,4 +2,7 @@ module git.hatter.ink/simple-ssh-server
 
 go 1.13
 
-require github.com/gliderlabs/ssh v0.3.3
+require (
+	github.com/creack/pty v1.1.13 // indirect
+	github.com/gliderlabs/ssh v0.3.3
+)
diff --git a/go.sum b/go.sum
index a141e53..2b7f83f 100644
--- a/go.sum
+++ b/go.sum
@@ -1,5 +1,7 @@
 github.com/anmitsu/go-shlex v0.0.0-20200514113438-38f4b401e2be h1:9AeTilPcZAjCFIImctFaOjnTIavg87rW78vTPkQqLI8=
 github.com/anmitsu/go-shlex v0.0.0-20200514113438-38f4b401e2be/go.mod h1:ySMOLuWl6zY27l47sB3qLNK6tF2fkHG55UZxx8oIVo4=
+github.com/creack/pty v1.1.13 h1:rTPnd/xocYRjutMfqide2zle1u96upp1gm6eUHKi7us=
+github.com/creack/pty v1.1.13/go.mod h1:MOBLtS5ELjhRRrroQr9kyvTxUAFNvYEK993ew/Vr4O4=
 github.com/gliderlabs/ssh v0.3.3 h1:mBQ8NiOgDkINJrZtoizkC3nDNYgSaWtxyem6S2XHBtA=
 github.com/gliderlabs/ssh v0.3.3/go.mod h1:ZSS+CUoKHDrqVakTfTWUlKSr9MtMFkC4UvtQKD7O914=
 golang.org/x/crypto v0.0.0-20210616213533-5ff15b29337e h1:gsTQYXdTw2Gq7RBsWvlQ91b+aEQ6bXFUngBGuR8sPpI=
diff --git a/main.go b/main.go
index ff61e96..5497ae6 100644
--- a/main.go
+++ b/main.go
@@ -4,8 +4,13 @@ import (
 	"fmt"
 	"io"
 	"log"
+	"os"
+	"os/exec"
 	"strings"
+	"syscall"
+	"unsafe"
 
+	"github.com/creack/pty"
 	"github.com/gliderlabs/ssh"
 	gossh "golang.org/x/crypto/ssh"
 )
@@ -19,6 +24,29 @@ func main() {
 		io.WriteString(s, fmt.Sprintf("\r\n%s\r\n", strings.Repeat("-", 88)))
 		s.Write(authorizedKey)
 		io.WriteString(s, fmt.Sprintf("\r\n%s\r\n", strings.Repeat("-", 88)))
+
+		cmd := exec.Command("/bin/bash")
+		ptyReq, winCh, isPty := s.Pty()
+		if isPty {
+			cmd.Env = append(cmd.Env, fmt.Sprintf("TERM=%s", ptyReq.Term))
+			f, err := pty.Start(cmd)
+			if err != nil {
+				panic(err)
+			}
+			go func() {
+				for win := range winCh {
+					setWinsize(f, win.Width, win.Height)
+				}
+			}()
+			go func() {
+				io.Copy(f, s) // stdin
+			}()
+			io.Copy(s, f) // stdout
+			cmd.Wait()
+		} else {
+			io.WriteString(s, "No PTY requested.\n")
+			s.Exit(1)
+		}
 	})
 
 	publicKeyOption := ssh.PublicKeyAuth(func(ctx ssh.Context, key ssh.PublicKey) bool {
@@ -32,3 +60,8 @@ func main() {
 	log.Println("Listening :222...")
 	log.Fatal(ssh.ListenAndServe(":2222", nil, publicKeyOption))
 }
+
+func setWinsize(f *os.File, w, h int) {
+	syscall.Syscall(syscall.SYS_IOCTL, f.Fd(), uintptr(syscall.TIOCSWINSZ),
+		uintptr(unsafe.Pointer(&struct{ h, w, x, y uint16 }{uint16(h), uint16(w), 0, 0})))
+}