aiops/backend/controller/webshell.go

package controller

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

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

	"ezops/middleware"
	"ezops/model"
	"ezops/pkg/wshub"
)

// upgrader is kept for backward compatibility (k8s_pod.go/kubectl_shell.go reference it).
// The WSManager has its own internal upgrader; this global is aliased for handlers
// that haven't been migrated yet or need low-level access.
var upgrader = websocket.Upgrader{
	CheckOrigin: func(r *http.Request) bool {
		return true
	},
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
}

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

// SessionContext 活跃会话上下文
type SessionContext struct {
	SessionID uint
	Cancel    chan struct{}
}

// asciicastRecorder writes terminal output in asciicast v2 format
type asciicastRecorder struct {
	buf       strings.Builder
	startTime time.Time
	mu        sync.Mutex
}

func newAsciicastRecorder() *asciicastRecorder {
	r := &asciicastRecorder{startTime: time.Now()}
	r.writeHeader()
	return r
}

func (r *asciicastRecorder) writeHeader() {
	h, _ := json.Marshal(map[string]interface{}{
		"version":   2,
		"width":     80,
		"height":    24,
		"timestamp": r.startTime.Unix(),
		"env": map[string]string{
			"SHELL": "/bin/bash",
			"TERM":  "xterm-256color",
		},
	})
	r.buf.Write(h)
	r.buf.WriteByte('\n')
}

func (r *asciicastRecorder) WriteRow(data []byte) {
	if len(data) == 0 {
		return
	}
	r.mu.Lock()
	defer r.mu.Unlock()
	delta := time.Since(r.startTime).Seconds()
	row, _ := json.Marshal([]interface{}{delta, "o", string(data)})
	r.buf.Write(row)
	r.buf.WriteByte('\n')
}

func (r *asciicastRecorder) String() string { return r.buf.String() }
func (r *asciicastRecorder) Len() int       { return r.buf.Len() }

// WebShellHandler WebShell处理器
type WebShellHandler struct {
	db       *gorm.DB
	auditHub *AuditHub
	wsMgr    *wshub.WSManager
	clients  sync.Map // sessionID → *SessionContext
}

func NewWebShellHandler(db *gorm.DB, auditHub *AuditHub, wsMgr *wshub.WSManager) *WebShellHandler {
	return &WebShellHandler{db: db, auditHub: auditHub, wsMgr: wsMgr}
}

// KillSession 强制终止指定终端会话
func (h *WebShellHandler) KillSession(sessionID uint) bool {
	if val, ok := h.clients.Load(sessionID); ok {
		ctx := val.(*SessionContext)
		close(ctx.Cancel)
		h.clients.Delete(sessionID)
		log.Printf("[WebShell] Kill signal sent session=%d", sessionID)
		return true
	}
	return false
}

// formatSSHError 将 Go SSH 错误转换为友好的中文提示
func formatSSHError(err error) string {
	errStr := err.Error()
	switch {
	case strings.Contains(errStr, "i/o timeout"):
		return "Connection timeout (target may be behind firewall or unreachable)"
	case strings.Contains(errStr, "connection refused"):
		return "Connection refused (SSH service not running or wrong port)"
	case strings.Contains(errStr, "no route to host"):
		return "No route to host (IP address unreachable)"
	case strings.Contains(errStr, "authentication failed") || strings.Contains(errStr, "unable to authenticate"):
		return "Authentication failed (wrong username or password)"
	case strings.Contains(errStr, "handshake"):
		return "SSH handshake failed (target may not be an SSH service)"
	case strings.Contains(errStr, "dns") || strings.Contains(errStr, "no such host"):
		return "DNS resolution failed (invalid IP address)"
	case strings.Contains(errStr, "ssh: must specify"):
		return "Authentication method error (no password or key configured)"
	default:
		return errStr
	}
}

// contains 字符串包含检查 (used by agent_install.go)
func contains(s, substr string) bool {
	return strings.Contains(s, substr)
}

