hadrian/proxychain/codegen/codegen.go

package main

import (
	"fmt"
	"go/ast"
	"go/parser"
	"go/token"
	"io"
	"io/fs"

	//"io/fs"
	"os"
	"path/filepath"
	"strings"
	//"strings"
)

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("  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("  rsmModMap[\"%s\"] = func(params ...string) proxychain.ResponseModification {\n    return tx.%s(%s)\n  }\n", name, name, params)
	}
	return x
}

func responseModCodeGen(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 !shouldGenCodeFor(file) {
			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, responseModToFactoryMap(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"
	tx "ladder/proxychain/responsemodifiers"
)

type ResponseModifierFactory func(params ...string) proxychain.ResponseModification

var rsmModMap map[string]ResponseModifierFactory

func init() {
	rsmModMap = make(map[string]ResponseModifierFactory)

	%s
}`, strings.Join(factoryMaps, "\n"))
	//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 !shouldGenCodeFor(file) {
			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 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 {
		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)
	}
}