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, 56 insertions, 0 deletions

diff --git a/libmproxy/flow.py b/libmproxy/flow.py
index 13b5acbc..d174b41f 100644
--- a/libmproxy/flow.py
+++ b/libmproxy/flow.py
@@ -323,6 +323,18 @@ class Flow:
                 self.response.ack()
             self.intercepting = False
 
+    def replace(self, pattern, repl, count=0, flags=0):
+        """
+            Replaces a regular expression pattern with repl in all parts of the
+            flow . Returns the number of replacements made. 
+        """
+        c = self.request.replace(pattern, repl, count, flags)
+        if self.response:
+            c += self.response.replace(pattern, repl, count, flags)
+        if self.error:
+            c += self.error.replace(pattern, repl, count, flags)
+        return c
+
 
 class State:
     def __init__(self):
diff --git a/libmproxy/proxy.py b/libmproxy/proxy.py
index a7cc31e8..c5e5971a 100644
--- a/libmproxy/proxy.py
+++ b/libmproxy/proxy.py
@@ -280,6 +280,16 @@ class Request(controller.Msg):
         else:
             return self.FMT_PROXY % (self.method, self.scheme, self.host, self.port, self.path, str(headers), content)
 
+    def replace(self, pattern, repl, count=0, flags=0):
+        """
+            Replaces a regular expression pattern with repl in both the headers
+            and the body of the request. Returns the number of replacements
+            made. 
+        """
+        self.content, c = re.subn(pattern, repl, self.content, count, flags)
+        c += self.headers.replace(pattern, repl, count, flags)
+        return c
+
 
 class Response(controller.Msg):
     FMT = '%s\r\n%s\r\n%s'
@@ -406,6 +416,16 @@ class Response(controller.Msg):
         data = (proto, str(headers), content)
         return self.FMT%data
 
+    def replace(self, pattern, repl, count=0, flags=0):
+        """
+            Replaces a regular expression pattern with repl in both the headers
+            and the body of the response. Returns the number of replacements
+            made. 
+        """
+        self.content, c = re.subn(pattern, repl, self.content, count, flags)
+        c += self.headers.replace(pattern, repl, count, flags)
+        return c
+
 
 class ClientDisconnect(controller.Msg):
     def __init__(self, client_conn):
@@ -473,6 +493,15 @@ class Error(controller.Msg):
     def __eq__(self, other):
         return self.get_state() == other.get_state()
 
+    def replace(self, pattern, repl, count=0, flags=0):
+        """
+            Replaces a regular expression pattern with repl in both the headers
+            and the body of the request. Returns the number of replacements
+            made. 
+        """
+        self.msg, c = re.subn(pattern, repl, self.msg, count, flags)
+        return c
+
 
 class FileLike:
     def __init__(self, o):
diff --git a/libmproxy/utils.py b/libmproxy/utils.py
index 4d8c7c41..9ac9c0b8 100644
--- a/libmproxy/utils.py
+++ b/libmproxy/utils.py
@@ -279,6 +279,21 @@ class Headers:
                     ret.append([name, value])
         self.lst = ret
 
+    def replace(self, pattern, repl, count=0, flags=0):
+        """
+            Replaces a regular expression pattern with repl in both header keys
+            and values. Returns the number of replacements made. 
+        """
+        nlst, count = [], 0
+        for i in self.lst:
+            k, c = re.subn(pattern, repl, i[0], count, flags)
+            count += c
+            v, c = re.subn(pattern, repl, i[1], count, flags)
+            count += c
+            nlst.append([k, v])
+        self.lst = nlst
+        return count
+
 
 def pretty_size(size):
     suffixes = [