mirror of
https://github.com/Grasscutters/mitmproxy.git
synced 2024-11-30 03:14:22 +00:00
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"
This commit is contained in:
parent
8d662e6636
commit
d51b8cab0c
@ -199,6 +199,34 @@ class ODictCaseless(ODict):
|
|||||||
return s.lower()
|
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):
|
class HTTPMsg(controller.Msg):
|
||||||
def decode(self):
|
def decode(self):
|
||||||
"""
|
"""
|
||||||
|
@ -996,6 +996,29 @@ class uODictCaseless(libpry.AutoTree):
|
|||||||
assert len(self.od) == 1
|
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 = [
|
tests = [
|
||||||
uStickyCookieState(),
|
uStickyCookieState(),
|
||||||
@ -1012,4 +1035,5 @@ tests = [
|
|||||||
uClientConnect(),
|
uClientConnect(),
|
||||||
uODict(),
|
uODict(),
|
||||||
uODictCaseless(),
|
uODictCaseless(),
|
||||||
|
udecoded()
|
||||||
]
|
]
|
||||||
|
Loading…
Reference in New Issue
Block a user