waf.go

package service

import (
	"bytes"
	"context"
	"crypto/tls"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"os"
	"strconv"
	"strings"
	"sync"
	"time"

	jsoniter "github.com/json-iterator/go"
	"github.com/olivere/elastic/v7"
	"github.com/rs/zerolog/log"
	"gitlab.com/tensorsecurity-rd/waf-console/internal/model"
	"gitlab.com/tensorsecurity-rd/waf-console/internal/utils"
	"gitlab.com/tensorsecurity-rd/waf-console/pkg/apis/waf.security.io/v1alpha1"
	"gorm.io/gorm"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/sets"
)

type wafService struct {
	clusterClientManager *utils.ClusterClientManager
	db                   *gorm.DB
	gatewayUrl           string
	elasticClient        *elastic.Client
}

func NewWafService(clusterClientManager *utils.ClusterClientManager, db *gorm.DB, gatewayUrl string, elasticClient *elastic.Client) Service {
	return &wafService{clusterClientManager: clusterClientManager, db: db, gatewayUrl: gatewayUrl, elasticClient: elasticClient}
}

func (s *wafService) GetWaf(ctx context.Context, regionCode, namespace, gatewayName string) (*WafService, error) {
	wafService := &model.WafService{}
	err := s.db.Model(&model.WafService{}).Where("gateway_name = ? AND region_code = ? AND namespace = ?", gatewayName, regionCode, namespace).First(wafService).Error
	if err != nil {
		if err == gorm.ErrRecordNotFound {
			// Create new WAF service record if not found
			wafService = &model.WafService{
				RegionCode:  regionCode,
				Namespace:   namespace,
				GatewayName: gatewayName,
				Mode:        string(WafModeAlert),
			}
			if err := s.db.Create(wafService).Error; err != nil {
				return nil, fmt.Errorf("failed to create WAF service: %v", err)
			}
		} else {
			return nil, fmt.Errorf("failed to query WAF service: %v", err)
		}
	}
	listenerWafs, err := s.ListListenerWafStatus(ctx, &GatewateInfo{
		GatewayName: gatewayName,
		Namespace:   namespace,
		RegionCode:  regionCode,
	})
	if err != nil {
		return nil, fmt.Errorf("failed to list listener WAF status: %v", err)
	}
	listeners := []string{}
	for _, listener := range listenerWafs {
		hosts := strings.Join(listener.Hosts, "@")
		listeners = append(listeners, fmt.Sprintf("%s-%d", hosts, listener.Port))
	}

	return &WafService{
		GatewayName: wafService.GatewayName,
		Mode:        wafService.Mode,
		RuleNum:     wafService.RuleNum,
		AttackNum:   wafService.AttackNum,
		Listeners:   listeners,
	}, nil
}

func (s *wafService) ListWafs(ctx context.Context) ([]WafService, error) {
	var wafs []WafService
	if err := s.db.Model(&model.WafService{}).Find(&wafs).Error; err != nil {
		return nil, err
	}
	for i, waf := range wafs {
		wafs[i].Name = waf.GatewayName
	}
	return wafs, nil
}

func (s *wafService) GetWafGatewayInfo(ctx context.Context, req *GetWafGatewayInfoReq) (*WafService, error) {
	wafService := &model.WafService{}
	err := s.db.Model(&model.WafService{}).Where("gateway_name = ? AND namespace = ? AND region_code = ?", req.GatewayName, req.Namespace, req.RegionCode).First(wafService).Error
	if err != nil {
		if err == gorm.ErrRecordNotFound {
			httpRequst := http.Request{
				Method: http.MethodPost,
				URL:    &url.URL{Scheme: "https", Host: "console.tensorsecurity.com", Path: "/api/v1/waf/gateway"},
				Header: http.Header{
					"Cookie": []string{req.Cookie},
				},
				Body: io.NopCloser(strings.NewReader(fmt.Sprintf(`{"gateway_name": "%s", "namespace": "%s", "region_code": "%s"}`, req.GatewayName, req.Namespace, req.RegionCode))),
			}
			resp, err := http.DefaultClient.Do(&httpRequst)
			if err != nil {
				return nil, fmt.Errorf("failed to get WAF service: %v", err)
			}
			defer resp.Body.Close()
			body, err := io.ReadAll(resp.Body)
			if err != nil {
				return nil, fmt.Errorf("failed to read WAF service: %v", err)
			}
			var wafService model.WafService
			err = json.Unmarshal(body, &wafService)
			if err != nil {
				return nil, fmt.Errorf("failed to unmarshal WAF service: %v", err)
			}
			wafService.ID = 0
			wafService.RuleCategoryStatus = nil
			wafService.RuleNum = 0
			wafService.AttackNum = 0
			// wafService.Host = model.HostList([]string{"*"})
			wafService.Mode = string(WafModeAlert)
			err = s.db.Create(wafService).Error
			if err != nil {
				return nil, fmt.Errorf("failed to create WAF service: %v", err)
			}
		} else {
			return nil, fmt.Errorf("failed to query WAF service: %v", err)
		}
	}
	return &WafService{
		GatewayName: wafService.GatewayName,
		Mode:        wafService.Mode,
		RuleNum:     wafService.RuleNum,
		AttackNum:   wafService.AttackNum,
	}, nil
}

