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/store/store.go | 109 ++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 109 insertions(+) (limited to 'internal/store') diff --git a/internal/store/store.go b/internal/store/store.go index a3c6b03..211162e 100644 --- a/internal/store/store.go +++ b/internal/store/store.go @@ -12,6 +12,8 @@ import ( "time" _ "modernc.org/sqlite" + + "mitmux/internal/rules" ) const schema = ` @@ -36,6 +38,18 @@ CREATE VIRTUAL TABLE IF NOT EXISTS history_fts USING fts5( method, host, path, request_text, response_text, tokenize = 'unicode61 remove_diacritics 2' ); + +CREATE TABLE IF NOT EXISTS rules ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + enabled INTEGER NOT NULL DEFAULT 1, + name TEXT NOT NULL DEFAULT '', + scope TEXT NOT NULL, + part TEXT NOT NULL, + match TEXT NOT NULL, + replace TEXT NOT NULL, + is_regex INTEGER NOT NULL DEFAULT 0, + position INTEGER NOT NULL DEFAULT 0 +); ` // Store is a handle to the history database. Safe for concurrent use. @@ -312,6 +326,101 @@ func prepareFTSQuery(q string) string { return strings.Join(fields, " ") } +// ListRules returns every match-and-replace rule, ordered for application. +func (s *Store) ListRules() ([]rules.Rule, error) { + rows, err := s.db.Query( + `SELECT id, enabled, name, scope, part, match, replace, is_regex, position + FROM rules ORDER BY position, id`, + ) + if err != nil { + return nil, fmt.Errorf("list rules: %w", err) + } + defer rows.Close() + + var out []rules.Rule + for rows.Next() { + var r rules.Rule + var enabled, isRegex int + if err := rows.Scan(&r.ID, &enabled, &r.Name, &r.Scope, &r.Part, &r.Match, &r.Replace, &isRegex, &r.Position); err != nil { + return nil, fmt.Errorf("scan rule row: %w", err) + } + r.Enabled = enabled != 0 + r.IsRegex = isRegex != 0 + out = append(out, r) + } + return out, rows.Err() +} + +// EnabledRules returns enabled rules for scope ("request" or +// "response"), ordered for application. +func (s *Store) EnabledRules(scope string) ([]rules.Rule, error) { + rows, err := s.db.Query( + `SELECT id, enabled, name, scope, part, match, replace, is_regex, position + FROM rules WHERE enabled = 1 AND scope = ? ORDER BY position, id`, + scope, + ) + if err != nil { + return nil, fmt.Errorf("enabled rules: %w", err) + } + defer rows.Close() + + var out []rules.Rule + for rows.Next() { + var r rules.Rule + var enabled, isRegex int + if err := rows.Scan(&r.ID, &enabled, &r.Name, &r.Scope, &r.Part, &r.Match, &r.Replace, &isRegex, &r.Position); err != nil { + return nil, fmt.Errorf("scan rule row: %w", err) + } + r.Enabled = enabled != 0 + r.IsRegex = isRegex != 0 + out = append(out, r) + } + return out, rows.Err() +} + +// AddRule stores r and returns its assigned ID. +func (s *Store) AddRule(r rules.Rule) (int64, error) { + res, err := s.db.Exec( + `INSERT INTO rules (enabled, name, scope, part, match, replace, is_regex, position) + VALUES (?, ?, ?, ?, ?, ?, ?, ?)`, + boolToInt(r.Enabled), r.Name, r.Scope, r.Part, r.Match, r.Replace, boolToInt(r.IsRegex), r.Position, + ) + if err != nil { + return 0, fmt.Errorf("add rule: %w", err) + } + return res.LastInsertId() +} + +// UpdateRule replaces the stored rule with the same ID as r. +func (s *Store) UpdateRule(r rules.Rule) error { + _, err := s.db.Exec( + `UPDATE rules SET enabled = ?, name = ?, scope = ?, part = ?, match = ?, replace = ?, is_regex = ?, position = ? + WHERE id = ?`, + boolToInt(r.Enabled), r.Name, r.Scope, r.Part, r.Match, r.Replace, boolToInt(r.IsRegex), r.Position, r.ID, + ) + if err != nil { + return fmt.Errorf("update rule %d: %w", r.ID, err) + } + return nil +} + +// SetRuleEnabled toggles a rule without touching its other fields. +func (s *Store) SetRuleEnabled(id int64, enabled bool) error { + _, err := s.db.Exec(`UPDATE rules SET enabled = ? WHERE id = ?`, boolToInt(enabled), id) + if err != nil { + return fmt.Errorf("set rule %d enabled: %w", id, err) + } + return nil +} + +// DeleteRule removes a rule. +func (s *Store) DeleteRule(id int64) error { + if _, err := s.db.Exec(`DELETE FROM rules WHERE id = ?`, id); err != nil { + return fmt.Errorf("delete rule %d: %w", id, err) + } + return nil +} + func boolToInt(b bool) int { if b { return 1 -- cgit v1.2.3