go_gameshield/internal/repository/webforwarding.go

package repository

import (
	"context"
	"errors"
	"fmt"
	"github.com/go-nunu/nunu-layout-advanced/internal/model"
	"github.com/qiniu/qmgo"
	"go.mongodb.org/mongo-driver/bson"
	"go.mongodb.org/mongo-driver/bson/primitive"
	"go.mongodb.org/mongo-driver/mongo"
	"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
	GetWebForwardingWafWebIdById(ctx context.Context, id int) (int, error)
	GetWebForwardingPortCountByHostId(ctx context.Context, hostId int) (int64, error)
	GetWebForwardingDomainCountByHostId(ctx context.Context, hostId int) (int64, []string, error)
	AddWebForwardingIps(ctx context.Context, req model.WebForwardingRule) (primitive.ObjectID, error)
	EditWebForwardingIps(ctx context.Context, req model.WebForwardingRule) error
	GetWebForwardingByID(ctx context.Context, webId int) (*model.WebForwardingRule, error)
}

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

type webForwardingRepository struct {
	*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 {
	if err := r.db.WithContext(ctx).Updates(req).Error; err != nil {
		return err
	}
	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) GetWebForwardingWafWebIdById(ctx context.Context, id int) (int, error) {
	var WafWebId int

	if err := r.db.Model(&model.WebForwarding{}).WithContext(ctx).Where("id = ?", id).Select("waf_web_id").Find(&WafWebId).Error; err != nil {
		return 0, err
	}
	return WafWebId, 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
}


// 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 req.AccessRule != "" {
		updateData["access_rule"] = req.AccessRule
	}

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

	if len(req.AllowIpList) > 0 {
		updateData["allow_ip_list"] = req.AllowIpList
	}

	if len(req.DenyIpList) > 0 {
		updateData["deny_ip_list"] = req.DenyIpList
	}

	// 始终更新更新时间
	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) GetWebForwardingByID(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, fmt.Errorf("记录不存在")
		}
		// 其他错误
		return nil, fmt.Errorf("查询MongoDB失败: %w", err)
	}

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