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

package handlers

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

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

	"ezops/internal/models"
)

// KubectlShellHandler kubectl WebSocket 终端处理（PTY 模式）
type KubectlShellHandler struct {
	db *gorm.DB
}

// NewKubectlShellHandler creates handler
func NewKubectlShellHandler(db *gorm.DB) *KubectlShellHandler {
	return &KubectlShellHandler{db: db}
}

// 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 models.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
	}

	conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
	if err != nil {
		log.Printf("[KubectlShell] WebSocket upgrade failed: %v", err)
		return
	}
	defer conn.Close()

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

	// 启动 bash -i（TERM=dumb，无 PS1 ANSI 码，不回显混乱）
	// 注意：不加 --noediting，否则无 TTY 时 bash 不回显输入
	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,
		Setctty: 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()
		// Kill the process group so child processes (kubectl, etc.) are cleaned up
		if cmd.Process != nil {
			syscall.Kill(-cmd.Process.Pid, syscall.SIGKILL)
		}
		cmd.Wait()
	}()

	// Welcome message (already on PTY so bash will show PS1 right after)
	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 (single goroutine via io.Copy)
	go func() {
		defer cancel() // unblock main goroutine when PTY closes
		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 (read JSON messages, write raw to PTY)
	go func() {
		defer ptmx.Close() // close PTY to stop bash and unblock Copy goroutine
		for {
			msgType, data, err := conn.ReadMessage()
			if err != nil {
				if !websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseNormalClosure) {
					// Normal close — client disconnected
				} 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 {
					// Binary data (rare) — pass through raw
					if msgType == websocket.BinaryMessage {
						ptmx.Write(data)
					}
					continue
				}
				switch msg.Status {
				case 0:
					// Keyboard input — write raw bytes to PTY
					if _, err := ptmx.Write([]byte(msg.Data)); err != nil {
						log.Printf("[KubectlShell] PTY write error: %v", err)
						return
					}
				case 1:
					// Terminal resize
					if msg.Cols > 0 && msg.Rows > 0 {
						if err := pty.Setsize(ptmx, &pty.Winsize{
							Cols: uint16(msg.Cols),
							Rows: uint16(msg.Rows),
						}); err != nil {
							log.Printf("[KubectlShell] PTY resize error: %v", err)
						}
					}
				}
			}
		}
	}()

	// Block until PTY closes or context cancelled
	<-ctx.Done()
	log.Printf("[KubectlShell] session ended for cluster %s", cluster.Name)
}

func init() {
	// ensure kubectl is available
	if _, err := exec.LookPath("kubectl"); err != nil {
		log.Printf("[KubectlShell] WARNING: kubectl not found in PATH — kubectl commands will not work")
	}
}