From d51b8cab0c0d1352865155865dfd258f66103ffe Mon Sep 17 00:00:00 2001 From: Aldo Cortesi Date: Fri, 16 Mar 2012 11:12:56 +1300 Subject: Add a decoded context manager. This simplifies a common chore when modifying traffic - decoding the object, modifying it, then re-encoding it with the same encoding afterwards. You can now simply say: with flow.decoded(request): request.content = "bar" --- libmproxy/flow.py | 28 ++++++++++++++++++++++++++++ 1 file changed, 28 insertions(+) (limited to 'libmproxy/flow.py') 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): """ -- cgit v1.2.3