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

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