Add CLI support for custom headers and HTTP method overrides

commands/run.go

@@ -2,6 +2,8 @@ package commands
 
 import (
 	"math"
+	"net/http"
+	"strings"
 	"sync"
 	"time"
 
@@ -18,7 +20,9 @@ var (
 	interval    time.Duration
 	size        string
 	tor         string
+	headers     string
 	url         string
+	method      string
 )
 
 const defaultInterval = 10 * time.Second
@@ -35,6 +39,8 @@ func (*Run) SetFlags(flags *pflag.FlagSet) {
 	flags.StringVarP(&size, "payload-size", "p", "1MB", "Random generated payload with the given size.")
 	flags.StringVarP(&tor, "tor", "t", "", "TOR endpoint (either socks5://1.1.1.1:1234, or 1.1.1.1:1234).")
 	flags.StringVarP(&url, "url", "u", "", "Target URL to send the attack to.")
+	flags.StringVarP(&headers, "headers", "H", "", "HTTP headers in HEADER=VALUE format, separated by commas.")
+	flags.StringVarP(&method, "method", "m", http.MethodPost, "HTTP method to use when sending requests.")
 }
 
 // GetRequiredFlags returns the server required flags.
@@ -74,13 +80,15 @@ func (*Run) Run() RunCmd {
 
 		waitgroup.Add(int(concurrents))
 
+		parsedHeaders := parseHeaders(headers)
+
 		for range concurrents {
 			go func() {
 				if isize > math.MaxInt64 {
 					return
 				}
 
-				req := request.NewRequest(int64(isize), url, interval)
+				req := request.NewRequest(int64(isize), url, interval, parsedHeaders, method)
 				if tor != "" {
 					req.WithTor(tor)
 				}
@@ -97,6 +105,38 @@ func (*Run) Run() RunCmd {
 	}
 }
 
+func parseHeaders(raw string) map[string]string {
+	if raw == "" {
+		return nil
+	}
+
+	headersMap := make(map[string]string)
+
+	entries := strings.Split(raw, ",")
+	for _, entry := range entries {
+		entry = strings.TrimSpace(entry)
+		if entry == "" {
+			continue
+		}
+
+		parts := strings.SplitN(entry, "=", 2)
+		if len(parts) != 2 {
+			panic("invalid header format: " + entry)
+		}
+
+		key := strings.TrimSpace(parts[0])
+		value := strings.TrimSpace(parts[1])
+
+		if key == "" {
+			panic("header name cannot be empty")
+		}
+
+		headersMap[key] = value
+	}
+
+	return headersMap
+}
+
 func newRun() command {
 	return &Run{}
 }

request/request.go

@@ -8,6 +8,7 @@ import (
 	"io"
 	"net/http"
 	"net/url"
+	"strings"
 	"time"
 
 	"github.com/darkweak/rudy/logger"
@@ -27,13 +28,23 @@ type Request interface {
 }
 
 // NewRequest creates the request.
-func NewRequest(size int64, u string, delay time.Duration) Request {
-	req, _ := http.NewRequestWithContext(context.Background(), http.MethodPost, u, nil)
+func NewRequest(size int64, u string, delay time.Duration, headers map[string]string, method string) Request {
+	if method == "" {
+		method = http.MethodPost
+	} else {
+		method = strings.ToUpper(method)
+	}
+
+	req, _ := http.NewRequestWithContext(context.Background(), method, u, nil)
 	req.ProtoMajor = 1
 	req.ProtoMinor = 1
 	req.TransferEncoding = []string{"chunked"}
 	req.Header = make(map[string][]string)
 
+	for key, value := range headers {
+		req.Header.Set(key, value)
+	}
+
 	return &request{
 		client:      http.DefaultClient,
 		delay:       delay,