Merge pull request #2120 from ujjwal96/replace-from-file

Replace from file through replacement editor #1458

docs/features/replacements.rst

@@ -48,14 +48,14 @@ In practice, it's pretty common for the replacement literal to be long and
 complex. For instance, it might be an XSS exploit that weighs in at hundreds or
 thousands of characters. To cope with this, there's a variation of the
 replacement hook specifier that lets you load the replacement text from a file.
-So, you might start **mitmdump** as follows:
+To specify a file as replacement, prefix the file path with ``@``.
+You might start **mitmdump** as follows:
 
->>> mitmdump --replace-from-file :~q:foo:~/xss-exploit
+>>> mitmdump --replacements :~q:foo:@~/xss-exploit
 
 This will load the replacement text from the file ``~/xss-exploit``.
 
-Both the ``--replace`` and ``--replace-from-file`` flags can be passed multiple
-times.
+The ``--replacements`` flag can be passed multiple times.
 
 
 Interactively
@@ -66,7 +66,6 @@ replacement hooks using a built-in editor. The context-sensitive help (:kbd:`?`)
 complete usage information.
 
 ================== =======================
-command-line       ``--replace``,
-                   ``--replace-from-file``
+command-line       ``--replacements``
 mitmproxy shortcut :kbd:`O` then :kbd:`R`
 ================== =======================

mitmproxy/addons/__init__.py

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

mitmproxy/addons/replace.py

@@ -1,3 +1,4 @@
+import os
 import re
 
 from mitmproxy import exceptions
@@ -42,7 +43,7 @@ def parse_hook(s):
     return patt, a, b
 
 
-class _ReplaceBase:
+class Replace:
     def __init__(self):
         self.lst = []
 
@@ -51,12 +52,12 @@ class _ReplaceBase:
             .replacements is a list of tuples (fpat, rex, s):
 
             fpatt: a string specifying a filter pattern.
-            rex: a regular expression, as bytes.
-            s: the replacement string, as bytes
+            rex: a regular expression, as string.
+            s: the replacement string
         """
-        if self.optionName in updated:
+        if "replacements" in updated:
             lst = []
-            for rep in getattr(options, self.optionName):
+            for rep in options.replacements:
                 fpatt, rex, s = parse_hook(rep)
 
                 flt = flowfilter.parse(fpatt)
@@ -65,11 +66,16 @@ class _ReplaceBase:
                         "Invalid filter pattern: %s" % fpatt
                     )
                 try:
+                    # We should ideally escape here before trying to compile
                     re.compile(rex)
                 except re.error as e:
                     raise exceptions.OptionsError(
                         "Invalid regular expression: %s - %s" % (rex, str(e))
                     )
+                if s.startswith("@") and not os.path.isfile(s[1:]):
+                    raise exceptions.OptionsError(
+                        "Invalid file path: {}".format(s[1:])
+                    )
                 lst.append((rex, s, flt))
             self.lst = lst
 
@@ -89,21 +95,13 @@ class _ReplaceBase:
         if not flow.reply.has_message:
             self.execute(flow)
 
-
-class Replace(_ReplaceBase):
-    optionName = "replacements"
-
     def replace(self, obj, rex, s):
+        if s.startswith("@"):
+            s = os.path.expanduser(s[1:])
+            try:
+                with open(s, "rb") as f:
+                    s = f.read()
+            except IOError:
+                ctx.log.warn("Could not read replacement file: %s" % s)
+                return
         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

@@ -120,13 +120,6 @@ class Options(optmanager.OptManager):
             the separator can be any character.
             """
         )
