aiops/backend/controller/rdp_file.go

package controller

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

	"ezops/model"
	"ezops/pkg/crypto"
	"ezops/pkg/rdpfilestage"
	"ezops/pkg/winsmb"

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

// RDPFileHandler 通过 SMB 管理共享实现 RDP 虚拟 Download 盘文件传输 // Logan.Li
type RDPFileHandler struct {
	db    *gorm.DB
	crypt *crypto.SecretBox
}

func NewRDPFileHandler(db *gorm.DB, masterKey string) *RDPFileHandler {
	return &RDPFileHandler{db: db, crypt: crypto.NewSecretBox(masterKey)}
}

func (h *RDPFileHandler) loadAuth(serverID uint) (winsmb.Auth, error) {
	var server model.Server
	if err := h.db.First(&server, serverID).Error; err != nil {
		return winsmb.Auth{}, err
	}
	normalizeServerRemoteAccess(&server)
	if !isRDPServer(server) {
		return winsmb.Auth{}, errors.New("该资产未配置为 RDP/Windows")
	}
	cred, err := loadRDPCredential(h.crypt, server)
	if err != nil {
		return winsmb.Auth{}, err
	}
	return winsmb.Auth{
		Host:     server.IP,
		Domain:   cred.Domain,
		Username: cred.Username,
		Password: cred.Password,
	}, nil
}

func (h *RDPFileHandler) ListFiles(c *gin.Context) {
	serverID, err := strconv.ParseUint(c.Param("id"), 10, 32)
	if err != nil {
		c.JSON(http.StatusBadRequest, model.Response{Code: 400, Message: "无效的服务器ID"})
		return
	}
	auth, err := h.loadAuth(uint(serverID))
	if err != nil {
		c.JSON(http.StatusBadRequest, model.Response{Code: 400, Message: err.Error()})
		return
	}
	entries, err := winsmb.ListDownload(auth)
	if err != nil {
		c.JSON(http.StatusBadGateway, model.Response{Code: 502, Message: err.Error()})
		return
	}
	type item struct {
		Name    string `json:"name"`
		Size    int64  `json:"size"`
		IsDir   bool   `json:"is_dir"`
		ModTime string `json:"mod_time"`
	}
	list := make([]item, 0, len(entries))
	for _, ent := range entries {
		if ent.IsDir {
			continue
		}
		list = append(list, item{
			Name:    ent.Name,
			Size:    ent.Size,
			IsDir:   ent.IsDir,
			ModTime: ent.ModTime.Format(time.RFC3339),
		})
	}
	c.JSON(http.StatusOK, model.Response{
		Code: 200,
		Data: gin.H{
			"root":  winsmb.DownloadFullPath,
			"list":  list,
			"total": len(list),
		},
	})
}

