实现WAF统计数据接口

internal/rpc/services/service_server.go

@@ -14,6 +14,7 @@ import (
 	"github.com/iwind/TeaGo/logs"
 	"github.com/iwind/TeaGo/maps"
 	"github.com/iwind/TeaGo/types"
+	timeutil "github.com/iwind/TeaGo/utils/time"
 )
 
 type ServerService struct {
@@ -1247,10 +1248,14 @@ func (this *ServerService) UploadServerHTTPRequestStat(ctx context.Context, req
 
 	var tx = this.NullTx()
 
-	// 全局
 	month := req.Month
-	if len(month) != 6 {
-		return nil, errors.New("invalid month '" + month + "'")
+	if len(month) == 0 {
+		month = timeutil.Format("Ym")
+	}
+
+	day := req.Day
+	if len(day) == 0 {
+		day = timeutil.Format("Ymd")
 	}
 
 	// 区域
@@ -1389,5 +1394,22 @@ func (this *ServerService) UploadServerHTTPRequestStat(ctx context.Context, req
 		}
 	}
 
+	// 防火墙
+	for _, result := range req.HttpFirewallRuleGroups {
+		err := func() error {
+			if result.HttpFirewallRuleGroupId <= 0 {
+				return nil
+			}
+			key := fmt.Sprintf("%d@%d@%s@%s", result.ServerId, result.HttpFirewallRuleGroupId, result.Action, day)
+			serverStatLocker.Lock()
+			serverHTTPFirewallRuleGroupStatMap[key] += result.Count
+			serverStatLocker.Unlock()
+			return nil
+		}()
+		if err != nil {
+			return nil, err
+		}
+	}
+
 	return this.Success()
 }

internal/rpc/services/service_server_.go

@@ -12,12 +12,13 @@ import (
 )
 
 // HTTP请求统计缓存队列
-var serverHTTPCountryStatMap = map[string]int64{}  // serverId@countryId@month => count
-var serverHTTPProvinceStatMap = map[string]int64{} // serverId@provinceId@month => count
-var serverHTTPCityStatMap = map[string]int64{}     // serverId@cityId@month => count
-var serverHTTPProviderStatMap = map[string]int64{} // serverId@providerId@month => count
-var serverHTTPSystemStatMap = map[string]int64{}   // serverId@systemId@version@month => count
-var serverHTTPBrowserStatMap = map[string]int64{}  // serverId@browserId@version@month => count
+var serverHTTPCountryStatMap = map[string]int64{}           // serverId@countryId@month => count
+var serverHTTPProvinceStatMap = map[string]int64{}          // serverId@provinceId@month => count
+var serverHTTPCityStatMap = map[string]int64{}              // serverId@cityId@month => count
+var serverHTTPProviderStatMap = map[string]int64{}          // serverId@providerId@month => count
+var serverHTTPSystemStatMap = map[string]int64{}            // serverId@systemId@version@month => count
+var serverHTTPBrowserStatMap = map[string]int64{}           // serverId@browserId@version@month => count
+var serverHTTPFirewallRuleGroupStatMap = map[string]int64{} // serverId@firewallRuleGroupId@action@day => count
 var serverStatLocker = sync.Mutex{}
 
 func init() {
@@ -29,7 +30,7 @@ func init() {
 			var duration = 30 * time.Minute
 			if Tea.IsTesting() {
 				// 测试条件下缩短时间，以便进行观察
-				duration = 1 * time.Minute
+				duration = 10 * time.Second
 			}
 			ticker := time.NewTicker(duration)
 			for range ticker.C {
@@ -151,5 +152,24 @@ func (this *ServerService) dumpServerHTTPStats() error {
 		}
 	}
 
+	// 防火墙
+	{
+		serverStatLocker.Lock()
+		m := serverHTTPFirewallRuleGroupStatMap
+		serverHTTPFirewallRuleGroupStatMap = map[string]int64{}
+		serverStatLocker.Unlock()
+		for k, count := range m {
+			pieces := strings.Split(k, "@")
+			if len(pieces) != 4 {
+				continue
+			}
+			err := stats.SharedServerHTTPFirewallDailyStatDAO.IncreaseDailyCount(nil, types.Int64(pieces[0]), types.Int64(pieces[1]), pieces[2], pieces[3], count)
+			if err != nil {
+				return err
+			}
+		}
+	}
+
+
 	return nil
 }

internal/rpc/services/service_server_http_firewall_daily_stat.go

@@ -0,0 +1,159 @@
+package services
+
+import (
+	"context"
+	"github.com/TeaOSLab/EdgeAPI/internal/db/models"
+	"github.com/TeaOSLab/EdgeAPI/internal/db/models/stats"
+	"github.com/TeaOSLab/EdgeAPI/internal/utils"
+	"github.com/TeaOSLab/EdgeCommon/pkg/rpc/pb"
+	"github.com/iwind/TeaGo/types"
+	timeutil "github.com/iwind/TeaGo/utils/time"
+	"time"
+)
+
+// WAF统计
+type ServerHTTPFirewallDailyStatService struct {
+	BaseService
+}
+
+// 组合Dashboard
+func (this *ServerHTTPFirewallDailyStatService) ComposeServerHTTPFirewallDashboard(ctx context.Context, req *pb.ComposeServerHTTPFirewallDashboardRequest) (*pb.ComposeServerHTTPFirewallDashboardResponse, error) {
+	_, userId, err := this.ValidateAdminAndUser(ctx, 0, 0)
+	if err != nil {
+		return nil, err
+	}
+
+	if userId > 0 {
+		if req.UserId > 0 && req.UserId != userId {
+			return nil, this.PermissionError()
+		}
+		if req.ServerId > 0 {
+			err = models.SharedServerDAO.CheckUserServer(nil, userId, req.ServerId)
+			if err != nil {
+				return nil, err
+			}
+		}
+	} else {
+		userId = req.UserId
+	}
+
+	day := req.Day
+	if len(day) != 8 {
+		day = timeutil.Format("Ymd")
+	}
+
+	date := time.Date(types.Int(day[:4]), time.Month(types.Int(day[4:6])), types.Int(day[6:]), 0, 0, 0, 0, time.Local)
+	var w = types.Int(timeutil.Format("w", date))
+	if w == 0 {
+		w = 7
+	}
+	weekFrom := timeutil.Format("Ymd", date.AddDate(0, 0, -w+1))
+	weekTo := timeutil.Format("Ymd", date.AddDate(0, 0, -w+7))
+
+	var d = types.Int(timeutil.Format("d"))
+	monthFrom := timeutil.Format("Ymd", date.AddDate(0, 0, -d+1))
+	monthTo := timeutil.Format("Ymd", date.AddDate(0, 1, -d))
+
+	var tx = this.NullTx()
+
+	countDailyLog, err := stats.SharedServerHTTPFirewallDailyStatDAO.SumDailyCount(tx, userId, req.ServerId, "log", day, day)
+	if err != nil {
+		return nil, err
+	}
+
+	countDailyBlock, err := stats.SharedServerHTTPFirewallDailyStatDAO.SumDailyCount(tx, userId, req.ServerId, "block", day, day)
+	if err != nil {
+		return nil, err
+	}
+
+	countDailyCaptcha, err := stats.SharedServerHTTPFirewallDailyStatDAO.SumDailyCount(tx, userId, req.ServerId, "captcha", day, day)
+	if err != nil {
+		return nil, err
+	}
+
+	countWeeklyBlock, err := stats.SharedServerHTTPFirewallDailyStatDAO.SumDailyCount(tx, userId, req.ServerId, "block", weekFrom, weekTo)
+	if err != nil {
+		return nil, err
+	}
+
+	countMonthlyBlock, err := stats.SharedServerHTTPFirewallDailyStatDAO.SumDailyCount(tx, userId, req.ServerId, "block", monthFrom, monthTo)
+	if err != nil {
+		return nil, err
+	}
+
+	resp := &pb.ComposeServerHTTPFirewallDashboardResponse{
+		CountDailyLog:     countDailyLog,
+		CountDailyBlock:   countDailyBlock,
+		CountDailyCaptcha: countDailyCaptcha,
+		CountWeeklyBlock:  countWeeklyBlock,
+		CountMonthlyBlock: countMonthlyBlock,
+	}
+
+	// 规则分组
+	groupStats, err := stats.SharedServerHTTPFirewallDailyStatDAO.GroupDailyCount(tx, userId, req.ServerId, monthFrom, monthTo, 0, 10)
+	if err != nil {
+		return nil, err
+	}
+	for _, stat := range groupStats {
+		ruleGroupName, err := models.SharedHTTPFirewallRuleGroupDAO.FindHTTPFirewallRuleGroupName(tx, int64(stat.HttpFirewallRuleGroupId))
+		if err != nil {
+			return nil, err
+		}
+		if len(ruleGroupName) == 0 {
+			continue
+		}
+
+		resp.HttpFirewallRuleGroups = append(resp.HttpFirewallRuleGroups, &pb.ComposeServerHTTPFirewallDashboardResponse_HTTPFirewallRuleGroupStat{
+			HttpFirewallRuleGroup: &pb.HTTPFirewallRuleGroup{Id: int64(stat.HttpFirewallRuleGroupId), Name: ruleGroupName},
+			Count:                 int64(stat.Count),
+		})
+	}
+
+	// 每日趋势
+	dayBefore := timeutil.Format("Ymd", date.AddDate(0, 0, -14))
+	days, err := utils.RangeDays(dayBefore, day)
+	if err != nil {
+		return nil, err
+	}
+	{
+		statList, err := stats.SharedServerHTTPFirewallDailyStatDAO.FindDailyStats(tx, userId, req.ServerId, "log", dayBefore, day)
+		if err != nil {
+			return nil, err
+		}
+		m := map[string]int64{} // day => count
+		for _, stat := range statList {
+			m[stat.Day] = int64(stat.Count)
+		}
+		for _, day := range days {
+			resp.LogDailyStats = append(resp.LogDailyStats, &pb.ComposeServerHTTPFirewallDashboardResponse_DailyStat{Day: day, Count: m[day]})
+		}
+	}
+	{
+		statList, err := stats.SharedServerHTTPFirewallDailyStatDAO.FindDailyStats(tx, userId, req.ServerId, "block", dayBefore, day)
+		if err != nil {
+			return nil, err
+		}
+		m := map[string]int64{} // day => count
+		for _, stat := range statList {
+			m[stat.Day] = int64(stat.Count)
+		}
+		for _, day := range days {
+			resp.BlockDailyStats = append(resp.BlockDailyStats, &pb.ComposeServerHTTPFirewallDashboardResponse_DailyStat{Day: day, Count: m[day]})
+		}
+	}
+	{
+		statList, err := stats.SharedServerHTTPFirewallDailyStatDAO.FindDailyStats(tx, userId, req.ServerId, "captcha", dayBefore, day)
+		if err != nil {
+			return nil, err
+		}
+		m := map[string]int64{} // day => count
+		for _, stat := range statList {
+			m[stat.Day] = int64(stat.Count)
+		}
+		for _, day := range days {
+			resp.CaptchaDailyStats = append(resp.CaptchaDailyStats, &pb.ComposeServerHTTPFirewallDashboardResponse_DailyStat{Day: day, Count: m[day]})
+		}
+	}
+
+	return resp, nil
+}

internal/rpc/services/service_server_http_firewall_daily_stat_test.go

@@ -0,0 +1,20 @@
+package services
+
+import (
+	rpcutils "github.com/TeaOSLab/EdgeAPI/internal/rpc/utils"
+	"github.com/TeaOSLab/EdgeCommon/pkg/rpc/pb"
+	"github.com/iwind/TeaGo/dbs"
+	"github.com/iwind/TeaGo/logs"
+	"testing"
+)
+
+func TestServerHTTPFirewallDailyStatService_ComposeServerHTTPFirewallDashboard(t *testing.T) {
+	dbs.NotifyReady()
+
+	service := new(ServerHTTPFirewallDailyStatService)
+	resp, err := service.ComposeServerHTTPFirewallDashboard(rpcutils.NewMockAdminNodeContext(1), &pb.ComposeServerHTTPFirewallDashboardRequest{})
+	if err != nil {
+		t.Fatal(err)
+	}
+	logs.PrintAsJSON(resp, t)
+}

internal/rpc/utils/mock_node_admin.go

@@ -0,0 +1,13 @@
+package rpcutils
+
+import "context"
+
+type MockAdminNodeContext struct {
+	context.Context
+
+	AdminId int64
+}
+
+func NewMockAdminNodeContext(adminId int64) context.Context {
+	return &MockAdminNodeContext{AdminId: adminId}
+}

internal/rpc/utils/utils.go

@@ -57,9 +57,18 @@ func ValidateRequest(ctx context.Context, userTypes ...UserType) (userType UserT
 	}
 
 	// 是否是模拟测试
-	mockCtx, isMock := ctx.(*MockNodeContext)
-	if isMock {
-		return UserTypeNode, mockCtx.NodeId, nil
+	{
+		mockCtx, isMock := ctx.(*MockNodeContext)
+		if isMock {
+			return UserTypeNode, mockCtx.NodeId, nil
+		}
+	}
+
+	{
+		mockCtx, isMock := ctx.(*MockAdminNodeContext)
+		if isMock {
+			return UserTypeAdmin, mockCtx.AdminId, nil
+		}
 	}
 
 	md, ok := metadata.FromIncomingContext(ctx)