Merge pull request #833 from zbuc/contentview_scripts

Contentview scripts

7 files changed, 198 insertions, 28 deletions

diff --git a/examples/custom_contentviews.py b/examples/custom_contentviews.py
new file mode 100644
index 00000000..1a2bcb1e
--- /dev/null
+++ b/examples/custom_contentviews.py
@@ -0,0 +1,66 @@
+import string
+from libmproxy import script, flow, utils
+import libmproxy.contentviews as cv
+from netlib.http import Headers
+import lxml.html
+import lxml.etree
+
+
+class ViewPigLatin(cv.View):
+    name = "pig_latin_HTML"
+    prompt = ("pig latin HTML", "l")
+    content_types = ["text/html"]
+
+    def __call__(self, data, **metadata):
+        if utils.isXML(data):
+            parser = lxml.etree.HTMLParser(
+                strip_cdata=True,
+                remove_blank_text=True
+            )
+            d = lxml.html.fromstring(data, parser=parser)
+            docinfo = d.getroottree().docinfo
+
+            def piglify(src):
+                words = string.split(src)
+                ret = ''
+                for word in words:
+                    idx = -1
+                    while word[idx] in string.punctuation and (idx * -1) != len(word): idx -= 1
+                    if word[0].lower() in 'aeiou':
+                        if idx == -1: ret += word[0:] + "hay"
+                        else: ret += word[0:len(word)+idx+1] + "hay" + word[idx+1:]
+                    else:
+                        if idx == -1: ret += word[1:] + word[0] + "ay"
+                        else: ret += word[1:len(word)+idx+1] + word[0] + "ay" + word[idx+1:]
+                    ret += ' '
+                return ret.strip()
+
+            def recurse(root):
+                if hasattr(root, 'text') and root.text:
+                        root.text = piglify(root.text)
+                if hasattr(root, 'tail') and root.tail:
+                        root.tail = piglify(root.tail)
+
+                if len(root):
+                    for child in root:
+                        recurse(child)
+
+            recurse(d)
+
+            s = lxml.etree.tostring(
+                d,
+                pretty_print=True,
+                doctype=docinfo.doctype
+            )
+            return "HTML", cv.format_text(s)
+
+
+pig_view = ViewPigLatin()
+
+
+def start(context, argv):
+    context.add_contentview(pig_view)
+
+
+def stop(context):
+    context.remove_contentview(pig_view)
diff --git a/libmproxy/contentviews.py b/libmproxy/contentviews.py
index 9af08033..2f46ccca 100644
--- a/libmproxy/contentviews.py
+++ b/libmproxy/contentviews.py
@@ -479,34 +479,9 @@ class ViewWBXML(View):
             return None
 
 
-views = [
-    ViewAuto(),
-    ViewRaw(),
-    ViewHex(),
-    ViewJSON(),
-    ViewXML(),
-    ViewWBXML(),
-    ViewHTML(),
-    ViewHTMLOutline(),
-    ViewJavaScript(),
-    ViewCSS(),
-    ViewURLEncoded(),
-    ViewMultipart(),
-    ViewImage(),
-]
-if pyamf:
-    views.append(ViewAMF())
-
-if ViewProtobuf.is_available():
-    views.append(ViewProtobuf())
-
+views = []
 content_types_map = {}
-for i in views:
-    for ct in i.content_types:
-        l = content_types_map.setdefault(ct, [])
-        l.append(i)
-
-view_prompts = [i.prompt for i in views]
+view_prompts = []
 
 
 def get_by_shortcut(c):
@@ -515,6 +490,58 @@ def get_by_shortcut(c):
             return i
 
 
+def add(view):
+    # TODO: auto-select a different name (append an integer?)
+    for i in views:
+        if i.name == view.name:
+            raise ContentViewException("Duplicate view: " + view.name)
+
+    # TODO: the UI should auto-prompt for a replacement shortcut
+    for prompt in view_prompts:
+        if prompt[1] == view.prompt[1]:
+            raise ContentViewException("Duplicate view shortcut: " + view.prompt[1])
+
+    views.append(view)
+
+    for ct in view.content_types:
+        l = content_types_map.setdefault(ct, [])
+        l.append(view)
+
+    view_prompts.append(view.prompt)
+
+
+def remove(view):
+    for ct in view.content_types:
+        l = content_types_map.setdefault(ct, [])
+        l.remove(view)
+
+        if not len(l):
+            del content_types_map[ct]
+
+    view_prompts.remove(view.prompt)
+    views.remove(view)
+
+
+add(ViewAuto())
+add(ViewRaw())
+add(ViewHex())
+add(ViewJSON())
+add(ViewXML())
+add(ViewWBXML())
+add(ViewHTML())
+add(ViewHTMLOutline())
+add(ViewJavaScript())
+add(ViewCSS())
+add(ViewURLEncoded())
+add(ViewMultipart())
+add(ViewImage())
+
+if pyamf:
+    add(ViewAMF())
+
+if ViewProtobuf.is_available():
+    add(ViewProtobuf())
+
 def get(name):
     for i in views:
         if i.name == name:
