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) 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 } 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 } 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 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) 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, fmt.Errorf("记录不存在") } // 其他错误 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 }