waf.go

package controller

import (
	"context"
	"errors"
	"fmt"
	"strconv"
	"strings"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/olivere/elastic/v7"
	"github.com/rs/zerolog/log"
	"gitlab.com/tensorsecurity-rd/waf-console/internal/service"
	"gitlab.com/tensorsecurity-rd/waf-console/internal/utils"
	"gorm.io/gorm"
)

type WafController struct {
	service     service.Service
	simpleProxy *service.SimpleProxy
}

func NewWafController(clusterClientManager *utils.ClusterClientManager, db *gorm.DB, gatewayUrl string, elasticClient *elastic.Client, regionUrlMap map[string]string) *WafController {
	return &WafController{
		service:     service.NewWafService(clusterClientManager, db, gatewayUrl, elasticClient),
		simpleProxy: service.NewSimpleProxy(regionUrlMap),
	}
}

func (c *WafController) Waf(ctx *gin.Context) {
	ctx1, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	gatewayName := ctx.Query("gateway_name")
	regionCode := ctx.Query("region_code")
	namespace := ctx.Query("namespace")

	waf, err := c.service.GetWaf(ctx1, regionCode, namespace, gatewayName)
	if err != nil {
		// logging.Get().Err(err).Msgf("get waf")
		utils.AssembleResponse(ctx, nil, err)
		return
	}
	count, err := c.simpleProxy.CountAttackLogs(ctx1, regionCode, waf.ID)
	if err != nil {
		utils.AssembleResponse(ctx, nil, err)
		return
	}
	waf.AttackNum = int(count)
	resp := &utils.SingleRespData{
		Item: waf,
	}
	utils.AssembleResponse(ctx, resp, nil)
}

func (c *WafController) ListWafs(ctx *gin.Context) {
	ctx1, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	wafs, err := c.service.ListWafs(ctx1)
	if err != nil {
		utils.AssembleResponse(ctx, nil, err)
		return
	}
	utils.AssembleResponse(ctx, wafs, nil)
}

func (c *WafController) ListWafsV2(ctx *gin.Context) {
	ctx1, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	wafs, err := c.service.ListWafs(ctx1)
	if err != nil {
		utils.AssembleResponse(ctx, nil, err)
		return
	}

	type wafService struct {
		Name      string `json:"name"`
		Namespace string `json:"namespace"`
		ID        uint32 `json:"id"`
	}

	type resp struct {
		Items []wafService `json:"items"`
	}

	items := []wafService{}
	for _, waf := range wafs {
		items = append(items, wafService{
			Name: waf.GatewayName,
			ID:   waf.ID,
		})
	}

	utils.AssembleResponse(ctx, resp{Items: items}, nil)
}

func (c *WafController) GetWafGatewayInfo(ctx *gin.Context) {
	ctx1, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	cookie := ctx.Request.Header.Get("Cookie")
	if cookie == "" {
		utils.AssembleResponse(ctx, nil, errors.New("cookie is required"))
		return
	}

	var req service.GetWafGatewayInfoReq
	if err := ctx.BindJSON(&req); err != nil {
		utils.AssembleResponse(ctx, nil, err)
		return
	}

	info, err := c.service.GetWafGatewayInfo(ctx1, &req)
	if err != nil {
		utils.AssembleResponse(ctx, nil, err)
		return
	}
	utils.AssembleResponse(ctx, info, nil)
}

func (c *WafController) CreateWaf(ctx *gin.Context) {
	ctx1, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	var req []service.CreateWafReq
	if err := ctx.BindJSON(&req); err != nil {
		utils.AssembleResponse(ctx, nil, err)
		return
	}
	for _, r := range req {
		_, err := c.service.CreateWaf(ctx1, &r)
		if err != nil {
			utils.AssembleResponse(ctx, nil, err)
			return
		}
	}
	utils.AssembleResponse(ctx, nil, nil)
}

// `
// curl -X POST http://127.0.0.1:8080/v1/api/waf \
//      -H "Content-Type: application/json" \
//      -d '[
//            {
//              "region_code": "cn-east-1",
//              "gatewayname": "istio-ingressgateway",
// 			 "namespace": "istio-system",
// 			 "port": 80,
// 			 "host": ["*"]
//            }
//          ]'
// `

func (c *WafController) UpdateMode(ctx *gin.Context) {
	ctx1, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	var req service.UpdateModeReq
	if err := ctx.BindJSON(&req); err != nil {
		utils.AssembleResponse(ctx, nil, err)
		return
	}

	_, err := c.service.UpdateMode(ctx1, &req)
	if err != nil {
		utils.AssembleResponse(ctx, nil, err)
		return
	}
	utils.AssembleResponse(ctx, nil, nil)
}

