Repository: mitmproxy/mitmproxy

Issue:
Add max replacements/count argument to flow.replace
##### Steps to reproduce the problem:

``` python
 regex1 = r'\d+\#\d+\#\d+\#\d+\#(\d+)\#\d+\#\d+\#\d+\#\d+\#'
 re.compile(regex1)
 flow.response.replace(regex1, r'444#5555#66#77#\1#0#0#0#0#', 10)
```

![notepad _2016-08-20_11-17-57](https://cloud.githubusercontent.com/assets/15959479/17833044/07dee18a-66c8-11e6-8bf7-dfb50aef2ee7.png)
##### What is the expected behavior?

`flow.response.replace(regex1, r'444#5555#66#77#\1#0#0#0#0#', 10)`

The **10** used should be Max Count or Max Replace
##### What went wrong?

In  `flow.response.replace(old, new[, maxreplace])`

is no longer working, all my older scripts are now broken; this used to work fine.
##### Any other comments? What have you tried so far?

I have tried changing my replace to a variable, changing the long regex and replace with something very small and simple.

Work Great on my other servers that run MITMProxy 0.15 and i think 0.16 was working too.

*When i updated to 0.17 or 0.18 it then no longer works and my scripts are broken.

---

Mitmproxy Version:  0.17 and 0.18
Operating System:  Ubuntu 14.04 and 16.04

<!-- Please use the mitmproxy forums (https://discourse.mitmproxy.org/) for support/how-to questions. Thanks! :) -->



Discussion:
I think the problem is not with replace, but in your script. The syntax has changed a bit.

``` python
import re

# Notice that the context parameter is no longer needed
def response(flow):
    flow.response.replace("this", "that")
```

I've tested that this works with latest code from master!

@dufferzafar 

Thank you for the reply,

Sorry, but that did not help :(

Thank you for letting me know that it is no longer needed, but this did not allow;
flow.response.replace("this", "that", **COUNT**) to work.

the normal replace works, but the "COUNT" Does not; this is my issue with version 0.17 and 0.18
- on all the older version eg. 0.15 and 0.16; COUNT works great!

for example, there are a 100 Zip codes and you only want to change the first one; it would be something like:
flow.response.replace("[0-9]{5}", "12345", 1) = Replace the First Occurrence of a 5 digit number.
Now Only flow.response.replace("[0-9]{5}", "12345") Works; Replacing ALL 5 digit numbers.

This is my issue, I hope this help explain it better... :)

Thank you again.

We'd be happy to accept a PR that reintroduces the count argument. Other than that, `response.content = response.content.replace("foo","bar")` may just be what you are looking for anyways.

@mhils 

Thank you for your reply.

I would love if someone would create a Pull Request to reintroduce the COUNT argument.

I was not sure if this is/was a BUG or if it was taken out, changed for some reason; No idea why this would be removed.

I was thinking of using re.sub() or some other python commands to have it's MAX COUNT function work within the .py script, but was not able to succeed.

For now I'm stuck using Version 0.16, but hopefully down the line I could use the latest Versions.

I'm open to other suggestion or examples on a workaround for Replace the FIRST Occurrence of a Regex in the BODY of a Response.

Thank you all.


Resolving patch:
diff --git a/netlib/http/headers.py b/netlib/http/headers.py
--- a/netlib/http/headers.py
+++ b/netlib/http/headers.py
@@ -158,7 +158,7 @@ def items(self, multi=False):
         else:
             return super(Headers, self).items()
 
-    def replace(self, pattern, repl, flags=0):
+    def replace(self, pattern, repl, flags=0, count=0):
         """
         Replaces a regular expression pattern with repl in each "name: value"
         header line.
@@ -172,10 +172,10 @@ def replace(self, pattern, repl, flags=0):
             repl = strutils.escaped_str_to_bytes(repl)
         pattern = re.compile(pattern, flags)
         replacements = 0
-
+        flag_count = count > 0
         fields = []
         for name, value in self.fields:
-            line, n = pattern.subn(repl, name + b": " + value)
+            line, n = pattern.subn(repl, name + b": " + value, count=count)
             try:
                 name, value = line.split(b": ", 1)
             except ValueError:
@@ -184,6 +184,10 @@ def replace(self, pattern, repl, flags=0):
                 pass
             else:
                 replacements += n
+                if flag_count:
+                    count -= n
+                    if count == 0:
+                        break
             fields.append((name, value))
         self.fields = tuple(fields)
         return replacements
diff --git a/netlib/http/message.py b/netlib/http/message.py
--- a/netlib/http/message.py
+++ b/netlib/http/message.py
@@ -260,7 +260,7 @@ def encode(self, e):
         if "content-encoding" not in self.headers:
             raise ValueError("Invalid content encoding {}".format(repr(e)))
 
-    def replace(self, pattern, repl, flags=0):
+    def replace(self, pattern, repl, flags=0, count=0):
         """
         Replaces a regular expression pattern with repl in both the headers
         and the body of the message. Encoded body will be decoded
@@ -276,9 +276,9 @@ def replace(self, pattern, repl, flags=0):
         replacements = 0
         if self.content:
             self.content, replacements = re.subn(
-                pattern, repl, self.content, flags=flags
+                pattern, repl, self.content, flags=flags, count=count
             )
-        replacements += self.headers.replace(pattern, repl, flags)
+        replacements += self.headers.replace(pattern, repl, flags=flags, count=count)
         return replacements
 
     # Legacy
diff --git a/netlib/http/request.py b/netlib/http/request.py
--- a/netlib/http/request.py
+++ b/netlib/http/request.py
@@ -80,7 +80,7 @@ def __repr__(self):
             self.method, hostport, path
         )
 
-    def replace(self, pattern, repl, flags=0):
+    def replace(self, pattern, repl, flags=0, count=0):
         """
             Replaces a regular expression pattern with repl in the headers, the
             request path and the body of the request. Encoded content will be
@@ -94,9 +94,9 @@ def replace(self, pattern, repl, flags=0):
         if isinstance(repl, six.text_type):
             repl = strutils.escaped_str_to_bytes(repl)
 
-        c = super(Request, self).replace(pattern, repl, flags)
+        c = super(Request, self).replace(pattern, repl, flags, count)
         self.path, pc = re.subn(
-            pattern, repl, self.data.path, flags=flags
+            pattern, repl, self.data.path, flags=flags, count=count
         )
         c += pc
         return c


Test patch:
diff --git a/test/netlib/http/test_headers.py b/test/netlib/http/test_headers.py
--- a/test/netlib/http/test_headers.py
+++ b/test/netlib/http/test_headers.py
@@ -75,6 +75,11 @@ def test_replace_remove_spacer(self):
         assert replacements == 0
         assert headers["Host"] == "example.com"
 
+    def test_replace_with_count(self):
+        headers = Headers(Host="foobarfoo.com", Accept="foo/bar")
+        replacements = headers.replace("foo", "bar", count=1)
+        assert replacements == 1
+
 
 def test_parse_content_type():
     p = parse_content_type
diff --git a/test/netlib/http/test_message.py b/test/netlib/http/test_message.py
--- a/test/netlib/http/test_message.py
+++ b/test/netlib/http/test_message.py
@@ -99,6 +99,16 @@ def test_timestamp_end(self):
     def test_http_version(self):
         _test_decoded_attr(tresp(), "http_version")
 
+    def test_replace(self):
+        r = tresp()
+        r.content = b"foofootoo"
+        r.replace(b"foo", "gg")
+        assert r.content == b"ggggtoo"
+
+        r.content = b"foofootoo"
+        r.replace(b"foo", "gg", count=1)
+        assert r.content == b"ggfootoo"
+
 
 class TestMessageContentEncoding(object):
     def test_simple(self):
diff --git a/test/netlib/http/test_request.py b/test/netlib/http/test_request.py
--- a/test/netlib/http/test_request.py
+++ b/test/netlib/http/test_request.py
@@ -26,6 +26,16 @@ def test_repr(self):
         request.host = None
         assert repr(request) == "Request(GET /path)"
 
+    def replace(self):
+        r = treq()
+        r.path = b"foobarfoo"
+        r.replace(b"foo", "bar")
+        assert r.path == b"barbarbar"
+
+        r.path = b"foobarfoo"
+        r.replace(b"foo", "bar", count=1)
+        assert r.path == b"barbarfoo"
+
     def test_first_line_format(self):
         _test_passthrough_attr(treq(), "first_line_format")