From bc715d910176c6390c6acbcde46f57861ca310bb Mon Sep 17 00:00:00 2001 From: Kevin Pham Date: Tue, 5 Dec 2023 01:51:30 -0600 Subject: [PATCH] implement proper marshaller/unmarshaller for rulesets in yaml format --- proxychain/codegen/README.md | 16 +++ proxychain/ruleset/rule.go | 87 +++++++++++++--- proxychain/ruleset/rule_test.go | 136 +++++++++++++++---------- proxychain/ruleset/rule_utils.go | 59 ----------- proxychain/ruleset/rule_utils_test.go | 138 -------------------------- 5 files changed, 170 insertions(+), 266 deletions(-) create mode 100644 proxychain/codegen/README.md delete mode 100644 proxychain/ruleset/rule_utils.go delete mode 100644 proxychain/ruleset/rule_utils_test.go diff --git a/proxychain/codegen/README.md b/proxychain/codegen/README.md new file mode 100644 index 0000000..2fbe814 --- /dev/null +++ b/proxychain/codegen/README.md @@ -0,0 +1,16 @@ +## TLDR +- If you create, delete or rename any request/response modifier, run `go run codegen.go`, so that ruleset unmarshaling will work properly. + +## Overview + +The `codegen.go` file is a utility for the rulesets that automatically generates Go code that maps functional options names found in response/request modifiers to corresponding factory functions. This generation is crucial for the serialization of rulesets from JSON or YAML into functional options suitable for use in proxychains. The tool processes Go files containing modifier functions and generates the necessary mappings. + +- The generated mappings will be written in `proxychain/ruleset/rule_reqmod_types.gen.go` and `proxychain/ruleset/rule_resmod_types.gen.go`. +- These files are used in UnmarshalJSON and UnmarshalYAML methods of the rule type, found in `proxychain/ruleset/rule.go` + + +## Usage +```sh +go run codegen.go +``` + diff --git a/proxychain/ruleset/rule.go b/proxychain/ruleset/rule.go index 343ea5a..0d15c69 100644 --- a/proxychain/ruleset/rule.go +++ b/proxychain/ruleset/rule.go @@ -3,8 +3,8 @@ package ruleset_v2 import ( "encoding/json" "fmt" + "gopkg.in/yaml.v3" "ladder/proxychain" - // _ "gopkg.in/yaml.v3" ) type Rule struct { @@ -17,22 +17,22 @@ type Rule struct { // internal represenation of ResponseModifications type _rsm struct { - Name string `json:"name"` - Params []string `json:"params"` + Name string `json:"name" yaml:"name"` + Params []string `json:"params" yaml:"params"` } // internal represenation of RequestModifications type _rqm struct { - Name string `json:"name"` - Params []string `json:"params"` + Name string `json:"name" yaml:"name"` + Params []string `json:"params" yaml:"params"` } // implement type encoding/json/Marshaler func (rule *Rule) UnmarshalJSON(data []byte) error { type Aux struct { Domains []string `json:"domains"` - RequestModifications []_rqm `json:"request_modifications"` - ResponseModifications []_rsm `json:"response_modifications"` + RequestModifications []_rqm `json:"requestmodifications"` + ResponseModifications []_rsm `json:"responsemodifications"` } aux := &Aux{} @@ -65,16 +65,75 @@ func (rule *Rule) UnmarshalJSON(data []byte) error { return nil } -func (r *Rule) MarshalJSON() ([]byte, error) { +func (rule *Rule) MarshalJSON() ([]byte, error) { aux := struct { Domains []string `json:"domains"` - RequestModifications []_rqm `json:"request_modifications"` - ResponseModifications []_rsm `json:"response_modifications"` + RequestModifications []_rqm `json:"requestmodifications"` + ResponseModifications []_rsm `json:"responsemodifications"` }{ - Domains: r.Domains, - RequestModifications: r._rqms, - ResponseModifications: r._rsms, + Domains: rule.Domains, + RequestModifications: rule._rqms, + ResponseModifications: rule._rsms, } - return json.Marshal(aux) + return json.MarshalIndent(aux, "", " ") +} + +// ============================================================ +// YAML + +// implement type yaml marshaller +func (rule *Rule) UnmarshalYAML(unmarshal func(interface{}) error) error { + + type Aux struct { + Domains []string `yaml:"domains"` + RequestModifications []_rqm `yaml:"requestmodifications"` + ResponseModifications []_rsm `yaml:"responsemodifications"` + } + + var aux Aux + if err := unmarshal(&aux); err != nil { + return err + } + + rule.Domains = aux.Domains + rule._rqms = aux.RequestModifications + rule._rsms = aux.ResponseModifications + + // convert requestModification function call string into actual functional option + for _, rqm := range aux.RequestModifications { + f, exists := rqmModMap[rqm.Name] + if !exists { + return fmt.Errorf("Rule::UnmarshalYAML => requestModifier '%s' does not exist, please check spelling", rqm.Name) + } + rule.RequestModifications = append(rule.RequestModifications, f(rqm.Params...)) + } + + // convert responseModification function call string into actual functional option + for _, rsm := range aux.ResponseModifications { + f, exists := rsmModMap[rsm.Name] + if !exists { + return fmt.Errorf("Rule::UnmarshalYAML => responseModifier '%s' does not exist, please check spelling", rsm.Name) + } + rule.ResponseModifications = append(rule.ResponseModifications, f(rsm.Params...)) + } + + return nil +} + +func (rule *Rule) MarshalYAML() (interface{}, error) { + + type Aux struct { + Domains []string `yaml:"domains"` + RequestModifications []_rqm `yaml:"requestmodifications"` + ResponseModifications []_rsm `yaml:"responsemodifications"` + } + + aux := &Aux{ + Domains: rule.Domains, + RequestModifications: rule._rqms, + ResponseModifications: rule._rsms, + } + + return yaml.Marshal(aux) } diff --git a/proxychain/ruleset/rule_test.go b/proxychain/ruleset/rule_test.go index 5207a6c..98f7bc9 100644 --- a/proxychain/ruleset/rule_test.go +++ b/proxychain/ruleset/rule_test.go @@ -3,53 +3,34 @@ package ruleset_v2 import ( "encoding/json" "fmt" - //yaml "gopkg.in/yaml.v3" + "gopkg.in/yaml.v3" "testing" ) -func TestRuleUnmarshalJSON(t *testing.T) { - ruleJSON := `{ - "domains": [ - "example.com", - "www.example.com" - ], - "response_modifications": [ - { - "name": "APIContent", - "params": [] - }, - { - "name": "SetContentSecurityPolicy", - "params": ["foobar"] - }, - { - "name": "SetIncomingCookie", - "params": ["authorization-bearer", "hunter2"] - } - ], - "request_modifications": [ - { - "name": "ForwardRequestHeaders", - "params": [] - } - ] -}` - - //fmt.Println(ruleJSON) +// unmarshalRule is a helper function to unmarshal a Rule from a JSON string. +func unmarshalRule(t *testing.T, ruleJSON string) *Rule { rule := &Rule{} err := json.Unmarshal([]byte(ruleJSON), rule) if err != nil { - t.Errorf("expected no error in Unmarshal, got '%s'", err) - return + t.Fatalf("expected no error in Unmarshal, got '%s'", err) } + return rule +} + +func TestRuleUnmarshalJSON(t *testing.T) { + ruleJSON := `{ + "domains": ["example.com", "www.example.com"], + "responsemodifications": [{"name": "APIContent", "params": []}, {"name": "SetContentSecurityPolicy", "params": ["foobar"]}, {"name": "SetIncomingCookie", "params": ["authorization-bearer", "hunter2"]}], + "requestmodifications": [{"name": "ForwardRequestHeaders", "params": []}] + }` + + rule := unmarshalRule(t, ruleJSON) if len(rule.Domains) != 2 { t.Errorf("expected number of domains to be 2") - 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)) @@ -57,8 +38,17 @@ func TestRuleUnmarshalJSON(t *testing.T) { if len(rule.RequestModifications) != 1 { t.Errorf("expected number of RequestModifications to be 1, got %d", len(rule.RequestModifications)) } +} + +func TestRuleMarshalJSON(t *testing.T) { + ruleJSON := `{ + "domains": ["example.com", "www.example.com"], + "responsemodifications": [{"name": "APIContent", "params": []}, {"name": "SetContentSecurityPolicy", "params": ["foobar"]}, {"name": "SetIncomingCookie", "params": ["authorization-bearer", "hunter2"]}], + "requestmodifications": [{"name": "ForwardRequestHeaders", "params": []}] + }` + + rule := unmarshalRule(t, ruleJSON) - // test marshal jsonRule, err := json.Marshal(rule) if err != nil { t.Errorf("expected no error marshalling rule to json, got '%s'", err.Error()) @@ -66,34 +56,45 @@ func TestRuleUnmarshalJSON(t *testing.T) { 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") -` +// =============================================== +// unmarshalYAMLRule is a helper function to unmarshal a Rule from a YAML string. +func unmarshalYAMLRule(t *testing.T, ruleYAML string) *Rule { rule := &Rule{} err := yaml.Unmarshal([]byte(ruleYAML), rule) if err != nil { - t.Errorf("expected no error in Unmarshal, got '%s'", err) - return + t.Fatalf("expected no error in Unmarshal, got '%s'", err) } + return rule +} + +func TestRuleUnmarshalYAML(t *testing.T) { + ruleYAML := ` +domains: +- example.com +- www.example.com +responsemodifications: +- name: APIContent + params: [] +- name: SetContentSecurityPolicy + params: + - foobar +- name: SetIncomingCookie + params: + - authorization-bearer + - hunter2 +requestmodifications: +- name: ForwardRequestHeaders + params: [] +` + + rule := unmarshalYAMLRule(t, ruleYAML) if len(rule.Domains) != 2 { - t.Errorf("expected number of domains to be 2, got %d", len(rule.Domains)) - return + t.Errorf("expected number of domains to be 2") } 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)) @@ -101,11 +102,36 @@ response_modifications: if len(rule.RequestModifications) != 1 { t.Errorf("expected number of RequestModifications to be 1, got %d", len(rule.RequestModifications)) } + fmt.Println(rule.RequestModifications[0].Name) +} + +func TestRuleMarshalYAML(t *testing.T) { + ruleYAML := ` +domains: +- example.com +- www.example.com +responsemodifications: +- name: APIContent + params: [] +- name: SetContentSecurityPolicy + params: + - foobar +- name: SetIncomingCookie + params: + - authorization-bearer + - hunter2 +requestmodifications: +- name: ForwardRequestHeaders + params: [] +` + + rule := unmarshalYAMLRule(t, ruleYAML) 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)) + if yamlRule == nil { + t.Errorf("expected marshalling rule to yaml to not be nil") + } } -*/ diff --git a/proxychain/ruleset/rule_utils.go b/proxychain/ruleset/rule_utils.go deleted file mode 100644 index 1a283ea..0000000 --- a/proxychain/ruleset/rule_utils.go +++ /dev/null @@ -1,59 +0,0 @@ -package ruleset_v2 - -import ( - "errors" - "strings" -) - -// parseFuncCall takes a string that look like foo("bar", "baz") and breaks it down -// into funcName = "foo" and params = []string{"bar", "baz"}] -func parseFuncCall(funcCall string) (funcName string, params []string, err error) { - // Splitting the input string into two parts: functionName and parameters - parts := strings.SplitN(funcCall, "(", 2) - if len(parts) != 2 { - return "", nil, errors.New("invalid function call format") - } - - // get function name - funcName = strings.TrimSpace(parts[0]) - - // Removing the closing parenthesis from the parameters part - paramsPart := strings.TrimSuffix(parts[1], ")") - if len(paramsPart) == 0 { - // No parameters - return funcName, []string{}, nil - } - - inQuote := false - inEscape := false - param := "" - for _, r := range paramsPart { - switch { - case inQuote && !inEscape && r == '\\': - inEscape = true - continue - case inEscape && inQuote && r == '"': - param += string(r) - inEscape = false - continue - case inEscape: - param += string(r) - inEscape = false - continue - case r == '"': - inQuote = !inQuote - if !inQuote { - params = append(params, param) - param = "" - } - continue - case !inQuote && r == ',': - continue - case inQuote: - param += string(r) - continue - } - } - - return funcName, params, nil -} diff --git a/proxychain/ruleset/rule_utils_test.go b/proxychain/ruleset/rule_utils_test.go deleted file mode 100644 index 55f23ed..0000000 --- a/proxychain/ruleset/rule_utils_test.go +++ /dev/null @@ -1,138 +0,0 @@ -package ruleset_v2 - -import ( - "errors" - "reflect" - "testing" -) - -func TestParseFuncCall(t *testing.T) { - testCases := []struct { - name string - input string - expected struct { - funcName string - params []string - err error - } - }{ - { - name: "Normal case, one param", - input: `one("baz")`, - expected: struct { - funcName string - params []string - err error - }{funcName: "one", params: []string{"baz"}, err: nil}, - }, - { - name: "Normal case, one param, extra space in function call", - input: `two("baz" )`, - expected: struct { - funcName string - params []string - err error - }{funcName: "two", params: []string{"baz"}, err: nil}, - }, - { - name: "Normal case, one param, extra space in param", - input: `three("baz ")`, - expected: struct { - funcName string - params []string - err error - }{funcName: "three", params: []string{"baz "}, err: nil}, - }, - { - name: "Space in front of function", - input: ` three("baz")`, - expected: struct { - funcName string - params []string - err error - }{funcName: "three", params: []string{"baz"}, err: nil}, - }, - { - name: "Normal case, two params", - input: `foobar("baz", "qux")`, - expected: struct { - funcName string - params []string - err error - }{funcName: "foobar", params: []string{"baz", "qux"}, err: nil}, - }, - { - name: "Normal case, two params, no spaces between param comma", - input: `foobar("baz","qux")`, - expected: struct { - funcName string - params []string - err error - }{funcName: "foobar", params: []string{"baz", "qux"}, err: nil}, - }, - { - name: "Escaped parenthesis", - input: `testFunc("hello\(world", "anotherParam")`, - expected: struct { - funcName string - params []string - err error - }{funcName: "testFunc", params: []string{`hello(world`, "anotherParam"}, err: nil}, - }, - { - name: "Escaped quote", - input: `testFunc("hello\"world", "anotherParam")`, - expected: struct { - funcName string - params []string - err error - }{funcName: "testFunc", params: []string{`hello"world`, "anotherParam"}, err: nil}, - }, - { - name: "Two Escaped quote", - input: `testFunc("hello: \"world\"", "anotherParam")`, - expected: struct { - funcName string - params []string - err error - }{funcName: "testFunc", params: []string{`hello: "world"`, "anotherParam"}, err: nil}, - }, - { - name: "No parameters", - input: `emptyFunc()`, - expected: struct { - funcName string - params []string - err error - }{funcName: "emptyFunc", params: []string{}, err: nil}, - }, - { - name: "Invalid format", - input: `invalidFunc`, - expected: struct { - funcName string - params []string - err error - }{funcName: "", params: nil, err: errors.New("invalid function call format")}, - }, - { - name: "Invalid format 2", - input: `invalidFunc "foo", "bar"`, - expected: struct { - funcName string - params []string - err error - }{funcName: "", params: nil, err: errors.New("invalid function call format")}, - }, - } - - for _, tc := range testCases { - t.Run(tc.name, func(t *testing.T) { - funcName, params, err := parseFuncCall(tc.input) - if funcName != tc.expected.funcName || !reflect.DeepEqual(params, tc.expected.params) || (err != nil && tc.expected.err != nil && err.Error() != tc.expected.err.Error()) { - //if funcName != tc.expected.funcName || (err != nil && tc.expected.err != nil && err.Error() != tc.expected.err.Error()) { - t.Errorf("Test %s failed: got (%s, %v, %v), want (%s, %v, %v)", tc.name, funcName, params, err, tc.expected.funcName, tc.expected.params, tc.expected.err) - } - }) - } -}