func (s *wafService) getRulesForService(req *CreateWafReq) ([]v1alpha1.Rule, error) {
	rules := []v1alpha1.Rule{}
	ruleCategories := []model.WafRuleCategory{}
	if err := s.db.Model(&model.WafRuleCategory{}).Where("status = ?", 0).Find(&ruleCategories).Error; err != nil {
		return nil, fmt.Errorf("failed to get rule categories: %v", err)
	}

	// Get existing WAF service config if any
	wafService := &model.WafService{}
	err := s.db.Model(&model.WafService{}).Where("gateway_name = ? AND namespace = ? AND region_code = ?", req.GatewayName, req.Namespace, req.RegionCode).First(wafService).Error
	if err != nil {
		if err == gorm.ErrRecordNotFound {
			// Create new WAF service record if not found
			wafService = &model.WafService{
				RegionCode:  req.RegionCode,
				Namespace:   req.Namespace,
				GatewayName: req.GatewayName,
				Mode:        string(WafModeAlert),
			}
			if err := s.db.Create(wafService).Error; err != nil {
				return nil, fmt.Errorf("failed to create WAF service: %v", err)
			}
		} else {
			return nil, fmt.Errorf("failed to query WAF service: %v", err)
		}
	}

	// Determine which rule categories to enable
	var enabledCategories []model.WafRuleCategory

	if wafService.RuleCategoryStatus != nil && len(wafService.RuleCategoryStatus.CategoryID) > 0 {
		// Only include categories not already enabled
		for _, category := range ruleCategories {
			for _, id := range wafService.RuleCategoryStatus.CategoryID {
				if id == category.CategoryID {
					enabledCategories = append(enabledCategories, category)
					continue
				}
			}
		}
	} else {
		// Enable all categories if none specified
		enabledCategories = ruleCategories
	}

	for _, category := range enabledCategories {
		for _, rule := range category.Rules {
			rules = append(rules, v1alpha1.Rule{
				ID:          rule.ID,
				Level:       rule.Level,
				Name:        rule.Name,
				Type:        rule.Type,
				Description: rule.Description,
				Expr:        rule.Expr,
				Mode:        rule.Mode,
			})
		}
	}

	return rules, nil
}

func (s *wafService) CreateWaf(ctx context.Context, req *CreateWafReq) (*WafService, error) {
	var errMsg string
	var status int = 1 // Success by default

	name := fmt.Sprintf("%s-%d", req.GatewayName, req.Port)

	defer func() {
		_ = s.addListenerHistory(ctx, name, req.ListenerName, req.GatewayName, req.Namespace, req.RegionCode, errMsg, status, model.OperationCreate)
	}()

	// Create the WAF service resource
	service := &v1alpha1.Service{
		ObjectMeta: metav1.ObjectMeta{
			Name:      name,
			Namespace: req.Namespace,
			Labels: map[string]string{
				"apigateway_name": req.GatewayName,
			},
		},
		Spec: v1alpha1.ServiceSpec{
			HostNames:   req.Host,
			ServiceName: req.ListenerName,
			Port:        req.Port,
			Workload: v1alpha1.WorkloadRef{
				Kind:       req.GatewayName,
				Name:       req.ListenerName,
				Namespace:  req.Namespace,
				ClusterKey: req.RegionCode,
			},
			Uri: &v1alpha1.StringMatch{
				Prefix: "/",
			},
			LogConfig: &v1alpha1.LogConfig{
				Enable: 1,
				Level:  "info",
			},
			Mode:      string(req.Mode),
			ServiceID: req.ServiceID,
		},
	}

	rules, err := s.getRulesForService(req)
	if err != nil {
		status = 0 // Failure
		errMsg = fmt.Sprintf("failed to get rules for service: %v", err)
		return nil, fmt.Errorf("%s", errMsg)
	}
	service.Spec.Rules = rules

	if len(service.Spec.Rules) == 0 {
		status = 0 // Failure
		errMsg = "cannot create WAF service with no rules"
		return nil, fmt.Errorf("%s", errMsg)
	}

	// Create the WAF service in Kubernetes
	client := s.clusterClientManager.GetClient(req.RegionCode)
	if client == nil {
		status = 0 // Failure
		errMsg = fmt.Sprintf("failed to get cluster client for region %s", req.RegionCode)
		return nil, fmt.Errorf("%s", errMsg)
	}
	if _, err := client.Versioned.WafV1alpha1().Services(req.Namespace).Create(ctx, service, metav1.CreateOptions{}); err != nil {
		status = 0 // Failure
		errMsg = fmt.Sprintf("failed to create WAF service: %v", err)
		return nil, fmt.Errorf("%s", errMsg)
	}

	return &WafService{
		GatewayName: req.GatewayName,
		Mode:        service.Spec.Mode,
		RuleNum:     len(service.Spec.Rules),
		AttackNum:   0,
	}, nil
}

func (s *wafService) DeleteListenerWaf(ctx context.Context, req *DeleteListenerReq) error {
	var errMsg string
	var status int = 1 // Success by default

	name := fmt.Sprintf("%s-%d", req.GatewayName, req.Port)

	defer func() {
		_ = s.addListenerHistory(ctx, name, req.ListenerName, req.GatewayName, req.Namespace, req.RegionCode, errMsg, status, model.OperationDelete)
	}()

	client := s.clusterClientManager.GetClient(req.RegionCode)
	if client == nil {
		status = 0 // Failure
		errMsg = fmt.Sprintf("failed to get cluster client for region %s", req.RegionCode)
		return fmt.Errorf("%s", errMsg)
	}

	if err := client.Versioned.WafV1alpha1().Services(req.Namespace).Delete(ctx, name, metav1.DeleteOptions{}); err != nil {
		status = 0 // Failure
		errMsg = fmt.Sprintf("failed to delete WAF service: %v", err)
		return fmt.Errorf("%s", errMsg)
	}

	return nil
}

