Permalink
Browse files

replace() methods now decode and re-encode contents before substitution.

  • Loading branch information...
1 parent d51b8ca commit d138af72171d833659cfb53edc80eade121ca836 @cortesi cortesi committed Mar 15, 2012
Showing with 41 additions and 9 deletions.
  1. +23 −9 libmproxy/flow.py
  2. +18 −0 test/test_flow.py
View
@@ -176,8 +176,11 @@ def match_re(self, expr):
def replace(self, pattern, repl, *args, **kwargs):
"""
- Replaces a regular expression pattern with repl in both keys
- and values. Returns the number of replacements made.
+ Replaces a regular expression pattern with repl in both keys and
+ values. Encoded content will be decoded before replacement, and
+ re-encoded afterwards.
+
+ Returns the number of replacements made.
"""
nlst, count = [], 0
for i in self.lst:
@@ -475,10 +478,13 @@ def _assemble(self, _proxy = False):
def replace(self, pattern, repl, *args, **kwargs):
"""
Replaces a regular expression pattern with repl in both the headers
- and the body of the request. Returns the number of replacements
- made.
+ and the body of the request. Encoded content will be decoded before
+ replacement, and re-encoded afterwards.
+
+ Returns the number of replacements made.
"""
- self.content, c = re.subn(pattern, repl, self.content, *args, **kwargs)
+ with decoded(self):
+ self.content, c = re.subn(pattern, repl, self.content, *args, **kwargs)
self.path, pc = re.subn(pattern, repl, self.path, *args, **kwargs)
c += pc
c += self.headers.replace(pattern, repl, *args, **kwargs)
@@ -630,10 +636,13 @@ def _assemble(self):
def replace(self, pattern, repl, *args, **kwargs):
"""
Replaces a regular expression pattern with repl in both the headers
- and the body of the response. Returns the number of replacements
- made.
+ and the body of the response. Encoded content will be decoded
+ before replacement, and re-encoded afterwards.
+
+ Returns the number of replacements made.
"""
- self.content, c = re.subn(pattern, repl, self.content, *args, **kwargs)
+ with decoded(self):
+ self.content, c = re.subn(pattern, repl, self.content, *args, **kwargs)
c += self.headers.replace(pattern, repl, *args, **kwargs)
return c
@@ -753,6 +762,8 @@ def replace(self, pattern, repl, *args, **kwargs):
Replaces a regular expression pattern with repl in both the headers
and the body of the request. Returns the number of replacements
made.
+
+ FIXME: Is replace useful on an Error object??
"""
self.msg, c = re.subn(pattern, repl, self.msg, *args, **kwargs)
return c
@@ -1062,7 +1073,10 @@ def accept_intercept(self):
def replace(self, pattern, repl, *args, **kwargs):
"""
Replaces a regular expression pattern with repl in all parts of the
- flow . Returns the number of replacements made.
+ flow. Encoded content will be decoded before replacement, and
+ re-encoded afterwards.
+
+ Returns the number of replacements made.
"""
c = self.request.replace(pattern, repl, *args, **kwargs)
if self.response:
View
@@ -278,6 +278,24 @@ def test_replace(self):
f.replace("error", "bar")
assert f.error.msg == "bar"
+ def test_replace_encoded(self):
+ f = tutils.tflow_full()
+ f.request.content = "afoob"
+ f.request.encode("gzip")
+ f.response.content = "afoob"
+ f.response.encode("gzip")
+
+ f.replace("foo", "bar")
+
+ assert f.request.content != "abarb"
+ f.request.decode()
+ assert f.request.content == "abarb"
+
+ assert f.response.content != "abarb"
+ f.response.decode()
+ assert f.response.content == "abarb"
+
+
class uState(libpry.AutoTree):
def test_backup(self):

0 comments on commit d138af7

Please sign in to comment.