DWS/hadrian
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

idiomatize (?) ruleset package and lint

Browse Source
This commit is contained in:
Damian Bednarczyk
2023-11-22 21:26:44 -06:00
parent 394eaf9805
commit dc69af9f38
7 changed files with 124 additions and 57 deletions
Download Patch File Download Diff File Expand all files Collapse all files

52
pkg/ruleset/ruleset.go
View File

@@ -1,6 +1,7 @@
package ruleset
import (
"compress/gzip"
"errors"
"fmt"
"io"
@@ -11,8 +12,6 @@ import (
"regexp"
"strings"
"compress/gzip"
"gopkg.in/yaml.v3"
)
@@ -41,7 +40,7 @@ type Rule struct {
GoogleCache bool `yaml:"googleCache,omitempty"`
RegexRules []Regex `yaml:"regexRules,omitempty"`
UrlMods struct {
URLMods struct {
Domain []Regex `yaml:"domain,omitempty"`
Path []Regex `yaml:"path,omitempty"`
Query []KV `yaml:"query,omitempty"`
@@ -55,6 +54,8 @@ type Rule struct {
} `yaml:"injections,omitempty"`
}
var remoteRegex = regexp.MustCompile(`^https?:\/\/(www\.)?[-a-zA-Z0-9@:%._\+~#=]{1,256}\.[a-zA-Z0-9()]{1,6}\b([-a-zA-Z0-9()!@:%_\+.~#?&\/\/=]*)`)
// NewRulesetFromEnv creates a new RuleSet based on the RULESET environment variable.
// It logs a warning and returns an empty RuleSet if the RULESET environment variable is not set.
// If the RULESET is set but the rules cannot be loaded, it panics.
@@ -64,10 +65,12 @@ func NewRulesetFromEnv() RuleSet {
log.Printf("WARN: No ruleset specified. Set the `RULESET` environment variable to load one for a better success rate.")
return RuleSet{}
}
ruleSet, err := NewRuleset(rulesPath)
if err != nil {
log.Println(err)
}
return ruleSet
}
@@ -75,16 +78,17 @@ func NewRulesetFromEnv() RuleSet {
// It supports loading rules from both local file paths and remote URLs.
// Returns a RuleSet and an error if any issues occur during loading.
func NewRuleset(rulePaths string) (RuleSet, error) {
ruleSet := RuleSet{}
errs := []error{}
var ruleSet RuleSet
var errs []error
rp := strings.Split(rulePaths, ";")
var remoteRegex = regexp.MustCompile(`^https?:\/\/(www\.)?[-a-zA-Z0-9@:%._\+~#=]{1,256}\.[a-zA-Z0-9()]{1,6}\b([-a-zA-Z0-9()!@:%_\+.~#?&\/\/=]*)`)
for _, rule := range rp {
rulePath := strings.Trim(rule, " ")
var err error
rulePath := strings.Trim(rule, " ")
isRemote := remoteRegex.MatchString(rulePath)
if isRemote {
err = ruleSet.loadRulesFromRemoteFile(rulePath)
} else {
@@ -94,6 +98,7 @@ func NewRuleset(rulePaths string) (RuleSet, error) {
if err != nil {
e := fmt.Errorf("WARN: failed to load ruleset from '%s'", rulePath)
errs = append(errs, errors.Join(e, err))
continue
}
}
@@ -101,6 +106,7 @@ func NewRuleset(rulePaths string) (RuleSet, error) {
if len(errs) != 0 {
e := fmt.Errorf("WARN: failed to load %d rulesets", len(rp))
errs = append(errs, e)
// panic if the user specified a local ruleset, but it wasn't found on disk
// don't fail silently
for _, err := range errs {
@@ -109,10 +115,13 @@ func NewRuleset(rulePaths string) (RuleSet, error) {
panic(errors.Join(e, err))
}
}
// else, bubble up any errors, such as syntax or remote host issues
return ruleSet, errors.Join(errs...)
}
ruleSet.PrintStats()
return ruleSet, nil
}
@@ -146,13 +155,16 @@ func (rs *RuleSet) loadRulesFromLocalDir(path string) error {
log.Printf("WARN: failed to load directory ruleset '%s': %s, skipping", path, err)
return nil
}
log.Printf("INFO: loaded ruleset %s\n", path)
return nil
})
if err != nil {
return err
}
return nil
}
@@ -167,42 +179,51 @@ func (rs *RuleSet) loadRulesFromLocalFile(path string) error {
var r RuleSet
err = yaml.Unmarshal(yamlFile, &r)
if err != nil {
e := fmt.Errorf("failed to load rules from local file, possible syntax error in '%s'", path)
ee := errors.Join(e, err)
if _, ok := os.LookupEnv("DEBUG"); ok {
debugPrintRule(string(yamlFile), ee)
}
return ee
}
*rs = append(*rs, r...)
return nil
}
// loadRulesFromRemoteFile loads rules from a remote URL.
// It supports plain and gzip compressed content.
// Returns an error if there's an issue accessing the URL or if there's a syntax error in the YAML.
func (rs *RuleSet) loadRulesFromRemoteFile(rulesUrl string) error {
func (rs *RuleSet) loadRulesFromRemoteFile(rulesURL string) error {
var r RuleSet
resp, err := http.Get(rulesUrl)
resp, err := http.Get(rulesURL)
if err != nil {
e := fmt.Errorf("failed to load rules from remote url '%s'", rulesUrl)
e := fmt.Errorf("failed to load rules from remote url '%s'", rulesURL)
return errors.Join(e, err)
}
defer resp.Body.Close()
if resp.StatusCode >= 400 {
e := fmt.Errorf("failed to load rules from remote url (%s) on '%s'", resp.Status, rulesUrl)
e := fmt.Errorf("failed to load rules from remote url (%s) on '%s'", resp.Status, rulesURL)
return errors.Join(e, err)
}
var reader io.Reader
isGzip := strings.HasSuffix(rulesUrl, ".gz") || strings.HasSuffix(rulesUrl, ".gzip") || resp.Header.Get("content-encoding") == "gzip"
isGzip := strings.HasSuffix(rulesURL, ".gz") || strings.HasSuffix(rulesURL, ".gzip") || resp.Header.Get("content-encoding") == "gzip"
if isGzip {
reader, err = gzip.NewReader(resp.Body)
if err != nil {
return fmt.Errorf("failed to create gzip reader for URL '%s' with status code '%s': %w", rulesUrl, resp.Status, err)
return fmt.Errorf("failed to create gzip reader for URL '%s' with status code '%s': %w", rulesURL, resp.Status, err)
}
} else {
reader = resp.Body
@@ -211,12 +232,14 @@ func (rs *RuleSet) loadRulesFromRemoteFile(rulesUrl string) error {
err = yaml.NewDecoder(reader).Decode(&r)
if err != nil {
e := fmt.Errorf("failed to load rules from remote url '%s' with status code '%s' and possible syntax error", rulesUrl, resp.Status)
e := fmt.Errorf("failed to load rules from remote url '%s' with status code '%s' and possible syntax error", rulesURL, resp.Status)
ee := errors.Join(e, err)
return ee
}
*rs = append(*rs, r...)
return nil
}
@@ -228,6 +251,7 @@ func (rs *RuleSet) Yaml() (string, error) {
if err != nil {
return "", err
}
return string(y), nil
}

30
pkg/ruleset/ruleset_test.go
View File

@@ -33,6 +33,7 @@ func TestLoadRulesFromRemoteFile(t *testing.T) {
c.SendString(validYAML)
return nil
})
app.Get("/invalid-config.yml", func(c *fiber.Ctx) error {
c.SendString(invalidYAML)
return nil
@@ -40,10 +41,12 @@ func TestLoadRulesFromRemoteFile(t *testing.T) {
app.Get("/valid-config.gz", func(c *fiber.Ctx) error {
c.Set("Content-Type", "application/octet-stream")
rs, err := loadRuleFromString(validYAML)
if err != nil {
t.Errorf("failed to load valid yaml from string: %s", err.Error())
}
s, err := rs.GzipYaml()
if err != nil {
t.Errorf("failed to load gzip serialize yaml: %s", err.Error())
@@ -70,15 +73,18 @@ func TestLoadRulesFromRemoteFile(t *testing.T) {
if err != nil {
t.Errorf("failed to load plaintext ruleset from http server: %s", err.Error())
}
assert.Equal(t, rs[0].Domain, "example.com")
rs, err = NewRuleset("http://127.0.0.1:9999/valid-config.gz")
if err != nil {
t.Errorf("failed to load gzipped ruleset from http server: %s", err.Error())
}
assert.Equal(t, rs[0].Domain, "example.com")
os.Setenv("RULESET", "http://127.0.0.1:9999/valid-config.gz")
rs = NewRulesetFromEnv()
if !assert.Equal(t, rs[0].Domain, "example.com") {
t.Error("expected no errors loading ruleset from gzip url using environment variable, but got one")
@@ -88,10 +94,14 @@ func TestLoadRulesFromRemoteFile(t *testing.T) {
func loadRuleFromString(yaml string) (RuleSet, error) {
// Create a temporary file and load it
tmpFile, _ := os.CreateTemp("", "ruleset*.yaml")
defer os.Remove(tmpFile.Name())
tmpFile.WriteString(yaml)
rs := RuleSet{}
err := rs.loadRulesFromLocalFile(tmpFile.Name())
return rs, err
}
@@ -101,6 +111,7 @@ func TestLoadRulesFromLocalFile(t *testing.T) {
if err != nil {
t.Errorf("Failed to load rules from valid YAML: %s", err)
}
assert.Equal(t, rs[0].Domain, "example.com")
assert.Equal(t, rs[0].RegexRules[0].Match, "^http:")
assert.Equal(t, rs[0].RegexRules[0].Replace, "https:")
@@ -118,30 +129,39 @@ func TestLoadRulesFromLocalDir(t *testing.T) {
if err != nil {
t.Fatalf("Failed to create temporary directory: %s", err)
}
defer os.RemoveAll(baseDir)
// Create a nested subdirectory
nestedDir := filepath.Join(baseDir, "nested")
err = os.Mkdir(nestedDir, 0755)
err = os.Mkdir(nestedDir, 0o755)
if err != nil {
t.Fatalf("Failed to create nested directory: %s", err)
}
// Create a nested subdirectory
nestedTwiceDir := filepath.Join(nestedDir, "nestedTwice")
err = os.Mkdir(nestedTwiceDir, 0755)
err = os.Mkdir(nestedTwiceDir, 0o755)
if err != nil {
t.Fatalf("Failed to create twice-nested directory: %s", err)
}
testCases := []string{"test.yaml", "test2.yaml", "test-3.yaml", "test 4.yaml", "1987.test.yaml.yml", "foobar.example.com.yaml", "foobar.com.yml"}
for _, fileName := range testCases {
filePath := filepath.Join(nestedDir, "2x-"+fileName)
os.WriteFile(filePath, []byte(validYAML), 0644)
os.WriteFile(filePath, []byte(validYAML), 0o644)
filePath = filepath.Join(nestedDir, fileName)
os.WriteFile(filePath, []byte(validYAML), 0644)
os.WriteFile(filePath, []byte(validYAML), 0o644)
filePath = filepath.Join(baseDir, "base-"+fileName)
os.WriteFile(filePath, []byte(validYAML), 0644)
os.WriteFile(filePath, []byte(validYAML), 0o644)
}
rs := RuleSet{}
err = rs.loadRulesFromLocalDir(baseDir)
assert.NoError(t, err)
assert.Equal(t, rs.Count(), len(testCases)*3)