func (s *wafService) UpdateMode(ctx context.Context, req *UpdateModeReq) (*WafService, error) {
	// Check if WAF service exists
	wafService := &model.WafService{}
	err := s.db.Model(&model.WafService{}).Where("gateway_name = ?", req.GatewayName).First(wafService).Error
	if err != nil {
		if err == gorm.ErrRecordNotFound {
			// Create new WAF service record if not found
			wafService = &model.WafService{
				RegionCode:  req.RegionCode,
				Namespace:   req.Namespace,
				GatewayName: req.GatewayName,
				Mode:        string(req.Mode),
			}
			if err := s.db.Create(wafService).Error; err != nil {
				return nil, fmt.Errorf("failed to create WAF service: %v", err)
			}
		} else {
			return nil, fmt.Errorf("failed to query WAF service: %v", err)
		}
	} else {
		// Update mode if service exists
		if err := s.db.Model(wafService).Update("mode", string(req.Mode)).Error; err != nil {
			return nil, fmt.Errorf("failed to update WAF service mode: %v", err)
		}
	}
	// Update mode for each listener
	client := s.clusterClientManager.GetClient(req.RegionCode)
	if client == nil {
		return nil, fmt.Errorf("failed to get cluster client for region %s", req.RegionCode)
	}
	listenerList, err := client.Versioned.WafV1alpha1().Services(req.Namespace).List(ctx, metav1.ListOptions{LabelSelector: fmt.Sprintf("apigateway_name=%s", req.GatewayName)})
	if err != nil {
		return nil, fmt.Errorf("failed to get listener list: %v", err)
	}
	var wg sync.WaitGroup
	for _, listener := range listenerList.Items {
		wg.Add(1)
		listener := listener // Create new variable for goroutine
		listener.Spec.Mode = string(req.Mode)
		go func() {
			defer wg.Done()
			log.Info().Msgf("update WAF service mode: %v", listener.Name)
			_, err := client.Versioned.WafV1alpha1().Services(req.Namespace).Update(ctx, &listener, metav1.UpdateOptions{})
			if err != nil {
				log.Error().Msgf("failed to update WAF service mode: %v", err)
			}
		}()
	}
	wg.Wait()
	return &WafService{
		GatewayName: req.GatewayName,
		Mode:        string(req.Mode),
	}, nil
}

// func (s *wafService) GetRuleCategories(ctx context.Context) ([]WafRuleCategory, error) {
// 	var categories []WafRuleCategory
// 	err := s.db.Table("waf_rule_categories").Find(&categories).Error
// 	if err != nil {
// 		return nil, err
// 	}
// 	return categories, nil
// }

func (s *wafService) GetRules(ctx context.Context, categoryID string) ([]WafRule, error) {
	var rules []WafRule
	err := s.db.Table("waf_rules").Where("category_id = ?", categoryID).Find(&rules).Error
	if err != nil {
		return nil, err
	}
	return rules, nil
}

func (s *wafService) GetRule(ctx context.Context, ruleID int) (*WafRule, error) {
	var rule WafRule
	err := s.db.Table("waf_rules").Where("id = ?", ruleID).First(&rule).Error
	if err != nil {
		return nil, err
	}
	return &rule, nil
}

func (s *wafService) SaveRuleCategoryToDB(ctx context.Context) error {
	var categories []WafRuleCategory
	jsonFile, err := os.ReadFile("rules/waf-rules.json")
	if err != nil {
		return fmt.Errorf("error reading yaml file: %v", err)
	}

	// err = yaml.Unmarshal(yamlFile, &categories)
	err = json.Unmarshal(jsonFile, &categories)
	if err != nil {
		return fmt.Errorf("error unmarshaling yaml: %v", err)
	}

	for _, category := range categories {
		rules := []model.WafRule{}
		for _, rule := range category.Rules {
			rules = append(rules, model.WafRule{
				ID:          rule.ID,
				CategoryID:  category.CategoryID,
				Level:       rule.Level,
				Name:        rule.Name,
				Type:        rule.Type,
				Description: rule.Description,
				Expr:        rule.Expr,
				Mode:        rule.Mode,
			})
		}
		model := model.WafRuleCategory{
			CategoryID:    category.CategoryID,
			Status:        category.Status,
			CategoryEN:    category.Category.EN,
			CategoryZH:    category.Category.Zh,
			DescriptionEN: category.Description.EN,
			DescriptionZH: category.Description.Zh,
			Rules:         model.RuleList(rules),
		}
		err = s.db.Table("waf_rule_categories").Create(&model).Error
		if err != nil {
			return err
		}
	}

	return nil
}

func (s *wafService) DeleteListener(ctx context.Context, req *DeleteListenerReq) error {
	listener := &model.GatewayListener{}
	err := s.db.Model(&model.GatewayListener{}).Where("gateway_name = ? AND namespace = ? AND region_code = ?", req.GatewayName, req.Namespace, req.RegionCode).First(listener).Error
	if err != nil {
		return err
	}

	err = s.db.Model(&model.GatewayListener{}).Where("gateway_name = ? AND namespace = ? AND region_code = ?", req.GatewayName, req.Namespace, req.RegionCode).Delete(listener).Error
	if err != nil {
		return err
	}

	return nil
}

func (s *wafService) getServiceID(ctx context.Context, gatewayName, namespace, regionCode string) (uint32, error) {
	service := &model.WafService{}
	err := s.db.WithContext(ctx).Model(&model.WafService{}).Where("gateway_name = ? AND namespace = ? AND region_code = ?", gatewayName, namespace, regionCode).First(service).Error
	if err != nil {
		return 0, err
	}

	return uint32(service.ID), nil
}

func (s *wafService) EnableListenerWaf(ctx context.Context, req *EnableListenerWafReq) error {
	if req.Enable {
		log.Info().Msgf("Create WAF for listener %s", req.GatewayName)
		serviceID, err := s.getServiceID(ctx, req.GatewayName, req.Namespace, req.RegionCode)
		if err != nil {
			if err == gorm.ErrRecordNotFound {
				service := &model.WafService{
					GatewayName: req.GatewayName,
					Namespace:   req.Namespace,
					RegionCode:  req.RegionCode,
					Mode:        string(req.Mode),
				}
				if err := s.db.Model(&model.WafService{}).Create(&service).Error; err != nil {
					return err
				}
				serviceID = uint32(service.ID)
			} else {
				return err
			}
		}

		_, err = s.CreateWaf(ctx, &CreateWafReq{
			GatewateInfo: GatewateInfo{
				GatewayName: req.GatewayName,
				Namespace:   req.Namespace,
				RegionCode:  req.RegionCode,
			},
			Port:         uint32(req.Port),
			Host:         req.Hosts,
			Mode:         req.Mode,
			ListenerName: req.ListenerName,
			ServiceID:    serviceID,
		})
		if err != nil {
			return err
		}
	} else {
		log.Info().Msgf("Delete WAF for listener %s", req.GatewayName)
		err := s.DeleteListenerWaf(ctx, &DeleteListenerReq{
			GatewateInfo: GatewateInfo{
				GatewayName: req.GatewayName,
				Namespace:   req.Namespace,
				RegionCode:  req.RegionCode,
			},
			Port:         req.Port,
			ListenerName: req.ListenerName,
		})
		if err != nil {
			return err
		}
	}
	return nil
}

