diff options
| -rw-r--r-- | libmproxy/flow.py | 28 | ||||
| -rw-r--r-- | test/test_flow.py | 24 |
2 files changed, 52 insertions, 0 deletions
diff --git a/libmproxy/flow.py b/libmproxy/flow.py index 450fef30..dbc0c109 100644 --- a/libmproxy/flow.py +++ b/libmproxy/flow.py @@ -199,6 +199,34 @@ class ODictCaseless(ODict): return s.lower() +class decoded(object): + """ + + A context manager that decodes a request, response or error, and then + re-encodes it with the same encoding after execution of the block. + + Example: + + with decoded(request): + request.content = request.content.replace("foo", "bar") + """ + def __init__(self, o): + self.o = o + ce = o.headers["content-encoding"] + if ce and ce[0] in encoding.ENCODINGS: + self.ce = ce[0] + else: + self.ce = None + + def __enter__(self): + if self.ce: + self.o.decode() + + def __exit__(self, type, value, tb): + if self.ce: + self.o.encode(self.ce) + + class HTTPMsg(controller.Msg): def decode(self): """ diff --git a/test/test_flow.py b/test/test_flow.py index b6818960..56303881 100644 --- a/test/test_flow.py +++ b/test/test_flow.py @@ -996,6 +996,29 @@ class uODictCaseless(libpry.AutoTree): assert len(self.od) == 1 +class udecoded(libpry.AutoTree): + def test_del(self): + r = tutils.treq() + assert r.content == "content" + assert not r.headers["content-encoding"] + r.encode("gzip") + assert r.headers["content-encoding"] + assert r.content != "content" + with flow.decoded(r): + assert not r.headers["content-encoding"] + assert r.content == "content" + assert r.headers["content-encoding"] + assert r.content != "content" + + with flow.decoded(r): + r.content = "foo" + + assert r.content != "foo" + r.decode() + assert r.content == "foo" + + + tests = [ uStickyCookieState(), @@ -1012,4 +1035,5 @@ tests = [ uClientConnect(), uODict(), uODictCaseless(), + udecoded() ] |