From c2443f27ef5a844f045c038c7689d217d1dbf0c4 Mon Sep 17 00:00:00 2001
From: srdusr <99972264+srdusr@users.noreply.github.com>
Date: Mon, 23 Sep 2024 21:33:00 +0200
Subject: Match-and-replace: header rewrite rules

Implements build-order step 6, scoped to headers only for this pass -
see PLAN.md for why bodies are a separate problem (request-body capture
currently depends on streaming straight through, which a body-rewriting
rule would have to interrupt; deciding what "exact" means for a
rule-modified request needs its own pass, not a rushed add-on to this
one).

internal/rules: Rule type and ApplyHeaders, which serializes a Header
map to a raw "Name: value\r\n" block, runs enabled rules' match/replace
over that text, and reparses it - operating on text rather than
per-value substitution is what lets a rule add or remove a header, not
just rewrite one, matching how Burp's header match/replace works.
Invalid rule output (bad regex, unparseable result) leaves the header
map untouched rather than corrupting the request.

internal/store: rules table + CRUD. internal/proxy: forward() fetches
enabled rules for each scope and applies them to outReq.Header /
resp.Header, positioned so the existing capture/history pipeline is
untouched - request_raw keeps showing what the client actually sent and
response_raw what the origin actually sent, while the wire itself
reflects the rules. Deliberate split: match-and-replace transforms
traffic, it doesn't rewrite the audit trail. internal/ipc gains
rules_list/rules_save/rules_delete/rules_toggle. cmd/mitmux gains a
rules view ('m' from history) with add/edit/delete/toggle and a small
form (name, match, replace, scope, regex).

Verified live against real external traffic, not just local echoes:
a request-scope rule rewriting User-Agent, confirmed via httpbin.org's
own header echo that the origin received the rewritten value while curl
sent the real one; a response-scope rule rewriting the Server header,
confirmed the client actually received the rewritten value; disabling a
rule confirmed via a follow-up request that it stops applying; and
throughout, history continued showing the pre-rule original on both
sides, confirming the capture/transform split holds.
---
 internal/rules/rules.go | 97 +++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 97 insertions(+)
 create mode 100644 internal/rules/rules.go

(limited to 'internal/rules/rules.go')

diff --git a/internal/rules/rules.go b/internal/rules/rules.go
new file mode 100644
index 0000000..1f50ec4
--- /dev/null
+++ b/internal/rules/rules.go
@@ -0,0 +1,97 @@
+// Package rules implements match-and-replace: user-defined rules that
+// rewrite request/response headers as they pass through the proxy.
+// Deliberately headers-only for now - see ApplyHeaders for why bodies
+// are a separate, harder problem (noted as a follow-up in PLAN.md).
+package rules
+
+import (
+	"bufio"
+	"net/http"
+	"net/textproto"
+	"regexp"
+	"sort"
+	"strings"
+)
+
+// Rule is one match-and-replace rule.
+type Rule struct {
+	ID      int64
+	Enabled bool
+	Name    string
+	Scope   string // "request" or "response"
+	Part    string // "header" (only part supported so far)
+	Match   string
+	Replace string
+	IsRegex bool
+	// Position orders rule application (ascending) when several rules
+	// could touch the same text.
+	Position int
+}
+
+// ApplyHeaders rewrites h in place by serializing it to a raw
+// "Name: value\r\n" block, running every enabled rule with Part=="header"
+// over that text (in Position order), and reparsing the result. Working
+// on the raw text rather than per-value substitution is what lets a rule
+// add or remove a header entirely, not just rewrite an existing value -
+// matching how Burp's header match/replace works. If a rule's output
+// doesn't parse back as valid headers, ApplyHeaders returns h unchanged
+// rather than risk sending something corrupted.
+func ApplyHeaders(h http.Header, rs []Rule) http.Header {
+	keys := make([]string, 0, len(h))
+	for k := range h {
+		keys = append(keys, k)
+	}
+	sort.Strings(keys)
+
+	var block strings.Builder
+	for _, k := range keys {
+		for _, v := range h[k] {
+			block.WriteString(k)
+			block.WriteString(": ")
+			block.WriteString(v)
+			block.WriteString("\r\n")
+		}
+	}
+	text := block.String()
+
+	changed := false
+	for _, r := range sortedByPosition(rs) {
+		if !r.Enabled || r.Part != "header" {
+			continue
+		}
+		if next, ok := apply(text, r); ok {
+			text, changed = next, true
+		}
+	}
+	if !changed {
+		return h
+	}
+
+	tp := textproto.NewReader(bufio.NewReader(strings.NewReader(text + "\r\n")))
+	mh, err := tp.ReadMIMEHeader()
+	if err != nil {
+		return h
+	}
+	return http.Header(mh)
+}
+
+func sortedByPosition(rs []Rule) []Rule {
+	out := make([]Rule, len(rs))
+	copy(out, rs)
+	sort.SliceStable(out, func(i, j int) bool { return out[i].Position < out[j].Position })
+	return out
+}
+
+func apply(text string, r Rule) (string, bool) {
+	if r.IsRegex {
+		re, err := regexp.Compile(r.Match)
+		if err != nil {
+			return text, false
+		}
+		return re.ReplaceAllString(text, r.Replace), true
+	}
+	if r.Match == "" {
+		return text, false
+	}
+	return strings.ReplaceAll(text, r.Match, r.Replace), true
+}
-- 
cgit v1.2.3