func getGatewayNameFromCrn(crn string) string {
	// crn:ucs::apigateway:lf-tst7:214613666997:instance/testaaa
	parts := strings.Split(crn, "/")
	return parts[len(parts)-1]
}

func (s *wafService) listListenerFromApiGateway(ctx context.Context, apiGatewayCrn string, regionCode string, cookie string) ([]GatewayRespListenerData, error) {
	body, err := json.Marshal(map[string]string{
		"apigateway_crn": apiGatewayCrn,
		"region_code":    regionCode,
	})
	if err != nil {
		return nil, fmt.Errorf("failed to marshal request body: %v", err)
	}
	request, err := http.NewRequestWithContext(ctx, "POST", "https://csm.console.test.tg.unicom.local/apigatewaymng/listener/lf-tst7/list_listeners", bytes.NewBuffer(body))
	if err != nil {
		return nil, fmt.Errorf("failed to create request: %v", err)
	}
	request.Header.Set("Cookie", cookie)
	// Create custom transport with TLS config
	tr := &http.Transport{
		TLSClientConfig: &tls.Config{
			InsecureSkipVerify: true, // Skip certificate verification for test environment
		},
	}
	client := &http.Client{Transport: tr}
	resp, err := client.Do(request)
	if err != nil {
		return nil, fmt.Errorf("failed to get listener list: %v", err)
	}
	defer resp.Body.Close()

	log.Info().Msgf("resp: %v", resp)
	// Parse response
	var response GatewayListenerResponseList

	if err := json.NewDecoder(resp.Body).Decode(&response); err != nil {
		return nil, fmt.Errorf("failed to parse listener list: %v", err)
	}
	log.Info().Msgf("response: %v", response)
	return response.Data, nil
}

func (s *wafService) EnableGatewayWaf(ctx context.Context, req *EnableGatewayWafReq) error {
	if req.Enable {
		listeners, err := s.listListenerFromApiGateway(ctx, req.ApiGatewayCrn, req.RegionCode, req.Cookie)
		if err != nil {
			return fmt.Errorf("failed to get listener list: %v", err)
		}
		log.Info().Msgf("listeners: %v", listeners)
		// Create WAF for each listener
		for _, listener := range listeners {
			gatewayName := getGatewayNameFromCrn(listener.ApiGatewayCrn)
			namespace := fmt.Sprintf("%s-%s", listener.CreateAccountName, listener.CreateAccountID)
			if _, err := s.CreateWaf(ctx, &CreateWafReq{
				GatewateInfo: GatewateInfo{
					GatewayName: gatewayName,
					Namespace:   namespace,
					RegionCode:  req.RegionCode,
				},
				Port: uint32(listener.Port),
				Host: listener.Hosts,
			}); err != nil {
				return fmt.Errorf("failed to create WAF for listener %d: %v", listener.Port, err)
			}
		}
	} else {
		s.DeleteGatewayWaf(ctx, &GatewateInfo{
			GatewayName: req.GatewayName,
			Namespace:   req.Namespace,
			RegionCode:  req.RegionCode,
		})
	}
	return nil
}

func (s *wafService) DeleteGatewayWaf(ctx context.Context, req *GatewateInfo) error {
	client := s.clusterClientManager.GetClient(req.RegionCode)
	if client == nil {
		return fmt.Errorf("failed to get cluster client")
	}
	labelSelector := fmt.Sprintf("apigateway_name=%s", req.GatewayName)
	if err := client.Versioned.WafV1alpha1().Services(req.Namespace).DeleteCollection(ctx, metav1.DeleteOptions{}, metav1.ListOptions{LabelSelector: labelSelector}); err != nil {
		return fmt.Errorf("failed to delete WAF service: %v", err)
	}
	return nil
}

func (s *wafService) UpdateRule(ctx context.Context, req *RuleRequest) error {
	wafService := &model.WafService{}
	err := s.db.Model(&model.WafService{}).Where("gateway_name = ?", req.GatewayName).First(wafService).Error
	if err != nil {
		if err == gorm.ErrRecordNotFound {
			// Create new WAF service record if not found
			wafService = &model.WafService{
				RegionCode:  req.RegionCode,
				Namespace:   req.Namespace,
				GatewayName: req.GatewayName,
				Mode:        string(WafModeAlert),
				RuleCategoryStatus: &model.RuleCategoryStatus{
					CategoryID: req.CategoryID,
					Status:     req.Status,
				},
			}
			if err := s.db.Create(wafService).Error; err != nil {
				return fmt.Errorf("failed to create WAF service: %v", err)
			}
		} else {
			return fmt.Errorf("failed to query WAF service: %v", err)
		}
	} else {
		// Update mode if service exists
		if err := s.db.Model(wafService).Update("rule_category_status", model.RuleCategoryStatus{
			CategoryID: req.CategoryID,
			Status:     req.Status,
		}).Error; err != nil {
			return fmt.Errorf("failed to update WAF service mode: %v", err)
		}
	}
	return nil
}