func (h *RDPFileHandler) UploadFile(c *gin.Context) {
	serverID, err := strconv.ParseUint(c.Param("id"), 10, 32)
	if err != nil {
		c.JSON(http.StatusBadRequest, model.Response{Code: 400, Message: "无效的服务器ID"})
		return
	}
	fileHeader, err := c.FormFile("file")
	if err != nil {
		c.JSON(http.StatusBadRequest, model.Response{Code: 400, Message: "请选择文件"})
		return
	}
	if fileHeader.Size > 100<<20 {
		c.JSON(http.StatusBadRequest, model.Response{Code: 400, Message: "单文件不能超过 100MB"})
		return
	}
	f, err := fileHeader.Open()
	if err != nil {
		c.JSON(http.StatusBadRequest, model.Response{Code: 400, Message: "读取上传文件失败"})
		return
	}
	defer f.Close()
	data, err := io.ReadAll(f)
	if err != nil {
		c.JSON(http.StatusInternalServerError, model.Response{Code: 500, Message: "读取上传文件失败"})
		return
	}
	auth, err := h.loadAuth(uint(serverID))
	if err != nil {
		c.JSON(http.StatusBadRequest, model.Response{Code: 400, Message: err.Error()})
		return
	}
	if err := winsmb.Upload(auth, fileHeader.Filename, data); err != nil {
		c.JSON(http.StatusBadGateway, model.Response{Code: 502, Message: err.Error()})
		return
	}
	c.JSON(http.StatusOK, model.Response{
		Code:    200,
		Message: "上传成功",
		Data: gin.H{
			"name": fileHeader.Filename,
			"path": winsmb.DownloadFullPath + `\` + fileHeader.Filename,
			"size": fileHeader.Size,
		},
	})
}

func (h *RDPFileHandler) DownloadFile(c *gin.Context) {
	serverID, err := strconv.ParseUint(c.Param("id"), 10, 32)
	if err != nil {
		c.JSON(http.StatusBadRequest, model.Response{Code: 400, Message: "无效的服务器ID"})
		return
	}
	name := c.Query("name")
	if name == "" {
		c.JSON(http.StatusBadRequest, model.Response{Code: 400, Message: "缺少文件名"})
		return
	}
	auth, err := h.loadAuth(uint(serverID))
	if err != nil {
		c.JSON(http.StatusBadRequest, model.Response{Code: 400, Message: err.Error()})
		return
	}
	data, err := winsmb.Download(auth, name)
	if err != nil {
		c.JSON(http.StatusBadGateway, model.Response{Code: 502, Message: err.Error()})
		return
	}
	c.Header("Content-Disposition", "attachment; filename=\""+name+"\"")
	c.Data(http.StatusOK, "application/octet-stream", data)
}

// StageUpload SMB 不可用时：文件暂存 AIOps，返回可由 Windows 拉取的一次性 URL // Logan.Li
func (h *RDPFileHandler) StageUpload(c *gin.Context) {
	serverID, err := strconv.ParseUint(c.Param("id"), 10, 32)
	if err != nil {
		c.JSON(http.StatusBadRequest, model.Response{Code: 400, Message: "无效的服务器ID"})
		return
	}
	fileHeader, err := c.FormFile("file")
	if err != nil {
		c.JSON(http.StatusBadRequest, model.Response{Code: 400, Message: "请选择文件"})
		return
	}
	if fileHeader.Size > 50<<20 {
		c.JSON(http.StatusBadRequest, model.Response{Code: 400, Message: "回退模式单文件不能超过 50MB"})
		return
	}
	f, err := fileHeader.Open()
	if err != nil {
		c.JSON(http.StatusBadRequest, model.Response{Code: 400, Message: "读取上传文件失败"})
		return
	}
	defer f.Close()
	data, err := io.ReadAll(f)
	if err != nil {
		c.JSON(http.StatusInternalServerError, model.Response{Code: 500, Message: "读取上传文件失败"})
		return
	}
	if _, err := h.loadAuth(uint(serverID)); err != nil {
		c.JSON(http.StatusBadRequest, model.Response{Code: 400, Message: err.Error()})
		return
	}
	token := rdpfilestage.Put(fileHeader.Filename, data, 30*time.Minute)
	base := strings.TrimRight(publicHTTPBase(c), "/")
	downloadURL := fmt.Sprintf("%s/api/v1/rdp-staged/%s", base, token)
	remotePath := winsmb.DownloadFullPath + `\` + fileHeader.Filename
	deployCmd := fmt.Sprintf(
		`powershell -NoProfile -WindowStyle Hidden -Command "$d='%s';if(!(Test-Path $d)){ni -ItemType Directory -Path $d -Force|Out-Null};iwr -UseBasicParsing -Uri '%s' -OutFile (Join-Path $d '%s')"`,
		strings.ReplaceAll(winsmb.DownloadFullPath, `'`, `''`),
		downloadURL,
		strings.ReplaceAll(fileHeader.Filename, `'`, `''`),
	)
	c.JSON(http.StatusOK, model.Response{
		Code: 200,
		Data: gin.H{
			"token":          token,
			"download_url":   downloadURL,
			"remote_path":    remotePath,
			"deploy_command": deployCmd,
			"mode":           "stage",
		},
	})
}

// ServeStaged 供 Windows 主机拉取暂存文件（token 即密钥）// Logan.Li
func (h *RDPFileHandler) ServeStaged(c *gin.Context) {
	token := c.Param("token")
	item, ok := rdpfilestage.Get(token)
	if !ok {
		c.JSON(http.StatusNotFound, model.Response{Code: 404, Message: "链接已过期或不存在"})
		return
	}
	c.Header("Content-Disposition", "attachment; filename=\""+item.Name+"\"")
	c.Data(http.StatusOK, "application/octet-stream", item.Data)
}

func publicHTTPBase(c *gin.Context) string {
	scheme := "http"
	if c.Request.TLS != nil || strings.EqualFold(c.GetHeader("X-Forwarded-Proto"), "https") {
		scheme = "https"
	}
	host := c.Request.Host
	if fwd := strings.TrimSpace(c.GetHeader("X-Forwarded-Host")); fwd != "" {
		host = strings.Split(fwd, ",")[0]
	}
	return scheme + "://" + host
}