add YAML serialization for ruleset

go.mod

@@ -39,6 +39,7 @@ require (
 	golang.org/x/crypto v0.16.0 // indirect
 	golang.org/x/exp v0.0.0-20231127185646-65229373498e // indirect
 	golang.org/x/text v0.14.0 // indirect
+	gopkg.in/yaml.v2 v2.4.0 // indirect
 )
 
 require (

go.sum

@@ -153,6 +153,8 @@ golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8T
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/check.v1 v1.0.0-20200902074654-038fdea0a05b h1:QRR6H1YWRnHb4Y/HeNFCTJLFVxaq6wH4YuVdsUOr75U=
 gopkg.in/check.v1 v1.0.0-20200902074654-038fdea0a05b/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
+gopkg.in/yaml.v2 v2.4.0 h1:D8xgwECY7CYvx+Y2n4sBz93Jn9JRvxdiyyo8CTfuKaY=
+gopkg.in/yaml.v2 v2.4.0/go.mod h1:RDklbk79AGWmwhnvt/jBztapEOGDOx6ZbXqjP6csGnQ=
 gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
 gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
 gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=

proxychain/codegen/codegen.go

@@ -6,6 +6,7 @@ import (
 	"go/parser"
 	"go/token"
 	"io"
+	"io/fs"
 
 	//"io/fs"
 	"os"
@@ -42,7 +43,7 @@ func responseModCodeGen(dir string) (code string, err error) {
 
 	factoryMaps := []string{}
 	for _, file := range files {
-		if filepath.Ext(file.Name()) != ".go" {
+		if !shouldGenCodeFor(file) {
 			continue
 		}
 
@@ -114,7 +115,7 @@ func requestModCodeGen(dir string) (code string, err error) {
 
 	factoryMaps := []string{}
 	for _, file := range files {
-		if filepath.Ext(file.Name()) != ".go" {
+		if !shouldGenCodeFor(file) {
 			continue
 		}
 
@@ -158,6 +159,19 @@ func init() {
 	return code, nil
 }
 
+func shouldGenCodeFor(file fs.DirEntry) bool {
+	if file.IsDir() {
+		return false
+	}
+	if filepath.Ext(file.Name()) != ".go" {
+		return false
+	}
+	if strings.HasSuffix(file.Name(), "_test.go") {
+		return false
+	}
+	return true
+}
+
 func main() {
 	rqmCode, err := requestModCodeGen("../requestmodifiers/")
 	if err != nil {

proxychain/requestmodifiers/modify_domain_with_regex.go

@@ -1,13 +1,18 @@
 package requestmodifiers
 
 import (
+	"fmt"
 	"regexp"
 
 	"ladder/proxychain"
 )
 
-func ModifyDomainWithRegex(match regexp.Regexp, replacement string) proxychain.RequestModification {
+func ModifyDomainWithRegex(matchRegex string, replacement string) proxychain.RequestModification {
+	match, err := regexp.Compile(matchRegex)
 	return func(px *proxychain.ProxyChain) error {
+		if err != nil {
+			return fmt.Errorf("RequestModification :: ModifyDomainWithRegex error => invalid match regex: %s - %s", matchRegex, err.Error())
+		}
 		px.Request.URL.Host = match.ReplaceAllString(px.Request.URL.Host, replacement)
 		return nil
 	}

proxychain/requestmodifiers/modify_path_with_regex.go

@@ -1,13 +1,17 @@
 package requestmodifiers
 
 import (
-	"regexp"
+	"fmt"
-
 	"ladder/proxychain"
+	"regexp"
 )
 
-func ModifyPathWithRegex(match regexp.Regexp, replacement string) proxychain.RequestModification {
+func ModifyPathWithRegex(matchRegex string, replacement string) proxychain.RequestModification {
+	match, err := regexp.Compile(matchRegex)
 	return func(px *proxychain.ProxyChain) error {
+		if err != nil {
+			return fmt.Errorf("RequestModification :: ModifyPathWithRegex error => invalid match regex: %s - %s", matchRegex, err.Error())
+		}
 		px.Request.URL.Path = match.ReplaceAllString(px.Request.URL.Path, replacement)
 		return nil
 	}

proxychain/ruleset/rule.go

@@ -4,6 +4,11 @@ import (
 	"encoding/json"
 	"fmt"
 	"ladder/proxychain"
+	"reflect"
+	"runtime"
+	"strings"
+
+	_ "gopkg.in/yaml.v3"
 )
 
 type Rule struct {
@@ -29,19 +34,6 @@ func (rule *Rule) UnmarshalJSON(data []byte) error {
 	rule.Domains = aux.Domains
 
 	// convert requestModification function call string into actual functional option
-	for _, resModStr := range aux.RequestModifications {
-		name, params, err := parseFuncCall(resModStr)
-		if err != nil {
-			return fmt.Errorf("Rule::UnmarshalJSON invalid function call syntax => '%s'", err)
-		}
-		f, exists := rsmModMap[name]
-		if !exists {
-			return fmt.Errorf("Rule::UnmarshalJSON => responseModifer '%s' does not exist, please check spelling", err)
-		}
-		rule.ResponseModifications = append(rule.ResponseModifications, f(params...))
-	}
-
-	// convert responseModification function call string into actual functional option
 	for _, rqmModStr := range aux.RequestModifications {
 		name, params, err := parseFuncCall(rqmModStr)
 		if err != nil {
@@ -49,14 +41,152 @@ func (rule *Rule) UnmarshalJSON(data []byte) error {
 		}
 		f, exists := rqmModMap[name]
 		if !exists {
-			return fmt.Errorf("Rule::UnmarshalJSON => requestModifier '%s' does not exist, please check spelling", err)
+			return fmt.Errorf("Rule::UnmarshalJSON => requestModifier '%s' does not exist, please check spelling", name)
 		}
 		rule.RequestModifications = append(rule.RequestModifications, f(params...))
 	}
 
+	// convert responseModification function call string into actual functional option
+	for _, rsmModStr := range aux.ResponseModifications {
+		name, params, err := parseFuncCall(rsmModStr)
+		if err != nil {
+			return fmt.Errorf("Rule::UnmarshalJSON invalid function call syntax => '%s'", err)
+		}
+		f, exists := rsmModMap[name]
+		if !exists {
+			return fmt.Errorf("Rule::UnmarshalJSON => responseModifier '%s' does not exist, please check spelling", name)
+		}
+		rule.ResponseModifications = append(rule.ResponseModifications, f(params...))
+	}
+
 	return nil
 }
 
+// not fully possible to go from rule to JSON rule because
+// reflection cannot get the parameters of the functional options
+// of requestmodifiers and responsemodifiers
 func (r *Rule) MarshalJSON() ([]byte, error) {
-	return []byte{}, nil
+	type Aux struct {
+		Domains               []string `json:"domains"`
+		RequestModifications  []string `json:"request_modifications"`
+		ResponseModifications []string `json:"response_modifications"`
+	}
+	aux := &Aux{}
+	aux.Domains = r.Domains
+
+	for _, rqmMod := range r.RequestModifications {
+		fnName := getFunctionName(rqmMod)
+		aux.RequestModifications = append(aux.RequestModifications, fnName)
+	}
+
+	for _, rsmMod := range r.ResponseModifications {
+		fnName := getFunctionName(rsmMod)
+		aux.ResponseModifications = append(aux.ResponseModifications, fnName)
+	}
+
+	return json.Marshal(aux)
+}
+
+// getFunctionName returns the name of the function
+func getFunctionName(i interface{}) string {
+	// Get the value of the interface
+	val := reflect.ValueOf(i)
+
+	// Ensure it's a function
+	if val.Kind() != reflect.Func {
+		return "Not a function"
+	}
+
+	// Get the pointer to the function
+	ptr := val.Pointer()
+
+	// Get the function details from runtime
+	funcForPc := runtime.FuncForPC(ptr)
+
+	if funcForPc == nil {
+		return "Unknown"
+	}
+
+	// Return the name of the function
+	return extractShortName(funcForPc.Name())
+}
+
+// extractShortName extracts the short function name from the full name
+func extractShortName(fullName string) string {
+	parts := strings.Split(fullName, ".")
+	if len(parts) > 0 {
+		// Assuming the function name is always the second last part
+		return parts[len(parts)-2]
+	}
+	return ""
+}
+
+// == YAML
+// UnmarshalYAML implements the yaml.Unmarshaler interface for Rule
+func (rule *Rule) UnmarshalYAML(unmarshal func(interface{}) error) error {
+	type Aux struct {
+		Domains               []string `yaml:"domains"`
+		RequestModifications  []string `yaml:"request_modifications"`
+		ResponseModifications []string `yaml:"response_modifications"`
+	}
+
+	aux := &Aux{}
+	if err := unmarshal(aux); err != nil {
+		return err
+	}
+
+	rule.Domains = aux.Domains
+
+	// Process requestModifications
+	for _, rqmModStr := range aux.RequestModifications {
+		name, params, err := parseFuncCall(rqmModStr)
+		if err != nil {
+			return fmt.Errorf("Rule::UnmarshalYAML invalid function call syntax => '%s'", err)
+		}
+		f, exists := rqmModMap[name]
+		if !exists {
+			return fmt.Errorf("Rule::UnmarshalYAML => requestModifier '%s' does not exist, please check spelling", name)
+		}
+		rule.RequestModifications = append(rule.RequestModifications, f(params...))
+	}
+
+	// Process responseModifications
+	for _, rsmModStr := range aux.ResponseModifications {
+		name, params, err := parseFuncCall(rsmModStr)
+		if err != nil {
+			return fmt.Errorf("Rule::UnmarshalYAML invalid function call syntax => '%s'", err)
+		}
+		f, exists := rsmModMap[name]
+		if !exists {
+			return fmt.Errorf("Rule::UnmarshalYAML => responseModifier '%s' does not exist, please check spelling", name)
+		}
+		rule.ResponseModifications = append(rule.ResponseModifications, f(params...))
+	}
+
+	return nil
+}
+
+func (r *Rule) MarshalYAML() (interface{}, error) {
+	type Aux struct {
+		Domains               []string `yaml:"domains"`
+		RequestModifications  []string `yaml:"request_modifications"`
+		ResponseModifications []string `yaml:"response_modifications"`
+	}
+
+	aux := &Aux{
+		Domains: r.Domains,
+	}
+
+	for _, rqmMod := range r.RequestModifications {
+		// Assuming getFunctionName returns a string representation of the function
+		fnName := getFunctionName(rqmMod)
+		aux.RequestModifications = append(aux.RequestModifications, fnName)
+	}
+
+	for _, rsmMod := range r.ResponseModifications {
+		fnName := getFunctionName(rsmMod)
+		aux.ResponseModifications = append(aux.ResponseModifications, fnName)
+	}
+
+	return aux, nil
 }

proxychain/ruleset/rule_resmod_types.gen.go

@@ -20,14 +20,6 @@ func init() {
     return tx.APIContent()
   }
 
-  rsmModMap["TestCreateAPIErrReader"] = func(params ...string) proxychain.ResponseModification {
-    return tx.TestCreateAPIErrReader(params[0])
-  }
-
-  rsmModMap["TestCreateAPIErrReader2"] = func(params ...string) proxychain.ResponseModification {
-    return tx.TestCreateAPIErrReader2(params[0])
-  }
-
   rsmModMap["BlockElementRemoval"] = func(params ...string) proxychain.ResponseModification {
     return tx.BlockElementRemoval(params[0])
   }

proxychain/ruleset/rule_test.go

@@ -3,7 +3,7 @@ package ruleset_v2
 import (
 	"encoding/json"
 	"fmt"
-	//"io"
+	yaml "gopkg.in/yaml.v3"
 	"testing"
 )
 
@@ -13,12 +13,14 @@ func TestRuleUnmarshalJSON(t *testing.T) {
         "example.com",
         "www.example.com"
     ],
-    "response_modifiers": [
+    "request_modifications": [
-        "APIContent()",
+    	"SpoofUserAgent(\"googlebot\")"
-        "SetContentSecurityPolicy(\"foobar\")",
-        "SetIncomingCookie(\"authorization-bearer\", \"hunter2\")"
     ],
-    "response_modifiers": []
+    "response_modifications": [
+      "APIContent()",
+      "SetContentSecurityPolicy(\"foobar\")",
+      "SetIncomingCookie(\"authorization-bearer\", \"hunter2\")"
+    ]
 }`
 
 	//fmt.Println(ruleJSON)
@@ -37,9 +39,59 @@ func TestRuleUnmarshalJSON(t *testing.T) {
 		t.Errorf("expected domain to be example.com")
 		return
 	}
-	if len(rule.ResponseModifications) == 3 {
+	if len(rule.ResponseModifications) != 3 {
-		t.Errorf("expected number of ResponseModifications to be 3")
+		t.Errorf("expected number of ResponseModifications to be 3, got %d", len(rule.ResponseModifications))
+	}
+	if len(rule.RequestModifications) != 1 {
+		t.Errorf("expected number of RequestModifications to be 1, got %d", len(rule.RequestModifications))
 	}
-	fmt.Println(rule.ResponseModifications)
 
+	// test marshal
+	jsonRule, err := json.Marshal(rule)
+	if err != nil {
+		t.Errorf("expected no error marshalling rule to json, got '%s'", err.Error())
+	}
+	fmt.Println(string(jsonRule))
+}
+
+func TestRuleUnmarshalYAML(t *testing.T) {
+	ruleYAML := `
+domains:
+  - example.com
+  - www.example.com
+request_modifications:
+  - SpoofUserAgent("googlebot")
+response_modifications:
+  - APIContent()
+  - SetContentSecurityPolicy("foobar")
+  - SetIncomingCookie("authorization-bearer", "hunter2")
+`
+
+	rule := &Rule{}
+	err := yaml.Unmarshal([]byte(ruleYAML), rule)
+	if err != nil {
+		t.Errorf("expected no error in Unmarshal, got '%s'", err)
+		return
+	}
+
+	if len(rule.Domains) != 2 {
+		t.Errorf("expected number of domains to be 2, got %d", len(rule.Domains))
+		return
+	}
+	if !(rule.Domains[0] == "example.com" || rule.Domains[1] == "example.com") {
+		t.Errorf("expected domain to be example.com")
+		return
+	}
+	if len(rule.ResponseModifications) != 3 {
+		t.Errorf("expected number of ResponseModifications to be 3, got %d", len(rule.ResponseModifications))
+	}
+	if len(rule.RequestModifications) != 1 {
+		t.Errorf("expected number of RequestModifications to be 1, got %d", len(rule.RequestModifications))
+	}
+
+	yamlRule, err := yaml.Marshal(rule)
+	if err != nil {
+		t.Errorf("expected no error marshalling rule to yaml, got '%s'", err.Error())
+	}
+	fmt.Println(string(yamlRule))
 }