waf.go

package service

import (
	"context"
	"fmt"
	"os"

	"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) {
	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,
			},
		},
	}
	_, err := s.client.WafV1alpha1().Services(req.Namespace).Create(context.Background(), service, metav1.CreateOptions{})
	if err != nil {
		return nil, 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("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),
		}
		s.db.Table("waf_rule_categories").Create(&model)
	}

	return nil
}