Add SetHeaders, analogous to ReplaceHooks, with a graphical editor in mitmproxy (H shortcut).

SetHeaders defines headers that are set on flows, based on a matching pattern.
Existing headers are over-ridden.

1 files changed, 62 insertions, 5 deletions

diff --git a/test/test_flow.py b/test/test_flow.py
index eccd11f4..47a09360 100644
--- a/test/test_flow.py
+++ b/test/test_flow.py
@@ -1017,7 +1017,16 @@ def test_replacehooks():
     h = flow.ReplaceHooks()
     h.add("~q", "foo", "bar")
     assert h.lst
-    h.remove("~q", "foo", "bar")
+
+    h.set(
+        [
+            (".*", "one", "two"),
+            (".*", "three", "four"),
+        ]
+    )
+    assert h.count() == 2
+
+    h.clear()
     assert not h.lst
 
     h.add("~q", "foo", "bar")
@@ -1026,10 +1035,6 @@ def test_replacehooks():
     v = h.get_specs()
     assert v == [('~q', 'foo', 'bar'), ('~s', 'foo', 'bar')]
     assert h.count() == 2
-    h.remove("~q", "foo", "bar")
-    assert h.count() == 1
-    h.remove("~q", "foo", "bar")
-    assert h.count() == 1
     h.clear()
     assert h.count() == 0
 
@@ -1056,3 +1061,55 @@ def test_replacehooks():
     assert not h.add("~", "foo", "bar")
     assert not h.add("foo", "*", "bar")
 
+
+def test_setheaders():
+    h = flow.SetHeaders()
+    h.add("~q", "foo", "bar")
+    assert h.lst
+
+    h.set(
+        [
+            (".*", "one", "two"),
+            (".*", "three", "four"),
+        ]
+    )
+    assert h.count() == 2
+
+    h.clear()
+    assert not h.lst
+
+    h.add("~q", "foo", "bar")
+    h.add("~s", "foo", "bar")
+
+    v = h.get_specs()
+    assert v == [('~q', 'foo', 'bar'), ('~s', 'foo', 'bar')]
+    assert h.count() == 2
+    h.clear()
+    assert h.count() == 0
+
+    f = tutils.tflow()
+    f.request.content = "foo"
+    h.add("~s", "foo", "bar")
+    h.run(f)
+    assert f.request.content == "foo"
+
+
+    h.clear()
+    h.add("~s", "one", "two")
+    h.add("~s", "one", "three")
+    f = tutils.tflow_full()
+    f.request.headers["one"] = ["xxx"]
+    f.response.headers["one"] = ["xxx"]
+    h.run(f)
+    assert f.request.headers["one"] == ["xxx"]
+    assert f.response.headers["one"] == ["two", "three"]
+
+    h.clear()
+    h.add("~q", "one", "two")
+    h.add("~q", "one", "three")
+    f = tutils.tflow()
+    f.request.headers["one"] = ["xxx"]
+    h.run(f)
+    assert f.request.headers["one"] == ["two", "three"]
+
+    assert not h.add("~", "foo", "bar")