aiops/backend/internal/api/handlers/webshell.go

package handlers

import (
	"encoding/json"
	"log"
	"net/http"
	"strconv"
	"sync"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"
	"golang.org/x/crypto/ssh"
	"gorm.io/gorm"

	"ezops/internal/models"
)

var upgrader = websocket.Upgrader{
	CheckOrigin: func(r *http.Request) bool {
		return true
	},
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
}

// WSMessage WebSocket消息结构 // Logan.Li
type WSMessage struct {
	Status int    `json:"status"` // 0=data, 1=resize
	Data   string `json:"data"`
	Cols   int    `json:"cols"`
	Rows   int    `json:"rows"`
}

// WebShellHandler WebShell处理器
type WebShellHandler struct {
	db      *gorm.DB
	clients sync.Map
}

// NewWebShellHandler 创建WebShell处理器
func NewWebShellHandler(db *gorm.DB) *WebShellHandler {
	return &WebShellHandler{db: db}
}

// Terminal WebSocket终端连接 — 真实SSH // Logan.Li
func (h *WebShellHandler) Terminal(c *gin.Context) {
	serverID := c.Query("id")
	if serverID == "" {
		c.JSON(http.StatusBadRequest, gin.H{"error": "missing server id"})
		return
	}

	var server models.Server
	if err := h.db.First(&server, serverID).Error; err != nil {
		c.JSON(http.StatusNotFound, gin.H{"error": "server not found"})
		return
	}

	// 先升级 WebSocket — 这样前端能看到连接过程
	conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
	if err != nil {
		log.Printf("[WebShell] WebSocket升级失败: %v", err)
		return
	}
	defer conn.Close()

	conn.WriteJSON(WSMessage{
		Status: 0,
		Data:   "\r\n\x1b[36m正在连接 " + server.Hostname + " (" + server.IP + ")...\x1b[0m\r\n",
	})

	// 构建 SSH 配置
	sshConfig := &ssh.ClientConfig{
		User:            server.Username,
		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
		Timeout:         10 * time.Second,
	}

	switch server.AuthType {
	case "key":
		signer, err := ssh.ParsePrivateKey([]byte(server.SSHKey))
		if err != nil {
			sshConfig.Auth = []ssh.AuthMethod{ssh.Password(server.Password)}
		} else {
			sshConfig.Auth = []ssh.AuthMethod{ssh.PublicKeys(signer)}
		}
	default:
		sshConfig.Auth = []ssh.AuthMethod{
			ssh.Password(server.Password),
		}
	}

	// 连接 SSH（带10秒超时）
	addr := server.IP + ":" + strconv.Itoa(server.Port)
	sshClient, err := ssh.Dial("tcp", addr, sshConfig)
	if err != nil {
		log.Printf("[WebShell] SSH连接失败 %s: %v", addr, err)
		conn.WriteJSON(WSMessage{
			Status: 0,
			Data:   "\r\n\x1b[31mSSH连接失败: " + err.Error() + "\x1b[0m\r\n\x1b[33m提示: 请检查服务器IP/端口/密码是否正确，目标主机是否在线\x1b[0m\r\n",
		})
		conn.Close()
		return
	}
	defer sshClient.Close()

	log.Printf("[WebShell] SSH已连接 %s (%s)", server.Hostname, addr)

	// 创建 SSH session
	session, err := sshClient.NewSession()
	if err != nil {
		log.Printf("[WebShell] 创建session失败: %v", err)
		conn.WriteJSON(WSMessage{Status: 0, Data: "\r\n\x1b[31m创建SSH会话失败\x1b[0m\r\n"})
		return
	}
	defer session.Close()

	// 请求 PTY
	modes := ssh.TerminalModes{
		ssh.ECHO:          1,
		ssh.TTY_OP_ISPEED: 14400,
		ssh.TTY_OP_OSPEED: 14400,
	}
	if err := session.RequestPty("xterm-256color", 24, 80, modes); err != nil {
		log.Printf("[WebShell] PTY请求失败: %v", err)
		conn.WriteJSON(WSMessage{Status: 0, Data: "\r\n\x1b[31mPTY请求失败\x1b[0m\r\n"})
		return
	}

	// 获取管道
	stdinPipe, err := session.StdinPipe()
	if err != nil {
		return
	}
	stdoutPipe, err := session.StdoutPipe()
	if err != nil {
		return
	}

	// 启动 shell
	if err := session.Shell(); err != nil {
		log.Printf("[WebShell] Shell启动失败: %v", err)
		return
	}

	// 发送欢迎消息
	conn.WriteJSON(WSMessage{
		Status: 0,
		Data:   "\r\n\x1b[32mEZOps WebShell\x1b[0m — 已连接 " + server.Hostname + " (" + server.IP + ")\r\n\r\n",
	})

	done := make(chan struct{})

	// goroutine: SSH stdout -> WebSocket
	go func() {
		defer close(done)
		buf := make([]byte, 8192)
		for {
			n, err := stdoutPipe.Read(buf)
			if err != nil {
				conn.WriteJSON(WSMessage{Status: 0, Data: "\r\n\x1b[31m[远程连接已断开]\x1b[0m\r\n"})
				break
			}
			if n > 0 {
				conn.WriteJSON(WSMessage{Status: 0, Data: string(buf[:n])})
			}
		}
	}()

	// 主循环: WebSocket -> SSH stdin
	for {
		msgType, data, err := conn.ReadMessage()
		if err != nil {
			break
		}

		if msgType == websocket.TextMessage {
			var msg WSMessage
			if err := json.Unmarshal(data, &msg); err != nil {
				continue
			}

			switch msg.Status {
			case 0: // 数据输入
				stdinPipe.Write([]byte(msg.Data))
			case 1: // resize
				if msg.Cols > 0 && msg.Rows > 0 {
					session.WindowChange(msg.Rows, msg.Cols)
				}
			}
		}
	}

	// 等待 stdout goroutine 结束
	<-done
	stdinPipe.Close()
	log.Printf("[WebShell] 断开 %s (%s)", server.Hostname, addr)
}