package main

import (
	"fmt"
	"io"
	"log"
	"net"
	"os"
	"os/exec"
	"strings"

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

func main() {
	// 生成SSH密钥对
	privateKeyBytes, err := os.ReadFile("../.ssh/id_rsa")
	if err != nil {
		log.Fatal("无法读取私钥文件：", err)
	}

	privateKey, err := ssh.ParsePrivateKey(privateKeyBytes)
	if err != nil {
		log.Fatal("无法解析私钥：", err)
	}

	// 配置SSH服务器
	config := &ssh.ServerConfig{
		PasswordCallback: func(conn ssh.ConnMetadata, password []byte) (*ssh.Permissions, error) {
			// 这里可以进行密码验证逻辑，例如检查用户名和密码是否匹配
			if conn.User() == "admin" && string(password) == "123456" {
				return nil, nil
			}
			return nil, fmt.Errorf("密码错误")
		},
		PublicKeyCallback: func(conn ssh.ConnMetadata, key ssh.PublicKey) (*ssh.Permissions, error) {
			// 这里可以进行公钥验证逻辑，例如检查用户的公钥是否在允许的公钥列表中
			return nil, fmt.Errorf("公钥验证失败")
		},
	}

	config.AddHostKey(privateKey)

	// 启动SSH服务器
	listener, err := net.Listen("tcp", "192.168.2.119:2222")
	if err != nil {
		log.Fatal("无法监听端口：", err)
	}

	log.Println("SSH服务器已启动，监听端口：2222")

	for {
		conn, err := listener.Accept()
		if err != nil {
			log.Fatal("无法接受连接：", err)
		}

		go handleSSHConnection(conn, config)
	}
}

func handleSSHConnection(conn net.Conn, config *ssh.ServerConfig) {
	// SSH握手
	sshConn, chans, reqs, err := ssh.NewServerConn(conn, config)
	if err != nil {
		log.Println("SSH握手失败：", err)
		return
	}

	log.Printf("SSH连接已建立，客户端版本：%s，用户：%s", sshConn.ClientVersion(), sshConn.User())

	// 处理SSH请求
	go ssh.DiscardRequests(reqs)

	// 处理SSH通道
	for newChannel := range chans {
		if newChannel.ChannelType() != "session" {
			newChannel.Reject(ssh.UnknownChannelType, "不支持的通道类型")
			continue
		}

		channel, requests, err := newChannel.Accept()
		if err != nil {
			log.Println("无法接受通道：", err)
			continue
		}

		go handleSSHChannel(channel, requests)
	}
}

func handleSSHChannel(channel ssh.Channel, requests <-chan *ssh.Request) {
	for req := range requests {
		switch req.Type {
		case "exec":
			// 执行远程命令
			cmd := exec.Command("cmd", "/c", strings.TrimSpace(string(req.Payload)))
			cmd.Stdin = channel
			cmd.Stdout = channel
			cmd.Stderr = channel.Stderr()
			log.Println("cmd:", strings.TrimSpace(string(req.Payload)))

			err := cmd.Run()
			if err != nil {
				log.Println("无法执行命令：", err)
			}

			channel.SendRequest("exit-status", false, []byte{0, 0, 0, 0})
			channel.Close()
		case "shell":
			// // 处理交互式shell会话
			// // 在这里添加您的处理逻辑，例如启动一个shell进程并将其连接到channel
			// // 请注意，处理交互式shell会话需要更复杂的逻辑，您可能需要使用类似于pty包来处理终端相关的操作
			// channel.Write([]byte("交互式shell会话已启动\n"))
			// channel.Close()
			handleSSHShell(channel)
		default:
			req.Reply(false, nil)
		}
	}
}

func handleSSHShell(channel ssh.Channel) {
	defer channel.Close()

	// 检查通道是否支持终端
	term := terminal.NewTerminal(channel, "> ")

	// 启动交互式shell会话
	for {
		line, err := term.ReadLine()
		if err != nil {
			if err == io.EOF {
				break
			}
			log.Println("Failed to read line: ", err)
			break
		}

		// 在这里处理输入的命令
		// 您可以根据需要添加更多的逻辑
		if strings.TrimSpace(line) == "exit" {
			break
		}

		response := fmt.Sprintf("You entered: %s\n", line)
		term.Write([]byte(response))
	}
}

// func main() {
// 	// 读取 SSH 私钥文件
// 	privateKeyBytes, err := os.ReadFile("../.ssh/private_key.pem")
// 	if err != nil {
// 		log.Fatalf("Failed to load private key: %v", err)
// 	}

// 	// 解析私钥
// 	privateKey, err := ssh.ParsePrivateKey(privateKeyBytes)
// 	if err != nil {
// 		log.Fatalf("Failed to parse private key: %v", err)
// 	}

// 	// 配置 SSH 服务器
// 	config := &ssh.ServerConfig{
// 		PasswordCallback: func(conn ssh.ConnMetadata, password []byte) (*ssh.Permissions, error) {
// 			// 在此处验证密码
// 			if string(password) == "123456" {
// 				return nil, nil
// 			}
// 			return nil, fmt.Errorf("password rejected for %s", conn.User())
// 		},
// 	}

// 	// 添加私钥到配置
// 	config.AddHostKey(privateKey)

// 	// 启动 SSH 服务器
// 	listener, err := net.Listen("tcp", "0.0.0.0:2222")
// 	if err != nil {
// 		log.Fatalf("Failed to listen on 2222: %v", err)
// 	}

// 	log.Println("SSH server listening on 0.0.0.0:2222")

// 	for {
// 		conn, err := listener.Accept()
// 		if err != nil {
// 			log.Fatalf("Failed to accept incoming connection: %v", err)
// 		}

// 		// 处理 SSH 连接
// 		go func() {
// 			sshConn, chans, reqs, err := ssh.NewServerConn(conn, config)
// 			if err != nil {
// 				log.Printf("Failed to handshake: %v", err)
// 				return
// 			}

// 			log.Printf("SSH connection established from %s", sshConn.RemoteAddr())

// 			// 处理 SSH 请求
// 			go ssh.DiscardRequests(reqs)

// 			// 处理 SSH 通道
// 			for newChannel := range chans {
// 				if newChannel.ChannelType() != "session" {
// 					newChannel.Reject(ssh.UnknownChannelType, "unknown channel type")
// 					continue
// 				}

// 				channel, requests, err := newChannel.Accept()
// 				if err != nil {
// 					log.Printf("Failed to accept channel: %v", err)
// 					continue
// 				}

// 				// 处理 SSH 会话
// 				go func(in <-chan *ssh.Request) {
// 					for req := range in {
// 						log.Printf("Received out-of-band request: %v", string(req))
// 					}
// 				}(requests)

// 				// 示例：执行远程命令
// 				go func() {
// 					defer channel.Close()

// 					command := "echo 'Hello, SSH!'"
// 					output, err := execCommand(command)
// 					if err != nil {
// 						log.Printf("Failed to execute command: %v", err)
// 						return
// 					}

// 					_, err = channel.Write([]byte(output))
// 					if err != nil {
// 						log.Printf("Failed to write output: %v", err)
// 						return
// 					}
// 				}()
// 			}
// 		}()
// 	}
// }

// func execCommand(command string) (string, error) {
// 	// 在此处执行远程命令并返回结果
// 	return "hello ssh", nil
// }