Fix autocalibration-strategy merging, add tests (#732)

CHANGELOG.md

@@ -2,6 +2,7 @@
 - master
   - New
   - Changed
+    - Fix a bug in autocalibration strategy merging, when two files have the same strategy key
   
 - v2.1.0
   - New

main.go

@@ -144,7 +144,7 @@ func ParseFlags(opts *ffuf.ConfigOptions) *ffuf.ConfigOptions {
 
 	opts.General.AutoCalibrationStrings = autocalibrationstrings
 	if len(autocalibrationstrategies) > 0 {
-		opts.General.AutoCalibrationStrategies = []string {}
+		opts.General.AutoCalibrationStrategies = []string{}
 		for _, strategy := range autocalibrationstrategies {
 			opts.General.AutoCalibrationStrategies = append(opts.General.AutoCalibrationStrategies, strings.Split(strategy, ",")...)
 		}

pkg/ffuf/autocalibration.go

@@ -1,14 +1,14 @@
 package ffuf
 
 import (
+	"encoding/json"
 	"fmt"
 	"log"
 	"math/rand"
+	"os"
+	"path/filepath"
 	"strconv"
 	"time"
-	"encoding/json"
-	"path/filepath"
-	"os"
 )
 
 type AutocalibrationStrategy map[string][]string
@@ -44,9 +44,9 @@ func (j *Job) autoCalibrationStrings() map[string][]string {
 }
 
 func setupDefaultAutocalibrationStrategies() error {
-	basic_strategy := AutocalibrationStrategy {
+	basic_strategy := AutocalibrationStrategy{
-		"basic_admin": []string{"admin"+RandomString(16), "admin"+RandomString(8)},
+		"basic_admin":  []string{"admin" + RandomString(16), "admin" + RandomString(8)},
-		"htaccess": []string{".htaccess"+RandomString(16), ".htaccess"+RandomString(8)},
+		"htaccess":     []string{".htaccess" + RandomString(16), ".htaccess" + RandomString(8)},
 		"basic_random": []string{RandomString(16), RandomString(8)},
 	}
 	basic_strategy_json, err := json.Marshal(basic_strategy)
@@ -54,12 +54,12 @@ func setupDefaultAutocalibrationStrategies() error {
 		return err
 	}
 
-	advanced_strategy := AutocalibrationStrategy {
+	advanced_strategy := AutocalibrationStrategy{
-		"basic_admin": []string{"admin"+RandomString(16), "admin"+RandomString(8)},
+		"basic_admin":  []string{"admin" + RandomString(16), "admin" + RandomString(8)},
-		"htaccess": []string{".htaccess"+RandomString(16), ".htaccess"+RandomString(8)},
+		"htaccess":     []string{".htaccess" + RandomString(16), ".htaccess" + RandomString(8)},
 		"basic_random": []string{RandomString(16), RandomString(8)},
-		"admin_dir": []string{"admin"+RandomString(16)+"/", "admin"+RandomString(8)+"/"},
+		"admin_dir":    []string{"admin" + RandomString(16) + "/", "admin" + RandomString(8) + "/"},
-		"random_dir": []string{RandomString(16)+"/", RandomString(8)+"/"},
+		"random_dir":   []string{RandomString(16) + "/", RandomString(8) + "/"},
 	}
 	advanced_strategy_json, err := json.Marshal(advanced_strategy)
 	if err != nil {

pkg/ffuf/autocalibration_test.go

@@ -0,0 +1,108 @@
+package ffuf
+
+import (
+	"encoding/json"
+	"os"
+	"path/filepath"
+	"testing"
+)
+
+// NullOutput is a dummy output provider that does nothing
+type NullOutput struct {
+	Results []Result
+}
+
+func NewNullOutput() *NullOutput                             { return &NullOutput{} }
+func (o *NullOutput) Banner()                                {}
+func (o *NullOutput) Finalize() error                        { return nil }
+func (o *NullOutput) Progress(status Progress)               {}
+func (o *NullOutput) Info(infostring string)                 {}
+func (o *NullOutput) Error(errstring string)                 {}
+func (o *NullOutput) Raw(output string)                      {}
+func (o *NullOutput) Warning(warnstring string)              {}
+func (o *NullOutput) Result(resp Response)                   {}
+func (o *NullOutput) PrintResult(res Result)                 {}
+func (o *NullOutput) SaveFile(filename, format string) error { return nil }
+func (o *NullOutput) GetCurrentResults() []Result            { return o.Results }
+func (o *NullOutput) SetCurrentResults(results []Result)     { o.Results = results }
+func (o *NullOutput) Reset()                                 {}
+func (o *NullOutput) Cycle()                                 {}
+
+func TestAutoCalibrationStrings(t *testing.T) {
+	// Create a temporary directory for the test
+	tmpDir, err := os.MkdirTemp("", "ffuf-test")
+	AUTOCALIBDIR = tmpDir
+	if err != nil {
+		t.Fatalf("Failed to create temporary directory: %v", err)
+	}
+	defer os.RemoveAll(tmpDir)
+
+	// Create a test strategy file
+	strategy := AutocalibrationStrategy{
+		"test": {"foo", "bar"},
+	}
+	strategyJSON, err := json.Marshal(strategy)
+	if err != nil {
+		t.Fatalf("Failed to marshal strategy to JSON: %v", err)
+	}
+	strategyFile := filepath.Join(tmpDir, "test.json")
+	err = os.WriteFile(strategyFile, strategyJSON, 0644)
+	if err != nil {
+		t.Fatalf("Failed to write strategy file: %v", err)
+	}
+
+	// Create a test job with the strategy
+	job := &Job{
+		Config: &Config{
+			AutoCalibrationStrategies: []string{"test"},
+		},
+		Output: NewNullOutput(),
+	}
+	cInputs := job.autoCalibrationStrings()
+
+	// Verify that the custom strategy was added
+	if len(cInputs["custom"]) != 0 {
+		t.Errorf("Expected custom strategy to be empty, but got %v", cInputs["custom"])
+	}
+
+	// Verify that the test strategy was added
+	expected := []string{"foo", "bar"}
+	if len(cInputs["test"]) != len(expected) {
+		t.Errorf("Expected test strategy to have %d inputs, but got %d", len(expected), len(cInputs["test"]))
+	}
+	for i, input := range cInputs["test"] {
+		if input != expected[i] {
+			t.Errorf("Expected test strategy input %d to be %q, but got %q", i, expected[i], input)
+		}
+	}
+
+	// Verify that a missing strategy is skipped
+	job = &Job{
+		Config: &Config{
+			AutoCalibrationStrategies: []string{"missing"},
+		},
+		Output: NewNullOutput(),
+	}
+	cInputs = job.autoCalibrationStrings()
+	if len(cInputs) != 0 {
+		t.Errorf("Expected missing strategy to be skipped, but got %v", cInputs)
+	}
+
+	// Verify that a malformed strategy is skipped
+	malformedStrategy := []byte(`{"test": "foo"}`)
+	malformedFile := filepath.Join(tmpDir, "malformed.json")
+	err = os.WriteFile(malformedFile, malformedStrategy, 0644)
+	if err != nil {
+		t.Fatalf("Failed to write malformed strategy file: %v", err)
+	}
+	job = &Job{
+		Config: &Config{
+			AutoCalibrationStrategies: []string{"malformed"},
+		},
+		Output: NewNullOutput(),
+	}
+	cInputs = job.autoCalibrationStrings()
+	if len(cInputs) != 0 {
+		t.Errorf("Expected malformed strategy to be skipped, but got %v", cInputs)
+	}
+}

pkg/ffuf/util.go

@@ -131,7 +131,16 @@ func mergeMaps(m1 map[string][]string, m2 map[string][]string) map[string][]stri
 		merged[k] = v
 	}
 	for key, value := range m2 {
+		if _, ok := merged[key]; !ok {
+			// Key not found, add it
 			merged[key] = value
+			continue
+		}
+		for _, entry := range value {
+			if !StrInSlice(entry, merged[key]) {
+				merged[key] = append(merged[key], entry)
+			}
+		}
 	}
 	return merged
 }