func (s *wafService) ListListenerWafStatus(ctx context.Context, req *GatewateInfo) ([]*GatewayListener, error) {
	client := s.clusterClientManager.GetClient(req.RegionCode)
	if client == nil {
		return nil, fmt.Errorf("failed to get cluster client")
	}

	listenerList, err := client.Versioned.WafV1alpha1().Services(req.Namespace).List(ctx, metav1.ListOptions{LabelSelector: fmt.Sprintf("apigateway_name=%s", req.GatewayName)})
	if err != nil {
		return nil, fmt.Errorf("failed to get listener list: %v", err)
	}

	listenerStatusList := []*GatewayListener{}
	for _, listener := range listenerList.Items {
		n := strings.LastIndex(listener.Name, "-")
		if n == -1 {
			return nil, fmt.Errorf("failed to get listener port: %v", listener.Name)
		}
		listenerPort := listener.Name[n+1:]
		listenerPortInt, err := strconv.Atoi(listenerPort)
		if err != nil {
			return nil, fmt.Errorf("failed to parse listener port: %v", err)
		}

		// hosts := strings.Join(listener.Spec.HostNames, "@")
		// log.Info().Msgf("hosts: %v", hosts)
		listenerStatusList = append(listenerStatusList, &GatewayListener{
			GatewayName: req.GatewayName,
			Namespace:   req.Namespace,
			RegionCode:  req.RegionCode,
			Port:        listenerPortInt,
			Hosts:       listener.Spec.HostNames,
		})
	}

	// for _, port := range portList {
	// 	listenerStatusList = append(listenerStatusList, &GatewayListener{
	// 		GatewayName: req.GatewayName,
	// 		Namespace:   req.Namespace,
	// 		RegionCode:  req.RegionCode,
	// 		Port:        port,
	// 		Enable:      true,
	// 	})
	// }

	return listenerStatusList, nil
}

func (s *wafService) EnableListenerWafs(ctx context.Context, req *EnableListenerWafsReq) error {

	client := s.clusterClientManager.GetClient(req.RegionCode)
	if client == nil {
		return fmt.Errorf("failed to get cluster client")
	}

	listenerList, err := client.Versioned.WafV1alpha1().Services(req.Namespace).List(ctx, metav1.ListOptions{LabelSelector: fmt.Sprintf("apigateway_name=%s", req.GatewayName)})
	if err != nil {
		log.Error().Msgf("failed to get listener list: %v", err)
		return err
	}

	portList := []int{}
	for _, listener := range listenerList.Items {
		n := strings.LastIndex(listener.Name, "-")
		if n == -1 {
			return fmt.Errorf("failed to get listener port: %v", listener.Name)
		}
		listenerPort := listener.Name[n+1:]
		listenerPortInt, err := strconv.Atoi(listenerPort)
		if err != nil {
			return fmt.Errorf("failed to parse listener port: %v", err)
		}
		portList = append(portList, listenerPortInt)
	}
	currentPortSet := sets.NewInt(portList...)

	desiredPortSet := sets.NewInt()
	wafMap := map[int]ListenerWaf{}
	for _, listener := range req.Listeners {
		// get port from listener.HostsAndPort, like hosts1@127.0.0.1@abc.com-8080
		index := strings.LastIndex(listener.HostsAndPort, "-")
		if index == -1 {
			return fmt.Errorf("failed to get listener port: %v", listener)
		}
		port := listener.HostsAndPort[index+1:]
		portInt, err := strconv.Atoi(port)
		if err != nil {
			return fmt.Errorf("failed to parse listener port: %v", err)
		}
		desiredPortSet.Insert(portInt)
		log.Info().Msgf("listener: %v", listener.Name)

		hosts := strings.Split(listener.HostsAndPort[:index], "@")
		wafMap[portInt] = ListenerWaf{
			Hosts:        hosts,
			HostsAndPort: listener.HostsAndPort,
			Name:         listener.Name,
		}
	}

	// enable WAF for ports that are in the desired port set but not in the current port set
	addingPortSet := desiredPortSet.Difference(currentPortSet)

	// Get mode from waf_services table
	wafService := &model.WafService{}
	err = s.db.Model(&model.WafService{}).Where("gateway_name = ?", req.GatewayName).First(wafService).Error
	if err != nil {
		if err == gorm.ErrRecordNotFound {
			return fmt.Errorf("waf service not found for gateway %s", req.GatewayName)
		}
		return fmt.Errorf("failed to query waf service: %v", err)
	}
	mode := WafMode(wafService.Mode)

	for _, port := range addingPortSet.List() {
		err := s.EnableListenerWaf(ctx, &EnableListenerWafReq{
			GatewateInfo: GatewateInfo{
				GatewayName: req.GatewayName,
				Namespace:   req.Namespace,
				RegionCode:  req.RegionCode,
			},
			Port:         port,
			Hosts:        wafMap[port].Hosts,
			Enable:       true,
			Mode:         mode,
			ListenerName: wafMap[port].Name,
		})
		if err != nil {
			return fmt.Errorf("failed to enable listener WAF: %v", err)
		}
	}

	// delete WAF for ports that are not in the desired port set
	deletingPortSet := currentPortSet.Difference(desiredPortSet)
	for _, port := range deletingPortSet.List() {
		err := s.DeleteListenerWaf(ctx, &DeleteListenerReq{
			GatewateInfo: GatewateInfo{
				GatewayName: req.GatewayName,
				Namespace:   req.Namespace,
				RegionCode:  req.RegionCode,
			},
			Port: port,
		})
		if err != nil {
			return fmt.Errorf("failed to delete listener WAF: %v", err)
		}
	}
	return nil
}

func generateListToken(sort []interface{}) (string, error) {
	var token []string

	for _, v := range sort {
		switch v.(type) {
		case string:
			token = append(token, v.(string))
		case float64:
			token = append(token, strconv.FormatInt(int64(v.(float64)), 10))
		default:
			// json-iterator 和 encoding/json 的Decode，返回的都是 encoding/json 的 Number
			t, ok := jsoniter.CastJsonNumber(v)
			if !ok {
				return "", fmt.Errorf("unsupported sort field. value: %v, type: %T", v, v)
			}
			token = append(token, t)
		}
	}

	return strings.Join(token, ","), nil
}

