Add -header flag to attack sub command

attack_test.go

@@ -1,8 +1,10 @@
 package main
 
 import (
+	"flag"
 	"io/ioutil"
 	"log"
+	"net/http"
 	"strings"
 	"testing"
 	"time"
@@ -14,48 +16,48 @@ func init() {
 }
 
 func TestRateValidation(t *testing.T) {
-	rate, duration, targetsf, ordering, output := defaultArguments()
+	rate, duration, targetsf, ordering, output, header := defaultArguments()
 	rate = 0
 
-	err := attack(rate, duration, targetsf, ordering, output)
+	err := attack(rate, duration, targetsf, ordering, output, header)
 	if err == nil || (err != nil && !strings.HasPrefix(err.Error(), errRatePrefix)) {
 		t.Errorf("Rate 0 shouldn't be valid: %s", err)
 	}
 }
 
 func TestDurationValidation(t *testing.T) {
-	rate, duration, targetsf, ordering, output := defaultArguments()
+	rate, duration, targetsf, ordering, output, header := defaultArguments()
 	duration = 0
 
-	err := attack(rate, duration, targetsf, ordering, output)
+	err := attack(rate, duration, targetsf, ordering, output, header)
 	if err == nil || (err != nil && !strings.HasPrefix(err.Error(), errDurationPrefix)) {
 		t.Errorf("Duration 0 shouldn't be valid: %s", err)
 	}
 }
 
 func TestTargetsValidation(t *testing.T) {
-	rate, duration, goodFile, ordering, output := defaultArguments()
+	rate, duration, goodFile, ordering, output, header := defaultArguments()
 
 	// Good case
-	err := attack(rate, duration, goodFile, ordering, output)
+	err := attack(rate, duration, goodFile, ordering, output, header)
 	if err != nil {
 		t.Errorf("Targets file `%s` should be valid: %s", goodFile, err)
 	}
 
 	// Bad case
 	badFile := "randomInexistingFile12345.txt"
-	err = attack(rate, duration, badFile, ordering, output)
+	err = attack(rate, duration, badFile, ordering, output, header)
 	if err == nil || (err != nil && !strings.HasPrefix(err.Error(), errTargetsFilePrefix)) {
 		t.Errorf("Targets file `%s` shouldn't be valid: %s", badFile, err)
 	}
 }
 
 func TestOrderingValidation(t *testing.T) {
-	rate, duration, targetsf, _, output := defaultArguments()
+	rate, duration, targetsf, _, output, header := defaultArguments()
 
 	// Good cases
 	for _, ordering := range []string{"random", "sequential"} {
-		err := attack(rate, duration, targetsf, ordering, output)
+		err := attack(rate, duration, targetsf, ordering, output, header)
 		if err != nil {
 			t.Errorf("Ordering `%s` should be valid: %s", ordering, err)
 		}
@@ -63,12 +65,36 @@ func TestOrderingValidation(t *testing.T) {
 
 	// Bad case
 	badOrdering := "lolcat"
-	err := attack(rate, duration, targetsf, badOrdering, output)
+	err := attack(rate, duration, targetsf, badOrdering, output, header)
 	if err == nil || (err != nil && !strings.HasPrefix(err.Error(), errOrderingPrefix)) {
 		t.Errorf("Ordering `%s` shouldn't be valid: %s", badOrdering, err)
 	}
 }
 
-func defaultArguments() (uint64, time.Duration, string, string, string) {
-	return uint64(1000), 5 * time.Millisecond, ".targets.txt", "random", "/dev/null"
+func TestHeadersParsing(t *testing.T) {
+	fs := flag.NewFlagSet("test", flag.ContinueOnError)
+	fs.SetOutput(ioutil.Discard)
+	hdrs := headers{Header: make(http.Header)}
+	fs.Var(hdrs, "H", "Header")
+	// Good case
+	good := []string{"-H", "Host: lolcathost"}
+	if err := fs.Parse(good); err != nil {
+		t.Errorf("%v should be a valid header", good[1])
+	}
+	// Bad cases
+	bad := [][]string{[]string{"-H", "Host:"}, []string{"-H", "Host"}}
+	for _, args := range bad {
+		if err := fs.Parse(args); err == nil {
+			t.Errorf("%v should not be a valid header", args[1])
+		}
+	}
+}
+
+func defaultArguments() (uint64, time.Duration, string, string, string, http.Header) {
+	return uint64(1000),
+		5 * time.Millisecond,
+		".targets.txt",
+		"random",
+		"/dev/null",
+		http.Header{}
 }