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

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
}