From 0b33765b4f64d080068e7ca9145a191d1ffc83e7 Mon Sep 17 00:00:00 2001 From: Kevin Pham Date: Tue, 5 Dec 2023 21:03:27 -0600 Subject: [PATCH] add support for rulesets in JSON format for frontend use --- handlers/proxy.go | 6 -- internal/cli/ruleset_merge.go | 2 +- proxychain/ruleset/rule.go | 7 +- proxychain/ruleset/ruleset.go | 109 ++++++++++++++++++++++++++--- proxychain/ruleset/ruleset_test.go | 45 ++++++++++-- 5 files changed, 146 insertions(+), 23 deletions(-) diff --git a/handlers/proxy.go b/handlers/proxy.go index 5377a89..2a9e40f 100644 --- a/handlers/proxy.go +++ b/handlers/proxy.go @@ -1,7 +1,6 @@ package handlers import ( - "fmt" "ladder/proxychain" rx "ladder/proxychain/requestmodifiers" tx "ladder/proxychain/responsemodifiers" @@ -61,12 +60,7 @@ func NewProxySiteHandler(opts *ProxyOptions) fiber.Handler { // load ruleset rule, exists := rs.GetRule(proxychain.Request.URL) - fmt.Println("============") - fmt.Println(proxychain.Request.URL) - fmt.Println(rs) - fmt.Println("============") if exists { - fmt.Println("===========EXISTS=") proxychain.AddOnceRequestModifications(rule.RequestModifications...) proxychain.AddOnceResponseModifications(rule.ResponseModifications...) } diff --git a/internal/cli/ruleset_merge.go b/internal/cli/ruleset_merge.go index 98524b4..2a7a75d 100644 --- a/internal/cli/ruleset_merge.go +++ b/internal/cli/ruleset_merge.go @@ -52,7 +52,7 @@ func HandleRulesetMerge(rulesetPath string, mergeRulesets bool, output *os.File) // Returns: // - An error if YAML conversion or file writing fails, otherwise nil. func yamlMerge(rs ruleset_v2.Ruleset, output io.Writer) error { - yaml, err := rs.Yaml() + yaml, err := rs.YAML() if err != nil { return err } diff --git a/proxychain/ruleset/rule.go b/proxychain/ruleset/rule.go index ccc8bc6..0e7568b 100644 --- a/proxychain/ruleset/rule.go +++ b/proxychain/ruleset/rule.go @@ -1,6 +1,7 @@ package ruleset_v2 import ( + "bytes" "encoding/json" "fmt" @@ -134,5 +135,9 @@ func (rule *Rule) MarshalYAML() (interface{}, error) { ResponseModifications: rule._rsms, } - return yaml.Marshal(aux) + var b bytes.Buffer + y := yaml.NewEncoder(&b) + y.SetIndent(2) + err := y.Encode(aux) + return b.String(), err } diff --git a/proxychain/ruleset/ruleset.go b/proxychain/ruleset/ruleset.go index 68056d2..cadb8fe 100644 --- a/proxychain/ruleset/ruleset.go +++ b/proxychain/ruleset/ruleset.go @@ -1,6 +1,7 @@ package ruleset_v2 import ( + "bytes" "compress/gzip" "errors" "fmt" @@ -12,13 +13,14 @@ import ( "path/filepath" "strings" + "encoding/json" "gopkg.in/yaml.v3" - //"encoding/json" ) type IRuleset interface { HasRule(url url.URL) bool GetRule(url url.URL) (rule Rule, exists bool) + YAML() (string, error) } type Ruleset struct { @@ -63,6 +65,7 @@ func (rs *Ruleset) MarshalYAML() (interface{}, error) { RequestModifications []_rqm `yaml:"requestmodifications"` ResponseModifications []_rsm `yaml:"responsemodifications"` } + type Aux struct { Rules []AuxRule `yaml:"rules"` } @@ -78,10 +81,68 @@ func (rs *Ruleset) MarshalYAML() (interface{}, error) { aux.Rules = append(aux.Rules, auxRule) } - out, err := yaml.Marshal(&aux) - return out, err + var b bytes.Buffer + y := yaml.NewEncoder(&b) + y.SetIndent(2) + err := y.Encode(&aux) + + return b.String(), err } +// ========================================================== + +func (rs *Ruleset) UnmarshalJSON(data []byte) error { + type AuxRuleset struct { + Rules []Rule `json:"rules"` + } + ar := &AuxRuleset{} + + if err := json.Unmarshal(data, ar); err != nil { + return err + } + + rs._rulemap = make(map[string]*Rule) + rs.Rules = ar.Rules + + for i, rule := range rs.Rules { + rulePtr := &rs.Rules[i] + for _, domain := range rule.Domains { + rs._rulemap[domain] = rulePtr + if !strings.HasPrefix(domain, "www.") { + rs._rulemap["www."+domain] = rulePtr + } + } + } + + return nil +} + +func (rs *Ruleset) MarshalJSON() ([]byte, error) { + type AuxRule struct { + Domains []string `json:"domains"` + RequestModifications []_rqm `json:"requestmodifications"` + ResponseModifications []_rsm `json:"responsemodifications"` + } + + type Aux struct { + Rules []AuxRule `json:"rules"` + } + + aux := Aux{} + for _, rule := range rs.Rules { + auxRule := AuxRule{ + Domains: rule.Domains, + RequestModifications: rule._rqms, + ResponseModifications: rule._rsms, + } + aux.Rules = append(aux.Rules, auxRule) + } + + return json.Marshal(aux) +} + +// =========================================================== + func (rs Ruleset) GetRule(url *url.URL) (rule *Rule, exists bool) { rule, exists = rs._rulemap[url.Hostname()] return rule, exists @@ -165,6 +226,9 @@ func (rs *Ruleset) loadRulesFromLocalDir(path string) error { // create a map of pointers to rules loaded above based on domain string keys // this way we don't have two copies of the rule in ruleset + if rs._rulemap == nil { + rs._rulemap = make(map[string]*Rule) + } for i, rule := range rs.Rules { rulePtr := &rs.Rules[i] for _, domain := range rule.Domains { @@ -185,17 +249,22 @@ func (rs *Ruleset) loadRulesFromLocalDir(path string) error { // loadRulesFromLocalFile loads rules from a local YAML file specified by the path. // Returns an error if the file cannot be read or if there's a syntax error in the YAML. func (rs *Ruleset) loadRulesFromLocalFile(path string) error { - yamlFile, err := os.ReadFile(path) + file, err := os.ReadFile(path) if err != nil { e := fmt.Errorf("failed to read rules from local file: '%s'", path) return errors.Join(e, err) } - err = yaml.Unmarshal(yamlFile, rs) + isJSON := strings.HasSuffix(path, ".json") + if isJSON { + err = json.Unmarshal(file, rs) + } else { + err = yaml.Unmarshal(file, rs) + } if err != nil { e := fmt.Errorf("failed to load rules from local file, possible syntax error in '%s' - %s", path, err) - debugPrintRule(string(yamlFile), e) + debugPrintRule(string(file), e) return e } @@ -232,7 +301,12 @@ func (rs *Ruleset) loadRulesFromRemoteFile(rulesURL string) error { reader = resp.Body } - err = yaml.NewDecoder(reader).Decode(&rs) + isJSON := strings.HasSuffix(rulesURL, ".json") || resp.Header.Get("content-type") == "application/json" + if isJSON { + err = json.NewDecoder(reader).Decode(&rs) + } else { + err = yaml.NewDecoder(reader).Decode(&rs) + } if err != nil { return fmt.Errorf("failed to load rules from remote url '%s' with status code '%s' and possible syntax error - %s", rulesURL, resp.Status, err) @@ -243,14 +317,29 @@ func (rs *Ruleset) loadRulesFromRemoteFile(rulesURL string) error { // ================= utility methods ========================== -// Yaml returns the ruleset as a Yaml string -func (rs *Ruleset) Yaml() (string, error) { +// YAML returns the ruleset as a Yaml string +func (rs *Ruleset) YAML() (string, error) { y, err := yaml.Marshal(rs) if err != nil { return "", err } - return string(y), nil + // for some reason, MarshalYAML seems to turn everything into a string block + // this is a workaround. Don't call yaml.Marshal directly, instead call this helper method. + x := strings.ReplaceAll(string(y), "\n ", "\n") + x = strings.Replace(x, "|\n", "", 1) + fmt.Println(x) + + return x, nil +} + +// YAML returns the ruleset as a JSON string +func (rs *Ruleset) JSON() (string, error) { + j, err := json.Marshal(rs) + if err != nil { + return "", err + } + return string(j), nil } // Domains extracts and returns a slice of all domains present in the RuleSet. diff --git a/proxychain/ruleset/ruleset_test.go b/proxychain/ruleset/ruleset_test.go index c7311bd..454edd0 100644 --- a/proxychain/ruleset/ruleset_test.go +++ b/proxychain/ruleset/ruleset_test.go @@ -1,20 +1,24 @@ package ruleset_v2 import ( + "encoding/json" "fmt" "net/url" "os" "path/filepath" + "strings" + + //"strings" "testing" "time" "github.com/gofiber/fiber/v2" "github.com/stretchr/testify/assert" + //"gopkg.in/yaml.v3" ) var ( - validYAML = ` -rules: + validYAML = `rules: - domains: - example.com - www.example.com @@ -51,6 +55,32 @@ rules: ` ) +func TestYAMLUnmarshal(t *testing.T) { + rs, err := loadRuleFromString(validYAML) + fmt.Println(validYAML) + assert.NoError(t, err, "expected no error loading valid yml") + yml, err := rs.YAML() + assert.NoError(t, err, "expected no error marshalling ruleset") + + rs2, err := loadRuleFromString(yml) + assert.NoError(t, err, "expected no error loading yaml marshalled -> unmarshalled -> marshalled ruleset") + assert.Equal(t, 1, rs2.Count(), "expected one rule to be returned after marshalled -> unmarshalled -> marshalled ruleset") +} + +func TestJSONUnmarshal(t *testing.T) { + rs, err := loadRuleFromString(validYAML) + assert.NoError(t, err, "expected no error loading valid yml") + j, err := json.Marshal(&rs) + assert.NoError(t, err, "expected no error marshalling ruleset to json") + + fmt.Println(string(j)) + + rs2, err := loadRuleFromString(string(j)) + assert.NoError(t, err, "expected no error loading JSON marshalled -> unmarshalled -> marshalled ruleset") + + assert.Equal(t, 1, rs2.Count(), "expected one rule to be returned after JSON marshalled -> unmarshalled -> marshalled ruleset") +} + func TestLoadRulesFromRemoteFile(t *testing.T) { app := fiber.New() defer app.Shutdown() @@ -103,13 +133,18 @@ func TestLoadRulesFromRemoteFile(t *testing.T) { } } -func loadRuleFromString(yaml string) (Ruleset, error) { +func loadRuleFromString(yamlOrJSON string) (Ruleset, error) { // Create a temporary file and load it - tmpFile, _ := os.CreateTemp("", "ruleset*.yaml") + var tmpFile *os.File + if strings.HasPrefix(yamlOrJSON, "{") { + tmpFile, _ = os.CreateTemp("", "ruleset*.json") + } else { + tmpFile, _ = os.CreateTemp("", "ruleset*.yaml") + } defer os.Remove(tmpFile.Name()) - tmpFile.WriteString(yaml) + tmpFile.WriteString(yamlOrJSON) rs := Ruleset{ _rulemap: map[string]*Rule{},