refactor wip

proxychain/responsemodifers/bypass_cors.go

@@ -0,0 +1,21 @@
+package responsemodifers
+
+import (
+	"ladder/proxychain"
+)
+
+// BypassCORS modifies response headers to prevent the browser
+// from enforcing any CORS restrictions. This should run at the end of the chain.
+func BypassCORS() proxychain.ResponseModification {
+	return func(chain *proxychain.ProxyChain) error {
+		chain.AddResponseModifications(
+			SetResponseHeader("Access-Control-Allow-Origin", "*"),
+			SetResponseHeader("Access-Control-Expose-Headers", "*"),
+			SetResponseHeader("Access-Control-Allow-Credentials", "true"),
+			SetResponseHeader("Access-Control-Allow-Methods", "GET, PUT, POST, DELETE, HEAD, OPTIONS, PATCH"),
+			SetResponseHeader("Access-Control-Allow-Headers", "*"),
+			DeleteResponseHeader("X-Frame-Options"),
+		)
+		return nil
+	}
+}

proxychain/responsemodifers/bypass_csp.go

@@ -0,0 +1,27 @@
+package responsemodifers
+
+import (
+	"ladder/proxychain"
+)
+
+// BypassContentSecurityPolicy modifies response headers to prevent the browser
+// from enforcing any CSP restrictions. This should run at the end of the chain.
+func BypassContentSecurityPolicy() proxychain.ResponseModification {
+	return func(chain *proxychain.ProxyChain) error {
+		chain.AddResponseModifications(
+			DeleteResponseHeader("Content-Security-Policy"),
+			DeleteResponseHeader("Content-Security-Policy-Report-Only"),
+			DeleteResponseHeader("X-Content-Security-Policy"),
+			DeleteResponseHeader("X-WebKit-CSP"),
+		)
+		return nil
+	}
+}
+
+// SetContentSecurityPolicy modifies response headers to a specific CSP
+func SetContentSecurityPolicy(csp string) proxychain.ResponseModification {
+	return func(chain *proxychain.ProxyChain) error {
+		chain.Response.Header.Set("Content-Security-Policy", csp)
+		return nil
+	}
+}

proxychain/responsemodifers/modify_incoming_cookies.go

@@ -0,0 +1,102 @@
+package responsemodifers
+
+import (
+	"fmt"
+	"ladder/proxychain"
+	"net/http"
+)
+
+// DeleteIncomingCookies prevents ALL cookies from being sent from the proxy server
+// back down to the client.
+func DeleteIncomingCookies(whitelist ...string) proxychain.ResponseModification {
+	return func(px *proxychain.ProxyChain) error {
+		px.Response.Header.Del("Set-Cookie")
+		return nil
+	}
+}
+
+// DeleteIncomingCookiesExcept prevents non-whitelisted cookies from being sent from the proxy server
+// to the client. Cookies whose names are in the whitelist are not removed.
+func DeleteIncomingCookiesExcept(whitelist ...string) proxychain.ResponseModification {
+	return func(px *proxychain.ProxyChain) error {
+		// Convert whitelist slice to a map for efficient lookups
+		whitelistMap := make(map[string]struct{})
+		for _, cookieName := range whitelist {
+			whitelistMap[cookieName] = struct{}{}
+		}
+
+		// If the response has no cookies, return early
+		if px.Response.Header == nil {
+			return nil
+		}
+
+		// Filter the cookies in the response
+		filteredCookies := []string{}
+		for _, cookieStr := range px.Response.Header["Set-Cookie"] {
+			cookie := parseCookie(cookieStr)
+			if _, found := whitelistMap[cookie.Name]; found {
+				filteredCookies = append(filteredCookies, cookieStr)
+			}
+		}
+
+		// Update the Set-Cookie header with the filtered cookies
+		if len(filteredCookies) > 0 {
+			px.Response.Header["Set-Cookie"] = filteredCookies
+		} else {
+			px.Response.Header.Del("Set-Cookie")
+		}
+
+		return nil
+	}
+}
+
+// parseCookie parses a cookie string and returns an http.Cookie object.
+func parseCookie(cookieStr string) *http.Cookie {
+	header := http.Header{}
+	header.Add("Set-Cookie", cookieStr)
+	request := http.Request{Header: header}
+	return request.Cookies()[0]
+}
+
+// SetIncomingCookies adds a raw cookie string being sent from the proxy server down to the client
+func SetIncomingCookies(cookies string) proxychain.ResponseModification {
+	return func(px *proxychain.ProxyChain) error {
+		px.Response.Header.Set("Set-Cookie", cookies)
+		return nil
+	}
+}
+
+// SetIncomingCookie modifies a specific cookie in the response from the proxy server to the client.
+func SetIncomingCookie(name string, val string) proxychain.ResponseModification {
+	return func(px *proxychain.ProxyChain) error {
+		if px.Response.Header == nil {
+			return nil
+		}
+
+		updatedCookies := []string{}
+		found := false
+
+		// Iterate over existing cookies and modify the one that matches the cookieName
+		for _, cookieStr := range px.Response.Header["Set-Cookie"] {
+			cookie := parseCookie(cookieStr)
+			if cookie.Name == name {
+				// Replace the cookie with the new value
+				updatedCookies = append(updatedCookies, fmt.Sprintf("%s=%s", name, val))
+				found = true
+			} else {
+				// Keep the cookie as is
+				updatedCookies = append(updatedCookies, cookieStr)
+			}
+		}
+
+		// If the specified cookie wasn't found, add it
+		if !found {
+			updatedCookies = append(updatedCookies, fmt.Sprintf("%s=%s", name, val))
+		}
+
+		// Update the Set-Cookie header
+		px.Response.Header["Set-Cookie"] = updatedCookies
+
+		return nil
+	}
+}

