From 4684617d2c8e17d84fd98b635fd5c799dd77b741 Mon Sep 17 00:00:00 2001
From: Maximilian Hils <git@maximilianhils.com>
Date: Sun, 12 Mar 2017 20:13:33 +0100
Subject: [PATCH] minor fixes

---
 mitmproxy/addons/replace.py                   | 17 +++----
 mitmproxy/tools/console/grideditor/editors.py |  2 +-
 test/mitmproxy/addons/test_replace.py         | 45 +++++++++----------
 3 files changed, 31 insertions(+), 33 deletions(-)

diff --git a/mitmproxy/addons/replace.py b/mitmproxy/addons/replace.py
index f66f60c35..d90d37ed2 100644
--- a/mitmproxy/addons/replace.py
+++ b/mitmproxy/addons/replace.py
@@ -46,19 +46,18 @@ def parse_hook(s):
 class Replace:
     def __init__(self):
         self.lst = []
-        self.optionName = "replacements"
 
     def configure(self, options, updated):
         """
             .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)
@@ -67,6 +66,7 @@ class Replace:
                         "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(
@@ -93,10 +93,11 @@ class Replace:
 
     def replace(self, obj, rex, s):
         if s.startswith("@"):
-            s = s.replace("@", "")
+            s = os.path.expanduser(s[1:])
             try:
-                s = open(os.path.expanduser(s), "rb").read()
-            except IOError as e:
+                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)
diff --git a/mitmproxy/tools/console/grideditor/editors.py b/mitmproxy/tools/console/grideditor/editors.py
index 503e9768d..0d9929aea 100644
--- a/mitmproxy/tools/console/grideditor/editors.py
+++ b/mitmproxy/tools/console/grideditor/editors.py
@@ -91,7 +91,7 @@ class ReplaceEditor(base.GridEditor):
             except re.error:
                 return "Invalid regular expression."
         elif col == 2:
-            if val.startswith("@") and not os.path.isfile(os.path.expanduser(val.replace("@", ""))):
+            if val.startswith("@") and not os.path.isfile(os.path.expanduser(val[1:])):
                 return "Invalid file path"
         return False
 
diff --git a/test/mitmproxy/addons/test_replace.py b/test/mitmproxy/addons/test_replace.py
index 79ef7baca..5d2680767 100644
--- a/test/mitmproxy/addons/test_replace.py
+++ b/test/mitmproxy/addons/test_replace.py
@@ -1,11 +1,11 @@
 import os.path
-import pytest
-from mitmproxy.test import tflow
-from mitmproxy.test import tutils
 
-from .. import tservers
+import pytest
+
 from mitmproxy.addons import replace
 from mitmproxy.test import taddons
+from mitmproxy.test import tflow
+from mitmproxy.test import tutils
 
 
 class TestReplace:
@@ -34,7 +34,7 @@ class TestReplace:
         with taddons.context() as tctx:
             tctx.configure(
                 r,
-                replacements = [
+                replacements=[
                     "/~q/foo/bar",
                     "/~s/foo/bar",
                 ]
@@ -49,25 +49,22 @@ 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:
@@ -80,7 +77,7 @@ class TestReplaceFile:
             with taddons.context() as tctx:
                 tctx.configure(
                     r,
-                    replacements = [
+                    replacements=[
                         "/~q/foo/@" + rp,
                         "/~s/foo/@" + rp,
                         "/~b nonexistent/nonexistent/@nonexistent",