aiops/backend/controller/log_datasource.go

package controller

import (
	"crypto/tls"
	"encoding/base64"
	"io"
	"net/http"
	"strings"
	"time"

	"ezops/model"

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

// LogDSHandler 日志数据源管理 // Logan.Li
type LogDSHandler struct {
	db *gorm.DB
}

func NewLogDSHandler(db *gorm.DB) *LogDSHandler {
	return &LogDSHandler{db: db}
}

// List GET /api/v1/log-datasources?type=loki
func (h *LogDSHandler) List(c *gin.Context) {
	q := h.db.Order("is_default DESC, id ASC")
	if t := c.Query("type"); t != "" {
		q = q.Where("type = ?", t)
	}
	var sources []model.LogDataSource
	if err := q.Find(&sources).Error; err != nil {
		c.JSON(500, model.Response{Code: 500, Message: err.Error()})
		return
	}
	if sources == nil {
		sources = []model.LogDataSource{}
	}
	c.JSON(200, model.Response{Code: 200, Data: sources})
}

// Get GET /api/v1/log-datasources/:id
func (h *LogDSHandler) Get(c *gin.Context) {
	var ds model.LogDataSource
	if err := h.db.First(&ds, c.Param("id")).Error; err != nil {
		c.JSON(404, model.Response{Code: 404, Message: "数据源不存在"})
		return
	}
	c.JSON(200, model.Response{Code: 200, Data: ds})
}

// Create POST /api/v1/log-datasources
func (h *LogDSHandler) Create(c *gin.Context) {
	var req model.CreateLogDSReq
	if err := c.ShouldBindJSON(&req); err != nil {
		c.JSON(400, model.Response{Code: 400, Message: err.Error()})
		return
	}
	req.Type = strings.ToLower(req.Type)
	if req.Type != "loki" && req.Type != "elasticsearch" {
		c.JSON(400, model.Response{Code: 400, Message: "type 必须为 loki 或 elasticsearch"})
		return
	}
	if req.IsDefault {
		h.db.Model(&model.LogDataSource{}).Where("type = ? AND is_default = ?", req.Type, true).Update("is_default", false)
	}
	ds := model.LogDataSource{
		Name: req.Name, Type: req.Type, URL: strings.TrimRight(req.URL, "/"),
		BasicAuth: req.BasicAuth, BasicAuthUser: req.BasicAuthUser, BasicAuthPass: req.BasicAuthPass,
		SkipTLSVerify: req.SkipTLSVerify, DefaultIndex: req.DefaultIndex,
		IsDefault: req.IsDefault, Status: 1,
	}
	if err := h.db.Create(&ds).Error; err != nil {
		c.JSON(500, model.Response{Code: 500, Message: err.Error()})
		return
	}
	c.JSON(200, model.Response{Code: 200, Message: "创建成功", Data: ds})
}

// Update PUT /api/v1/log-datasources/:id
func (h *LogDSHandler) Update(c *gin.Context) {
	var ds model.LogDataSource
	if err := h.db.First(&ds, c.Param("id")).Error; err != nil {
		c.JSON(404, model.Response{Code: 404, Message: "数据源不存在"})
		return
	}
	var req model.UpdateLogDSReq
	if err := c.ShouldBindJSON(&req); err != nil {
		c.JSON(400, model.Response{Code: 400, Message: err.Error()})
		return
	}
	updates := map[string]interface{}{}
	if req.Name != nil {
		updates["name"] = *req.Name
	}
	if req.Type != nil {
		updates["type"] = strings.ToLower(*req.Type)
	}
	if req.URL != nil {
		updates["url"] = strings.TrimRight(*req.URL, "/")
	}
	if req.BasicAuth != nil {
		updates["basic_auth"] = *req.BasicAuth
	}
	if req.BasicAuthUser != nil {
		updates["basic_auth_user"] = *req.BasicAuthUser
	}
	if req.BasicAuthPass != nil && *req.BasicAuthPass != "" {
		updates["basic_auth_pass"] = *req.BasicAuthPass
	}
	if req.SkipTLSVerify != nil {
		updates["skip_tls_verify"] = *req.SkipTLSVerify
	}
	if req.DefaultIndex != nil {
		updates["default_index"] = *req.DefaultIndex
	}
	if req.Status != nil {
		updates["status"] = *req.Status
	}
	if req.IsDefault != nil && *req.IsDefault {
		h.db.Model(&model.LogDataSource{}).Where("type = ? AND is_default = ?", ds.Type, true).Update("is_default", false)
		updates["is_default"] = true
	} else if req.IsDefault != nil {
		updates["is_default"] = false
	}
	if err := h.db.Model(&ds).Updates(updates).Error; err != nil {
		c.JSON(500, model.Response{Code: 500, Message: err.Error()})
		return
	}
	h.db.First(&ds, ds.ID)
	c.JSON(200, model.Response{Code: 200, Message: "更新成功", Data: ds})
}

// Delete DELETE /api/v1/log-datasources/:id
func (h *LogDSHandler) Delete(c *gin.Context) {
	if err := h.db.Delete(&model.LogDataSource{}, c.Param("id")).Error; err != nil {
		c.JSON(500, model.Response{Code: 500, Message: "删除失败"})
		return
	}
	c.JSON(200, model.Response{Code: 200, Message: "已删除"})
}

// SetDefault POST /api/v1/log-datasources/:id/default
func (h *LogDSHandler) SetDefault(c *gin.Context) {
	var ds model.LogDataSource
	if err := h.db.First(&ds, c.Param("id")).Error; err != nil {
		c.JSON(404, model.Response{Code: 404, Message: "数据源不存在"})
		return
	}
	h.db.Model(&model.LogDataSource{}).Where("type = ? AND is_default = ?", ds.Type, true).Update("is_default", false)
	h.db.Model(&ds).Update("is_default", true)
	c.JSON(200, model.Response{Code: 200, Message: "已设为默认"})
}

// TestConnection POST /api/v1/log-datasources/test
func (h *LogDSHandler) TestConnection(c *gin.Context) {
	var req model.TestLogDSReq
	if err := c.ShouldBindJSON(&req); err != nil {
		c.JSON(400, model.Response{Code: 400, Message: err.Error()})
		return
	}
	ok, msg := testLogDSConnection(req)
	if !ok {
		c.JSON(502, model.Response{Code: 502, Message: msg})
		return
	}
	c.JSON(200, model.Response{Code: 200, Message: "连接成功", Data: map[string]string{"status": "ok"}})
}

func testLogDSConnection(req model.TestLogDSReq) (bool, string) {
	client := newLogHTTPClient(req.SkipTLSVerify, 10*time.Second)
	url := strings.TrimRight(req.URL, "/")
	var testURL string
	switch strings.ToLower(req.Type) {
	case "loki":
		testURL = url + "/loki/api/v1/labels"
	case "elasticsearch":
		testURL = url + "/_cluster/health"
	default:
		return false, "未知数据源类型"
	}
	httpReq, err := http.NewRequest(http.MethodGet, testURL, nil)
	if err != nil {
		return false, err.Error()
	}
	applyLogDSAuth(httpReq, req.BasicAuth, req.BasicAuthUser, req.BasicAuthPass)
	resp, err := client.Do(httpReq)
	if err != nil {
		return false, "连接失败: " + err.Error()
	}
	defer resp.Body.Close()
	if resp.StatusCode >= 400 {
		body, _ := io.ReadAll(io.LimitReader(resp.Body, 512))
		return false, "HTTP " + resp.Status + ": " + string(body)
	}
	return true, ""
}

func loadLogDS(h *LogDSHandler, id uint) (*model.LogDataSource, error) {
	var ds model.LogDataSource
	if err := h.db.First(&ds, id).Error; err != nil {
		return nil, err
	}
	return &ds, nil
}

func newLogHTTPClient(skipTLS bool, timeout time.Duration) *http.Client {
	tr := &http.Transport{}
	if skipTLS {
		tr.TLSClientConfig = &tls.Config{InsecureSkipVerify: true} //nolint:gosec
	}
	return &http.Client{Timeout: timeout, Transport: tr}
}

func applyLogDSAuth(req *http.Request, basic bool, user, pass string) {
	if basic && user != "" {
		req.Header.Set("Authorization", "Basic "+base64.StdEncoding.EncodeToString([]byte(user+":"+pass)))
	}
}