proxychain/responsemodifers/modify_response_header.go

@@ -0,0 +1,21 @@
+package responsemodifers
+
+import (
+	"ladder/proxychain"
+)
+
+// SetResponseHeader modifies response headers from the upstream server
+func SetResponseHeader(key string, value string) proxychain.ResponseModification {
+	return func(px *proxychain.ProxyChain) error {
+		px.Context.Response().Header.Set(key, value)
+		return nil
+	}
+}
+
+// DeleteResponseHeader removes response headers from the upstream server
+func DeleteResponseHeader(key string) proxychain.ResponseModification {
+	return func(px *proxychain.ProxyChain) error {
+		px.Context.Response().Header.Del(key)
+		return nil
+	}
+}

proxychain/responsemodifers/rewrite_http_resource_urls.go

@@ -0,0 +1,114 @@
+package responsemodifers
+
+import (
+	"bytes"
+	"io"
+	"ladder/proxychain"
+	"net/url"
+	"strings"
+
+	"golang.org/x/net/html"
+)
+
+type HTMLResourceURLRewriter struct {
+	src    io.Reader
+	buffer *bytes.Buffer 	// buffer to temporarily hold rewritten output for the reader
+	proxyURL    *url.URL    // proxyURL is the URL of the proxy, not the upstream URL
+}
+
+func NewHTMLResourceURLRewriter(src io.Reader, proxyURL *url.URL) *HTMLResourceURLRewriter {
+	return &HTMLResourceURLRewriter{
+		src:    	 src,
+		buffer: 	 new(bytes.Buffer),
+		proxyURL:    proxyURL,
+	}
+}
+
+func rewriteToken(token *html.Token, baseURL *url.URL) {
+	attrsToRewrite := map[string]bool{"href": true, "src": true, "action": true, "srcset": true}
+	for i := range token.Attr {
+		attr := &token.Attr[i]
+		if attrsToRewrite[attr.Key] && strings.HasPrefix(attr.Val, "/") {
+			// Make URL absolute
+			attr.Val = "/https://" + baseURL.Host + attr.Val
+		}
+	}
+}
+
+func (r *HTMLResourceURLRewriter) Read(p []byte) (int, error) {
+	if r.buffer.Len() != 0 {
+		return r.buffer.Read(p)
+	}
+
+		tokenizer := html.NewTokenizer(r.src)
+		for {
+			tokenType := tokenizer.Next()
+			if tokenType == html.ErrorToken {
+				err := tokenizer.Err()
+				if err == io.EOF {
+					return 0, io.EOF // End of document
+				}
+				return 0, err // Actual error
+			}
+			token := tokenizer.Token()
+			if tokenType == html.StartTagToken || tokenType == html.SelfClosingTagToken {
+				rewriteToken(&token, r.url)
+			}
+			r.buffer.WriteString(token.String())
+			if r.buffer.Len() > 0 {
+				break
+			}
+		}
+	}
+}
+
+
+// RewriteHTMLResourceURLs updates src/href attributes in HTML content to route through the proxy.
+func RewriteHTMLResourceURLs() proxychain.ResponseModification {
+	return func(chain *proxychain.ProxyChain) error {
+		ct := chain.Response.Header.Get("content-type")
+		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)
+				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
+}