From 9f33201d70901c76f92347e73b1f350221782017 Mon Sep 17 00:00:00 2001
From: William Vu <wvu@vulncheck.com>
Date: Thu, 21 Nov 2024 21:53:55 -0600
Subject: [PATCH] Add protocol.MultipartAddPart

---
 protocol/httphelper.go | 19 +++++++++++++------
 1 file changed, 13 insertions(+), 6 deletions(-)

diff --git a/protocol/httphelper.go b/protocol/httphelper.go
index 6dbf7f4..fbbffe1 100644
--- a/protocol/httphelper.go
+++ b/protocol/httphelper.go
@@ -481,22 +481,29 @@ func MultipartAddField(writer *multipart.Writer, name string, value string) bool
 	return err == nil
 }
 
-func MultipartAddFile(writer *multipart.Writer, name, filename, ctype, value string) bool {
-	// CreateFormFile doesn't expose Content-Type
+func MultipartAddPart(writer *multipart.Writer, headers map[string]string, body string) bool {
 	h := make(textproto.MIMEHeader)
-	h.Set("Content-Disposition",
-		fmt.Sprintf(`form-data; name="%s"; filename="%s"`, name, filename))
-	h.Set("Content-Type", ctype)
+	for k, v := range headers {
+		h.Set(k, v)
+	}
 
 	fw, err := writer.CreatePart(h)
 	if err != nil {
 		return false
 	}
-	_, err = io.Copy(fw, strings.NewReader(value))
+	_, err = io.Copy(fw, strings.NewReader(body))
 
 	return err == nil
 }
 
+func MultipartAddFile(writer *multipart.Writer, name, filename, ctype, value string) bool {
+	// CreateFormFile doesn't expose Content-Type
+	return MultipartAddPart(writer, map[string]string{
+		"Content-Disposition": fmt.Sprintf(`form-data; name="%s"; filename="%s"`, name, filename),
+		"Content-Type":        ctype,
+	}, value)
+}
+
 // Provided an HTTP request, find the Set-Cookie headers, and extract
 // the value of the specified cookie. Example:.
 func GetSetCookieValue(resp *http.Response, name string) (string, bool) {