-        self.add_option(
-            "replacement_files", Sequence[str], [],
-            """
-            Replacement pattern, where the replacement clause is a path to a
-            file.
-            """
-        )
         self.add_option(
             "server_replay_use_headers", Sequence[str], [],
             "Request headers to be considered during replay."

mitmproxy/tools/cmdline.py

@@ -93,7 +93,6 @@ def common_options(parser, opts):
     # Replacements
     group = parser.add_argument_group("Replacements")
     opts.make_parser(group, "replacements", metavar="PATTERN", short="R")
-    opts.make_parser(group, "replacement_files", metavar="PATTERN")
 
     # Set headers
     group = parser.add_argument_group("Set Headers")

mitmproxy/tools/console/grideditor/editors.py

@@ -1,5 +1,8 @@
+import os
 import re
+
 import urwid
+
 from mitmproxy import exceptions
 from mitmproxy import flowfilter
 from mitmproxy.addons import script
@@ -87,6 +90,9 @@ class ReplaceEditor(base.GridEditor):
                 re.compile(val)
             except re.error:
                 return "Invalid regular expression."
+        elif col == 2:
+            if val.startswith("@") and not os.path.isfile(os.path.expanduser(val[1:])):
+                return "Invalid file path"
         return False
 
 

test/mitmproxy/addons/test_replace.py

@@ -1,6 +1,5 @@
 import pytest
 
-from .. import tservers
 from mitmproxy.addons import replace
 from mitmproxy.test import taddons
 from mitmproxy.test import tflow
@@ -32,7 +31,7 @@ class TestReplace:
         with taddons.context() as tctx:
             tctx.configure(
                 r,
-                replacements = [
+                replacements=[
                     "/~q/foo/bar",
                     "/~s/foo/bar",
                 ]
@@ -47,53 +46,57 @@ class TestReplace:
             r.response(f)
             assert f.response.content == b"bar"
 
-
-class TestUpstreamProxy(tservers.HTTPUpstreamProxyTest):
-    ssl = False
-
     def test_order(self):
-        sa = replace.Replace()
-        self.proxy.tmaster.addons.add(sa)
-
-        self.proxy.tmaster.options.replacements = [
-            "/~q/foo/bar",
-            "/~q/bar/baz",
-            "/~q/foo/oh noes!",
-            "/~s/baz/ORLY"
-        ]
-        p = self.pathoc()
-        with p.connect():
-            req = p.request("get:'%s/p/418:b\"foo\"'" % self.server.urlbase)
-        assert req.content == b"ORLY"
-        assert req.status_code == 418
+        r = replace.Replace()
+        with taddons.context() as tctx:
+            tctx.configure(
+                r,
+                replacements=[
+                    "/foo/bar",
+                    "/bar/baz",
+                    "/foo/oh noes!",
+                    "/bar/oh noes!",
+                ]
+            )
+            f = tflow.tflow()
+            f.request.content = b"foo"
+            r.request(f)
+            assert f.request.content == b"baz"
 
 
 class TestReplaceFile:
     def test_simple(self, tmpdir):
-        r = replace.ReplaceFile()
-        rp = tmpdir.join("replacement")
-        rp.write("bar")
+        r = replace.Replace()
         with taddons.context() as tctx:
+            tmpfile = tmpdir.join("replacement")
+            tmpfile.write("bar")
             tctx.configure(
                 r,
-                replacement_files = [
-                    "/~q/foo/" + str(rp),
-                    "/~s/foo/" + str(rp),
-                    "/~b nonexistent/nonexistent/nonexistent",
-                ]
+                replacements=["/~q/foo/@" + str(tmpfile)]
             )
             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"
+    def test_nonexistent(self, tmpdir):
+        r = replace.Replace()
+        with taddons.context() as tctx:
+            with pytest.raises(Exception, match="Invalid file path"):
+                tctx.configure(
+                    r,
+                    replacements=["/~q/foo/@nonexistent"]
+                )
 
+            tmpfile = tmpdir.join("replacement")
+            tmpfile.write("bar")
+            tctx.configure(
+                r,
+                replacements=["/~q/foo/@" + str(tmpfile)]
+            )
+            tmpfile.remove()
             f = tflow.tflow()
-            f.request.content = b"nonexistent"
+            f.request.content = b"foo"
             assert not tctx.master.event_log
             r.request(f)
             assert tctx.master.event_log