Add utility functions to search and replace strings in flows

This is a common task in pentesting scenarios. This commit adds the following
functions:

utils.Headers.replace
proxy.Request.replace
proxy.Response.replace
flow.Flow.replace

3 files changed, 46 insertions, 1 deletions

diff --git a/test/test_flow.py b/test/test_flow.py
index 365b6c58..ba276248 100644
--- a/test/test_flow.py
+++ b/test/test_flow.py
@@ -219,6 +219,21 @@ class uFlow(libpry.AutoTree):
         f = flow.Flow(None)
         f.request = tutils.treq()
 
+    def test_replace(self):
+        f = tutils.tflow_full()
+        f.request.headers["foo"] = ["foo"]
+        f.request.content = "afoob"
+
+        f.response.headers["foo"] = ["foo"]
+        f.response.content = "afoob"
+
+        assert f.replace("foo", "bar") == 6
+
+        assert f.request.headers["bar"] == ["bar"]
+        assert f.request.content == "abarb"
+        assert f.response.headers["bar"] == ["bar"]
+        assert f.response.content == "abarb"
+
 
 class uState(libpry.AutoTree):
     def test_backup(self):
diff --git a/test/test_proxy.py b/test/test_proxy.py
index 46235e97..a7627c1b 100644
--- a/test/test_proxy.py
+++ b/test/test_proxy.py
@@ -1,4 +1,4 @@
-import cStringIO, time
+import cStringIO, time, re
 import libpry
 from libmproxy import proxy, controller, utils, dump, script
 import email.utils
@@ -128,6 +128,14 @@ class uRequest(libpry.AutoTree):
         r.load_state(r2.get_state())
         assert not r.client_conn
 
+    def test_replace(self):
+        r = tutils.treq()
+        r.headers["Foo"] = ["fOo"]
+        r.content = "afoob"
+        assert r.replace("foo", "boo", flags=re.I) == 3
+        assert not "foo" in r.content
+        assert r.headers["boo"] == ["boo"]
+
 
 class uResponse(libpry.AutoTree):
     def test_simple(self):
@@ -185,6 +193,14 @@ class uResponse(libpry.AutoTree):
         resp.load_state(resp2.get_state())
         assert resp == resp2
 
+    def test_replace(self):
+        r = tutils.tresp()
+        r.headers["Foo"] = ["fOo"]
+        r.content = "afoob"
+        assert r.replace("foo", "boo", flags=re.I) == 3
+        assert not "foo" in r.content
+        assert r.headers["boo"] == ["boo"]
+
 
 class uError(libpry.AutoTree):
     def test_getset_state(self):
@@ -203,6 +219,11 @@ class uError(libpry.AutoTree):
         e3 = e.copy()
         assert e3 == e
 
+    def test_replace(self):
+        e = proxy.Error(None, "amoop")
+        e.replace("moo", "bar")
+        assert e.msg == "abarp"
+
 
 class uProxyError(libpry.AutoTree):
     def test_simple(self):
diff --git a/test/test_utils.py b/test/test_utils.py
index 2ff951d4..e22ec039 100644
--- a/test/test_utils.py
+++ b/test/test_utils.py
@@ -137,6 +137,15 @@ class uHeaders(libpry.AutoTree):
         del self.hd["foo"]
         assert len(self.hd.lst) == 1
 
+    def test_replace(self):
+        self.hd.add("one", "two")
+        self.hd.add("two", "one")
+        assert self.hd.replace("one", "vun") == 2
+        assert self.hd.lst == [
+            ["vun", "two"],
+            ["two", "vun"],
+        ]
+
 
 class uisStringLike(libpry.AutoTree):
     def test_all(self):