func (c *WafController) UpdateRule(ctx *gin.Context) {
	ctx1, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	var req service.RuleRequest
	if err := ctx.BindJSON(&req); err != nil {
		utils.AssembleResponse(ctx, nil, err)
		return
	}

	err := c.service.UpdateRule(ctx1, &req)
	if err != nil {
		utils.AssembleResponse(ctx, nil, err)
		return
	}
	utils.AssembleResponse(ctx, nil, nil)
}

func (c *WafController) ListRules(ctx *gin.Context) {
	ctx1, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	regionCode := ctx.Query("region_code")
	namespace := ctx.Query("namespace")
	gatewayName := ctx.Query("gateway_name")
	language := ctx.Request.Header.Get("Accept-Language")
	name := ctx.Query("name")
	if language == "" {
		language = "zh"
	}

	rules, err := c.service.ListRules(ctx1, regionCode, namespace, gatewayName, language, name)
	if err != nil {
		utils.AssembleResponse(ctx, nil, err)
		return
	}
	type resp struct {
		Items []service.RuleGroupResp `json:"items"`
	}
	items := []service.RuleGroupResp{}
	for _, rule := range rules {
		items = append(items, service.RuleGroupResp{
			CategoryID:  rule.CategoryID,
			Status:      rule.Status,
			Category:    rule.Category,
			Description: rule.Description,
		})
	}

	utils.AssembleResponse(ctx, resp{Items: items}, nil)
}

func (c *WafController) SaveRuleCategoryToDB(ctx *gin.Context) {
	ctx1, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	err := c.service.SaveRuleCategoryToDB(ctx1)
	if err != nil {
		utils.AssembleResponse(ctx, nil, err)
		return
	}
	utils.AssembleResponse(ctx, nil, nil)
}

func (c *WafController) EnableListenerWaf(ctx *gin.Context) {
	ctx1, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	var req service.EnableListenerWafReq
	if err := ctx.BindJSON(&req); err != nil {
		utils.AssembleResponse(ctx, nil, err)
		return
	}

	err := c.service.EnableListenerWaf(ctx1, &req)
	if err != nil {
		utils.AssembleResponse(ctx, nil, err)
		return
	}
	utils.AssembleResponse(ctx, nil, nil)
}

func (c *WafController) EnableGatewayWaf(ctx *gin.Context) {
	ctx1, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	var req service.EnableGatewayWafReq
	if err := ctx.BindJSON(&req); err != nil {
		utils.AssembleResponse(ctx, nil, err)
		return
	}
	req.Cookie = ctx.Request.Header.Get("Cookie")

	err := c.service.EnableGatewayWaf(ctx1, &req)
	if err != nil {
		utils.AssembleResponse(ctx, nil, err)
		return
	}
	utils.AssembleResponse(ctx, nil, nil)
}

func (c *WafController) EnableListenerWafs(ctx *gin.Context) {
	ctx1, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	var req service.EnableListenerWafsReq
	if err := ctx.BindJSON(&req); err != nil {
		utils.AssembleResponse(ctx, nil, err)
		return
	}

	err := c.service.EnableListenerWafs(ctx1, &req)
	if err != nil {
		utils.AssembleResponse(ctx, nil, err)
		return
	}
	utils.AssembleResponse(ctx, nil, nil)
}

func (c *WafController) DeleteListenerWaf(ctx *gin.Context) {
	ctx1, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	regionCode := ctx.Param("region_code")
	namespace := ctx.Param("namespace")
	gatewayName := ctx.Param("gateway_name")
	port, err := strconv.Atoi(ctx.Param("port"))
	if err != nil {
		utils.AssembleResponse(ctx, nil, err)
		return
	}

	err = c.service.DeleteListenerWaf(ctx1, &service.DeleteListenerReq{
		GatewateInfo: service.GatewateInfo{
			RegionCode:  regionCode,
			Namespace:   namespace,
			GatewayName: gatewayName,
		},
		Port: port,
	})
	if err != nil {
		utils.AssembleResponse(ctx, nil, err)
		return
	}
	utils.AssembleResponse(ctx, nil, nil)
}

func (c *WafController) DeleteGatewayWaf(ctx *gin.Context) {
	ctx1, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	regionCode := ctx.Param("region_code")
	namespace := ctx.Param("namespace")
	gatewayName := ctx.Param("gateway_name")

	err := c.service.DeleteGatewayWaf(ctx1, &service.GatewateInfo{
		RegionCode:  regionCode,
		Namespace:   namespace,
		GatewayName: gatewayName,
	})
	if err != nil {
		utils.AssembleResponse(ctx, nil, err)
		return
	}
	utils.AssembleResponse(ctx, nil, nil)
}

