github.com/TeaOSLab/EdgeNode@v1.3.8/internal/waf/rule_group.go (about)

     1  package waf
     2  
     3  import (
     4  	"fmt"
     5  	"github.com/TeaOSLab/EdgeNode/internal/waf/requests"
     6  )
     7  
     8  // rule group
     9  type RuleGroup struct {
    10  	Id          int64      `yaml:"id" json:"id"`
    11  	IsOn        bool       `yaml:"isOn" json:"isOn"`
    12  	Name        string     `yaml:"name" json:"name"` // such as SQL Injection
    13  	Description string     `yaml:"description" json:"description"`
    14  	Code        string     `yaml:"code" json:"code"` // identify the group
    15  	RuleSets    []*RuleSet `yaml:"ruleSets" json:"ruleSets"`
    16  	IsInbound   bool       `yaml:"isInbound" json:"isInbound"`
    17  
    18  	hasRuleSets bool
    19  }
    20  
    21  func NewRuleGroup() *RuleGroup {
    22  	return &RuleGroup{
    23  		IsOn: true,
    24  	}
    25  }
    26  
    27  func (this *RuleGroup) Init(waf *WAF) error {
    28  	this.hasRuleSets = len(this.RuleSets) > 0
    29  
    30  	if this.hasRuleSets {
    31  		for _, set := range this.RuleSets {
    32  			err := set.Init(waf)
    33  			if err != nil {
    34  				return fmt.Errorf("init set '%d' failed: %w", set.Id, err)
    35  			}
    36  		}
    37  	}
    38  	return nil
    39  }
    40  
    41  func (this *RuleGroup) AddRuleSet(ruleSet *RuleSet) {
    42  	this.RuleSets = append(this.RuleSets, ruleSet)
    43  }
    44  
    45  func (this *RuleGroup) FindRuleSet(id int64) *RuleSet {
    46  	for _, ruleSet := range this.RuleSets {
    47  		if ruleSet.Id == id {
    48  			return ruleSet
    49  		}
    50  	}
    51  	return nil
    52  }
    53  
    54  func (this *RuleGroup) FindRuleSetWithCode(code string) *RuleSet {
    55  	if len(code) == 0 {
    56  		return nil
    57  	}
    58  	for _, ruleSet := range this.RuleSets {
    59  		if ruleSet.Code == code {
    60  			return ruleSet
    61  		}
    62  	}
    63  	return nil
    64  }
    65  
    66  func (this *RuleGroup) RemoveRuleSet(id int64) {
    67  	result := []*RuleSet{}
    68  	for _, ruleSet := range this.RuleSets {
    69  		if ruleSet.Id == id {
    70  			continue
    71  		}
    72  		result = append(result, ruleSet)
    73  	}
    74  	this.RuleSets = result
    75  }
    76  
    77  func (this *RuleGroup) MatchRequest(req requests.Request) (b bool, hasRequestBody bool, resultSet *RuleSet, err error) {
    78  	if !this.hasRuleSets {
    79  		return
    80  	}
    81  	for _, set := range this.RuleSets {
    82  		if !set.IsOn {
    83  			continue
    84  		}
    85  		b, hasRequestBody, err = set.MatchRequest(req)
    86  		if err != nil {
    87  			return false, hasRequestBody, nil, err
    88  		}
    89  		if b {
    90  			return true, hasRequestBody, set, nil
    91  		}
    92  	}
    93  	return
    94  }
    95  
    96  func (this *RuleGroup) MatchResponse(req requests.Request, resp *requests.Response) (b bool, hasRequestBody bool, resultSet *RuleSet, err error) {
    97  	if !this.hasRuleSets {
    98  		return
    99  	}
   100  	for _, set := range this.RuleSets {
   101  		if !set.IsOn {
   102  			continue
   103  		}
   104  		b, hasRequestBody, err = set.MatchResponse(req, resp)
   105  		if err != nil {
   106  			return false, hasRequestBody, nil, err
   107  		}
   108  		if b {
   109  			return true, hasRequestBody, set, nil
   110  		}
   111  	}
   112  	return
   113  }
   114  
   115  func (this *RuleGroup) MoveRuleSet(fromIndex int, toIndex int) {
   116  	if fromIndex < 0 || fromIndex >= len(this.RuleSets) {
   117  		return
   118  	}
   119  	if toIndex < 0 || toIndex >= len(this.RuleSets) {
   120  		return
   121  	}
   122  	if fromIndex == toIndex {
   123  		return
   124  	}
   125  
   126  	location := this.RuleSets[fromIndex]
   127  	result := []*RuleSet{}
   128  	for i := 0; i < len(this.RuleSets); i++ {
   129  		if i == fromIndex {
   130  			continue
   131  		}
   132  		if fromIndex > toIndex && i == toIndex {
   133  			result = append(result, location)
   134  		}
   135  		result = append(result, this.RuleSets[i])
   136  		if fromIndex < toIndex && i == toIndex {
   137  			result = append(result, location)
   138  		}
   139  	}
   140  
   141  	this.RuleSets = result
   142  }