From 547cf61a7d878ce597fd9b4336331d532acef6d1 Mon Sep 17 00:00:00 2001 From: Kevin Pham Date: Sun, 3 Dec 2023 21:32:03 -0600 Subject: [PATCH] begin work on refactor of ruleset and functional options serializer for proxychain response/request modifiers --- pkg/ruleset_v2/rule.go | 49 +++++++++ pkg/ruleset_v2/rule_reqmod_types.gen.go | 1 + pkg/ruleset_v2/rule_resmod_types.gen.go | 24 +++++ pkg/ruleset_v2/rule_test.go | 45 ++++++++ pkg/ruleset_v2/rule_utils.go | 59 ++++++++++ pkg/ruleset_v2/rule_utils_test.go | 138 ++++++++++++++++++++++++ pkg/ruleset_v2/ruleset.go | 32 ++++++ 7 files changed, 348 insertions(+) create mode 100644 pkg/ruleset_v2/rule.go create mode 100644 pkg/ruleset_v2/rule_reqmod_types.gen.go create mode 100644 pkg/ruleset_v2/rule_resmod_types.gen.go create mode 100644 pkg/ruleset_v2/rule_test.go create mode 100644 pkg/ruleset_v2/rule_utils.go create mode 100644 pkg/ruleset_v2/rule_utils_test.go create mode 100644 pkg/ruleset_v2/ruleset.go diff --git a/pkg/ruleset_v2/rule.go b/pkg/ruleset_v2/rule.go new file mode 100644 index 0000000..63f48ba --- /dev/null +++ b/pkg/ruleset_v2/rule.go @@ -0,0 +1,49 @@ +package ruleset_v2 + +import ( + "encoding/json" + "fmt" + "ladder/proxychain" +) + +type Rule struct { + Domains []string + RequestModifications []proxychain.RequestModification + ResponseModifications []proxychain.ResponseModification +} + +// implement type encoding/json/Marshaler +func (rule *Rule) UnmarshalJSON(data []byte) error { + type Aux struct { + Domains []string `json:"domains"` + RequestModifications []string `json:"request_modifications"` + ResponseModifications []string `json:"response_modifications"` + } + + aux := &Aux{} + if err := json.Unmarshal(data, aux); err != nil { + return err + } + + //fmt.Println(aux.Domains) + 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 := resModMap[name] + if !exists { + return fmt.Errorf("Rule::UnmarshalJSON => requestModifier '%s' does not exist, please check spelling", err) + } + rule.ResponseModifications = append(rule.ResponseModifications, f(params...)) + } + + return nil +} + +func (r *Rule) MarshalJSON() ([]byte, error) { + return []byte{}, nil +} diff --git a/pkg/ruleset_v2/rule_reqmod_types.gen.go b/pkg/ruleset_v2/rule_reqmod_types.gen.go new file mode 100644 index 0000000..11b5a3b --- /dev/null +++ b/pkg/ruleset_v2/rule_reqmod_types.gen.go @@ -0,0 +1 @@ +package ruleset_v2 diff --git a/pkg/ruleset_v2/rule_resmod_types.gen.go b/pkg/ruleset_v2/rule_resmod_types.gen.go new file mode 100644 index 0000000..d578020 --- /dev/null +++ b/pkg/ruleset_v2/rule_resmod_types.gen.go @@ -0,0 +1,24 @@ +package ruleset_v2 + +import ( + "ladder/proxychain" + rx "ladder/proxychain/responsemodifiers" +) + +type ResponseModifierFactory func(params ...string) proxychain.ResponseModification + +var resModMap map[string]ResponseModifierFactory + +// TODO: create codegen using AST parsing of exported methods in ladder/proxychain/responsemodifiers/*.go +func init() { + resModMap = make(map[string]ResponseModifierFactory) + resModMap["APIContent"] = func(_ ...string) proxychain.ResponseModification { + return rx.APIContent() + } + resModMap["SetContentSecurityPolicy"] = func(params ...string) proxychain.ResponseModification { + return rx.SetContentSecurityPolicy(params[0]) + } + resModMap["SetIncomingCookie"] = func(params ...string) proxychain.ResponseModification { + return rx.SetIncomingCookie(params[0], params[1]) + } +} diff --git a/pkg/ruleset_v2/rule_test.go b/pkg/ruleset_v2/rule_test.go new file mode 100644 index 0000000..63b34a1 --- /dev/null +++ b/pkg/ruleset_v2/rule_test.go @@ -0,0 +1,45 @@ +package ruleset_v2 + +import ( + "encoding/json" + "fmt" + //"io" + "testing" +) + +func TestRuleUnmarshalJSON(t *testing.T) { + ruleJSON := `{ + "domains": [ + "example.com", + "www.example.com" + ], + "response_modifiers": [ + "APIContent()", + "SetContentSecurityPolicy(\"foobar\")", + "SetIncomingCookie(\"authorization-bearer\", \"hunter2\")" + ], + "response_modifiers": [] +}` + + //fmt.Println(ruleJSON) + rule := &Rule{} + err := json.Unmarshal([]byte(ruleJSON), 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") + 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") + } + fmt.Println(rule.ResponseModifications) + +} diff --git a/pkg/ruleset_v2/rule_utils.go b/pkg/ruleset_v2/rule_utils.go new file mode 100644 index 0000000..1a283ea --- /dev/null +++ b/pkg/ruleset_v2/rule_utils.go @@ -0,0 +1,59 @@ +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/pkg/ruleset_v2/rule_utils_test.go b/pkg/ruleset_v2/rule_utils_test.go new file mode 100644 index 0000000..55f23ed --- /dev/null +++ b/pkg/ruleset_v2/rule_utils_test.go @@ -0,0 +1,138 @@ +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) + } + }) + } +} diff --git a/pkg/ruleset_v2/ruleset.go b/pkg/ruleset_v2/ruleset.go new file mode 100644 index 0000000..499f0ad --- /dev/null +++ b/pkg/ruleset_v2/ruleset.go @@ -0,0 +1,32 @@ +package ruleset_v2 + +import ( + "net/url" +) + +type IRuleset interface { + HasRule(url url.URL) bool + GetRule(url url.URL) (rule Rule, exists bool) +} + +type Ruleset struct { + rulesetPath string + rules map[string]Rule +} + +func (rs Ruleset) GetRule(url url.URL) (rule Rule, exists bool) { + rule, exists = rs.rules[url.Hostname()] + return rule, exists +} + +func (rs Ruleset) HasRule(url url.URL) bool { + _, exists := rs.GetRule(url) + return exists +} + +func NewRuleset(path string) (Ruleset, error) { + rs := Ruleset{ + rulesetPath: path, + } + return rs, nil +}