aiops/backend/controller/kubectl_shell.go

package controller

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"os/exec"
	"strconv"
	"syscall"

	"ezops/middleware"
	"ezops/model"
	"ezops/pkg/toolinstaller"
	"ezops/pkg/wshub"

	"github.com/creack/pty"
	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"
	"gorm.io/gorm"
)

// KubectlShellHandler kubectl WebSocket 终端处理（PTY 模式）
type KubectlShellHandler struct {
	db            *gorm.DB
	toolInstaller *toolinstaller.ToolInstaller
	wsMgr         *wshub.WSManager
}

// NewKubectlShellHandler creates handler
func NewKubectlShellHandler(db *gorm.DB, ti *toolinstaller.ToolInstaller, wsMgr *wshub.WSManager) *KubectlShellHandler {
	return &KubectlShellHandler{db: db, toolInstaller: ti, wsMgr: wsMgr}
}

// wsWriter wraps a websocket.Conn as an io.Writer (each write = one WSMessage)
type wsWriter struct {
	conn *websocket.Conn
}

func (w wsWriter) Write(p []byte) (int, error) {
	err := w.conn.WriteJSON(WSMessage{Status: 0, Data: string(p)})
	if err != nil {
		return 0, err
	}
	return len(p), nil
}

// Terminal starts a PTY-backed bash session with the cluster's kubeconfig injected
func (h *KubectlShellHandler) Terminal(c *gin.Context) {
	clusterID := c.Query("id")
	if clusterID == "" {
		c.JSON(http.StatusBadRequest, gin.H{"error": "missing cluster id"})
		return
	}
	id, err := strconv.Atoi(clusterID)
	if err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": "invalid cluster id"})
		return
	}

	var cluster model.KubernetesCluster
	if err := h.db.First(&cluster, id).Error; err != nil {
		c.JSON(http.StatusNotFound, gin.H{"error": "cluster not found"})
		return
	}
	if cluster.KubeConfig == "" {
		c.JSON(http.StatusBadRequest, gin.H{"error": "no kubeconfig configured"})
		return
	}

	userID := middleware.GetUserID(c)
	username := middleware.GetUsername(c)

	conn, info, err := h.wsMgr.Upgrade(c, wshub.TypeKubectl, userID, username)
	if err != nil {
		log.Printf("[KubectlShell] WebSocket upgrade failed: %v", err)
		return
	}
	defer h.wsMgr.Unregister(info)
	h.wsMgr.StartHeartbeat(conn, info)

	// Write kubeconfig to a temp file (isolated per session)
	tmpFile, err := os.CreateTemp("", "kubeconfig-*.yaml")
	if err != nil {
		conn.WriteJSON(WSMessage{Status: 0, Data: "\r\n\x1b[31mFailed to create temp kubeconfig\x1b[0m\r\n"})
		return
	}
	defer os.Remove(tmpFile.Name())
	if _, err := tmpFile.WriteString(cluster.KubeConfig); err != nil {
		conn.WriteJSON(WSMessage{Status: 0, Data: "\r\n\x1b[31mFailed to write kubeconfig\x1b[0m\r\n"})
		return
	}
	tmpFile.Close()

	// Lazy-load kubectl on first K8s terminal connection (uses ToolInstaller)
	if err := h.toolInstaller.EnsureAvailable("kubectl"); err != nil {
		log.Printf("[KubectlShell] kubectl not available: %v", err)
		conn.WriteJSON(WSMessage{Status: 0, Data: fmt.Sprintf("\r\n\x1b[31mkubectl 不可用: %v\x1b[0m\r\n", err)})
		// Continue — maybe user just wants bash without kubectl
	}

	// 启动 bash -i
	cmd := exec.CommandContext(c.Request.Context(), "bash", "--norc", "-i")
	cmd.Env = append(os.Environ(),
		"KUBECONFIG="+tmpFile.Name(),
		"HOME=/root",
		"TERM=xterm-256color",
		fmt.Sprintf("PS1=\x1b[32m[%s]\x1b[0m $ ", cluster.Code),
		"COLUMNS=120",
		"LINES=30",
	)
	cmd.SysProcAttr = &syscall.SysProcAttr{
		Setsid: true,
	}

	ptmx, err := pty.StartWithSize(cmd, &pty.Winsize{Cols: 120, Rows: 30})
	if err != nil {
		conn.WriteJSON(WSMessage{Status: 0, Data: "\r\n\x1b[31mFailed to start shell (bash not available?)\x1b[0m\r\n"})
		return
	}
	defer func() {
		ptmx.Close()
		if cmd.Process != nil {
			syscall.Kill(-cmd.Process.Pid, syscall.SIGKILL)
		}
		cmd.Wait()
	}()

	// Welcome message
	welcome := fmt.Sprintf("\r\n\x1b[32m=== kubectl shell — %s (%s) ===\x1b[0m\r\n", cluster.Name, cluster.Code)
	welcome += "\x1b[36mKUBECONFIG 已预配 | 输入 'exit' 退出\x1b[0m\r\n\r\n"
	conn.WriteJSON(WSMessage{Status: 0, Data: welcome})

	// Use a cancellable context for the two goroutines
	ctx, cancel := context.WithCancel(c.Request.Context())
	defer cancel()

	// PTY → WebSocket
	go func() {
		defer cancel()
		n, err := io.Copy(wsWriter{conn: conn}, ptmx)
		if err != nil && err != io.EOF {
			log.Printf("[KubectlShell] PTY copy ended: %v (bytes=%d)", err, n)
		}
	}()

	// WebSocket → PTY
	go func() {
		defer ptmx.Close()
		for {
			msgType, data, err := conn.ReadMessage()
			if err != nil {
				if !websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseNormalClosure) {
				} else {
					log.Printf("[KubectlShell] WS read error: %v", err)
				}
				return
			}
			if msgType == websocket.TextMessage || msgType == websocket.BinaryMessage {
				var msg WSMessage
				if err := json.Unmarshal(data, &msg); err != nil {
					if msgType == websocket.BinaryMessage {
						ptmx.Write(data)
					}
					continue
				}
				switch msg.Status {
				case 0:
					if _, err := ptmx.Write([]byte(msg.Data)); err != nil {
						log.Printf("[KubectlShell] PTY write error: %v", err)
						return
					}
				case 1:
					if msg.Cols > 0 && msg.Rows > 0 {
						pty.Setsize(ptmx, &pty.Winsize{
							Cols: uint16(msg.Cols),
							Rows: uint16(msg.Rows),
						})
					}
				}
			}
		}
	}()

	<-ctx.Done()
	log.Printf("[KubectlShell] session ended for cluster %s", cluster.Name)
}