func (c *WafController) ListAttackLogs(ctx *gin.Context) {
	var filter service.AttackLogFilter

	id := ctx.Query("serviceId")
	filter.ServiceId, _ = strconv.ParseInt(id, 10, 64)
	filter.Limit, filter.Offset, _ = getLimitAndOffset(ctx)
	filter.Cluster = ctx.Query("region_code")
	filter.AttackUrl = ctx.Query("attackAddr")
	filter.AttackIp = ctx.Query("attackIp")
	filter.AttackType = ctx.Query("attackType")
	filter.AttackApp = ctx.Query("attackApp")
	filter.AttackListener = ctx.Query("attackListener")
	filter.Token = ctx.Query("token")
	filter.Action = ctx.Query("action")
	stime := ctx.Query("startTime")
	filter.StartTime, _ = strconv.ParseInt(stime, 10, 64)
	etime := ctx.Query("endTime")
	filter.EndTime, _ = strconv.ParseInt(etime, 10, 64)

	ctx1, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	logs, token, err := c.service.ListAttackLogs(ctx1, &filter)
	if err != nil {
		utils.AssembleResponse(ctx, nil, err)
		return
	}
	log.Info().Interface("logs", logs).Msg("logs")

	type resp struct {
		Items []service.AttackLog `json:"items"`
		Token string              `json:"token"`
	}

	utils.AssembleResponse(ctx, resp{
		Items: logs,
		Token: token,
	}, nil)
}

func (c *WafController) CountAttackLogs(ctx *gin.Context) {
	ctx1, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	regionCode := ctx.Query("region_code")
	serviceID := ctx.Query("service_id")
	serviceIDUint, err := strconv.ParseUint(serviceID, 10, 32)
	if err != nil {
		utils.AssembleResponse(ctx, nil, err)
		return
	}

	count, err := c.service.CountAttackLogs(ctx1, regionCode, uint32(serviceIDUint))
	if err != nil {
		utils.AssembleResponse(ctx, nil, err)
		return
	}
	utils.AssembleResponse(ctx, count, nil)
}

// getLimitAndOffset extracts pagination parameters from the context
// Returns limit, offset, and error
func getLimitAndOffset(ctx *gin.Context) (int, int, error) {
	limit := 10 // default limit
	offset := 0 // default offset

	// Get limit from query parameters
	if limitStr := ctx.Query("limit"); limitStr != "" {
		parsedLimit, err := strconv.Atoi(limitStr)
		if err != nil {
			return 0, 0, fmt.Errorf("invalid limit parameter: %v", err)
		}
		if parsedLimit > 0 {
			limit = parsedLimit
		}
	}

	// Get offset from query parameters
	if offsetStr := ctx.Query("offset"); offsetStr != "" {
		parsedOffset, err := strconv.Atoi(offsetStr)
		if err != nil {
			return 0, 0, fmt.Errorf("invalid offset parameter: %v", err)
		}
		if parsedOffset >= 0 {
			offset = parsedOffset
		}
	}

	return limit, offset, nil
}

func (c *WafController) GetAttackLogDetails(ctx *gin.Context) {
	ctx1, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	uuid := ctx.Query("uuid")

	details, err := c.service.GetAttackLogDetails(ctx1, uuid)
	if err != nil {
		utils.AssembleResponse(ctx, nil, err)
		return
	}
	resp := &utils.SingleRespData{
		Item: details,
	}
	utils.AssembleResponse(ctx, resp, nil)
}

func (c *WafController) GetAttackLogRsp(ctx *gin.Context) {
	ctx1, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	uuid := ctx.Query("uuid")

	length := ctx.Query("length")
	lengthUint, err := strconv.ParseUint(length, 10, 32)
	if err != nil {
		utils.AssembleResponse(ctx, nil, err)
		return
	}

	rsp, err := c.service.GetAttackLogRsp(ctx1, uuid, uint32(lengthUint))
	if err != nil {
		utils.AssembleResponse(ctx, nil, err)
		return
	}
	resp := &utils.SingleRespData{
		Item: rsp,
	}
	utils.AssembleResponse(ctx, resp, nil)
}

func (c *WafController) CreateBlackWhiteList(ctx *gin.Context) {
	ctx1, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	var req service.MatcherExpr
	if err := ctx.BindJSON(&req); err != nil {
		utils.AssembleResponse(ctx, nil, err)
		return
	}

	err := c.service.CreateBlackWhiteList(ctx1, &req)
	if err != nil {
		utils.AssembleResponse(ctx, nil, err)
		return
	}
	utils.AssembleResponse(ctx, nil, nil)
}

