idiomatize (?) ruleset package and lint

This commit is contained in:
Damian Bednarczyk
2023-11-22 21:26:44 -06:00
parent 394eaf9805
commit dc69af9f38
7 changed files with 124 additions and 57 deletions

View File

@@ -3,10 +3,10 @@ package cli
import (
"fmt"
"io"
"io/fs"
"ladder/pkg/ruleset"
"os"
"ladder/pkg/ruleset"
"golang.org/x/term"
)
@@ -14,32 +14,38 @@ import (
// Exits the program with an error message if the ruleset path is not provided or if loading the ruleset fails.
//
// Parameters:
// - rulesetPath: A pointer to a string specifying the path to the ruleset file.
// - mergeRulesets: A pointer to a boolean indicating if a merge operation should be performed.
// - mergeRulesetsGzip: A pointer to a boolean indicating if the merge should be in Gzip format.
// - mergeRulesetsOutput: A pointer to a string specifying the output file path. If empty, the output is printed to stdout.
// - rulesetPath: Specifies the path to the ruleset file.
// - mergeRulesets: Indicates if a merge operation should be performed.
// - useGzip: Indicates if the merged rulesets should be gzip-ped.
// - output: Specifies the output file. If nil, stdout will be used.
//
// Returns:
// - An error if the ruleset loading or merging process fails, otherwise nil.
func HandleRulesetMerge(rulesetPath *string, mergeRulesets *bool, mergeRulesetsGzip *bool, mergeRulesetsOutput *string) error {
if *rulesetPath == "" {
*rulesetPath = os.Getenv("RULESET")
func HandleRulesetMerge(rulesetPath string, mergeRulesets bool, useGzip bool, output *os.File) error {
if !mergeRulesets {
return nil
}
if *rulesetPath == "" {
fmt.Println("ERROR: no ruleset provided. Try again with --ruleset <ruleset.yaml>")
if rulesetPath == "" {
rulesetPath = os.Getenv("RULESET")
}
if rulesetPath == "" {
fmt.Println("error: no ruleset provided. Try again with --ruleset <ruleset.yaml>")
os.Exit(1)
}
rs, err := ruleset.NewRuleset(*rulesetPath)
rs, err := ruleset.NewRuleset(rulesetPath)
if err != nil {
fmt.Println(err)
os.Exit(1)
}
if *mergeRulesetsGzip {
return gzipMerge(rs, mergeRulesetsOutput)
if useGzip {
return gzipMerge(rs, output)
}
return yamlMerge(rs, mergeRulesetsOutput)
return yamlMerge(rs, output)
}
// gzipMerge takes a RuleSet and an optional output file path pointer. It compresses the RuleSet into Gzip format.
@@ -48,33 +54,33 @@ func HandleRulesetMerge(rulesetPath *string, mergeRulesets *bool, mergeRulesetsG
//
// Parameters:
// - rs: The ruleset.RuleSet to be compressed.
// - mergeRulesetsOutput: A pointer to a string specifying the output file path. If empty, the output is directed to stdout.
// - output: The output for the gzip data. If nil, stdout will be used.
//
// Returns:
// - An error if compression or file writing fails, otherwise nil.
func gzipMerge(rs ruleset.RuleSet, mergeRulesetsOutput *string) error {
func gzipMerge(rs ruleset.RuleSet, output io.Writer) error {
gzip, err := rs.GzipYaml()
if err != nil {
return err
}
if *mergeRulesetsOutput != "" {
out, err := os.Create(*mergeRulesetsOutput)
defer out.Close()
_, err = io.Copy(out, gzip)
if output != nil {
_, err = io.Copy(output, gzip)
if err != nil {
return err
}
}
if term.IsTerminal(int(os.Stdout.Fd())) {
println("WARNING: binary output can mess up your terminal. Use '--merge-rulesets-output <ruleset.gz>' or pipe it to a file.")
println("warning: binary output can mess up your terminal. Use '--merge-rulesets-output <ruleset.gz>' or pipe it to a file.")
os.Exit(1)
}
_, err = io.Copy(os.Stdout, gzip)
if err != nil {
return err
}
return nil
}
@@ -83,23 +89,25 @@ func gzipMerge(rs ruleset.RuleSet, mergeRulesetsOutput *string) error {
//
// Parameters:
// - rs: The ruleset.RuleSet to be converted to YAML.
// - mergeRulesetsOutput: A pointer to a string specifying the output file path. If empty, the output is printed to stdout.
// - output: The output for the merged data. If nil, stdout will be used.
//
// Returns:
// - An error if YAML conversion or file writing fails, otherwise nil.
func yamlMerge(rs ruleset.RuleSet, mergeRulesetsOutput *string) error {
func yamlMerge(rs ruleset.RuleSet, output io.Writer) error {
yaml, err := rs.Yaml()
if err != nil {
return err
}
if *mergeRulesetsOutput == "" {
fmt.Printf(yaml)
if output == nil {
fmt.Println(yaml)
os.Exit(0)
}
err = os.WriteFile(*mergeRulesetsOutput, []byte(yaml), fs.FileMode(os.O_RDWR))
_, err = io.WriteString(output, yaml)
if err != nil {
return fmt.Errorf("ERROR: failed to write merged YAML ruleset to '%s'\n", *mergeRulesetsOutput)
return fmt.Errorf("failed to write merged YAML ruleset: %v", err)
}
return nil
}

View File

@@ -80,7 +80,6 @@ func extractUrl(c *fiber.Ctx) (string, error) {
// default behavior:
// eg: https://localhost:8080/https://realsite.com/images/foobar.jpg -> https://realsite.com/images/foobar.jpg
return urlQuery.String(), nil
}
func ProxySite(rulesetPath string) fiber.Handler {
@@ -121,18 +120,18 @@ func modifyURL(uri string, rule ruleset.Rule) (string, error) {
return "", err
}
for _, urlMod := range rule.UrlMods.Domain {
for _, urlMod := range rule.URLMods.Domain {
re := regexp.MustCompile(urlMod.Match)
newUrl.Host = re.ReplaceAllString(newUrl.Host, urlMod.Replace)
}
for _, urlMod := range rule.UrlMods.Path {
for _, urlMod := range rule.URLMods.Path {
re := regexp.MustCompile(urlMod.Match)
newUrl.Path = re.ReplaceAllString(newUrl.Path, urlMod.Replace)
}
v := newUrl.Query()
for _, query := range rule.UrlMods.Query {
for _, query := range rule.URLMods.Query {
if query.Value == "" {
v.Del(query.Key)
continue
@@ -223,11 +222,11 @@ func fetchSite(urlpath string, queries map[string]string) (string, *http.Request
}
if rule.Headers.CSP != "" {
//log.Println(rule.Headers.CSP)
// log.Println(rule.Headers.CSP)
resp.Header.Set("Content-Security-Policy", rule.Headers.CSP)
}
//log.Print("rule", rule) TODO: Add a debug mode to print the rule
// log.Print("rule", rule) TODO: Add a debug mode to print the rule
body := rewriteHtml(bodyB, u, rule)
return body, req, resp, nil
}

View File

@@ -2,12 +2,13 @@
package handlers
import (
"ladder/pkg/ruleset"
"net/http"
"net/http/httptest"
"net/url"
"testing"
"ladder/pkg/ruleset"
"github.com/gofiber/fiber/v2"
"github.com/stretchr/testify/assert"
)