package service import ( "context" "fmt" "os" "slices" "gitlab.com/tensorsecurity-rd/waf-console/internal/model" "gitlab.com/tensorsecurity-rd/waf-console/pkg/apis/waf.security.io/v1alpha1" "gitlab.com/tensorsecurity-rd/waf-console/pkg/generated/clientset/versioned" "gopkg.in/yaml.v3" "gorm.io/gorm" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" ) type wafService struct { client *versioned.Clientset db *gorm.DB } func NewWafService(client *versioned.Clientset, db *gorm.DB) Service { return &wafService{client: client, db: db} } func (s *wafService) GetWaf(ctx context.Context, gatewayName string) (*Waf, error) { waf := &Waf{ GatewayName: gatewayName, Mode: "block", RuleNum: 100, AttackNum: 100, } return waf, nil } func (s *wafService) CreateWaf(ctx context.Context, req *CreateWafReq) (*Waf, error) { // Create the WAF service resource service := &v1alpha1.Service{ ObjectMeta: metav1.ObjectMeta{ Name: req.GatewayName, Namespace: req.Namespace, }, Spec: v1alpha1.ServiceSpec{ HostNames: req.Host, ServiceName: req.GatewayName, Port: req.Port, Workload: v1alpha1.WorkloadRef{ Kind: "Deployment", Name: req.GatewayName, Namespace: req.Namespace, }, }, } // Get enabled rule categories from DB var ruleCategories []model.WafRuleCategory if err := s.db.Model(&model.WafRuleCategory{}).Where("status = ?", 1).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 = ?", req.GatewayName).First(wafService).Error if err != nil && err != gorm.ErrRecordNotFound { return nil, fmt.Errorf("failed to get WAF service: %v", err) } // Determine which rule categories to enable var enabledCategories []model.WafRuleCategory if len(wafService.RuleCategoryStatus.CategoryID) > 0 { // Only include categories not already enabled for _, category := range ruleCategories { if !slices.Contains(wafService.RuleCategoryStatus.CategoryID, category.CategoryID) { enabledCategories = append(enabledCategories, category) } } } else { // Enable all categories if none specified enabledCategories = ruleCategories } // Add rules from enabled categories for _, category := range enabledCategories { for _, rule := range category.Rules { service.Spec.Rules = append(service.Spec.Rules, v1alpha1.Rule{ ID: rule.ID, Level: rule.Level, Name: rule.Name, Type: rule.Type, Description: rule.Description, Expr: rule.Expr, Mode: rule.Mode, }) } } // Create the WAF service in Kubernetes if _, err := s.client.WafV1alpha1().Services(req.Namespace).Create(ctx, service, metav1.CreateOptions{}); err != nil { return nil, fmt.Errorf("failed to create WAF service: %v", err) } return nil, nil } func (s *wafService) UpdateMode(ctx context.Context, req *UpdateModeReq) (*Waf, error) { return nil, 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 yamlFile, err := os.ReadFile("rules/waf-rules.yaml") if err != nil { return fmt.Errorf("error reading yaml file: %v", err) } err = yaml.Unmarshal(yamlFile, &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.Catagory.EN, CategoryZH: category.Catagory.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 }