Merge pull request #833 from zbuc/contentview_scripts

Contentview scripts

3 files changed, 64 insertions, 28 deletions

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:
     """