Merge pull request #1973 from cortesi/replace

Revamp replacement hooks

mitmproxy/addons/__init__.py

@@ -28,6 +28,7 @@ def default_addons():
         onboarding.Onboarding(),
         proxyauth.ProxyAuth(),
         replace.Replace(),
+        replace.ReplaceFile(),
         script.ScriptLoader(),
         serverplayback.ServerPlayback(),
         setheaders.SetHeaders(),

mitmproxy/addons/replace.py

@@ -2,9 +2,47 @@ import re
 
 from mitmproxy import exceptions
 from mitmproxy import flowfilter
+from mitmproxy import ctx
 
 
-class Replace:
+def parse_hook(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 _ReplaceBase:
     def __init__(self):
         self.lst = []
 
@@ -16,9 +54,14 @@ class Replace:
             rex: a regular expression, as bytes.
             s: the replacement string, as bytes
         """
-        if "replacements" in updated:
+        if self.optionName in updated:
             lst = []
-            for fpatt, rex, s in options.replacements:
+            for rep in getattr(options, self.optionName):
+                if isinstance(rep, str):
+                    fpatt, rex, s = parse_hook(rep)
+                else:
+                    fpatt, rex, s = rep
+
                 flt = flowfilter.parse(fpatt)
                 if not flt:
                     raise exceptions.OptionsError(
@@ -37,9 +80,9 @@ class Replace:
         for rex, s, flt in self.lst:
             if flt(f):
                 if f.response:
-                    f.response.replace(rex, s, flags=re.DOTALL)
+                    self.replace(f.response, rex, s)
                 else:
-                    f.request.replace(rex, s, flags=re.DOTALL)
+                    self.replace(f.request, rex, s)
 
     def request(self, flow):
         if not flow.reply.has_message:
@@ -48,3 +91,22 @@ class Replace:
     def response(self, flow):
         if not flow.reply.has_message:
             self.execute(flow)
+
+
+class Replace(_ReplaceBase):
+    optionName = "replacements"
+
+    def replace(self, obj, rex, s):
+        obj.replace(rex, s, flags=re.DOTALL)
+
+
+class ReplaceFile(_ReplaceBase):
+    optionName = "replacement_files"
+
+    def replace(self, obj, rex, s):
+        try:
+            v = open(s, "rb").read()
+        except IOError as e:
+            ctx.log.warn("Could not read replacement file: %s" % s)
+            return
+        obj.replace(rex, v, flags=re.DOTALL)

mitmproxy/options.py

@@ -1,4 +1,4 @@
-from typing import Tuple, Optional, Sequence
+from typing import Tuple, Optional, Sequence, Union
 
 from mitmproxy import optmanager
 
@@ -38,7 +38,8 @@ class Options(optmanager.OptManager):
         rfile: Optional[str] = None,
         scripts: Sequence[str] = [],
         showhost: bool = False,
-        replacements: Sequence[Tuple[str, str, str]] = [],
+        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]] = [],
         server_replay: Sequence[str] = [],
@@ -124,6 +125,7 @@ class Options(optmanager.OptManager):
         self.scripts = scripts
         self.showhost = showhost
         self.replacements = replacements
+        self.replacement_files = replacement_files
         self.server_replay_use_headers = server_replay_use_headers
         self.setheaders = setheaders
         self.server_replay = server_replay

mitmproxy/tools/cmdline.py

@@ -1,5 +1,4 @@
 import argparse
-import re
 import os
 
 from mitmproxy import exceptions
@@ -41,40 +40,6 @@ def _parse_hook(s):
     return patt, a, b
 
 
-def parse_replace_hook(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/
-
-        Checks that pattern and regex are both well-formed. Raises
-        ParseException on error.
-    """
-    patt, regex, replacement = _parse_hook(s)
-    try:
-        re.compile(regex)
-    except re.error as e:
-        raise ParseException("Malformed replacement regex: %s" % str(e))
-    return patt, regex, replacement
-
-
 def parse_setheader(s):
     """
         Returns a (pattern, header, value) tuple.
@@ -116,26 +81,6 @@ def get_common_options(args):
     if stream_large_bodies:
         stream_large_bodies = human.parse_size(stream_large_bodies)
 
-    reps = []
-    for i in args.replace or []:
-        try:
-            p = parse_replace_hook(i)
-        except ParseException as e:
-            raise exceptions.OptionsError(e)
-        reps.append(p)
-    for i in args.replace_file or []:
-        try:
-            patt, rex, path = parse_replace_hook(i)
-        except ParseException as e:
-            raise exceptions.OptionsError(e)
-        try:
-            v = open(path, "rb").read()
-        except IOError as e:
-            raise exceptions.OptionsError(
-                "Could not read replace file: %s" % path
-            )
-        reps.append((patt, rex, v))
-
     setheaders = []
     for i in args.setheader or []:
         try:
@@ -224,7 +169,8 @@ def get_common_options(args):
         refresh_server_playback=not args.norefresh,
         server_replay_use_headers=args.server_replay_use_headers,
         rfile=args.rfile,
-        replacements=reps,
+        replacements=args.replacements,
+        replacement_files=args.replacement_files,
         setheaders=setheaders,
         server_replay=args.server_replay,
         scripts=args.scripts,
@@ -676,13 +622,13 @@ def replacements(parser):
     )
     group.add_argument(
         "--replace",
-        action="append", type=str, dest="replace",
+        action="append", type=str, dest="replacements",
         metavar="PATTERN",
         help="Replacement pattern."
     )
     group.add_argument(
         "--replace-from-file",
-        action="append", type=str, dest="replace_file",
+        action="append", type=str, dest="replacement_files",
         metavar="PATH",
         help="""
             Replacement pattern, where the replacement clause is a path to a

mitmproxy/tools/console/options.py

@@ -6,6 +6,8 @@ from mitmproxy.tools.console import grideditor
 from mitmproxy.tools.console import select
 from mitmproxy.tools.console import signals
 
+from mitmproxy.addons import replace
+
 footer = [
     ('heading_key', "enter/space"), ":toggle ",
     ('heading_key', "C"), ":clear all ",
@@ -215,10 +217,16 @@ class Options(urwid.WidgetWrap):
         )
 
     def replacepatterns(self):
+        data = []
+        for d in self.master.options.replacements:
+            if isinstance(d, str):
+                data.append(replace.parse_hook(d))
+            else:
+                data.append(d)
         self.master.view_grideditor(
             grideditor.ReplaceEditor(
                 self.master,
-                self.master.options.replacements,
+                data,
                 self.master.options.setter("replacements")
             )
         )

test/mitmproxy/addons/test_replace.py

@@ -1,57 +1,59 @@
+import os.path
 from mitmproxy.test import tflow
 from mitmproxy.test import tutils
 
 from .. import tservers
 from mitmproxy.addons import replace
-from mitmproxy import master
-from mitmproxy import options
-from mitmproxy import proxy
+from mitmproxy.test import taddons
 
 
 class TestReplace:
+    def test_parse_hook(self):
+        x = replace.parse_hook("/foo/bar/voing")
+        assert x == ("foo", "bar", "voing")
+        x = replace.parse_hook("/foo/bar/vo/ing/")
+        assert x == ("foo", "bar", "vo/ing/")
+        x = replace.parse_hook("/bar/voing")
+        assert x == (".*", "bar", "voing")
+        tutils.raises("invalid replacement", replace.parse_hook, "/")
+
     def test_configure(self):
         r = replace.Replace()
-        updated = set(["replacements"])
-        r.configure(options.Options(
-            replacements=[("one", "two", "three")]
-        ), updated)
-        tutils.raises(
-            "invalid filter pattern",
-            r.configure,
-            options.Options(
+        with taddons.context() as tctx:
+            tctx.configure(r, replacements=[("one", "two", "three")])
+            tutils.raises(
+                "invalid filter pattern",
+                tctx.configure,
+                r,
                 replacements=[("~b", "two", "three")]
-            ),
-            updated
-        )
-        tutils.raises(
-            "invalid regular expression",
-            r.configure,
-            options.Options(
+            )
+            tutils.raises(
+                "invalid regular expression",
+                tctx.configure,
+                r,
                 replacements=[("foo", "+", "three")]
-            ),
-            updated
-        )
+            )
+            tctx.configure(r, replacements=["/a/b/c/"])
 
     def test_simple(self):
-        o = options.Options(
-            replacements = [
-                ("~q", "foo", "bar"),
-                ("~s", "foo", "bar"),
-            ]
-        )
-        m = master.Master(o, proxy.DummyServer())
-        sa = replace.Replace()
-        m.addons.add(sa)
+        r = replace.Replace()
+        with taddons.context() as tctx:
+            tctx.configure(
+                r,
+                replacements = [
+                    ("~q", "foo", "bar"),
+                    ("~s", "foo", "bar"),
+                ]
+            )
+            f = tflow.tflow()
+            f.request.content = b"foo"
+            r.request(f)
+            assert f.request.content == b"bar"
 
-        f = tflow.tflow()
-        f.request.content = b"foo"
-        m.request(f)
-        assert f.request.content == b"bar"
-
-        f = tflow.tflow(resp=True)
-        f.response.content = b"foo"
-        m.response(f)
-        assert f.response.content == b"bar"
+            f = tflow.tflow(resp=True)
+            f.response.content = b"foo"
+            r.response(f)
+            assert f.response.content == b"bar"
 
 
 class TestUpstreamProxy(tservers.HTTPUpstreamProxyTest):
@@ -72,3 +74,36 @@ class TestUpstreamProxy(tservers.HTTPUpstreamProxyTest):
             req = p.request("get:'%s/p/418:b\"foo\"'" % self.server.urlbase)
         assert req.content == b"ORLY"
         assert req.status_code == 418
+
+
+class TestReplaceFile:
+    def test_simple(self):
+        r = replace.ReplaceFile()
+        with tutils.tmpdir() as td:
+            rp = os.path.join(td, "replacement")
+            with open(rp, "w") as f:
+                f.write("bar")
+            with taddons.context() as tctx:
+                tctx.configure(
+                    r,
+                    replacement_files = [
+                        ("~q", "foo", rp),
+                        ("~s", "foo", rp),
+                        ("~b nonexistent", "nonexistent", "nonexistent"),
+                    ]
+                )
+                f = tflow.tflow()
+                f.request.content = b"foo"
+                r.request(f)
+                assert f.request.content == b"bar"
+
+                f = tflow.tflow(resp=True)
+                f.response.content = b"foo"
+                r.response(f)
+                assert f.response.content == b"bar"
+
+                f = tflow.tflow()
+                f.request.content = b"nonexistent"
+                assert not tctx.master.event_log
+                r.request(f)
+                assert tctx.master.event_log

test/mitmproxy/test_cmdline.py

@@ -3,38 +3,6 @@ from mitmproxy.tools import cmdline
 from mitmproxy.test import tutils
 
 
-def test_parse_replace_hook():
-    x = cmdline.parse_replace_hook("/foo/bar/voing")
-    assert x == ("foo", "bar", "voing")
-
-    x = cmdline.parse_replace_hook("/foo/bar/vo/ing/")
-    assert x == ("foo", "bar", "vo/ing/")
-
-    x = cmdline.parse_replace_hook("/bar/voing")
-    assert x == (".*", "bar", "voing")
-
-    tutils.raises(
-        cmdline.ParseException,
-        cmdline.parse_replace_hook,
-        "/foo"
-    )
-    tutils.raises(
-        "replacement regex",
-        cmdline.parse_replace_hook,
-        "patt/[/rep"
-    )
-    tutils.raises(
-        "filter pattern",
-        cmdline.parse_replace_hook,
-        "/~/foo/rep"
-    )
-    tutils.raises(
-        "empty clause",
-        cmdline.parse_replace_hook,
-        "//"
-    )
-
-
 def test_parse_setheaders():
     x = cmdline.parse_setheader("/foo/bar/voing")
     assert x == ("foo", "bar", "voing")
@@ -65,38 +33,6 @@ def test_common():
     )
     opts.setheader = []
 
-    opts.replace = ["/foo/bar/voing"]
-    v = cmdline.get_common_options(opts)
-    assert v["replacements"] == [("foo", "bar", "voing")]
-
-    opts.replace = ["//"]
-    tutils.raises(
-        "empty clause",
-        cmdline.get_common_options,
-        opts
-    )
-
-    opts.replace = []
-    opts.replace_file = [("/foo/bar/nonexistent")]
-    tutils.raises(
-        "could not read replace file",
-        cmdline.get_common_options,
-        opts
-    )
-
-    opts.replace_file = [("/~/bar/nonexistent")]
-    tutils.raises(
-        "filter pattern",
-        cmdline.get_common_options,
-        opts
-    )
-
-    p = tutils.test_data.path("mitmproxy/data/replace")
-    opts.replace_file = [("/foo/bar/%s" % p)]
-    v = cmdline.get_common_options(opts)["replacements"]
-    assert len(v) == 1
-    assert v[0][2].strip() == b"replacecontents"
-
 
 def test_mitmproxy():
     ap = cmdline.mitmproxy()