func (s *wafService) ListAttackLogs(ctx context.Context, req *AttackLogFilter) ([]AttackLog, string, error) {
	boolQuery := elastic.NewBoolQuery()
	if req.ServiceId != 0 {
		boolQuery.Must(elastic.NewTermQuery("service_id", req.ServiceId))
	}
	if req.Cluster != "" {
		boolQuery.Filter(elastic.NewTermQuery("cluster_key", req.Cluster))
	}
	if req.AttackUrl != "" {
		boolQuery.Filter(elastic.NewMatchPhraseQuery("attacked_url", req.AttackUrl).Slop(0))
	}
	if req.AttackIp != "" {
		boolQuery.Filter(elastic.NewMatchPhraseQuery("attack_ip", req.AttackIp).Slop(0))
	}
	if req.AttackApp != "" {
		boolQuery.Filter(elastic.NewMatchPhraseQuery("attacked_app", req.AttackApp).Slop(0))
	}
	if req.AttackType != "" {
		attackTypes := strings.Split(req.AttackType, ",")
		var terms []any
		for _, attackType := range attackTypes {
			terms = append(terms, attackType)
		}
		boolQuery.Filter(elastic.NewTermsQuery("attack_type", terms...))
	}
	if req.Action != "" {
		actions := strings.Split(req.Action, ",")
		var terms []interface{}
		for _, action := range actions {
			terms = append(terms, action)
		}
		boolQuery.Filter(elastic.NewTermsQuery("action", terms...))
	} else {
		boolQuery.Filter(elastic.NewBoolQuery().MustNot(elastic.NewTermQuery("action", "pass")))
	}

	hasStart := req.StartTime > 0
	hasEnd := req.EndTime > 0
	if hasStart || hasEnd {
		rangeQuery := elastic.NewRangeQuery("attack_time")
		if hasStart {
			rangeQuery.Gte(req.StartTime)
		}
		if hasEnd {
			rangeQuery.Lte(req.EndTime)
		}

		boolQuery.Filter(rangeQuery)
	}
	src, _ := boolQuery.Source()
	log.Info().Interface("src", src.(map[string]interface{})).Msg("find waf detections src")

	ss := s.elasticClient.Search("waf-detections*")
	if req.Token != "" {
		for _, t := range strings.Split(req.Token, ",") {
			ss.SearchAfter(t)
		}
	}

	log.Info().Interface("limit", req.Limit).Msg("limit")
	result, err := ss.Query(boolQuery).Size(req.Limit).
		SortBy(elastic.NewFieldSort("attack_time").Order(false),
			elastic.NewFieldSort("id.digit").Order(false)).
		Do(ctx)

	if err != nil {
		return nil, "", fmt.Errorf("failed to search waf detections: %v", err)
	}

	list := make([]model.WafDetection, len(result.Hits.Hits))
	endIdx := len(result.Hits.Hits) - 1
	pageToken := ""
	log.Info().Interface("res", result).Msg("list attack logs res")
	for i, hit := range result.Hits.Hits {
		log.Info().Interface("hit source", hit.Source).Msg("hit")
		wafDetection := model.WafDetection{}
		if err = json.Unmarshal(hit.Source, &wafDetection); err != nil {
			return nil, "", fmt.Errorf("failed to unmarshal waf detection: %v", err)
		}

		list[i] = wafDetection
		if i == endIdx {
			pageToken, err = generateListToken(hit.Sort)
			if err != nil {
				return nil, "", fmt.Errorf("failed to generate list token: %v", err)
			}
		}
	}

	attackLogs := make([]AttackLog, len(list))
	for i, wafDetection := range list {
		attackLogs[i] = AttackLog{
			Uuid:           wafDetection.ID,
			AttackTime:     wafDetection.AttackTime,
			AttackIp:       wafDetection.AttackIP,
			AttackedApp:    wafDetection.AttackedApp,
			AttackListener: wafDetection.ResName,
			AttackType:     wafDetection.AttackType,
			Action:         wafDetection.Action,
			ClusterKey:     wafDetection.ClusterKey,
			AttackedAddr:   wafDetection.AttackedURL,
		}
	}
	return attackLogs, pageToken, nil
}

func (s *wafService) ListRules(ctx context.Context, regionCode, namespace, gatewayName, language, name string) ([]RuleGroupResp, error) {
	ruleCategories := []model.WafRuleCategory{}
	err := s.db.Model(&model.WafRuleCategory{}).Find(&ruleCategories).Error
	if err != nil {
		return nil, fmt.Errorf("failed to get waf service: %v", err)
	}

	ruleGroupResp := []RuleGroupResp{}
	wafService := &model.WafService{}
	err = s.db.Model(&model.WafService{}).Where("gateway_name = ? and namespace = ? and region_code = ?", gatewayName, namespace, regionCode).First(wafService).Error
	if err != nil {
		return nil, fmt.Errorf("failed to get waf service: %v", err)
	}
	if wafService.RuleCategoryStatus != nil && wafService.RuleCategoryStatus.Status == 1 {
		for _, category := range ruleCategories {
			for _, categoryID := range wafService.RuleCategoryStatus.CategoryID {
				if category.CategoryID == categoryID {
					category.Status = 1
				}
			}
			if language == "en" {
				ruleGroupResp = append(ruleGroupResp, RuleGroupResp{
					CategoryID:  category.CategoryID,
					Status:      category.Status,
					Category:    category.CategoryEN,
					Description: category.DescriptionEN,
				})
			} else {
				ruleGroupResp = append(ruleGroupResp, RuleGroupResp{
					CategoryID:  category.CategoryID,
					Status:      category.Status,
					Category:    category.CategoryZH,
					Description: category.DescriptionZH,
				})
			}
		}
	} else {
		for _, category := range ruleCategories {
			if language == "en" {
				ruleGroupResp = append(ruleGroupResp, RuleGroupResp{
					CategoryID:  category.CategoryID,
					Status:      category.Status,
					Category:    category.CategoryEN,
					Description: category.DescriptionEN,
				})
			} else {
				ruleGroupResp = append(ruleGroupResp, RuleGroupResp{
					CategoryID:  category.CategoryID,
					Status:      category.Status,
					Category:    category.CategoryZH,
					Description: category.DescriptionZH,
				})
			}
		}
	}
	return ruleGroupResp, nil
}

func (s *wafService) getWafServiceMap(ctx context.Context, req *MatcherExpr) (map[string][]model.WafService, error) {
	svcMap := make(map[string][]model.WafService)
	wafServices := []model.WafService{}
	if req.Global {
		if err := s.db.WithContext(ctx).Model(&model.WafService{}).Find(&wafServices).Error; err != nil {
			return nil, err
		}
	} else {
		if err := s.db.WithContext(ctx).Model(&model.WafService{}).Where("id in ?", req.Scope).Find(&wafServices).Error; err != nil {
			return nil, err
		}
	}
	for _, wafService := range wafServices {
		svcMap[wafService.RegionCode] = append(svcMap[wafService.RegionCode], wafService)
	}
	return svcMap, nil
}

