aiops/backend/controller/rdp_session.go

package controller

import (
	"fmt"
	"net/http"
	"strconv"
	"strings"
	"time"

	"ezops/middleware"
	"ezops/model"

	"github.com/gin-gonic/gin"
	"gorm.io/gorm"
)

// RDPSessionHandler 管理 Windows RDP 图形会话审计记录 // Logan.Li
type RDPSessionHandler struct {
	db       *gorm.DB
	killFunc func(sessionID uint) bool
}

func NewRDPSessionHandler(db *gorm.DB, killFunc func(sessionID uint) bool) *RDPSessionHandler {
	return &RDPSessionHandler{db: db, killFunc: killFunc}
}

func (h *RDPSessionHandler) ListRDPSessions(c *gin.Context) {
	var req model.RDPSessionQuery
	if err := c.ShouldBindQuery(&req); err != nil {
		c.JSON(http.StatusBadRequest, model.Response{Code: 400, Message: err.Error()})
		return
	}
	if req.Page <= 0 {
		req.Page = 1
	}
	if req.PageSize <= 0 {
		req.PageSize = 20
	}

	query := h.db.Model(&model.RDPSession{})
	if req.ServerID > 0 {
		query = query.Where("server_id = ?", req.ServerID)
	}
	if req.Username != "" {
		query = query.Where("username LIKE ?", "%"+req.Username+"%")
	}
	if req.Status != "" {
		query = query.Where("status = ?", req.Status)
	}
	if kw := strings.TrimSpace(req.Keyword); kw != "" {
		like := "%" + kw + "%"
		query = query.Where("username LIKE ? OR server_name LIKE ? OR server_ip LIKE ? OR client_ip LIKE ?", like, like, like, like)
	}
	if req.StartDate != "" {
		if t, err := time.ParseInLocation("2006-01-02", req.StartDate, time.Local); err == nil {
			query = query.Where("start_time >= ?", t)
		}
	}
	if req.EndDate != "" {
		if t, err := time.ParseInLocation("2006-01-02", req.EndDate, time.Local); err == nil {
			query = query.Where("start_time < ?", t.Add(24*time.Hour))
		}
	}
	if req.HasRecord == "1" || strings.EqualFold(req.HasRecord, "true") {
		query = query.Where("record_size > 0")
	}

	var total int64
	if err := query.Count(&total).Error; err != nil {
		c.JSON(http.StatusInternalServerError, model.Response{Code: 500, Message: "查询失败"})
		return
	}
	var sessions []model.RDPSession
	if err := query.Order("id DESC").Offset((req.Page - 1) * req.PageSize).Limit(req.PageSize).Find(&sessions).Error; err != nil {
		c.JSON(http.StatusInternalServerError, model.Response{Code: 500, Message: "查询失败"})
		return
	}
	c.JSON(http.StatusOK, model.Response{Code: 200, Data: model.ListResponse{List: sessions, Total: total}})
}

func (h *RDPSessionHandler) GetRDPSession(c *gin.Context) {
	id, err := strconv.ParseUint(c.Param("id"), 10, 64)
	if err != nil {
		c.JSON(http.StatusBadRequest, model.Response{Code: 400, Message: "无效的会话ID"})
		return
	}
	var session model.RDPSession
	if err := h.db.First(&session, id).Error; err != nil {
		c.JSON(http.StatusNotFound, model.Response{Code: 404, Message: "会话不存在"})
		return
	}
	c.JSON(http.StatusOK, model.Response{Code: 200, Data: session})
}

func (h *RDPSessionHandler) GetRDPRecord(c *gin.Context) {
	id, err := strconv.ParseUint(c.Param("id"), 10, 64)
	if err != nil {
		c.JSON(http.StatusBadRequest, model.Response{Code: 400, Message: "无效的会话ID"})
		return
	}
	var session model.RDPSession
	if err := h.db.First(&session, id).Error; err != nil {
		c.JSON(http.StatusNotFound, model.Response{Code: 404, Message: "会话不存在"})
		return
	}
	if session.RecordData == "" {
		c.JSON(http.StatusNotFound, model.Response{Code: 404, Message: "无录制数据"})
		return
	}
	c.JSON(http.StatusOK, model.Response{Code: 200, Data: session.RecordData})
}

func (h *RDPSessionHandler) KillRDPSession(c *gin.Context) {
	id, err := strconv.ParseUint(c.Param("id"), 10, 64)
	if err != nil {
		c.JSON(http.StatusBadRequest, model.Response{Code: 400, Message: "无效的会话ID"})
		return
	}
	if h.killFunc != nil && h.killFunc(uint(id)) {
		c.JSON(http.StatusOK, model.Response{Code: 200, Message: "已终止"})
		return
	}
	now := time.Now()
	var session model.RDPSession
	if err := h.db.First(&session, id).Error; err != nil {
		c.JSON(http.StatusNotFound, model.Response{Code: 404, Message: "会话不存在"})
		return
	}
	duration := int(now.Sub(session.StartTime).Seconds())
	_ = h.db.Model(&session).Updates(map[string]interface{}{"status": "closed", "end_time": now, "duration": duration}).Error
	c.JSON(http.StatusOK, model.Response{Code: 200, Message: "已终止"})
}

func (h *RDPSessionHandler) DownloadRDPRecord(c *gin.Context) {
	id, err := strconv.ParseUint(c.Param("id"), 10, 64)
	if err != nil {
		c.JSON(http.StatusBadRequest, model.Response{Code: 400, Message: "无效的会话ID"})
		return
	}
	var session model.RDPSession
	if err := h.db.First(&session, id).Error; err != nil || session.RecordData == "" {
		c.JSON(http.StatusNotFound, model.Response{Code: 404, Message: "无录制数据"})
		return
	}
	c.Header("Content-Type", "application/json; charset=utf-8")
	c.Header("Content-Disposition", fmt.Sprintf(`attachment; filename="rdp-session-%d.json"`, session.ID))
	c.Header("Content-Length", fmt.Sprintf("%d", len(session.RecordData)))
	c.String(http.StatusOK, session.RecordData)
}

func (h *RDPSessionHandler) LogRDPReplayView(c *gin.Context) {
	id, err := strconv.ParseUint(c.Param("id"), 10, 64)
	if err != nil {
		c.JSON(http.StatusBadRequest, model.Response{Code: 400, Message: "无效的会话ID"})
		return
	}
	var session model.RDPSession
	if err := h.db.Select("id", "username", "server_name", "server_ip").First(&session, id).Error; err != nil {
		c.JSON(http.StatusOK, model.Response{Code: 404, Message: "会话不存在"})
		return
	}
	uid := middleware.GetUserID(c)
	log := model.AuditLog{
		UserID:     &uid,
		Username:   middleware.GetUsername(c),
		IP:         c.ClientIP(),
		UserAgent:  c.GetHeader("User-Agent"),
		Module:     "rdp_audit",
		Action:     "replay_view",
		Resource:   "rdp_session",
		ResourceID: strconv.FormatUint(id, 10),
		Request:    fmt.Sprintf("回放 Windows RDP 会话 #%d（%s → %s/%s）", id, session.Username, session.ServerName, session.ServerIP),
		Status:     200,
	}
	_ = h.db.Create(&log).Error
	c.JSON(http.StatusOK, model.Response{Code: 200, Message: "ok"})
}