func (c *WafController) UpdateBlackWhiteList(ctx *gin.Context) {
	ctx1, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	var req service.MatcherExpr
	if err := ctx.BindJSON(&req); err != nil {
		utils.AssembleResponse(ctx, nil, err)
		return
	}

	err := c.service.UpdateBlackWhiteList(ctx1, &req)
	if err != nil {
		utils.AssembleResponse(ctx, nil, err)
		return
	}
	utils.AssembleResponse(ctx, nil, nil)
}
func (c *WafController) EnableBlackWhiteList(ctx *gin.Context) {
	ctx1, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	var req service.MatcherExpr
	if err := ctx.BindJSON(&req); err != nil {
		utils.AssembleResponse(ctx, nil, err)
		return
	}

	err := c.service.EnableBlackWhiteList(ctx1, &req)
	if err != nil {
		utils.AssembleResponse(ctx, nil, err)
		return
	}
	utils.AssembleResponse(ctx, nil, nil)
}

func (c *WafController) DeleteBlackWhiteList(ctx *gin.Context) {
	ctx1, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	id := ctx.Param("id")
	idUint, err := strconv.ParseUint(id, 10, 32)
	if err != nil {
		utils.AssembleResponse(ctx, nil, err)
		return
	}
	log.Info().Interface("delete black white list", idUint).Msg("delete black white list")
	err = c.service.DeleteBlackWhiteList(ctx1, uint32(idUint))
	if err != nil {
		utils.AssembleResponse(ctx, nil, err)
		return
	}
	utils.AssembleResponse(ctx, nil, nil)
}

func (c *WafController) GetBlackWhiteLists(ctx *gin.Context) {
	ctx1, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	limit, offset, err := getLimitAndOffset(ctx)
	if err != nil {
		utils.AssembleResponse(ctx, nil, err)
		return
	}
	name := ctx.Query("name")
	status := ctx.Query("status")
	mode := ctx.Query("mode")
	expr := ctx.Query("expr")

	query := service.MatchExprQuery()

	if status != "" {
		statusList := strings.Split(status, ",")
		if len(statusList) == 1 && statusList[0] != "" {
			switch status {
			case "0":
				query.WithStatus(0)
			case "1":
				query.WithStatus(1)
			}
		}
	}

	if name != "" {
		query.WithFuzzyName(name)
	}

	if mode != "" {
		modes := strings.Split(mode, ",")
		query.WithInConditionCustom("mode", modes)
	}

	if expr != "" {
		query.WithFuzzyExpr(expr)
	}
	lists, total, err := c.service.GetBlackWhiteLists(ctx1, query, limit, offset)
	if err != nil {
		utils.AssembleResponse(ctx, nil, err)
		return
	}

	respData := utils.ListRespData{
		Items:        lists,
		TotalItems:   total,
		ItemsPerPage: limit,
	}
	utils.AssembleResponse(ctx, respData, nil)
}

func (c *WafController) ListListenerHistory(ctx *gin.Context) {
	ctx1, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	limit, offset, err := getLimitAndOffset(ctx)
	if err != nil {
		utils.AssembleResponse(ctx, nil, err)
		return
	}
	limit = 1000000
	name := ctx.Query("name")
	gatewayName := ctx.Query("gateway_name")
	listenerName := ctx.Query("listener_name")
	regionCode := ctx.Query("region_code")
	operation := ctx.Query("operation")
	status := ctx.Query("status")

	query := service.WafListenerHistoryQuery()
	if name != "" {
		query.WithFuzzyName(name)
	}
	if gatewayName != "" {
		query.WithFuzzyGatewayName(gatewayName)
	}
	if listenerName != "" {
		query.WithFuzzyListenerName(listenerName)
	}

	if regionCode != "" {
		query.WithFuzzyRegionCode(regionCode)
	}
	if operation != "" {
		query.WithOperation(operation)
	}
	if status != "" {
		statusInt, err := strconv.Atoi(status)
		if err != nil {
			utils.AssembleResponse(ctx, nil, err)
			return
		}
		query.WithStatus(int32(statusInt))
	}

	lists, total, err := c.service.ListListenerHistory(ctx1, query, limit, offset)
	if err != nil {
		utils.AssembleResponse(ctx, nil, err)
		return
	}

	respData := utils.ListRespData{
		Items:        lists,
		TotalItems:   total,
		ItemsPerPage: limit,
	}
	utils.AssembleResponse(ctx, respData, nil)
}

func (c *WafController) AttackClassesList(ctx *gin.Context) {
	ctx1, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	language := ctx.Request.Header.Get("Accept-Language")
	if language == "" {
		language = "zh"
	}

	classes := c.service.ListAttackClasses(ctx1, language)
	respData := utils.ListRespData{
		Items:        classes,
		TotalItems:   len(classes),
		ItemsPerPage: 10,
	}

	utils.AssembleResponse(ctx, respData, nil)
}