From 2709441d5ba203da7c685b53692e66e273d20058 Mon Sep 17 00:00:00 2001 From: Aldo Cortesi Date: Thu, 9 Feb 2012 16:40:31 +1300 Subject: [PATCH] Add get_query and set_query methods to Request. --- libmproxy/console/connview.py | 6 +++--- libmproxy/flow.py | 21 ++++++++++++++++++++- libmproxy/utils.py | 12 +++++++++++- test/test_flow.py | 33 ++++++++++++++++++++++++++++----- 4 files changed, 62 insertions(+), 10 deletions(-) diff --git a/libmproxy/console/connview.py b/libmproxy/console/connview.py index 7fe4c899e..d057a59f4 100644 --- a/libmproxy/console/connview.py +++ b/libmproxy/console/connview.py @@ -377,7 +377,7 @@ class ConnectionView(common.WWrap): conn = self.flow.response self.flow.backup() - if part == "b": + if part == "r": c = self.master.spawn_editor(conn.content or "") conn.content = c.rstrip("\n") elif part == "h": @@ -445,7 +445,7 @@ class ConnectionView(common.WWrap): "Edit request", ( ("header", "h"), - ("body", "b"), + ("raw body", "r"), ("url", "u"), ("method", "m"), ), @@ -458,7 +458,7 @@ class ConnectionView(common.WWrap): ("code", "c"), ("message", "m"), ("header", "h"), - ("body", "b"), + ("raw body", "r"), ), self.edit ) diff --git a/libmproxy/flow.py b/libmproxy/flow.py index 818d81390..e929440eb 100644 --- a/libmproxy/flow.py +++ b/libmproxy/flow.py @@ -2,7 +2,7 @@ This module provides more sophisticated flow tracking. These match requests with their responses, and provide filtering and interception facilities. """ -import hashlib, Cookie, cookielib, copy, re +import hashlib, Cookie, cookielib, copy, re, urlparse import time import tnetstring, filt, script, utils, encoding, proxy from email.utils import parsedate_tz, formatdate, mktime_tz @@ -312,6 +312,25 @@ class Request(HTTPMsg): host = "%s:%s"%(self.host, self.port) return host + def get_query(self): + """ + Gets the request query string. Returns a list of (key, value) + tuples. + """ + _, _, _, _, query, _ = urlparse.urlparse(self.get_url()) + if not query: + return [] + return utils.urldecode(query) + + def set_query(self, q): + """ + Takes a list of (key, value) tuples, and sets the request query + string. + """ + scheme, netloc, path, params, _, fragment = urlparse.urlparse(self.get_url()) + query = utils.urlencode(q) + self.set_url(urlparse.urlunparse([scheme, netloc, path, params, query, fragment])) + def get_url(self): """ Returns a URL string, constructed from the Request's URL compnents. diff --git a/libmproxy/utils.py b/libmproxy/utils.py index 4339ec6d4..c12ccc9bc 100644 --- a/libmproxy/utils.py +++ b/libmproxy/utils.py @@ -12,7 +12,7 @@ # # You should have received a copy of the GNU General Public License # along with this program. If not, see . -import re, os, subprocess, datetime, urlparse, string +import re, os, subprocess, datetime, urlparse, string, urllib import time, functools, cgi, textwrap, hashlib import json @@ -123,9 +123,19 @@ def pretty_json(s): def urldecode(s): + """ + Takes a urlencoded string and returns a list of (key, value) tuples. + """ return cgi.parse_qsl(s) +def urlencode(s): + """ + Takes a list of (key, value) tuples and returns a urlencoded string. + """ + return urllib.urlencode(s, False) + + def hexdump(s): """ Returns a set of typles: diff --git a/test/test_flow.py b/test/test_flow.py index 85355b73b..616a2b2ff 100644 --- a/test/test_flow.py +++ b/test/test_flow.py @@ -10,7 +10,7 @@ class uStickyCookieState(libpry.AutoTree): def _response(self, cookie, host): s = flow.StickyCookieState(filt.parse(".*")) f = tutils.tflow_full() - f.request.host = host + f.request.host = host f.response.headers["Set-Cookie"] = [cookie] s.handle_response(f) return s, f @@ -153,12 +153,12 @@ class uFlow(libpry.AutoTree): def test_getset_state(self): f = tutils.tflow() f.response = tutils.tresp(f.request) - state = f._get_state() + state = f._get_state() assert f._get_state() == flow.Flow._from_state(state)._get_state() f.response = None f.error = flow.Error(f.request, "error") - state = f._get_state() + state = f._get_state() assert f._get_state() == flow.Flow._from_state(state)._get_state() f2 = tutils.tflow() @@ -368,7 +368,7 @@ class uState(libpry.AutoTree): flows = c.view[:] c.clear() - + c.load_flows(flows) assert isinstance(c._flow_list[0], flow.Flow) @@ -473,7 +473,7 @@ class uFlowMaster(libpry.AutoTree): rx = tutils.tresp() assert not fm.handle_response(rx) - + dc = flow.ClientDisconnect(req.client_conn) fm.handle_clientdisconnect(dc) @@ -577,6 +577,29 @@ class uRequest(libpry.AutoTree): r2 = r.copy() assert r == r2 + def test_getset_query(self): + h = flow.Headers() + + r = flow.Request(None, "host", 22, "https", "GET", "/foo?x=y&a=b", h, "content") + q = r.get_query() + assert q == [("x", "y"), ("a", "b")] + + r = flow.Request(None, "host", 22, "https", "GET", "/", h, "content") + q = r.get_query() + assert not q + + r = flow.Request(None, "host", 22, "https", "GET", "/?adsfa", h, "content") + q = r.get_query() + assert not q + + r = flow.Request(None, "host", 22, "https", "GET", "/foo?x=y&a=b", h, "content") + assert r.get_query() + r.set_query([]) + assert not r.get_query() + qv = [("a", "b"), ("c", "d")] + r.set_query(qv) + assert r.get_query() == qv + def test_anticache(self): h = flow.Headers() r = flow.Request(None, "host", 22, "https", "GET", "/", h, "content")