improve codegen for JSON serialization of response/request modifiers

proxychain/codegen/codegen.go

@@ -5,6 +5,8 @@ import (
 	"go/ast"
 	"go/parser"
 	"go/token"
+	"io"
+
 	//"io/fs"
 	"os"
 	"path/filepath"
@@ -12,30 +14,27 @@ import (
 	//"strings"
 )
 
-func modToFactoryMap(fn *ast.FuncDecl) (modMap string) {
+func responseModToFactoryMap(fn *ast.FuncDecl) (modMap string) {
 	paramCount := len(fn.Type.Params.List)
 	name := fn.Name.Name
 	var x string
 	switch paramCount {
 	case 0:
-		x = fmt.Sprintf("  resModMap[\"%s\"] = func(_ ...string) proxychain.ResponseModification {\n    return rx.%s()\n  }\n", name, name)
+		x = fmt.Sprintf("  rsmModMap[\"%s\"] = func(_ ...string) proxychain.ResponseModification {\n    return tx.%s()\n  }\n", name, name)
 	default:
 		p := []string{}
 		for i := 0; i < paramCount; i++ {
 			p = append(p, fmt.Sprintf("params[%d]", i))
 		}
 		params := strings.Join(p, ", ")
-		x = fmt.Sprintf("  resModMap[\"%s\"] = func(params, ...string) proxychain.ResponseModification {\n    return rx.%s(%s)\n  }\n", name, name, params)
+		x = fmt.Sprintf("  rsmModMap[\"%s\"] = func(params ...string) proxychain.ResponseModification {\n    return tx.%s(%s)\n  }\n", name, name, params)
 	}
 	return x
 }
 
-func main() {
+func responseModCodeGen(dir string) (code string, err error) {
 	fset := token.NewFileSet()
 
-	// Directory containing your Go files
-	dir := "../requestmodifiers/"
-
 	files, err := os.ReadDir(dir)
 	if err != nil {
 		panic(err)
@@ -50,37 +49,143 @@ func main() {
 		// Parse each Go file
 		node, err := parser.ParseFile(fset, filepath.Join(dir, file.Name()), nil, parser.ParseComments)
 		if err != nil {
-			panic(err)
+			return "", err
 		}
 
 		ast.Inspect(node, func(n ast.Node) bool {
 			fn, ok := n.(*ast.FuncDecl)
 			if ok && fn.Recv == nil && fn.Name.IsExported() {
-				factoryMaps = append(factoryMaps, modToFactoryMap(fn))
+				factoryMaps = append(factoryMaps, responseModToFactoryMap(fn))
 			}
 			return true
 		})
 
 	}
 
-	code := fmt.Sprintf(`
+	code = fmt.Sprintf(`
 package ruleset_v2
+// DO NOT EDIT THIS FILE. It is automatically generated by ladder/proxychain/codegen/codegen.go
+// The purpose of this is serialization of rulesets from JSON or YAML into functional options suitable
+// for use in proxychains.
 
 import (
 	"ladder/proxychain"
-	rx "ladder/proxychain/responsemodifiers"
+	tx "ladder/proxychain/responsemodifiers"
 )
 
 type ResponseModifierFactory func(params ...string) proxychain.ResponseModification
 
-var resModMap map[string]ResponseModifierFactory
+var rsmModMap map[string]ResponseModifierFactory
 
-// TODO: create codegen using AST parsing of exported methods in ladder/proxychain/responsemodifiers/*.go
 func init() {
-	resModMap = make(map[string]ResponseModifierFactory)
+	rsmModMap = make(map[string]ResponseModifierFactory)
 
 	%s
 }`, strings.Join(factoryMaps, "\n"))
-	fmt.Println(code)
-
+	//fmt.Println(code)
+	return code, nil
+}
+
+func requestModToFactoryMap(fn *ast.FuncDecl) (modMap string) {
+	paramCount := len(fn.Type.Params.List)
+	name := fn.Name.Name
+	var x string
+	switch paramCount {
+	case 0:
+		x = fmt.Sprintf("  rqmModMap[\"%s\"] = func(_ ...string) proxychain.RequestModification {\n    return rx.%s()\n  }\n", name, name)
+	default:
+		p := []string{}
+		for i := 0; i < paramCount; i++ {
+			p = append(p, fmt.Sprintf("params[%d]", i))
+		}
+		params := strings.Join(p, ", ")
+		x = fmt.Sprintf("  rqmModMap[\"%s\"] = func(params ...string) proxychain.RequestModification {\n    return rx.%s(%s)\n  }\n", name, name, params)
+	}
+	return x
+}
+
+func requestModCodeGen(dir string) (code string, err error) {
+	fset := token.NewFileSet()
+
+	files, err := os.ReadDir(dir)
+	if err != nil {
+		panic(err)
+	}
+
+	factoryMaps := []string{}
+	for _, file := range files {
+		if filepath.Ext(file.Name()) != ".go" {
+			continue
+		}
+
+		// Parse each Go file
+		node, err := parser.ParseFile(fset, filepath.Join(dir, file.Name()), nil, parser.ParseComments)
+		if err != nil {
+			return "", err
+		}
+
+		ast.Inspect(node, func(n ast.Node) bool {
+			fn, ok := n.(*ast.FuncDecl)
+			if ok && fn.Recv == nil && fn.Name.IsExported() {
+				factoryMaps = append(factoryMaps, requestModToFactoryMap(fn))
+			}
+			return true
+		})
+
+	}
+
+	code = fmt.Sprintf(`
+package ruleset_v2
+// DO NOT EDIT THIS FILE. It is automatically generated by ladder/proxychain/codegen/codegen.go
+// The purpose of this is serialization of rulesets from JSON or YAML into functional options suitable
+// for use in proxychains.
+
+import (
+	"ladder/proxychain"
+	rx "ladder/proxychain/requestmodifiers"
+)
+
+type RequestModifierFactory func(params ...string) proxychain.RequestModification
+
+var rqmModMap map[string]RequestModifierFactory
+
+func init() {
+	rqmModMap = make(map[string]RequestModifierFactory)
+
+	%s
+}`, strings.Join(factoryMaps, "\n"))
+	//fmt.Println(code)
+	return code, nil
+}
+
+func main() {
+	rqmCode, err := requestModCodeGen("../requestmodifiers/")
+	if err != nil {
+		panic(err)
+	}
+	//fmt.Println(rqmCode)
+
+	fq, err := os.Create("../ruleset/rule_reqmod_types.gen.go")
+	if err != nil {
+		panic(err)
+	}
+	_, err = io.WriteString(fq, rqmCode)
+	if err != nil {
+		panic(err)
+	}
+
+	rsmCode, err := responseModCodeGen("../responsemodifiers/")
+	if err != nil {
+		panic(err)
+	}
+	//fmt.Println(rsmCode)
+
+	fs, err := os.Create("../ruleset/rule_resmod_types.gen.go")
+	if err != nil {
+		panic(err)
+	}
+	_, err = io.WriteString(fs, rsmCode)
+	if err != nil {
+		panic(err)
+	}
 }