diff --git a/libmproxy/flow.py b/libmproxy/flow.py
index 55a4dbcf..3343e694 100644
--- a/libmproxy/flow.py
+++ b/libmproxy/flow.py
@@ -9,7 +9,7 @@ import cookielib
 import os
 import re
 import urlparse
-
+import inspect
 
 from netlib import wsgi
 from netlib.exceptions import HttpException
@@ -21,6 +21,7 @@ from .proxy.config import HostMatcher
 from .protocol.http_replay import RequestReplayThread
 from .protocol import Kill
 from .models import ClientConnection, ServerConnection, HTTPResponse, HTTPFlow, HTTPRequest
+from . import contentviews as cv
 
 
 class AppRegistry:
diff --git a/libmproxy/script.py b/libmproxy/script.py
index 9d051c12..4da40c52 100644
--- a/libmproxy/script.py
+++ b/libmproxy/script.py
@@ -5,6 +5,8 @@ import threading
 import shlex
 import sys
 
+from . import contentviews as cv
+
 
 class ScriptError(Exception):
     pass
@@ -56,6 +58,12 @@ class ScriptContext:
     def app_registry(self):
         return self._master.apps
 
+    def add_contentview(self, view_obj):
+        cv.add(view_obj)
+
+    def remove_contentview(self, view_obj):
+        cv.remove(view_obj)
+
 
 class Script:
     """
diff --git a/test/test_contentview.py b/test/test_contentview.py
index 14adcd83..2a70b414 100644
--- a/test/test_contentview.py
+++ b/test/test_contentview.py
@@ -210,6 +210,21 @@ Larry
         assert "decoded gzip" in r[0]
         assert "Raw" in r[0]
 
+    def test_add_cv(self):
+        class TestContentView(cv.View):
+            name = "test"
+            prompt = ("t", "test")
+
+        tcv = TestContentView()
+        cv.add(tcv)
+
+        # repeated addition causes exception
+        tutils.raises(
+            ContentViewException,
+            cv.add,
+            tcv
+        )
+
 
 if pyamf:
     def test_view_amf_request():
diff --git a/test/test_custom_contentview.py b/test/test_custom_contentview.py
new file mode 100644
index 00000000..4b5a3e53
--- /dev/null
+++ b/test/test_custom_contentview.py
@@ -0,0 +1,52 @@
+from libmproxy import script, flow
+import libmproxy.contentviews as cv
+from netlib.http import Headers
+
+
+def test_custom_views():
+    class ViewNoop(cv.View):
+        name = "noop"
+        prompt = ("noop", "n")
+        content_types = ["text/none"]
+
+        def __call__(self, data, **metadata):
+            return "noop", cv.format_text(data)
+
+
+    view_obj = ViewNoop()
+
+    cv.add(view_obj)
+
+    assert cv.get("noop")
+
+    r = cv.get_content_view(
+        cv.get("noop"),
+        "[1, 2, 3]",
+        headers=Headers(
+            content_type="text/plain"
+        )
+    )
+    assert "noop" in r[0]
+
+    # now try content-type matching
+    r = cv.get_content_view(
+        cv.get("Auto"),
+        "[1, 2, 3]",
+        headers=Headers(
+            content_type="text/none"
+        )
+    )
+    assert "noop" in r[0]
+
+    # now try removing the custom view
+    cv.remove(view_obj)
+    r = cv.get_content_view(
+        cv.get("Auto"),
+        "[1, 2, 3]",
+        headers=Headers(
+            content_type="text/none"
+        )
+    )
+    assert "noop" not in r[0]
+
+    
diff --git a/test/test_script.py b/test/test_script.py
index 1b0e5a5b..8612d5f3 100644
--- a/test/test_script.py
+++ b/test/test_script.py
@@ -127,3 +127,4 @@ def test_command_parsing():
     absfilepath = os.path.normcase(tutils.test_data.path("scripts/a.py"))
     s = script.Script(absfilepath, fm)
     assert os.path.isfile(s.args[0])
+