add support for rulesets in JSON format for frontend use

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...)
 		}

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
 	}

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
 }

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.

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{},