wip

proxychain/requestmodifers/modify_outgoing_cookies.go

@@ -2,20 +2,76 @@ package requestmodifers
 
 import (
 	"ladder/proxychain"
+	"net/http"
 )
 
-// BlockOutgoingCookies prevents ALL cookies from being sent from the client
-// to the upstream proxy server.
-func BlockOutgoingCookies() proxychain.RequestModification {
+// SetOutgoingCookie modifes a specific cookie name
+// by modifying the request cookie headers going to the upstream server.
+// If the cookie name does not already exist, it is created.
+func SetOutgoingCookie(name string, val string) proxychain.RequestModification {
+	return func(chain *proxychain.ProxyChain) error {
+		cookies := chain.Request.Cookies()
+		hasCookie := false
+		for _, cookie := range cookies {
+			if cookie.Name != name {
+				continue
+			}
+			hasCookie = true
+			cookie.Value = val
+		}
+
+		if hasCookie {
+			return nil
+		}
+
+		chain.Request.AddCookie(&http.Cookie{
+			Domain: chain.Request.URL.Host,
+			Name:   name,
+			Value:  val,
+		})
+
+		return nil
+	}
+}
+
+// SetOutgoingCookies modifies a client request's cookie header
+// to a raw Cookie string, overwriting existing cookies
+func SetOutgoingCookies(cookies string) proxychain.RequestModification {
+	return func(chain *proxychain.ProxyChain) error {
+		chain.Request.Header.Set("Cookies", cookies)
+		return nil
+	}
+}
+
+// DeleteOutgoingCookie modifies the http request's cookies header to
+// delete a specific request cookie going to the upstream server.
+// If the cookie does not exist, it does not do anything.
+func DeleteOutgoingCookie(name string) proxychain.RequestModification {
+	return func(chain *proxychain.ProxyChain) error {
+		cookies := chain.Request.Cookies()
+		chain.Request.Header.Del("Cookies")
+
+		for _, cookie := range cookies {
+			if cookie.Name == name {
+				chain.Request.AddCookie(cookie)
+			}
+		}
+		return nil
+	}
+}
+
+// DeleteOutgoingCookies removes the cookie header entirely,
+// preventing any cookies from reaching the upstream server.
+func DeleteOutgoingCookies() proxychain.RequestModification {
 	return func(px *proxychain.ProxyChain) error {
 		px.Request.Header.Del("Cookie")
 		return nil
 	}
 }
 
-// BlockOutgoingCookiesExcept prevents non-whitelisted cookies from being sent from the client
+// DeleteOutGoingCookiesExcept prevents non-whitelisted cookies from being sent from the client
 // to the upstream proxy server. Cookies whose names are in the whitelist are not removed.
-func BlockOutgoingCookiesExcept(whitelist ...string) proxychain.RequestModification {
+func DeleteOutgoingCookiesExcept(whitelist ...string) proxychain.RequestModification {
 	return func(px *proxychain.ProxyChain) error {
 		// Convert whitelist slice to a map for efficient lookups
 		whitelistMap := make(map[string]struct{})

proxychain/responsemodifers/rewrite_http_resource_urls.go

@@ -50,19 +50,19 @@ func (r *HTMLResourceURLRewriter) Read(p []byte) (int, error) {
 			}
 			return 0, err // Actual error
 		}
+
 		token := tokenizer.Token()
 		if tokenType == html.StartTagToken || tokenType == html.SelfClosingTagToken {
-				rewriteToken(&token, r.url)
+			rewriteToken(&token, r.proxyURL)
 		}
 		r.buffer.WriteString(token.String())
 		if r.buffer.Len() > 0 {
 			break
 		}
 	}
-	}
+	return r.buffer.Read(p)
 }
 
-
 // RewriteHTMLResourceURLs updates src/href attributes in HTML content to route through the proxy.
 func RewriteHTMLResourceURLs() proxychain.ResponseModification {
 	return func(chain *proxychain.ProxyChain) error {
@@ -70,45 +70,8 @@ func RewriteHTMLResourceURLs() proxychain.ResponseModification {
 		if ct != "text/html" {
 			return nil
 		}
-
-		// parse dom
-		tokenizer := html.NewTokenizer(chain.Body)
-		var buffer bytes.Buffer
-
-		// traverse dom and proxify existing src/img resource links
-		for {
-			tokenType := tokenizer.Next()
-			switch tokenType {
-			case html.ErrorToken:
-				// End of the document, set the new body
-				chain.Body = io.ReaderFrom(buffer)
+		// TODO: implement chaining rewriter chaining method
+		// so we can compose multiple body rewriters together
 		return nil
-			case html.StartTagToken, html.SelfClosingTagToken:
-				token := tokenizer.Token()
-				// Rewrite the necessary attributes
-				token = rewriteToken(token, u)
-				buffer.WriteString(token.String())
-			case html.TextToken, html.CommentToken, html.DoctypeToken, html.EndTagToken:
-				// Write the token to the buffer as is
-				buffer.WriteString(tokenizer.Token().String())
-			}
-		}
 	}
 }
-
-// rewriteToken rewrites the tokens with URLs to point to the proxy server.
-func rewriteToken(token html.Token, u *url.URL) html.Token {
-	// Define attributes to rewrite, add more as needed such as "srcset"
-	rewriteAttrs := map[string]bool{"href": true, "src": true, "action": true, "srcset": true}
-
-	for i, attr := range token.Attr {
-		_, shouldRewrite := rewriteAttrs[attr.Key]
-		if shouldRewrite {
-			val := attr.Val
-			if strings.HasPrefix(val, "/") {
-				token.Attr[i].Val = "/https://" + u.Host + val
-			}
-		}
-	}
-	return token
-}