Merge pull request #47 from dxbednarczyk/main

Idiomatize(?) ruleset package and run lint

Author: ladddder

Makefile

@@ -1,6 +1,6 @@
 lint:
 	gofumpt -l -w .
-	golangci-lint run -c .golangci-lint.yaml
+	golangci-lint run -c .golangci-lint.yaml --fix
 
 	go mod tidy
 	go clean

cmd/main.go

@@ -29,6 +29,7 @@ func main() {
 	if os.Getenv("PORT") == "" {
 		portEnv = "8080"
 	}
+
 	port := parser.String("p", "port", &argparse.Options{
 		Required: false,
 		Default:  portEnv,
@@ -49,10 +50,12 @@ func main() {
 		Required: false,
 		Help:     "Compiles a directory of yaml files into a single ruleset.yaml. Requires --ruleset arg.",
 	})
+
 	mergeRulesetsGzip := parser.Flag("", "merge-rulesets-gzip", &argparse.Options{
 		Required: false,
 		Help:     "Compiles a directory of yaml files into a single ruleset.gz Requires --ruleset arg.",
 	})
+
 	mergeRulesetsOutput := parser.String("", "merge-rulesets-output", &argparse.Options{
 		Required: false,
 		Help:     "Specify output file for --merge-rulesets and --merge-rulesets-gzip. Requires --ruleset and --merge-rulesets args.",
@@ -65,7 +68,18 @@ func main() {
 
 	// utility cli flag to compile ruleset directory into single ruleset.yaml
 	if *mergeRulesets || *mergeRulesetsGzip {
-		err = cli.HandleRulesetMerge(ruleset, mergeRulesets, mergeRulesetsGzip, mergeRulesetsOutput)
+		output := os.Stdout
+
+		if *mergeRulesetsOutput != "" {
+			output, err = os.Create(*mergeRulesetsOutput)
+			
+			if err != nil {
+				fmt.Println(err)
+				os.Exit(1)
+			}
+		}
+
+		err = cli.HandleRulesetMerge(*ruleset, *mergeRulesets, *mergeRulesetsGzip, output)
 		if err != nil {
 			fmt.Println(err)
 			os.Exit(1)
@@ -87,6 +101,7 @@ func main() {
 	userpass := os.Getenv("USERPASS")
 	if userpass != "" {
 		userpass := strings.Split(userpass, ":")
+
 		app.Use(basicauth.New(basicauth.Config{
 			Users: map[string]string{
 				userpass[0]: userpass[1],
@@ -102,23 +117,28 @@ func main() {
 	if os.Getenv("NOLOGS") != "true" {
 		app.Use(func(c *fiber.Ctx) error {
 			log.Println(c.Method(), c.Path())
+
 			return c.Next()
 		})
 	}
 
 	app.Get("/", handlers.Form)
+
 	app.Get("/styles.css", func(c *fiber.Ctx) error {
 		cssData, err := cssData.ReadFile("styles.css")
 		if err != nil {
 			return c.Status(fiber.StatusInternalServerError).SendString("Internal Server Error")
 		}
+
 		c.Set("Content-Type", "text/css")
+
 		return c.Send(cssData)
 	})
-	app.Get("ruleset", handlers.Ruleset)
 
+	app.Get("ruleset", handlers.Ruleset)
 	app.Get("raw/*", handlers.Raw)
 	app.Get("api/*", handlers.Api)
 	app.Get("/*", handlers.ProxySite(*ruleset))
+
 	log.Fatal(app.Listen(":" + *port))
 }

handlers/cli/cli.go

@@ -3,43 +3,49 @@ package cli
 import (
 	"fmt"
 	"io"
-	"io/fs"
-	"ladder/pkg/ruleset"
 	"os"
 
+	"ladder/pkg/ruleset"
+
 	"golang.org/x/term"
 )
 
 // HandleRulesetMerge merges a set of ruleset files, specified by the rulesetPath or RULESET env variable, into either YAML or Gzip format.
 // 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
 }

handlers/proxy.go

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

handlers/proxy.test.go

@@ -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"
 )