func groupWafServicesByNamespace(wafServices []model.WafService) map[string][]model.WafService {
	svcMap := make(map[string][]model.WafService)
	for _, wafService := range wafServices {
		svcMap[wafService.Namespace] = append(svcMap[wafService.Namespace], wafService)
	}
	return svcMap
}

func (s *wafService) createConfigMap(ctx context.Context, req *MatcherExpr, regionCode string, wafSvc []model.WafService) error {
	client := s.clusterClientManager.GetClient(regionCode)
	if client == nil {
		return fmt.Errorf("failed to get cluster client")
	}

	gatewayNames := []string{}
	for _, wafSvc := range wafSvc {
		gatewayNames = append(gatewayNames, wafSvc.GatewayName)
	}
	scope := strings.Join(gatewayNames, ",")
	matchExpr := v1alpha1.MatchExpression{
		ID:     req.ID,
		Name:   req.Name,
		Scope:  scope,
		Mode:   req.Mode,
		Expr:   req.Expr,
		Status: req.Status,
	}
	matchExprJson, err := json.Marshal(matchExpr)
	if err != nil {
		return fmt.Errorf("failed to marshal match expression: %v", err)
	}
	name := fmt.Sprintf("waf-black-white-list-%d", req.ID)
	log.Info().Interface("name", name).Msg("create config map")
	_, err = client.Clientset.CoreV1().ConfigMaps(wafSvc[0].Namespace).Create(ctx, &corev1.ConfigMap{
		ObjectMeta: metav1.ObjectMeta{
			Name:   name,
			Labels: map[string]string{"waf.security.io/black-white-list": "true"},
		},
		Data: map[string]string{
			"match-expression": string(matchExprJson),
		},
	}, metav1.CreateOptions{})
	if err != nil {
		return fmt.Errorf("failed to create config map: %v", err)
	}
	return nil
}

func (s *wafService) updateConfigMap(ctx context.Context, req *MatcherExpr, regionCode string, wafSvc []model.WafService) error {
	client := s.clusterClientManager.GetClient(regionCode)
	if client == nil {
		return fmt.Errorf("failed to get cluster client")
	}
	name := fmt.Sprintf("waf-black-white-list-%d", req.ID)
	configMap, err := client.Clientset.CoreV1().ConfigMaps(wafSvc[0].Namespace).Get(ctx, name, metav1.GetOptions{})
	if err != nil {
		return fmt.Errorf("failed to get config map: %v", err)
	}

	gatewayNames := []string{}
	for _, wafSvc := range wafSvc {
		gatewayNames = append(gatewayNames, wafSvc.GatewayName)
	}
	matchExpr := v1alpha1.MatchExpression{
		ID:     req.ID,
		Name:   req.Name,
		Scope:  strings.Join(gatewayNames, ","),
		Mode:   req.Mode,
		Expr:   req.Expr,
		Status: req.Status,
	}

	matchExprJson, err := json.Marshal(matchExpr)
	if err != nil {
		return fmt.Errorf("failed to marshal match expression: %v", err)
	}

	newConfigMap := configMap.DeepCopy()
	newConfigMap.Data["match-expression"] = string(matchExprJson)
	_, err = client.Clientset.CoreV1().ConfigMaps(wafSvc[0].Namespace).Update(ctx, newConfigMap, metav1.UpdateOptions{})
	if err != nil {
		return fmt.Errorf("failed to update config map: %v", err)
	}
	return nil
}

func (s *wafService) CreateBlackWhiteList(ctx context.Context, req *MatcherExpr) error {
	tErr := s.db.Transaction(func(tx *gorm.DB) error {
		matcherExpr := model.MatcherExpr{
			Name:   req.Name,
			Scope:  req.Scope,
			Mode:   req.Mode,
			Expr:   req.Expr,
			Global: req.Global,
		}
		err := s.db.WithContext(ctx).Create(&matcherExpr).Error
		if err != nil {
			return err
		}

		req.ID = matcherExpr.ID
		svcMap, err := s.getWafServiceMap(ctx, req)
		if err != nil {
			return err
		}
		log.Info().Interface("svcMap", svcMap).Msg("svcMap")
		for regionCode, wafServices := range svcMap {
			err := s.createConfigMap(ctx, req, regionCode, wafServices)
			if err != nil {
				return err
			}
		}

		return nil
	})
	return tErr
}

func (s *wafService) UpdateBlackWhiteList(ctx context.Context, req *MatcherExpr) error {
	tErr := s.db.Transaction(func(tx *gorm.DB) error {
		matcherExpr := model.MatcherExpr{}
		err := tx.WithContext(ctx).Where("id = ?", req.ID).First(&matcherExpr).Error
		if err != nil {
			return err
		}

		matcherExpr.Name = req.Name
		matcherExpr.Scope = req.Scope
		matcherExpr.Mode = req.Mode
		matcherExpr.Expr = req.Expr
		matcherExpr.Global = req.Global
		err = tx.WithContext(ctx).Save(&matcherExpr).Error
		if err != nil {
			return err
		}

		svcMap, err := s.getWafServiceMap(ctx, &MatcherExpr{
			ID:     req.ID,
			Name:   matcherExpr.Name,
			Scope:  matcherExpr.Scope,
			Mode:   matcherExpr.Mode,
			Expr:   matcherExpr.Expr,
			Global: matcherExpr.Global,
		})
		if err != nil {
			return err
		}
		for regionCode, wafSvc := range svcMap {
			err = s.updateConfigMap(ctx, req, regionCode, wafSvc)
			if err != nil {
				return err
			}
		}

		return nil
	})
	return tErr
}

