From 9c985f2d20e4086881bde5ecc63c21a208393894 Mon Sep 17 00:00:00 2001
From: Aldo Cortesi <aldo@nullcube.com>
Date: Fri, 10 Feb 2012 14:27:39 +1300
Subject: Methods for getting and setting form urlencoded data on Request.

---
 libmproxy/console/connview.py |  5 +----
 libmproxy/flow.py             | 23 +++++++++++++++++++++++
 test/test_flow.py             | 12 ++++++++++++
 3 files changed, 36 insertions(+), 4 deletions(-)

diff --git a/libmproxy/console/connview.py b/libmproxy/console/connview.py
index b15f0c8f..38fecfa0 100644
--- a/libmproxy/console/connview.py
+++ b/libmproxy/console/connview.py
@@ -176,7 +176,7 @@ class ConnectionView(common.WWrap):
             if i[0].lower() == "content-type":
                 ctype = i[1]
                 break
-        if ctype and "x-www-form-urlencoded" in ctype:
+        if ctype and flow.HDR_FORM_URLENCODED in ctype:
             data = utils.urldecode(content)
             if data:
                 return self._view_conn_urlencoded(data)
@@ -221,9 +221,6 @@ class ConnectionView(common.WWrap):
                 txt.extend(self._view_conn_raw(content))
         return urwid.ListBox(txt)
 
-
-
-
     def _tab(self, content, active):
         if active:
             attr = "heading"
diff --git a/libmproxy/flow.py b/libmproxy/flow.py
index e929440e..85c48ce1 100644
--- a/libmproxy/flow.py
+++ b/libmproxy/flow.py
@@ -8,6 +8,8 @@ import tnetstring, filt, script, utils, encoding, proxy
 from email.utils import parsedate_tz, formatdate, mktime_tz
 import controller, version
 
+HDR_FORM_URLENCODED = "x-www-form-urlencoded"
+
 class RunException(Exception):
     def __init__(self, msg, returncode, errout):
         Exception.__init__(self, msg)
@@ -56,6 +58,8 @@ class Headers:
         return new
 
     def __setitem__(self, k, hdrs):
+        if isinstance(hdrs, basestring):
+            raise ValueError("Header values should be lists.")
         k = self._kconv(k)
         new = self._filter_lst(k, self.lst)
         for i in hdrs:
@@ -312,6 +316,25 @@ class Request(HTTPMsg):
             host = "%s:%s"%(self.host, self.port)
         return host
 
+    def get_form_urlencoded(self):
+        """
+            Retrieves the URL-encoded form data, returning a list of (key,
+            value) tuples. Returns an empty list if there is no data or the
+            content-type indicates non-form data.
+        """
+        hv = [i.lower() for i in self.headers["content-type"]]
+        if HDR_FORM_URLENCODED in hv:
+            return utils.urldecode(self.content)
+        return []
+
+    def set_form_urlencoded(self, data):
+        """
+            Sets the body to the URL-encoded form data, and adds the
+            appropriate content-type header.
+        """
+        self.headers["content-type"] = [HDR_FORM_URLENCODED]
+        self.content = utils.urlencode(data)
+
     def get_query(self):
         """
             Gets the request query string. Returns a list of (key, value)
diff --git a/test/test_flow.py b/test/test_flow.py
index 616a2b2f..e4cd6da2 100644
--- a/test/test_flow.py
+++ b/test/test_flow.py
@@ -577,6 +577,18 @@ class uRequest(libpry.AutoTree):
         r2 = r.copy()
         assert r == r2
 
+    def test_getset_form_urlencoded(self):
+        h = flow.Headers()
+        h["content-type"] = [flow.HDR_FORM_URLENCODED]
+        d = [("one", "two"), ("three", "four")]
+        r = flow.Request(None, "host", 22, "https", "GET", "/", h, utils.urlencode(d))
+        assert r.get_form_urlencoded() == d
+
+        d = [("x", "y")]
+        r.set_form_urlencoded(d)
+        assert r.get_form_urlencoded() == d
+
+
     def test_getset_query(self):
         h = flow.Headers()
 
-- 
cgit v1.2.3