package main

import (
	"fmt"
	"net"
	"os"
	"os/exec"

	"golang.org/x/crypto/ssh"
)

func shell(conn net.Conn, config *ssh.ServerConfig) {
	sshConn, chans, reqs, err := ssh.NewServerConn(conn, config)
	if err != nil {
		fmt.Println("不能创建连接:", err)
		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"
		}

		command := exec.Command(shell)
		command.Stdin = channel
		command.Stdout = channel
		command.Stderr = channel

		if err := command.Start(); err != nil {
			fmt.Println("Failed to start shell:", err)
			return
		}

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