From d138af72171d833659cfb53edc80eade121ca836 Mon Sep 17 00:00:00 2001
From: Aldo Cortesi <aldo@nullcube.com>
Date: Fri, 16 Mar 2012 11:24:18 +1300
Subject: replace() methods now decode and re-encode contents before
 substitution.

---
 libmproxy/flow.py | 32 +++++++++++++++++++++++---------
 test/test_flow.py | 18 ++++++++++++++++++
 2 files changed, 41 insertions(+), 9 deletions(-)

diff --git a/libmproxy/flow.py b/libmproxy/flow.py
index dbc0c109..1ae2bb8a 100644
--- a/libmproxy/flow.py
+++ b/libmproxy/flow.py
@@ -176,8 +176,11 @@ class ODict:
 
     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 @@ class Request(HTTPMsg):
     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 @@ class Response(HTTPMsg):
     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 @@ class Error(controller.Msg):
             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 @@ class Flow:
     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:
diff --git a/test/test_flow.py b/test/test_flow.py
index 56303881..ff35f899 100644
--- a/test/test_flow.py
+++ b/test/test_flow.py
@@ -278,6 +278,24 @@ class uFlow(libpry.AutoTree):
         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):
-- 
cgit v1.2.3