network-topology-discovery/internal/terminal/handler.go

package terminal

import (
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"sync"
	"time"

	"github.com/gorilla/websocket"
	"golang.org/x/crypto/ssh"
)

var upgrader = websocket.Upgrader{
	CheckOrigin: func(r *http.Request) bool {
		return true // 允许所有来源
	},
}

// TerminalSession 终端会话
type TerminalSession struct {
	sshClient  *ssh.Client
	sshSession *ssh.Session
	stdin      io.Writer
	stdout     io.Reader
	wsConn     *websocket.Conn
	done       chan struct{}
	mu         sync.Mutex
}

// ConnectSSH 建立SSH连接并创建交互式Shell
func ConnectSSH(host string, port int, username, password string) (*TerminalSession, error) {
	config := &ssh.ClientConfig{
		User:            username,
		Auth:            []ssh.AuthMethod{ssh.Password(password)},
		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
		Timeout:         10 * time.Second,
		Config: ssh.Config{
			Ciphers: []string{
				"aes128-ctr", "aes192-ctr", "aes256-ctr",
				"aes128-gcm@openssh.com", "aes256-gcm@openssh.com",
				"chacha20-poly1305@openssh.com",
				"aes128-cbc", "aes256-cbc",
			},
			KeyExchanges: []string{
				"curve25519-sha256", "curve25519-sha256@libssh.org",
				"ecdh-sha2-nistp256", "ecdh-sha2-nistp384", "ecdh-sha2-nistp521",
				"diffie-hellman-group14-sha256", "diffie-hellman-group16-sha512",
				"diffie-hellman-group14-sha1", "diffie-hellman-group1-sha1",
			},
		},
	}

	addr := fmt.Sprintf("%s:%d", host, port)
	client, err := ssh.Dial("tcp", addr, config)
	if err != nil {
		return nil, fmt.Errorf("SSH连接失败: %w", err)
	}

	session, err := client.NewSession()
	if err != nil {
		client.Close()
		return nil, fmt.Errorf("创建SSH会话失败: %w", err)
	}

	// 获取 stdin 管道
	stdin, err := session.StdinPipe()
	if err != nil {
		session.Close()
		client.Close()
		return nil, fmt.Errorf("获取stdin失败: %w", err)
	}

	// 获取 stdout 管道
	stdout, err := session.StdoutPipe()
	if err != nil {
		session.Close()
		client.Close()
		return nil, fmt.Errorf("获取stdout失败: %w", err)
	}

	// 也获取 stderr
	session.Stderr = io.Discard

	// 请求 PTY（xterm 终端）
	modes := ssh.TerminalModes{
		ssh.ECHO:          1,
		ssh.TTY_OP_ISPEED: 14400,
		ssh.TTY_OP_OSPEED: 14400,
	}
	if err := session.RequestPty("xterm", 40, 120, modes); err != nil {
		session.Close()
		client.Close()
		return nil, fmt.Errorf("请求PTY失败: %w", err)
	}

	// 启动 Shell
	if err := session.Shell(); err != nil {
		session.Close()
		client.Close()
		return nil, fmt.Errorf("启动Shell失败: %w", err)
	}

	return &TerminalSession{
		sshClient:  client,
		sshSession: session,
		stdin:      stdin,
		stdout:     stdout,
		done:       make(chan struct{}),
	}, nil
}

// HandleTerminal 处理WebSocket终端连接
func HandleTerminal(w http.ResponseWriter, r *http.Request, host string, port int, username, password string) {
	// 建立 SSH 连接
	session, err := ConnectSSH(host, port, username, password)
	if err != nil {
		log.Printf("[终端] SSH连接失败 %s: %v", host, err)
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	// 升级为 WebSocket
	wsConn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		session.Close()
		log.Printf("[终端] WebSocket升级失败: %v", err)
		return
	}
	session.wsConn = wsConn

	log.Printf("[终端] 已连接到 %s (%s)", host, username)

	// 启动 SSH -> WebSocket 的数据转发
	go session.sshToWs()
	// 启动 WebSocket -> SSH 的数据转发
	go session.wsToSsh()
	// 等待结束
	<-session.done

	session.Close()
	log.Printf("[终端] 已断开 %s", host)
}

// sshToWs 从SSH读取输出并转发到WebSocket
func (s *TerminalSession) sshToWs() {
	buf := make([]byte, 8192)
	for {
		select {
		case <-s.done:
			return
		default:
		}

		n, err := s.stdout.Read(buf)
		if err != nil {
			log.Printf("[终端] SSH读取结束: %v", err)
			s.closeDone()
			return
		}

		if n > 0 {
			s.mu.Lock()
			err := s.wsConn.WriteMessage(websocket.TextMessage, buf[:n])
			s.mu.Unlock()
			if err != nil {
				log.Printf("[终端] WebSocket写入失败: %v", err)
				s.closeDone()
				return
			}
		}
	}
}

// wsToSsh 从WebSocket读取输入并转发到SSH
func (s *TerminalSession) wsToSsh() {
	for {
		select {
		case <-s.done:
			return
		default:
		}

		_, message, err := s.wsConn.ReadMessage()
		if err != nil {
			log.Printf("[终端] WebSocket读取失败: %v", err)
			s.closeDone()
			return
		}

		if len(message) > 0 {
			// 解析JSON消息格式（xterm.js发送的）
			var msg map[string]interface{}
			if err := json.Unmarshal(message, &msg); err == nil {
				if input, ok := msg["input"].(string); ok {
					_, err := s.stdin.Write([]byte(input))
					if err != nil {
						log.Printf("[终端] SSH写入失败: %v", err)
						s.closeDone()
						return
					}
				}
				// 处理resize消息
				if msg["type"] == "resize" {
					if cols, ok := msg["cols"].(float64); ok {
						if rows, ok := msg["rows"].(float64); ok {
							_ = s.sshSession.WindowChange(int(rows), int(cols))
						}
					}
				}
			} else {
				// 原始二进制数据，直接写入
				_, err := s.stdin.Write(message)
				if err != nil {
					s.closeDone()
					return
				}
			}
		}
	}
}

// closeDone 安全关闭
func (s *TerminalSession) closeDone() {
	s.mu.Lock()
	defer s.mu.Unlock()
	select {
	case <-s.done:
		// 已经关闭
	default:
		close(s.done)
	}
}

// Close 关闭会话
func (s *TerminalSession) Close() {
	s.closeDone()
	if s.wsConn != nil {
		s.wsConn.Close()
	}
	if s.sshSession != nil {
		s.sshSession.Close()
	}
	if s.sshClient != nil {
		s.sshClient.Close()
	}
}