pycharm-ssh.go

// main
/**
 * A golang ssh server which supports vscode remote ssh plugin.
 *
 * ssh_config:
 *   Host upterm
 *    HostName 192.168.101.135
 *    Port 2224
 *
 * VSCode remote ssh plugin use this command to setup its server, and use dynamic port fowarding to connect to this server.
 * 	ssh -T -D 1234 192.168.101.135 -p 2224  bash
 *  sftp https://gist.github.com/artyom/1a7d46c9511dd93b0714baf7b9745605
 *  bandwidth https://pkg.go.dev/github.com/gerritjvv/tcpshaper/bandwidth
 */
package main

import (
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"strings"
	"syscall"
	"unsafe"

	"github.com/gliderlabs/ssh"
	"github.com/pkg/sftp"

	"github.com/kr/pty"
)

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})))
}

func SftpHandler(sess ssh.Session) {
	debugStream := ioutil.Discard
	serverOptions := []sftp.ServerOption{
		sftp.WithDebug(debugStream),
	}
	server, err := sftp.NewServer(
		sess,
		serverOptions...,
	)
	if err != nil {
		log.Printf("sftp server init error: %s\n", err)
		return
	}
	if err := server.Serve(); err == io.EOF {
		server.Close()
		fmt.Println("sftp client exited session.")
	} else if err != nil {
		fmt.Println("sftp server completed with error:", err)
	}

}

func main() {
	svr := ssh.Server{
		Addr: ":2224",
		LocalPortForwardingCallback: ssh.LocalPortForwardingCallback(func(ctx ssh.Context, dhost string, dport uint32) bool {
			log.Println("Accepted forward", dhost, dport)
			return true
		}),
		// ReversePortForwardingCallback: ssh.ReversePortForwardingCallback(func(ctx ssh.Context, host string, port uint32) bool {
		// 	log.Println("attempt to bind", host, port, "granted")
		// 	return true
		// }),
		// PasswordHandler: func(ctx ssh.Context, password string) bool {
		// 	return true
		// },

		Handler: ssh.Handler(func(s ssh.Session) {
			ptyReq, winCh, isPty := s.Pty()
			if isPty {
				fmt.Println("In PTY....")
				cmd := exec.Command("/bin/bash")
				cmd.Env = s.Environ()

				// cmd.Env = append(cmd.Env, fmt.Sprintf("TERM=%s", ptyReq.Term))
				// cmd.Env = append(cmd.Env, fmt.Sprintf("0=%s", "zsh"))
				// cmd.Env = append(cmd.Env, fmt.Sprintf("USER=%s", "corvo"))
				// cmd.Env = append(cmd.Env, fmt.Sprintf("PATH=%s", "PATH=/home/corvo/.cargo/bin:/usr/local/bin:/usr/bin:/bin:/usr/games"))
				// cmd.Env = append(cmd.Env, fmt.Sprintf("HOME=%s", "/home/corvo"))

				fmt.Println(cmd.Env, 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 {
				cmds := s.Command()
				var cmd *exec.Cmd

				// NOTICE: pycharm need session open
				// defer s.Close()

				if len(cmds) == 0 {
					cmds = []string{"bash"}
					cmd = exec.Command(cmds[0], cmds[1:]...)
				} else {
					if cmds[0] == "/usr/bin/zsh" || cmds[0] == "/bin/zsh" || cmds[0] == "/bin/bash" || cmds[0] == "/bin/sh" {
						fmt.Println(cmds[0], cmds[1], cmds[2])
						cmd = exec.Command(cmds[0], cmds[1:]...)
					} else {
						cmds = []string{"-c", strings.Join(cmds, " ")}
						cmd = exec.Command("/usr/bin/zsh", cmds...)
					}
				}

				// cmd.Env = append(cmd.Env, fmt.Sprintf("TERM=%s", ptyReq.Term))
				// cmd.Env = append(cmd.Env, fmt.Sprintf("SHELL=%s", "/bin/zsh"))
				// cmd.Env = append(cmd.Env, fmt.Sprintf("SHELL=%s", "/bin/bash"))
				// cmd.Env = append(cmd.Env, fmt.Sprintf("0=%s", "bash"))
				// cmd.Env = append(cmd.Env, fmt.Sprintf("USER=%s", "corvo"))
				// cmd.Env = append(cmd.Env, fmt.Sprintf("PATH=%s", "PATH=/home/corvo/.cargo/bin:/usr/local/bin:/usr/bin:/bin:/usr/games"))
				// cmd.Env = append(cmd.Env, fmt.Sprintf("HOME=%s", "/home/corvo"))

				cmd.Stdout = s
				cmd.Stderr = s
				cmd.Stdin = s
				err := cmd.Start()
				if err != nil {
					log.Printf("could not start command (%s)", err)
				}

				_, err = cmd.Process.Wait()
				if err != nil {
					log.Printf("failed to exit bash (%s)", err)
				}

				log.Printf("session closed: %s", cmds)
			}
		}),

		ChannelHandlers: map[string]ssh.ChannelHandler{
			"session":      ssh.DefaultSessionHandler,
			"direct-tcpip": ssh.DirectTCPIPHandler,
		},
		SubsystemHandlers: map[string]ssh.SubsystemHandler{
			"sftp": func(sess ssh.Session) {
				debugStream := ioutil.Discard
				serverOptions := []sftp.ServerOption{
					sftp.WithDebug(debugStream),
				}
				server, err := sftp.NewServer(
					sess,
					serverOptions...,
				)
				if err != nil {
					log.Printf("sftp server init error: %s\n", err)
					return
				}
				if err := server.Serve(); err == io.EOF {
					server.Close()
					fmt.Println("sftp client exited session.")
				} else if err != nil {
					fmt.Println("sftp server completed with error:", err)
				}
			},
		},
	}
	svr.SetOption(
		ssh.HostKeyFile("/home/corvo/.ssh/id_rsa_tencent_4096"),
	)
	log.Println("starting ssh server on port 2224...")
	log.Fatal(svr.ListenAndServe())
}