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

add YAML serialization for ruleset

Browse Source
This commit is contained in:
Kevin Pham
2023-12-04 19:55:50 -06:00
parent b39955025e
commit 6157d6543f
8 changed files with 238 additions and 38 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
proxychain/codegen/codegen.go
View File

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

7
proxychain/requestmodifiers/modify_domain_with_regex.go
View File

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

10
proxychain/requestmodifiers/modify_path_with_regex.go
View File

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

160
proxychain/ruleset/rule.go
View File

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

8
proxychain/ruleset/rule_resmod_types.gen.go
View File

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

70
proxychain/ruleset/rule_test.go
View File

@@ -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": [
"APIContent()",
"SetContentSecurityPolicy(\"foobar\")",
"SetIncomingCookie(\"authorization-bearer\", \"hunter2\")"
"request_modifications": [
"SpoofUserAgent(\"googlebot\")"
],
"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 {
t.Errorf("expected number of ResponseModifications to be 3")
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))
}
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))
}