go_gameshield/internal/repository/api/waf/webforwarding.go

package waf

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	v1 "github.com/go-nunu/nunu-layout-advanced/api/v1"
	"github.com/go-nunu/nunu-layout-advanced/internal/model"
	"github.com/go-nunu/nunu-layout-advanced/internal/repository"
	"github.com/qiniu/qmgo"
	"go.mongodb.org/mongo-driver/bson"
	"go.mongodb.org/mongo-driver/bson/primitive"
	"go.mongodb.org/mongo-driver/mongo"
	"gorm.io/gorm"
	"time"
)

type WebForwardingRepository interface {
	GetWebForwarding(ctx context.Context, id int64) (*model.WebForwarding, error)
	AddWebForwarding(ctx context.Context, req *model.WebForwarding) (int, error)
	EditWebForwarding(ctx context.Context, req *model.WebForwarding) error
	DeleteWebForwarding(ctx context.Context, id int64) error
	GetWebForwardingPortCountByHostId(ctx context.Context, hostId int) (int64, error)
	GetWebForwardingDomainCountByHostId(ctx context.Context, hostId int) (int64, []string, error)
	GetWebForwardingWafWebAllIds(ctx context.Context, hostId int) ([]int, error)
	AddWebForwardingIps(ctx context.Context, req model.WebForwardingRule) (primitive.ObjectID, error)
	EditWebForwardingIps(ctx context.Context, req model.WebForwardingRule) error
	GetWebForwardingIpsByID(ctx context.Context, webId int) (*model.WebForwardingRule, error)
	DeleteWebForwardingIpsById(ctx context.Context, webId int) error
	// 获取域名数量
	GetDomainCount(ctx context.Context, hostId int, domain string) (int, error)
	// 获取IP数量等于1的IP
	GetIpCountByIp(ctx context.Context, ips []string) ([]v1.IpCountResult, error)
	GetSslCertId (ctx context.Context, sslPocyID int64) ([]v1.SslCertsJSON, error)
	// 获取CDN的web配置的id
	GetWebConfigId(ctx context.Context, id int64) (int64, error)
	// 获取域名
	GetDomainByHostIdPort(ctx context.Context, hostId int64, port string) ([]v1.Domain, error)
	// 获取CDN的web配置的id
	GetWebId(ctx context.Context, serverId int64) (int64, error)
	GetWebAll(ctx context.Context, hostIds []int) ([]int, error)
}

func NewWebForwardingRepository(
	repository *repository.Repository,
) WebForwardingRepository {
	return &webForwardingRepository{
		Repository: repository,
	}
}

type webForwardingRepository struct {
	*repository.Repository
}

func (r *webForwardingRepository) GetWebForwarding(ctx context.Context, id int64) (*model.WebForwarding, error) {
	var webForwarding model.WebForwarding
	if err := r.Db.WithContext(ctx).Where("id = ?", id).First(&webForwarding).Error; err != nil {
		return nil, err
	}
	return &webForwarding, nil
}

func (r *webForwardingRepository) AddWebForwarding(ctx context.Context, req *model.WebForwarding) (int, error) {
	if err := r.Db.WithContext(ctx).Create(req).Error; err != nil {
		return 0, err
	}
	return req.Id, nil
}

func (r *webForwardingRepository) EditWebForwarding(ctx context.Context, req *model.WebForwarding) error {
	forceUpdateFields := map[string]interface{}{
		"domain": req.Domain,
		"is_https": req.IsHttps,
		"proxy" : req.Proxy,
		"cc" : req.Cc,
	}
	// 核心逻辑：
	// 1. Model(req): 定位要更新的记录。
	// 2. Updates(req): 先用 struct 更新。GORM 会自动忽略 req 中的零值字段。
	//    - 如果 req.Domain 是 "abc"，它会被更新。
	//    - 如果 req.Domain 是 ""，它会被忽略。
	//    - 如果 req.TargetURL 是 "xyz"，它会被更新。
	//    - 如果 req.TargetURL 是 ""，它会被忽略。
	// 3. Updates(forceUpdateFields): 接着用 map 更新。这会无视零值，强制更新 map 中指定的字段。
	//    - 它会用 req.Domain 的值（无论是 "abc" 还是 ""）覆盖上一步的结果。
	//
	// 最终效果：
	// - Domain 字段总能被正确更新（无论新值是不是 ""）。
	// - 其他字段遵循 GORM 的默认行为（非零值才更新）。
	// - 这一切都在一个 UPDATE 语句中完成。
	db := r.Db.WithContext(ctx).Model(req).Updates(req).Updates(forceUpdateFields)

	if db.Error != nil {
		return db.Error
	}
	return nil
}

func (r *webForwardingRepository) DeleteWebForwarding(ctx context.Context, id int64) error {
	if err := r.Db.WithContext(ctx).Where("id = ?", id).Delete(&model.WebForwarding{}).Error; err != nil {
		return err
	}
	return nil
}

