From c1eaa9f74c9cf3edc34ced3ab441aede4d0bed01 Mon Sep 17 00:00:00 2001 From: Stephen Altamirano Date: Tue, 26 Jul 2011 22:03:41 -0700 Subject: [PATCH 1/2] Adds encode and decode methods to Response objects --- libmproxy/console.py | 9 ++------- libmproxy/proxy.py | 27 ++++++++++++++++++++++++--- 2 files changed, 26 insertions(+), 10 deletions(-) diff --git a/libmproxy/console.py b/libmproxy/console.py index 438e3305c..67e64d379 100644 --- a/libmproxy/console.py +++ b/libmproxy/console.py @@ -543,8 +543,7 @@ class ConnectionView(WWrap): conn = self.flow.response e = conn.headers["content-encoding"] or ["identity"] if e[0] != "identity": - conn.content = encoding.decode(e[0], conn.content) - conn.headers["content-encoding"] = ["identity"] + conn.decode() else: self.master.prompt_onekey( "Select encoding: ", @@ -563,11 +562,7 @@ class ConnectionView(WWrap): "z": "gzip", "d": "deflate", } - conn.content = encoding.encode( - encoding_map[key], - conn.content - ) - conn.headers["content-encoding"] = [encoding_map[key]] + conn.encode(encoding_map[key]) self.master.refresh_connection(self.flow) diff --git a/libmproxy/proxy.py b/libmproxy/proxy.py index 8d2888a48..9f2ed6447 100644 --- a/libmproxy/proxy.py +++ b/libmproxy/proxy.py @@ -284,7 +284,7 @@ class Request(controller.Msg): """ Replaces a regular expression pattern with repl in both the headers and the body of the request. Returns the number of replacements - made. + made. """ self.content, c = re.subn(pattern, repl, self.content, count, flags) self.path, pc = re.subn(pattern, repl, self.path, count, flags) @@ -422,12 +422,33 @@ class Response(controller.Msg): """ Replaces a regular expression pattern with repl in both the headers and the body of the response. Returns the number of replacements - made. + made. """ self.content, c = re.subn(pattern, repl, self.content, count, flags) c += self.headers.replace(pattern, repl, count, flags) return c + def decode(self): + """ + Alters Response object, decoding its content based on the current + Content-Encoding header and changing Content-Encoding header to + 'identity'. + """ + self.content = encoding.decode( + (self.headers["content-encoding"] or ["identity"])[0], + self.content + ) + self.headers["content-encoding"] = ["identity"] + + def encode(self, e): + """ + Alters Response object, encoding its content with the specified + coding. This method should only be called on Responses with + Content-Encoding headers of 'identity'. + """ + self.content = encoding.encode(e, self.content) + self.headers["content-encoding"] = [e] + class ClientDisconnect(controller.Msg): def __init__(self, client_conn): @@ -501,7 +522,7 @@ class Error(controller.Msg): """ Replaces a regular expression pattern with repl in both the headers and the body of the request. Returns the number of replacements - made. + made. """ self.msg, c = re.subn(pattern, repl, self.msg, count, flags) return c From acd511f6767cdcc332851e44a62fbd843399446a Mon Sep 17 00:00:00 2001 From: Stephen Altamirano Date: Tue, 26 Jul 2011 22:38:42 -0700 Subject: [PATCH 2/2] Adds test method for Response encoding and decoding --- test/test_proxy.py | 25 +++++++++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/test/test_proxy.py b/test/test_proxy.py index dacdbcc3d..1f63f3b3e 100644 --- a/test/test_proxy.py +++ b/test/test_proxy.py @@ -203,6 +203,31 @@ class uResponse(libpry.AutoTree): assert not "foo" in r.content assert r.headers["boo"] == ["boo"] + def test_decodeencode(self): + r = tutils.tresp() + r.headers["content-encoding"] = ["identity"] + r.content = "falafel" + r.decode() + assert r.headers["content-encoding"] == ["identity"] + assert r.content == "falafel" + + r = tutils.tresp() + r.headers["content-encoding"] = ["identity"] + r.content = "falafel" + r.encode("identity") + assert r.headers["content-encoding"] == ["identity"] + assert r.content == "falafel" + + r = tutils.tresp() + r.headers["content-encoding"] = ["identity"] + r.content = "falafel" + r.encode("gzip") + assert r.headers["content-encoding"] == ["gzip"] + assert r.content != "falafel" + r.decode() + assert r.headers["content-encoding"] == ["identity"] + assert r.content == "falafel" + class uError(libpry.AutoTree): def test_getset_state(self):