package waf import ( "context" "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" "go.mongodb.org/mongo-driver/bson" "go.mongodb.org/mongo-driver/bson/primitive" "go.mongodb.org/mongo-driver/mongo" "time" ) type TcpforwardingRepository interface { GetTcpforwarding(ctx context.Context, id int64) (*model.Tcpforwarding, error) AddTcpforwarding(ctx context.Context, req *model.Tcpforwarding) (int, error) EditTcpforwarding(ctx context.Context, req *model.Tcpforwarding) error DeleteTcpforwarding(ctx context.Context, id int64) error GetTcpforwardingWafTcpIdById(ctx context.Context, id int) (int, error) GetTcpForwardingPortCountByHostId(ctx context.Context, hostId int) (int64, error) GetTcpForwardingAllIdsByID(ctx context.Context, hostId int) ([]int, error) AddTcpforwardingIps(ctx context.Context,req model.TcpForwardingRule) (primitive.ObjectID, error) EditTcpforwardingIps(ctx context.Context, req model.TcpForwardingRule) error GetTcpForwardingIpsByID(ctx context.Context, tcpId int) (*model.TcpForwardingRule, error) DeleteTcpForwardingIpsById(ctx context.Context, tcpId int) error // 获取IP数量等于1的IP GetIpCountByIp(ctx context.Context,ips []string) ([]v1.IpCountResult, error) // 获取端口数量 GetPortCount(ctx context.Context,hostId int64, port string) (int64, error) // 获取所有数据 GetTcpAll(ctx context.Context, hostIds []int) ([]int, error) } func NewTcpforwardingRepository( repository *repository.Repository, ) TcpforwardingRepository { return &tcpforwardingRepository{ Repository: repository, } } type tcpforwardingRepository struct { *repository.Repository } func (r *tcpforwardingRepository) GetTcpforwarding(ctx context.Context, id int64) (*model.Tcpforwarding, error) { var tcpforwarding model.Tcpforwarding if err := r.Db.Where("id = ?", id).First(&tcpforwarding).Error; err != nil { return nil, err } return &tcpforwarding, nil } func (r *tcpforwardingRepository) AddTcpforwarding(ctx context.Context, req *model.Tcpforwarding) (int, error) { if err := r.Db.Create(req).Error; err != nil { return 0, err } return req.Id, nil } func (r *tcpforwardingRepository) EditTcpforwarding(ctx context.Context, req *model.Tcpforwarding) error { data := map[string]interface{}{ "proxy" : req.Proxy, } if err := r.Db.Updates(req).Updates(data).Error; err != nil { return err } return nil } func (r *tcpforwardingRepository) DeleteTcpforwarding(ctx context.Context, id int64) error { if err := r.Db.Where("id = ?", id).Delete(&model.Tcpforwarding{}).Error; err != nil { return err } return nil } func (r *tcpforwardingRepository) GetTcpforwardingWafTcpIdById(ctx context.Context, id int) (int, error) { var WafTcpId int if err := r.Db.Model(&model.Tcpforwarding{}).Where("id = ?", id).Select("waf_tcp_id").Find(&WafTcpId).Error; err != nil { return 0, err } return WafTcpId, nil } func (r *tcpforwardingRepository) GetTcpForwardingPortCountByHostId(ctx context.Context, hostId int) (int64, error) { var count int64 if err := r.Db.Model(&model.Tcpforwarding{}).Where("host_id = ?", hostId).Count(&count).Error; err != nil { return 0, err } return count, nil } func (r *tcpforwardingRepository) GetTcpForwardingAllIdsByID(ctx context.Context, hostId int) ([]int, error) { var res []int if err := r.Db.WithContext(ctx).Model(&model.Tcpforwarding{}).Where("host_id = ?", hostId).Select("id").Find(&res).Error; err != nil { return nil, err } return res, nil } //mongodb 插入 func (r *tcpforwardingRepository) AddTcpforwardingIps(ctx context.Context,req model.TcpForwardingRule) (primitive.ObjectID, error) { collection := r.MongoDB.Collection("tcp_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 *tcpforwardingRepository) EditTcpforwardingIps(ctx context.Context, req model.TcpForwardingRule) error { collection := r.MongoDB.Collection("tcp_forwarding_rules") updateData := bson.M{} if req.Uid != 0 { updateData["uid"] = req.Uid } if req.HostId != 0 { updateData["host_id"] = req.HostId } if req.TcpId != 0 { updateData["tcp_id"] = req.TcpId } 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{"tcp_id": req.TcpId}, update) if err != nil { if errors.Is(err, mongo.ErrNoDocuments) { return fmt.Errorf("记录不存在") } return fmt.Errorf("更新MongoDB文档失败: %w", err) } return nil } func (r *tcpforwardingRepository) GetTcpForwardingIpsByID(ctx context.Context, tcpId int) (*model.TcpForwardingRule, error) { collection := r.MongoDB.Collection("tcp_forwarding_rules") var res model.TcpForwardingRule err := collection.Find(ctx, bson.M{"tcp_id": tcpId}).One(&res) if err != nil { if errors.Is(err, mongo.ErrNoDocuments) { return nil, nil } return nil, fmt.Errorf("查询MongoDB失败: %w", err) } return &res, nil } func (r *tcpforwardingRepository) DeleteTcpForwardingIpsById(ctx context.Context, tcpId int) error { collection := r.MongoDB.Collection("tcp_forwarding_rules") err := collection.Remove(ctx, bson.M{"tcp_id": tcpId}) if err != nil { if errors.Is(err, mongo.ErrNoDocuments) { return fmt.Errorf("记录不存在") } return fmt.Errorf("删除MongoDB文档失败: %w", err) } return nil } // 获取IP数量等于1的IP func (r *tcpforwardingRepository) GetIpCountByIp(ctx context.Context, ips []string) ([]v1.IpCountResult, error) { if len(ips) == 0 { return []v1.IpCountResult{}, nil } pipeline := []bson.M{ // 1. 展开 backend_list 数组。此时 backend_list 字段会变成 "ip:port" 字符串。 { "$unwind": "$backend_list", }, // 2. 添加新字段 extracted_ip,存放从 "ip:port" 中解析出的 IP。 { "$addFields": bson.M{ "extracted_ip": bson.M{ "$arrayElemAt": []interface{}{ // 直接在 backend_list 字符串上分割 bson.M{"$split": []string{"$backend_list", ":"}}, 0, }, }, }, }, // 3. 匹配我们关心的 IP { "$match": bson.M{ "extracted_ip": bson.M{"$in": ips}, }, }, // 4. 按解析出的 IP 地址进行分组和计数 { "$group": bson.M{ "_id": "$extracted_ip", "count": bson.M{"$sum": 1}, }, }, // 5. 格式化最终输出 { "$project": bson.M{ "_id": 0, "ip": "$_id", "count": 1, }, }, } var results []v1.IpCountResult err := r.MongoDB.Collection("tcp_forwarding_rules").Aggregate(ctx, pipeline).All(&results) if err != nil { return nil, fmt.Errorf("聚合查询 tcp_forwarding_rules 失败: %w", err) } return results, nil } func (r *tcpforwardingRepository) GetPortCount(ctx context.Context,hostId int64, port string) (int64, error) { var count int64 if err := r.Db.WithContext(ctx).Model(&model.Tcpforwarding{}).Where("host_id = ? AND port = ?", hostId, port).Count(&count).Error; err != nil { return 0, err } return count, nil } func (r *tcpforwardingRepository) GetTcpAll(ctx context.Context, hostIds []int) ([]int, error) { var res []int if err := r.Db.WithContext(ctx).Model(&model.Tcpforwarding{}).Where("host_id IN ?", hostIds).Select("cdn_web_id").Scan(&res).Error; err != nil { return nil, err } return res, nil }