func (r *webForwardingRepository) GetWebForwardingPortCountByHostId(ctx context.Context, hostId int) (int64, error) {
	var count int64
	if err := r.Db.Model(&model.WebForwarding{}).WithContext(ctx).Where("host_id = ?", hostId).Count(&count).Error; err != nil {
		return 0, err
	}
	return count, nil
}

func (r *webForwardingRepository) GetWebForwardingDomainCountByHostId(ctx context.Context, hostId int) (int64, []string, error) {
	var distinctDomains []string
	err := r.Db.Model(&model.WebForwarding{}).WithContext(ctx).
		Distinct().                                                           // 确保我们只获取唯一的 domain 值
		Where("host_id = ? AND domain IS NOT NULL AND domain != ''", hostId). // 额外添加 domain != '' 以排除空字符串
		Pluck("domain", &distinctDomains).Error

	if err != nil {
		return 0, nil, err
	}
	count := int64(len(distinctDomains))

	return count, distinctDomains, nil
}

func (r *webForwardingRepository) GetWebForwardingWafWebAllIds(ctx context.Context, hostId int) ([]int, error) {
	var ids []int
	if err := r.Db.Model(&model.WebForwarding{}).WithContext(ctx).Where("host_id = ?", hostId).Select("id").Find(&ids).Error; err != nil {
		return nil, err
	}

	return ids, nil
}

// mongodb 插入
func (r *webForwardingRepository) AddWebForwardingIps(ctx context.Context, req model.WebForwardingRule) (primitive.ObjectID, error) {
	collection := r.MongoDB.Collection("web_forwarding_rules")
	req.CreatedAt = time.Now()
	result, err := collection.InsertOne(ctx, req)
	if err != nil {
		return primitive.NilObjectID, fmt.Errorf("插入MongoDB失败: %w", err)
	}

	// 返回插入文档的ID
	return result.InsertedID.(primitive.ObjectID), nil
}

func (r *webForwardingRepository) EditWebForwardingIps(ctx context.Context, req model.WebForwardingRule) error {
	collection := r.MongoDB.Collection("web_forwarding_rules")
	updateData := bson.M{}

	if req.Uid != 0 {
		updateData["uid"] = req.Uid
	}

	if req.HostId != 0 {
		updateData["host_id"] = req.HostId
	}

	if req.WebId != 0 {
		updateData["web_id"] = req.WebId
	}

	if len(req.BackendList) > 0 {
		updateData["backend_list"] = req.BackendList
	}

	updateData["cdn_origin_ids"] = req.CdnOriginIds

	// 始终更新更新时间
	updateData["updated_at"] = time.Now()

	// 如果没有任何字段需要更新，则直接返回
	if len(updateData) == 0 {
		return nil
	}

	// 执行更新
	update := bson.M{"$set": updateData}
	err := collection.UpdateOne(ctx, bson.M{"web_id": req.WebId}, update)
	if err != nil {
		return fmt.Errorf("更新MongoDB文档失败: %w", err)
	}

	return nil
}

func (r *webForwardingRepository) GetWebForwardingIpsByID(ctx context.Context, webId int) (*model.WebForwardingRule, error) {

	// 获取集合
	collection := r.MongoDB.Collection("web_forwarding_rules")

	// 创建一个结构体来存储查询到的文档
	var rule model.WebForwardingRule

	// 使用 FindByID 方法来查找文档
	// FindByID 是 QMgo 封装的一个方便的方法，它内部会构建查询 _id = id
	err := collection.Find(ctx, qmgo.M{"web_id": webId}).One(&rule) // QMgo 的 FindOne 返回一个 QueryBuilder，接着调用 .One() 来执行查询并解码到 rule

	if err != nil {
		if errors.Is(err, mongo.ErrNoDocuments) {
			return nil, nil
		}
		// 其他错误
		return nil, fmt.Errorf("查询MongoDB失败: %w", err)
	}

	// 返回找到的文档
	return &rule, nil
}

func (r *webForwardingRepository) DeleteWebForwardingIpsById(ctx context.Context, webId int) error {

	collection := r.MongoDB.Collection("web_forwarding_rules")

	err := collection.Remove(ctx, bson.M{"web_id": webId})

	if err != nil {
		if errors.Is(err, mongo.ErrNoDocuments) {
			return fmt.Errorf("记录不存在")
		}
		return fmt.Errorf("删除MongoDB文档失败: %w", err)
	}
	return nil
}

// 获取域名数量
func (r *webForwardingRepository) GetDomainCount(ctx context.Context, hostId int, domain string) (int, error) {
	var count int64
	if err := r.Db.Model(&model.WebForwarding{}).WithContext(ctx).Where("host_id = ? AND domain = ?", hostId, domain).Count(&count).Error; err != nil {
		return 0, err
	}
	return int(count), nil
}

// 获取IP数量等于1的IP

