From bd4bbb01c104cd3b50a5c78f89b66e1c3a551c22 Mon Sep 17 00:00:00 2001 From: Roi Feng <37480123+Rayzggz@users.noreply.github.com> Date: Tue, 27 May 2025 13:33:28 -0400 Subject: [PATCH] feat: Use one yml file to config --- config_example/rules/CAPTCHA.yml | 4 - config_example/rules/HTTPFlood.yml | 4 - config_example/rules/Server.yml | 17 +++++ config_example/rules/VerifyBot.yml | 6 -- internal/config/config.go | 115 +++++++++++------------------ 5 files changed, 62 insertions(+), 84 deletions(-) delete mode 100644 config_example/rules/CAPTCHA.yml delete mode 100644 config_example/rules/HTTPFlood.yml create mode 100644 config_example/rules/Server.yml delete mode 100644 config_example/rules/VerifyBot.yml diff --git a/config_example/rules/CAPTCHA.yml b/config_example/rules/CAPTCHA.yml deleted file mode 100644 index d3bbb2c..0000000 --- a/config_example/rules/CAPTCHA.yml +++ /dev/null @@ -1,4 +0,0 @@ -secret_key: "0378b0f84c4310279918d71a5647ba5d" -captcha_validate_time: 600 -captcha_challenge_session_timeout: 120 -hcaptcha_secret: "" \ No newline at end of file diff --git a/config_example/rules/HTTPFlood.yml b/config_example/rules/HTTPFlood.yml deleted file mode 100644 index ed4f43b..0000000 --- a/config_example/rules/HTTPFlood.yml +++ /dev/null @@ -1,4 +0,0 @@ -HTTPFloodSpeedLimit: - - "150/10s" -HTTPFloodSameURILimit: - - "50/10s" \ No newline at end of file diff --git a/config_example/rules/Server.yml b/config_example/rules/Server.yml new file mode 100644 index 0000000..dae22ee --- /dev/null +++ b/config_example/rules/Server.yml @@ -0,0 +1,17 @@ +CAPTCHA: + secret_key: "0378b0f84c4310279918d71a5647ba5d" + captcha_validate_time: 600 + captcha_challenge_session_timeout: 120 + hcaptcha_secret: "" +HTTPFlood: + HTTPFloodSpeedLimit: + - "150/10s" + HTTPFloodSameURILimit: + - "50/10s" +VerifyBot: + verify_google_bot: true + verify_bing_bot: true + verify_baidu_bot: true + verify_yandex_bot: true + verify_sogou_bot: true + verify_apple_bot: true \ No newline at end of file diff --git a/config_example/rules/VerifyBot.yml b/config_example/rules/VerifyBot.yml deleted file mode 100644 index d20da9f..0000000 --- a/config_example/rules/VerifyBot.yml +++ /dev/null @@ -1,6 +0,0 @@ -verify_google_bot: true -verify_bing_bot: true -verify_baidu_bot: true -verify_yandex_bot: true -verify_sogou_bot: true -verify_apple_bot: true \ No newline at end of file diff --git a/internal/config/config.go b/internal/config/config.go index fd49fcb..d86e2a1 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -2,6 +2,7 @@ package config import ( "bufio" + "fmt" "gopkg.in/yaml.v3" "net" "os" @@ -74,6 +75,18 @@ type RuleSet struct { HTTPFloodRule *dataType.HTTPFloodRule } +// ruleSetWrapper +type ruleSetWrapper struct { + CAPTCHARule *dataType.CaptchaRule `yaml:"CAPTCHA"` + VerifyBotRule *dataType.VerifyBotRule `yaml:"VerifyBot"` + HTTPFloodRule httpFloodRuleWrapper `yaml:"HTTPFlood"` +} + +type httpFloodRuleWrapper struct { + HTTPFloodSpeedLimit []string `yaml:"HTTPFloodSpeedLimit"` + HTTPFloodSameURILimit []string `yaml:"HTTPFloodSameURILimit"` +} + // LoadRules Load all rules from the specified path func LoadRules(rulePath string) (*RuleSet, error) { rs := RuleSet{ @@ -110,50 +123,52 @@ func LoadRules(rulePath string) (*RuleSet, error) { return nil, err } - // Load CAPTCHA Rule - captchaFile := rulePath + "/CAPTCHA.yml" - if err := loadCAPTCHARule(captchaFile, rs.CAPTCHARule); err != nil { - return nil, err - } - - // Load Verify Bot Rule - verifyBotFile := rulePath + "/VerifyBot.yml" - if err := loadVerifyBotRule(verifyBotFile, rs.VerifyBotRule); err != nil { - return nil, err - } - - // Load HTTP Flood Rule - httpFloodFile := rulePath + "/HTTPFlood.yml" - if err := loadHTTPFloodRule(httpFloodFile, rs.HTTPFloodRule); err != nil { - return nil, err + YAMLFile := filepath.Join(rulePath, "rules.yml") + set, err := loadServerRules(YAMLFile, rs) + if err != nil { + return set, err } return &rs, nil } -func loadCAPTCHARule(file string, rule *dataType.CaptchaRule) error { - data, err := os.ReadFile(file) +func loadServerRules(YAMLFile string, rs RuleSet) (*RuleSet, error) { + yamlData, err := os.ReadFile(YAMLFile) if err != nil { - return err + if os.IsNotExist(err) { + return nil, fmt.Errorf("[ERROR] rules file %s does not exist: %w", YAMLFile, err) + } else { + return nil, fmt.Errorf("[ERROR] failed to read rules file %s: %w", YAMLFile, err) + } } - if err := yaml.Unmarshal(data, &rule); err != nil { - return err + var wrapper ruleSetWrapper + if err := yaml.Unmarshal(yamlData, &wrapper); err != nil { + return nil, fmt.Errorf("[ERROR] failed to parse rules file %s: %w", YAMLFile, err) } - return nil + *rs.CAPTCHARule = *wrapper.CAPTCHARule + *rs.VerifyBotRule = *wrapper.VerifyBotRule -} + rs.HTTPFloodRule.HTTPFloodSpeedLimit = make(map[int64]int64) + rs.HTTPFloodRule.HTTPFloodSameURILimit = make(map[int64]int64) -func loadVerifyBotRule(file string, rule *dataType.VerifyBotRule) error { - data, err := os.ReadFile(file) - if err != nil { - return err + for _, s := range wrapper.HTTPFloodRule.HTTPFloodSpeedLimit { + limit, seconds, err := utils.ParseRate(s) + if err != nil { + return nil, err + } + rs.HTTPFloodRule.HTTPFloodSpeedLimit[seconds] = limit } - if err := yaml.Unmarshal(data, &rule); err != nil { - return err + + for _, s := range wrapper.HTTPFloodRule.HTTPFloodSameURILimit { + limit, seconds, err := utils.ParseRate(s) + if err != nil { + return nil, err + } + rs.HTTPFloodRule.HTTPFloodSameURILimit[seconds] = limit } - return nil + return nil, nil } // loadIPRules read the IP rule file and insert the rules into the trie @@ -231,43 +246,3 @@ func loadURLRules(filePath string, list *dataType.URLRuleList) error { return scanner.Err() } - -func loadHTTPFloodRule(file string, rule *dataType.HTTPFloodRule) error { - data, err := os.ReadFile(file) - if err != nil { - return err - } - - type httpFloodRuleYAML struct { - HTTPFloodSpeedLimit []string `yaml:"HTTPFloodSpeedLimit"` - HTTPFloodSameURILimit []string `yaml:"HTTPFloodSameURILimit"` - } - - var ymlRule httpFloodRuleYAML - err = yaml.Unmarshal(data, &ymlRule) - if err != nil { - return err - } - - rule.HTTPFloodSpeedLimit = make(map[int64]int64) - rule.HTTPFloodSameURILimit = make(map[int64]int64) - - for _, s := range ymlRule.HTTPFloodSpeedLimit { - limit, seconds, err := utils.ParseRate(s) - if err != nil { - return err - } - rule.HTTPFloodSpeedLimit[seconds] = limit - } - - for _, s := range ymlRule.HTTPFloodSameURILimit { - limit, seconds, err := utils.ParseRate(s) - if err != nil { - return err - } - rule.HTTPFloodSameURILimit[seconds] = limit - } - - return nil - -}