DWS/hadrian
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

implement proper marshaller/unmarshaller for rulesets in yaml format

Browse Source
This commit is contained in:
Kevin Pham
2023-12-05 01:51:30 -06:00
parent f8621e72ee
commit bc715d9101
5 changed files with 170 additions and 266 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
proxychain/codegen/README.md Normal file
View File

@@ -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
```

87
proxychain/ruleset/rule.go
View File

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

136
proxychain/ruleset/rule_test.go
View File

@@ -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")
}
}
*/

59
proxychain/ruleset/rule_utils.go
View File

@@ -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
}

138
proxychain/ruleset/rule_utils_test.go
View File

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