func (r *webForwardingRepository) GetIpCountByIp(ctx context.Context, ips []string) ([]v1.IpCountResult, error) {
	if len(ips) == 0 {
		return []v1.IpCountResult{}, nil
	}

	pipeline := []bson.M{
		// 第 1 步: $unwind - 展开 backend_list 数组
		// 将包含多个 backend 对象的文档拆分成多条，每条只包含一个 backend 对象。
		{
			"$unwind": "$backend_list",
		},

		// 第 2 步: $addFields - 添加一个新字段用于存放解析出的 IP
		// 我们需要从 "ip:port" 格式的 addr 字段中把 ip 提取出来。
		// 使用 $split 操作符按 ":" 分割字符串，然后用 $arrayElemAt 取第一个元素。
		{
			"$addFields": bson.M{
				"extracted_ip": bson.M{
					"$arrayElemAt": []interface{}{
						bson.M{"$split": []string{"$backend_list.addr", ":"}},
						0,
					},
				},
			},
		},

		// 第 3 步: $match - 匹配我们关心的 IP
		// 在上一步创建的 extracted_ip 字段上进行匹配。
		{
			"$match": bson.M{
				"extracted_ip": bson.M{"$in": ips},
			},
		},

		// 第 4 步: $group - 按解析出的 IP 地址进行分组和计数
		{
			"$group": bson.M{
				"_id":   "$extracted_ip", // 使用我们新创建的 extracted_ip 字段作为分组依据
				"count": bson.M{"$sum": 1},
			},
		},

		// 第 5 步: $project - 格式化最终输出
		// 这个阶段和之前一样，只是为了让输出结果更清晰，并匹配 Go 结构体。
		{
			"$project": bson.M{
				"_id":   0,
				"ip":    "$_id",
				"count": 1,
			},
		},
	}

	var results []v1.IpCountResult
	// 使用 qmgo 执行聚合查询
	err := r.MongoDB.Collection("web_forwarding_rules").Aggregate(ctx, pipeline).All(&results)
	if err != nil {
		// 加上错误包装，方便调试
		return nil, fmt.Errorf("聚合查询 web_forwarding_rules 失败: %w", err)
	}

	return results, nil
}

func (r *webForwardingRepository) GetSslCertId (ctx context.Context, sslPolicyID int64) ([]v1.SslCertsJSON, error) {
	var certsJSON string

	// 2. 查询数据库，只获取 `certs` 字段的字符串内容
	// 使用 Scopes 来确保没有不必要的 ORDER BY 子句，或者直接用 Raw/Scan
	// 但在这里，用 .First(&certsJSON) 通常是安全的，因为目标是简单类型 string
	err := r.DBWithName(ctx, "cdn").WithContext(ctx).
		Table("cloud_ssl_policies").
		Select("certs").
		Where("id = ?", sslPolicyID).
		Row(). // 获取 sql.Row
		Scan(&certsJSON) // 将结果扫描到字符串变量中

	// 如果查询出错，或者没有找到记录 (sql.ErrNoRows)
	if err != nil {
		if err == gorm.ErrRecordNotFound {
			// 如果记录不存在是正常情况，可以返回一个空切片和nil错误
			return []v1.SslCertsJSON{}, nil
		}
		return nil, err
	}

	// 如果certs字段在数据库中可能是NULL或者空字符串，需要处理
	if certsJSON == "" {
		return []v1.SslCertsJSON{}, nil
	}

	// 3. 将JSON字符串反序列化到Go结构体切片中
	var res []v1.SslCertsJSON
	err = json.Unmarshal([]byte(certsJSON), &res)
	if err != nil {
		// 这里是真正的JSON格式转换错误
		return nil, err
	}

	return res, nil
}

// 获取CDN的web配置的id
func (r *webForwardingRepository) GetWebConfigId(ctx context.Context, id int64) (int64, error) {
	var webConfigId int64
	if err := r.DBWithName(ctx,"cdn").Table("cloud_servers").Where("id = ?", id).Select("webId").Scan(&webConfigId).Error; err != nil {
		return 0, err
	}
	return webConfigId, nil
}


func (r *webForwardingRepository) GetDomainByHostIdPort(ctx context.Context, hostId int64, port string) ([]v1.Domain, error) {
	var domains []v1.Domain
	if err := r.Db.WithContext(ctx).Model(&model.WebForwarding{}).Where("host_id = ? AND port = ?", hostId, port).Select("domain,id,is_https").Scan(&domains).Error; err != nil {
		return nil, err
	}
	return domains, nil

}

// 获取CDN的web配置的id
func (r *webForwardingRepository) GetWebId(ctx context.Context, serverId int64) (int64, error) {
	var webId int64
	if err := r.DBWithName(ctx,"cdn").Table("cloud_servers").Where("id = ?", serverId).Select("webId").Scan(&webId).Error; err != nil {
		return 0, err
	}
	return webId, nil

}

func (r *webForwardingRepository) GetWebAll(ctx context.Context, hostIds []int) ([]int, error) {
	var res []int
	if err := r.Db.Model(&model.WebForwarding{}).WithContext(ctx).Where("host_id IN ?", hostIds).Select("cdn_web_id").Scan(&res).Error; err != nil {
		return nil, err
	}

	return res, nil
}