Merge pull request #1974 from cortesi/setheaders

Regularise setheaders options

6 files changed, 73 insertions, 93 deletions

diff --git a/mitmproxy/addons/setheaders.py b/mitmproxy/addons/setheaders.py
index 601e7521..95cf9a09 100644
--- a/mitmproxy/addons/setheaders.py
+++ b/mitmproxy/addons/setheaders.py
@@ -2,6 +2,43 @@ from mitmproxy import exceptions
 from mitmproxy import flowfilter
 
 
+def parse_setheader(s):
+    """
+        Returns a (pattern, regex, replacement) tuple.
+
+        The general form for a replacement hook is as follows:
+
+            /patt/regex/replacement
+
+        The first character specifies the separator. Example:
+
+            :~q:foo:bar
+
+        If only two clauses are specified, the pattern is set to match
+        universally (i.e. ".*"). Example:
+
+            /foo/bar/
+
+        Clauses are parsed from left to right. Extra separators are taken to be
+        part of the final clause. For instance, the replacement clause below is
+        "foo/bar/":
+
+            /one/two/foo/bar/
+    """
+    sep, rem = s[0], s[1:]
+    parts = rem.split(sep, 2)
+    if len(parts) == 2:
+        patt = ".*"
+        a, b = parts
+    elif len(parts) == 3:
+        patt, a, b = parts
+    else:
+        raise exceptions.OptionsError(
+            "Invalid replacement specifier: %s" % s
+        )
+    return patt, a, b
+
+
 class SetHeaders:
     def __init__(self):
         self.lst = []
