From 3189d144a521fcc98695dd079fb3dd4304de2eee Mon Sep 17 00:00:00 2001
From: Sahn Lam <sahn@pobox.com>
Date: Thu, 16 Aug 2012 23:27:47 -0700
Subject: Optional AMF decoding support

If PyAMF is installed, enable AMF decoding.
---
 libmproxy/console/contentview.py |  28 ++++++++++++++++++++++++++--
 libmproxy/console/help.py        |   4 ++++
 libmproxy/utils.py               |  36 ++++++++++++++++++++++++++++++++++++
 test/data/test.amf               | Bin 0 -> 432 bytes
 test/test_console_contentview.py |  23 +++++++++++++++++++++++
 5 files changed, 89 insertions(+), 2 deletions(-)
 create mode 100644 test/data/test.amf

diff --git a/libmproxy/console/contentview.py b/libmproxy/console/contentview.py
index 2443c188..4efbb2b1 100644
--- a/libmproxy/console/contentview.py
+++ b/libmproxy/console/contentview.py
@@ -20,6 +20,7 @@ VIEW_RAW = 7
 VIEW_HEX = 8
 VIEW_HTML = 9
 VIEW_OUTLINE = 10
+VIEW_AMF = 11
 
 VIEW_NAMES = {
     VIEW_AUTO: "Auto",
@@ -36,7 +37,7 @@ VIEW_NAMES = {
 }
 
 
-VIEW_PROMPT = (
+VIEW_PROMPT = [
     ("auto detect", "a"),
     ("hex", "e"),
     ("html", "h"),
@@ -48,7 +49,7 @@ VIEW_PROMPT = (
     ("multipart", "m"),
     ("urlencoded", "u"),
     ("xml", "x"),
-)
+]
 
 VIEW_SHORTCUTS = {
     "a": VIEW_AUTO,
@@ -285,6 +286,10 @@ def view_image(hdrs, content, limit):
         )
     return "%s image"%img.format, fmt
 
+def view_amf(hdrs, content, limit):
+    s = utils.pretty_amf(content)
+    if s:
+        return "AMF", _view_text(s[:limit], len(s), limit)
 
 PRETTY_FUNCTION_MAP = {
     VIEW_XML: view_xml,
@@ -344,3 +349,22 @@ def get_content_view(viewmode, hdrItems, content, limit):
     else:
         msg.append(ret[0])
     return " ".join(msg), ret[1]
+
+
+#
+# Enable optional decoding methods at runtime
+#
+
+# AMF decoding requires pyamf
+try:
+    import pyamf
+
+    VIEW_SHORTCUTS["f"] = VIEW_AMF
+    VIEW_PROMPT.append(("amf", "f"))
+    VIEW_NAMES[VIEW_AMF] = "AMF"
+    CONTENT_TYPES_MAP["application/x-amf"] = VIEW_AMF
+    PRETTY_FUNCTION_MAP[VIEW_AMF] = view_amf
+except ImportError:
+    pass
+
+
diff --git a/libmproxy/console/help.py b/libmproxy/console/help.py
index 64428fcc..e278e763 100644
--- a/libmproxy/console/help.py
+++ b/libmproxy/console/help.py
@@ -83,6 +83,10 @@ class HelpView(urwid.ListBox):
                     common.highlight_key("xml", "x") +
                     [("text", ": XML")]
                 ),
+                (None,
+                    common.highlight_key("amf", "f") +
+                    [("text", ": AMF (requires PyAMF)")]
+                ),
             ("o", "toggle options:"),
                 (None,
                     common.highlight_key("anticache", "a") +
diff --git a/libmproxy/utils.py b/libmproxy/utils.py
index 871beb35..a89419a9 100644
--- a/libmproxy/utils.py
+++ b/libmproxy/utils.py
@@ -81,6 +81,42 @@ def pretty_json(s):
     return json.dumps(p, sort_keys=True, indent=4).split("\n")
 
 
+def pretty_amf(s):
+    """
+        Takes an AMF encoded string, decodes it and returns a nicely indented
+        string in JSON format.
+
+        Reqires pyamf module. The function returns None if pyamf is not
+        installed.
+    """
+    try:
+        import pyamf
+        from pyamf import remoting
+    except ImportError:
+        return None
+
+    envelope = remoting.decode(s)
+    if not envelope:
+        return None
+
+    data = {}
+    data['amfVersion'] = envelope.amfVersion
+    for target, message in iter(envelope):
+
+        one_message = {}
+
+        if hasattr(message, 'status'):
+            one_message['status'] = message.status
+
+        if hasattr(message, 'target'):
+            one_message['target'] = message.target
+
+        one_message['body'] = message.body
+        data[target] = one_message
+
+    return json.dumps(data, indent=4)
+
+
 def urldecode(s):
     """
         Takes a urlencoded string and returns a list of (key, value) tuples.
diff --git a/test/data/test.amf b/test/data/test.amf
new file mode 100644
index 00000000..c8fc261d
Binary files /dev/null and b/test/data/test.amf differ
diff --git a/test/test_console_contentview.py b/test/test_console_contentview.py
index fbb7e6d2..6f3958df 100644
--- a/test/test_console_contentview.py
+++ b/test/test_console_contentview.py
@@ -53,6 +53,20 @@ class TestContentView:
               )
         assert f is cv.view_xml
 
+        try:
+            import pyamf
+
+            f = cv.get_view_func(
+                    cv.VIEW_AUTO,
+                    flow.ODictCaseless(
+                        [["content-type", "application/x-amf"]],
+                    ),
+                    ""
+                  )
+            assert f is cv.view_amf
+        except ImportError:
+            pass
+
     def test_view_urlencoded(self):
         d = utils.urlencode([("one", "two"), ("three", "four")])
         assert cv.view_urlencoded([], d, 100)
@@ -111,6 +125,15 @@ class TestContentView:
 
         assert not cv.view_image([], "flibble", sys.maxint)
 
+    def test_view_amf(self):
+        try:
+            import pyamf
+
+            p = tutils.test_data.path("data/test.amf")
+            assert cv.view_amf([], file(p).read(), sys.maxint)
+        except ImportError:
+            pass
+
     def test_view_multipart(self):
         v = """
 --AaB03x
-- 
cgit v1.2.3