// Terminal WebSocket终端连接 — 真实SSH + asciicast v2录制 + 命令审计
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 model.Server
	if err := h.db.First(&server, serverID).Error; err != nil {
		c.JSON(http.StatusNotFound, gin.H{"error": "server not found"})
		return
	}

	// Extract user identity BEFORE upgrade (needed for WSManager registration).
	userID := middleware.GetUserID(c)
	username := middleware.GetUsername(c)

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

	var userIDPtr *uint
	if userID > 0 {
		userIDPtr = &userID
	}
	usernameStr := username

	termSession := model.TerminalSession{
		UserID:     userIDPtr,
		Username:   usernameStr,
		ServerID:   server.ID,
		ServerName: server.Hostname,
		ServerIP:   server.IP,
		ServerPort: server.Port,
		ClientIP:   c.ClientIP(),
		Status:     "active",
		StartTime:  time.Now(),
	}
	if err := h.db.Create(&termSession).Error; err != nil {
		log.Printf("[WebShell] Create session record failed: %v", err)
	}

	recorder := newAsciicastRecorder()
	const maxRecordSize = 1024 * 1024
	var cmdBuf strings.Builder

	finalizeSession := func(status string) {
		endTime := time.Now()
		duration := int(endTime.Sub(termSession.StartTime).Seconds())
		recData := recorder.String()
		recSize := recorder.Len()
		log.Printf("[WebShell] Finalize session=%d status=%s duration=%ds rec=%d bytes", termSession.ID, status, duration, recSize)
		// Raw SQL with error logging
		result := h.db.Exec("UPDATE sys_terminal_sessions SET status=?, duration=?, end_time=?, record_data=?, record_size=? WHERE id=?",
			status, duration, endTime, recData, recSize, termSession.ID)
		if result.Error != nil {
			log.Printf("[WebShell] ERROR finalize session=%d: %v", termSession.ID, result.Error)
		} else {
			log.Printf("[WebShell] Finalize OK session=%d rows_affected=%d", termSession.ID, result.RowsAffected)
		}
	}

	connectingMsg := "\r\n\x1b[36mConnecting to " + server.Hostname + " (" + server.IP + ":" + strconv.Itoa(server.Port) + ")...\x1b[0m\r\n"
	conn.WriteJSON(WSMessage{Status: 0, Data: connectingMsg})

	sshConfig := &ssh.ClientConfig{
		User:    server.Username,
		Timeout: 10 * time.Second,
		HostKeyCallback: func(hostname string, remote net.Addr, key ssh.PublicKey) error {
			log.Printf("[WebShell] TOFU: trusting host key for %s: %s", server.Hostname, ssh.FingerprintSHA256(key))
			return nil
		},
	}

	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)}
	}

	addr := server.IP + ":" + strconv.Itoa(server.Port)
	sshClient, err := ssh.Dial("tcp", addr, sshConfig)
	if err != nil {
		log.Printf("[WebShell] SSH connect failed %s: %v", addr, err)
		friendlyErr := formatSSHError(err)
		errMsg := "\r\n\x1b[31mConnection failed: " + friendlyErr + "\x1b[0m\r\n"
		conn.WriteJSON(WSMessage{Status: 0, Data: errMsg})
		recorder.WriteRow([]byte(errMsg))
		finalizeSession("failed")
		return
	}
	defer sshClient.Close()

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

	session, err := sshClient.NewSession()
	if err != nil {
		log.Printf("[WebShell] Create SSH session failed: %v", err)
		conn.WriteJSON(WSMessage{Status: 0, Data: "\r\n\x1b[31mFailed to create SSH session\x1b[0m\r\n"})
		finalizeSession("failed")
		return
	}
	defer session.Close()

	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 request failed: %v", err)
		conn.WriteJSON(WSMessage{Status: 0, Data: "\r\n\x1b[31mPTY request failed\x1b[0m\r\n"})
		finalizeSession("failed")
		return
	}

	stdinPipe, err := session.StdinPipe()
	if err != nil {
		finalizeSession("failed")
		return
	}
	stdoutPipe, err := session.StdoutPipe()
	if err != nil {
		finalizeSession("failed")
		return
	}

	if err := session.Shell(); err != nil {
		log.Printf("[WebShell] Shell start failed: %v", err)
		finalizeSession("failed")
		return
	}

	welcome := "\r\n\x1b[32mAIOps WebShell\x1b[0m \u2014 connected to " + server.Hostname + " (" + server.IP + ")\r\n\r\n"
	conn.WriteJSON(WSMessage{Status: 0, Data: welcome})
	recorder.WriteRow([]byte(welcome))

	sessCtx := &SessionContext{
		SessionID: termSession.ID,
		Cancel:    make(chan struct{}),
	}
	h.clients.Store(termSession.ID, sessCtx)
	log.Printf("[WebShell] Session registered session=%d", termSession.ID)
	defer func() {
		h.clients.Delete(termSession.ID)
		if h.auditHub != nil {
			h.auditHub.Cleanup(termSession.ID)
		}
		log.Printf("[WebShell] Session unregistered session=%d", termSession.ID)
	}()

	done := make(chan struct{})

	// goroutine 1: Cancel listener
	go func() {
		select {
		case <-sessCtx.Cancel:
			log.Printf("[WebShell] Received Kill signal session=%d, closing", termSession.ID)
			conn.Close()
		case <-done:
		}
	}()

	// goroutine 2: SSH stdout → WebSocket + asciicast recording + AuditHub
	go func() {
		defer close(done)
		buf := make([]byte, 8192)
		for {
			n, err := stdoutPipe.Read(buf)
			if err != nil {
				discMsg := "\r\n\x1b[31m[Remote connection closed]\x1b[0m\r\n"
				conn.WriteJSON(WSMessage{Status: 0, Data: discMsg})
				recorder.WriteRow([]byte(discMsg))
				if h.auditHub != nil {
					h.auditHub.Broadcast(termSession.ID, []byte(discMsg))
				}
				break
			}
			if n > 0 {
				data := buf[:n]
				conn.WriteJSON(WSMessage{Status: 0, Data: string(data)})
				if recorder.Len() < maxRecordSize {
					recorder.WriteRow(data)
				}
				if h.auditHub != nil {
					h.auditHub.Broadcast(termSession.ID, data)
				}
			}
		}
	}()

	// Main loop: WebSocket → SSH stdin (command capture + idle timeout)
	const idleTimeout = 15 * time.Minute
	for {
		conn.SetReadDeadline(time.Now().Add(idleTimeout))
		msgType, data, err := conn.ReadMessage()
		if err != nil {
			if netErr, ok := err.(net.Error); ok && netErr.Timeout() {
				log.Printf("[WebShell] Idle timeout 15min session=%d", termSession.ID)
			} else {
				log.Printf("[WebShell] WS read error session=%d: %v", termSession.ID, err)
			}
			break
		}

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

			switch msg.Status {
			case 0:
				if msg.Data != "" {
					stdinPipe.Write([]byte(msg.Data))

					// Command extraction from stdin
					cmdBuf.WriteString(msg.Data)
					if strings.Contains(msg.Data, "\r") || strings.Contains(msg.Data, "\n") {
						cmd := extractCommand(cmdBuf.String())
						if cmd != "" {
							recordCommand(h.db, termSession.ID, termSession.ServerID, usernameStr, cmd)
						}
						cmdBuf.Reset()
					}
					if cmdBuf.Len() > 8192 {
						cmdBuf.Reset()
					}
				}
			case 1: // resize
				if msg.Cols > 0 && msg.Rows > 0 {
					session.WindowChange(msg.Rows, msg.Cols)
				}
			}
		}
	}

	// Close SSH to unblock stdout goroutine, then wait for it
	sshClient.Close()
	<-done
	stdinPipe.Close()
	finalizeSession("closed")
	log.Printf("[WebShell] Disconnected %s (%s)", server.Hostname, addr)
}

// extractCommand extracts a clean command from raw terminal input
func extractCommand(input string) string {
	cleaned := strings.Map(func(r rune) rune {
		if r < 32 && r != '\t' {
			return -1
		}
		if r == 127 {
			return -1
		}
		return r
	}, input)
	cleaned = strings.TrimSpace(cleaned)
	if len(cleaned) < 2 || len(cleaned) > 500 {
		return ""
	}
	if strings.HasPrefix(cleaned, "\x1b[") {
		return ""
	}
	return cleaned
}

// recordCommand records a terminal command to sys_session_commands
func recordCommand(db *gorm.DB, sessionID, serverID uint, username, command string) {
	now := time.Now()
	cmd := model.SessionCommand{
		SessionID: sessionID,
		ServerID:  serverID,
		Username:  username,
		Command:   command,
		Timestamp: now,
	}
	if err := db.Create(&cmd).Error; err != nil {
		log.Printf("[WebShell] Command record failed: %v", err)
	} else {
		log.Printf("[WebShell] Command recorded session=%d cmd=%q", sessionID, command)
	}
}