@@ -16,7 +53,12 @@ class SetHeaders:
         """
         if "setheaders" in updated:
             self.lst = []
-            for fpatt, header, value in options.setheaders:
+            for shead in options.setheaders:
+                if isinstance(shead, str):
+                    fpatt, header, value = parse_setheader(shead)
+                else:
+                    fpatt, header, value = shead
+
                 flt = flowfilter.parse(fpatt)
                 if not flt:
                     raise exceptions.OptionsError(
diff --git a/mitmproxy/options.py b/mitmproxy/options.py
index 3b64cc6a..630d2964 100644
--- a/mitmproxy/options.py
+++ b/mitmproxy/options.py
@@ -41,7 +41,7 @@ class Options(optmanager.OptManager):
         replacements: Sequence[Union[Tuple[str, str, str], str]] = [],
         replacement_files: Sequence[Union[Tuple[str, str, str], str]] = [],
         server_replay_use_headers: Sequence[str] = [],
-        setheaders: Sequence[Tuple[str, str, str]] = [],
+        setheaders: Sequence[Union[Tuple[str, str, str], str]] = [],
         server_replay: Sequence[str] = [],
         stickycookie: Optional[str] = None,
         stickyauth: Optional[str] = None,
diff --git a/mitmproxy/tools/cmdline.py b/mitmproxy/tools/cmdline.py
index bb11b9c2..1c620fd6 100644
--- a/mitmproxy/tools/cmdline.py
+++ b/mitmproxy/tools/cmdline.py
@@ -2,7 +2,6 @@ import argparse
 import os
 
 from mitmproxy import exceptions
-from mitmproxy import flowfilter
 from mitmproxy import options
 from mitmproxy import platform
 from mitmproxy.utils import human
@@ -18,57 +17,6 @@ class ParseException(Exception):
     pass
 
 
-def _parse_hook(s):
-    sep, rem = s[0], s[1:]
-    parts = rem.split(sep, 2)
-    if len(parts) == 2:
-        patt = ".*"
-        a, b = parts
-    elif len(parts) == 3:
-        patt, a, b = parts
-    else:
-        raise ParseException(
-            "Malformed hook specifier - too few clauses: %s" % s
-        )
-
-    if not a:
-        raise ParseException("Empty clause: %s" % str(patt))
-
-    if not flowfilter.parse(patt):
-        raise ParseException("Malformed filter pattern: %s" % patt)
-
-    return patt, a, b
-
-
-def parse_setheader(s):
-    """
-        Returns a (pattern, header, value) tuple.
-
-        The general form for a replacement hook is as follows:
-
-            /patt/header/value
-
-        The first character specifies the separator. Example:
-
-            :~q:foo:bar
-
-        If only two clauses are specified, the pattern is set to match
-        universally (i.e. ".*"). Example:
-
-            /foo/bar/
-
-        Clauses are parsed from left to right. Extra separators are taken to be
-        part of the final clause. For instance, the value clause below is
-        "foo/bar/":
-
-            /one/two/foo/bar/
-
-        Checks that pattern and regex are both well-formed. Raises
-        ParseException on error.
-    """
-    return _parse_hook(s)
-
-
 def get_common_options(args):
     stickycookie, stickyauth = None, None
     if args.stickycookie_filt:
@@ -81,14 +29,6 @@ def get_common_options(args):
     if stream_large_bodies:
         stream_large_bodies = human.parse_size(stream_large_bodies)
 
-    setheaders = []
-    for i in args.setheader or []:
-        try:
-            p = parse_setheader(i)
-        except ParseException as e:
-            raise exceptions.OptionsError(e)
-        setheaders.append(p)
-
     if args.streamfile and args.streamfile[0] == args.rfile:
         if args.streamfile[1] == "wb":
             raise exceptions.OptionsError(
@@ -171,7 +111,7 @@ def get_common_options(args):
         rfile=args.rfile,
         replacements=args.replacements,
         replacement_files=args.replacement_files,
-        setheaders=setheaders,
+        setheaders=args.setheaders,
         server_replay=args.server_replay,
         scripts=args.scripts,
         stickycookie=stickycookie,
@@ -648,7 +588,7 @@ def set_headers(parser):
     )
     group.add_argument(
         "--setheader",
-        action="append", type=str, dest="setheader",
+        action="append", type=str, dest="setheaders",
         metavar="PATTERN",
         help="Header set pattern."
     )
@@ -708,8 +648,8 @@ def common_options(parser):
 
 
 def mitmproxy():
-    # Don't import mitmproxy.tools.console for mitmdump, urwid is not available on all
-    # platforms.
+    # Don't import mitmproxy.tools.console for mitmdump, urwid is not available
+    # on all platforms.
     from .console import palettes
 
     parser = argparse.ArgumentParser(usage="%(prog)s [options]")
diff --git a/mitmproxy/tools/console/options.py b/mitmproxy/tools/console/options.py
index e88006fe..4115bd18 100644
--- a/mitmproxy/tools/console/options.py
+++ b/mitmproxy/tools/console/options.py
@@ -7,6 +7,7 @@ from mitmproxy.tools.console import select
 from mitmproxy.tools.console import signals
 
 from mitmproxy.addons import replace
+from mitmproxy.addons import setheaders
 
 footer = [
     ('heading_key', "enter/space"), ":toggle ",
@@ -190,10 +191,16 @@ class Options(urwid.WidgetWrap):
         )
 
     def setheaders(self):
+        data = []
+        for d in self.master.options.setheaders:
+            if isinstance(d, str):
+                data.append(setheaders.parse_setheader(d))
+            else:
+                data.append(d)
         self.master.view_grideditor(
             grideditor.SetHeadersEditor(
                 self.master,
-                self.master.options.setheaders,
+                data,
                 self.master.options.setter("setheaders")
             )
         )
diff --git a/test/mitmproxy/addons/test_setheaders.py b/test/mitmproxy/addons/test_setheaders.py
index 34395ddf..a721c180 100644
--- a/test/mitmproxy/addons/test_setheaders.py
+++ b/test/mitmproxy/addons/test_setheaders.py
@@ -3,19 +3,28 @@ from mitmproxy.test import tutils
 from mitmproxy.test import taddons
 
 from mitmproxy.addons import setheaders
-from mitmproxy import options
 
 
 class TestSetHeaders:
+    def test_parse_setheaders(self):
+        x = setheaders.parse_setheader("/foo/bar/voing")
+        assert x == ("foo", "bar", "voing")
+        x = setheaders.parse_setheader("/foo/bar/vo/ing/")
+        assert x == ("foo", "bar", "vo/ing/")
+        x = setheaders.parse_setheader("/bar/voing")
+        assert x == (".*", "bar", "voing")
+        tutils.raises("invalid replacement", setheaders.parse_setheader, "/")
+
     def test_configure(self):
         sh = setheaders.SetHeaders()
-        o = options.Options(
-            setheaders = [("~b", "one", "two")]
-        )
-        tutils.raises(
-            "invalid setheader filter pattern",
-            sh.configure, o, o.keys()
-        )
+        with taddons.context() as tctx:
+            tutils.raises(
+                "invalid setheader filter pattern",
+                tctx.configure,
+                sh,
+                setheaders = [("~b", "one", "two")]
+            )
+            tctx.configure(sh, setheaders = ["/foo/bar/voing"])
 
     def test_setheaders(self):
         sh = setheaders.SetHeaders()
diff --git a/test/mitmproxy/test_cmdline.py b/test/mitmproxy/test_cmdline.py
index fe0373d1..96d5ae31 100644
--- a/test/mitmproxy/test_cmdline.py
+++ b/test/mitmproxy/test_cmdline.py
@@ -1,11 +1,5 @@
 import argparse
 from mitmproxy.tools import cmdline
-from mitmproxy.test import tutils
-
-
-def test_parse_setheaders():
-    x = cmdline.parse_setheader("/foo/bar/voing")
-    assert x == ("foo", "bar", "voing")
 
 
 def test_common():
@@ -21,18 +15,6 @@ def test_common():
     assert v["stickycookie"] == "foo"
     assert v["stickyauth"] == "foo"
 
-    opts.setheader = ["/foo/bar/voing"]
-    v = cmdline.get_common_options(opts)
-    assert v["setheaders"] == [("foo", "bar", "voing")]
-
-    opts.setheader = ["//"]
-    tutils.raises(
-        "empty clause",
-        cmdline.get_common_options,
-        opts
-    )
-    opts.setheader = []
-
 
 def test_mitmproxy():
     ap = cmdline.mitmproxy()