func (s *wafService) EnableBlackWhiteList(ctx context.Context, req *MatcherExpr) error {
	matcherExpr := model.MatcherExpr{}
	err := s.db.WithContext(ctx).Where("id = ?", req.ID).First(&matcherExpr).Error
	if err != nil {
		return err
	}
	// if status is the same, do nothing
	if req.Status == int32(matcherExpr.Status) {
		return nil
	}

	svcMap, err := s.getWafServiceMap(ctx, &MatcherExpr{
		ID:     req.ID,
		Name:   matcherExpr.Name,
		Scope:  matcherExpr.Scope,
		Mode:   matcherExpr.Mode,
		Expr:   matcherExpr.Expr,
		Global: matcherExpr.Global,
	})
	if err != nil {
		return err
	}

	if req.Status == 1 {
		matcherExpr.Status = 1
		for regionCode, wafServices := range svcMap {
			err := s.deleteConfigMap(ctx, req.ID, regionCode, wafServices)
			if err != nil {
				return err
			}
		}
	} else {
		for regionCode, wafServices := range svcMap {
			err := s.createConfigMap(ctx, req, regionCode, wafServices)
			if err != nil {
				return err
			}
		}
	}
	err = s.db.WithContext(ctx).Model(&model.MatcherExpr{}).Where("id = ?", req.ID).Update("status", req.Status).Error
	if err != nil {
		return err
	}
	return nil
}

func (s *wafService) deleteConfigMap(ctx context.Context, id uint32, regionCode string, wafSvc []model.WafService) error {
	client := s.clusterClientManager.GetClient(regionCode)
	if client == nil {
		return fmt.Errorf("failed to get cluster client")
	}
	name := fmt.Sprintf("waf-black-white-list-%d", id)
	return client.Clientset.CoreV1().ConfigMaps(wafSvc[0].Namespace).Delete(ctx, name, metav1.DeleteOptions{})
}

func (s *wafService) DeleteBlackWhiteList(ctx context.Context, ID uint32) error {
	log.Info().Interface("delete black white list", ID).Msg("delete black white list")
	matcherExpr := model.MatcherExpr{}
	err := s.db.WithContext(ctx).Where("id = ?", ID).First(&matcherExpr).Error
	if err != nil {
		return err
	}

	svcMap, err := s.getWafServiceMap(ctx, &MatcherExpr{
		ID:     ID,
		Name:   matcherExpr.Name,
		Scope:  matcherExpr.Scope,
		Mode:   matcherExpr.Mode,
		Expr:   matcherExpr.Expr,
		Global: matcherExpr.Global,
	})
	if err != nil {
		return err
	}
	log.Info().Interface("delete svcMap", svcMap).Msg("delete svcMap")
	for regionCode, wafServices := range svcMap {
		err := s.deleteConfigMap(ctx, ID, regionCode, wafServices)
		if err != nil {
			return err
		}
	}
	err = s.db.WithContext(ctx).Delete(&model.MatcherExpr{}, ID).Error
	if err != nil {
		return err
	}

	return nil
}

func GetLikeExpr(s string) string {
	sb := strings.Builder{}
	sb.WriteByte('%')
	sb.WriteString(s)
	sb.WriteByte('%')
	return sb.String()
}

func (s *wafService) GetBlackWhiteLists(ctx context.Context, query *MatchExprQueryOption, limit int, offset int) ([]MatcherExpr, int, error) {
	oneCtx, oneCancel := context.WithTimeout(ctx, 750*time.Millisecond)
	defer oneCancel()

	exprs := []model.MatcherExpr{}
	db := s.db.WithContext(oneCtx).Model(&model.MatcherExpr{})

	if len(query.whereEqCondition) > 0 {
		db = db.Where(query.whereEqCondition)
	}
	for column, val := range query.WhereLikeCondition {
		db = db.Where(fmt.Sprintf("%s LIKE ?", column), GetLikeExpr(val))
	}
	for column, val := range query.whereInCondition {
		db = db.Where(fmt.Sprintf("%s in ?", column), val)
	}
	var total int64
	err := db.Count(&total).Error
	if err != nil {
		return nil, 0, err
	}

	if limit > 0 && offset >= 0 {
		db = db.Offset(offset).Limit(limit)
	}
	err = db.Order("updated_at DESC").Find(&exprs).Error
	if err != nil {
		return nil, 0, err
	}
	exprsResp := []MatcherExpr{}
	for _, expr := range exprs {
		exprsResp = append(exprsResp, MatcherExpr{
			ID:     expr.ID,
			Name:   expr.Name,
			Scope:  expr.Scope,
			Mode:   expr.Mode,
			Expr:   expr.Expr,
			Status: int32(expr.Status),
			Global: expr.Global,
		})
	}
	return exprsResp, int(total), nil
}

func (s *wafService) ListListenerHistory(ctx context.Context, query *WafListenerHistoryOption, limit int, offset int) ([]model.WafListenerHistory, int, error) {
	listenerHistories := []model.WafListenerHistory{}
	db := s.db.WithContext(ctx).Model(&model.WafListenerHistory{})
	if len(query.WhereEqCondition) > 0 {
		db = db.Where(query.WhereEqCondition)
	}
	for _, val := range query.WhereLikeCondition {
		db = db.Where(fmt.Sprintf("%s LIKE ?", "name"), GetLikeExpr(val)).Or("%s LIKE ?", "listener_name").Or("%s LIKE ?", "gateway_name")
	}
	if limit > 0 && offset >= 0 {
		db = db.Offset(offset).Limit(limit)
	}
	err := db.Order("created_at DESC").Find(&listenerHistories).Error
	if err != nil {
		return nil, 0, err
	}
	var total int64
	err = db.Count(&total).Error
	if err != nil {
		return nil, 0, err
	}
	return listenerHistories, int(total), nil
}

func (s *wafService) addListenerHistory(ctx context.Context, name, listenerName, gatewayName, namespace, regionCode, description string, status int, operation model.Operation) error {
	listenerHistory := model.WafListenerHistory{
		Name:         name,
		GatewayName:  gatewayName,
		ListenerName: listenerName,
		Namespace:    namespace,
		RegionCode:   regionCode,
		Description:  description,
		Status:       model.Status(status),
		Operation:    operation,
	}
	err := s.db.WithContext(ctx).Create(&listenerHistory).Error
	if err != nil {
		return err
	}
	return nil
}