Add get_query and set_query methods to Request.

4 files changed, 62 insertions, 10 deletions

diff --git a/libmproxy/console/connview.py b/libmproxy/console/connview.py
index 7fe4c899..d057a59f 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 818d8139..e929440e 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 4339ec6d..c12ccc9b 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 <http://www.gnu.org/licenses/>.
-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 85355b73..616a2b2f 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")