diff options
| -rw-r--r-- | examples/custom_contentviews.py | 68 | ||||
| -rw-r--r-- | libmproxy/console/__init__.py | 4 | ||||
| -rw-r--r-- | libmproxy/console/grideditor.py | 2 | ||||
| -rw-r--r-- | libmproxy/contentviews.py | 81 | ||||
| -rw-r--r-- | libmproxy/exceptions.py | 4 | ||||
| -rw-r--r-- | libmproxy/flow.py | 12 | ||||
| -rw-r--r-- | libmproxy/script.py | 216 | ||||
| -rw-r--r-- | libmproxy/script/__init__.py | 11 | ||||
| -rw-r--r-- | libmproxy/script/concurrent.py | 62 | ||||
| -rw-r--r-- | libmproxy/script/script.py | 123 | ||||
| -rw-r--r-- | libmproxy/script/script_context.py | 59 | ||||
| -rw-r--r-- | setup.py | 42 | ||||
| -rw-r--r-- | test/test_contentview.py | 17 | ||||
| -rw-r--r-- | test/test_custom_contentview.py | 52 | ||||
| -rw-r--r-- | test/test_examples.py | 2 | ||||
| -rw-r--r-- | test/test_script.py | 26 |
16 files changed, 493 insertions, 288 deletions
diff --git a/examples/custom_contentviews.py b/examples/custom_contentviews.py new file mode 100644 index 00000000..17920e51 --- /dev/null +++ b/examples/custom_contentviews.py @@ -0,0 +1,68 @@ +import string +import lxml.html +import lxml.etree +from libmproxy import utils, contentviews + + +class ViewPigLatin(contentviews.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", contentviews.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/console/__init__.py b/libmproxy/console/__init__.py index e1b38c45..6f09a6b3 100644 --- a/libmproxy/console/__init__.py +++ b/libmproxy/console/__init__.py @@ -315,8 +315,8 @@ class ConsoleMaster(flow.FlowMaster): signals.add_event("Running script on flow: %s" % command, "debug") try: - s = script.Script(command, self) - except script.ScriptError as v: + s = script.Script(command, script.ScriptContext(self)) + except script.ScriptException as v: signals.status_message.send( message = "Error loading script." ) diff --git a/libmproxy/console/grideditor.py b/libmproxy/console/grideditor.py index d32ce5b4..237eea28 100644 --- a/libmproxy/console/grideditor.py +++ b/libmproxy/console/grideditor.py @@ -636,7 +636,7 @@ class ScriptEditor(GridEditor): def is_error(self, col, val): try: script.Script.parse_command(val) - except script.ScriptError as v: + except script.ScriptException as v: return str(v) 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/exceptions.py b/libmproxy/exceptions.py index 8f23bd92..e2bde980 100644 --- a/libmproxy/exceptions.py +++ b/libmproxy/exceptions.py @@ -48,3 +48,7 @@ class ContentViewException(ProxyException): class ReplayException(ProxyException): pass + + +class ScriptException(ProxyException): + pass
\ No newline at end of file diff --git a/libmproxy/flow.py b/libmproxy/flow.py index 6f57c1a6..f090d9c6 100644 --- a/libmproxy/flow.py +++ b/libmproxy/flow.py @@ -10,11 +10,9 @@ import os import re import urlparse - from netlib import wsgi from netlib.exceptions import HttpException from netlib.http import CONTENT_MISSING, Headers, http1 -import netlib.http from . import controller, tnetstring, filt, script, version from .onboarding import app from .proxy.config import HostMatcher @@ -642,7 +640,7 @@ class FlowMaster(controller.Master): self.stream = None self.apps = AppRegistry() - script.sig_script_change.connect(self.script_change) + script.script_change.connect(self.script_change) def start_app(self, host, port): self.apps.add( @@ -664,7 +662,7 @@ class FlowMaster(controller.Master): def unload_script(self, script_obj): try: script_obj.unload() - except script.ScriptError as e: + except script.ScriptException as e: self.add_event("Script error:\n" + str(e), "error") self.scripts.remove(script_obj) @@ -674,8 +672,8 @@ class FlowMaster(controller.Master): wrong. """ try: - s = script.Script(command, self) - except script.ScriptError as v: + s = script.Script(command, script.ScriptContext(self)) + except script.ScriptException as v: return v.args[0] self.scripts.append(s) @@ -683,7 +681,7 @@ class FlowMaster(controller.Master): if script_obj and not self.pause_scripts: try: script_obj.run(name, *args, **kwargs) - except script.ScriptError as e: + except script.ScriptException as e: self.add_event("Script error:\n" + str(e), "error") def run_script_hook(self, name, *args, **kwargs): diff --git a/libmproxy/script.py b/libmproxy/script.py deleted file mode 100644 index b1ea83c8..00000000 --- a/libmproxy/script.py +++ /dev/null @@ -1,216 +0,0 @@ -from __future__ import absolute_import -import os -import traceback -import threading -import shlex -import sys -from watchdog.observers import Observer -from watchdog.events import PatternMatchingEventHandler, FileModifiedEvent -import blinker - -sig_script_change = blinker.Signal() - -class ScriptError(Exception): - pass - - -class ScriptContext: - """ - The script context should be used to interact with the global mitmproxy state from within a - script. - """ - def __init__(self, master): - self._master = master - - def log(self, message, level="info"): - """ - Logs an event. - - By default, only events with level "error" get displayed. This can be controlled with the "-v" switch. - How log messages are handled depends on the front-end. mitmdump will print them to stdout, - mitmproxy sends output to the eventlog for display ("e" keyboard shortcut). - """ - self._master.add_event(message, level) - - def kill_flow(self, f): - """ - Kills a flow immediately. No further data will be sent to the client or the server. - """ - f.kill(self._master) - - def duplicate_flow(self, f): - """ - Returns a duplicate of the specified flow. The flow is also - injected into the current state, and is ready for editing, replay, - etc. - """ - self._master.pause_scripts = True - f = self._master.duplicate_flow(f) - self._master.pause_scripts = False - return f - - def replay_request(self, f): - """ - Replay the request on the current flow. The response will be added - to the flow object. - """ - return self._master.replay_request(f, block=True, run_scripthooks=False) - - @property - def app_registry(self): - return self._master.apps - - -class Script: - """ - Script object representing an inline script. - """ - - def __init__(self, command, master): - self.command = command - self.args = self.parse_command(command) - self.ctx = ScriptContext(master) - self.ns = None - self.load() - self.observe_scripts() - - @classmethod - def parse_command(cls, command): - if not command or not command.strip(): - raise ScriptError("Empty script command.") - if os.name == "nt": # Windows: escape all backslashes in the path. - backslashes = shlex.split(command, posix=False)[0].count("\\") - command = command.replace("\\", "\\\\", backslashes) - args = shlex.split(command) - args[0] = os.path.expanduser(args[0]) - if not os.path.exists(args[0]): - raise ScriptError( - ("Script file not found: %s.\r\n" - "If your script path contains spaces, " - "make sure to wrap it in additional quotes, e.g. -s \"'./foo bar/baz.py' --args\".") % - args[0]) - elif os.path.isdir(args[0]): - raise ScriptError("Not a file: %s" % args[0]) - return args - - def load(self): - """ - Loads an inline script. - - Returns: - The return value of self.run("start", ...) - - Raises: - ScriptError on failure - """ - if self.ns is not None: - self.unload() - script_dir = os.path.dirname(os.path.abspath(self.args[0])) - ns = {'__file__': os.path.abspath(self.args[0])} - sys.path.append(script_dir) - try: - execfile(self.args[0], ns, ns) - except Exception as e: - # Python 3: use exception chaining, https://www.python.org/dev/peps/pep-3134/ - raise ScriptError(traceback.format_exc(e)) - sys.path.pop() - self.ns = ns - return self.run("start", self.args) - - def unload(self): - ret = self.run("done") - self.ns = None - return ret - - def run(self, name, *args, **kwargs): - """ - Runs an inline script hook. - - Returns: - The return value of the method. - None, if the script does not provide the method. - - Raises: - ScriptError if there was an exception. - """ - f = self.ns.get(name) - if f: - try: - return f(self.ctx, *args, **kwargs) - except Exception as e: - raise ScriptError(traceback.format_exc(e)) - else: - return None - - def observe_scripts(self): - script_dir = os.path.dirname(self.args[0]) - event_handler = ScriptModified(self) - observer = Observer() - observer.schedule(event_handler, script_dir) - observer.start() - - -class ReplyProxy(object): - def __init__(self, original_reply, script_thread): - self.original_reply = original_reply - self.script_thread = script_thread - self._ignore_call = True - self.lock = threading.Lock() - - def __call__(self, *args, **kwargs): - with self.lock: - if self._ignore_call: - self.script_thread.start() - self._ignore_call = False - return - self.original_reply(*args, **kwargs) - - def __getattr__(self, k): - return getattr(self.original_reply, k) - - -def _handle_concurrent_reply(fn, o, *args, **kwargs): - # Make first call to o.reply a no op and start the script thread. - # We must not start the script thread before, as this may lead to a nasty race condition - # where the script thread replies a different response before the normal reply, which then gets swallowed. - - def run(): - fn(*args, **kwargs) - # If the script did not call .reply(), we have to do it now. - reply_proxy() - - script_thread = ScriptThread(target=run) - - reply_proxy = ReplyProxy(o.reply, script_thread) - o.reply = reply_proxy - - -class ScriptThread(threading.Thread): - name = "ScriptThread" - - -def concurrent(fn): - if fn.func_name in ( - "request", - "response", - "error", - "clientconnect", - "serverconnect", - "clientdisconnect", - "next_layer"): - def _concurrent(ctx, obj): - _handle_concurrent_reply(fn, obj, ctx, obj) - - return _concurrent - raise NotImplementedError( - "Concurrent decorator not supported for '%s' method." % fn.func_name) - - -class ScriptModified(PatternMatchingEventHandler): - - def __init__(self, script): - PatternMatchingEventHandler.__init__(self, ignore_directories=True, patterns=["*.py"]) - self.script = script - - def on_modified(self, event=FileModifiedEvent): - sig_script_change.send(self.script) diff --git a/libmproxy/script/__init__.py b/libmproxy/script/__init__.py new file mode 100644 index 00000000..0f487795 --- /dev/null +++ b/libmproxy/script/__init__.py @@ -0,0 +1,11 @@ +from .script import Script, script_change +from .script_context import ScriptContext +from .concurrent import concurrent +from ..exceptions import ScriptException + +__all__ = [ + "Script", "script_change", + "ScriptContext", + "concurrent", + "ScriptException" +]
\ No newline at end of file diff --git a/libmproxy/script/concurrent.py b/libmproxy/script/concurrent.py new file mode 100644 index 00000000..bee2d43b --- /dev/null +++ b/libmproxy/script/concurrent.py @@ -0,0 +1,62 @@ +""" +This module provides a @concurrent decorator primitive to +offload computations from mitmproxy's main master thread. +""" +from __future__ import absolute_import, print_function, division +import threading + + +class ReplyProxy(object): + def __init__(self, original_reply, script_thread): + self.original_reply = original_reply + self.script_thread = script_thread + self._ignore_call = True + self.lock = threading.Lock() + + def __call__(self, *args, **kwargs): + with self.lock: + if self._ignore_call: + self.script_thread.start() + self._ignore_call = False + return + self.original_reply(*args, **kwargs) + + def __getattr__(self, k): + return getattr(self.original_reply, k) + + +def _handle_concurrent_reply(fn, o, *args, **kwargs): + # Make first call to o.reply a no op and start the script thread. + # We must not start the script thread before, as this may lead to a nasty race condition + # where the script thread replies a different response before the normal reply, which then gets swallowed. + + def run(): + fn(*args, **kwargs) + # If the script did not call .reply(), we have to do it now. + reply_proxy() + + script_thread = ScriptThread(target=run) + + reply_proxy = ReplyProxy(o.reply, script_thread) + o.reply = reply_proxy + + +class ScriptThread(threading.Thread): + name = "ScriptThread" + + +def concurrent(fn): + if fn.func_name in ( + "request", + "response", + "error", + "clientconnect", + "serverconnect", + "clientdisconnect", + "next_layer"): + def _concurrent(ctx, obj): + _handle_concurrent_reply(fn, obj, ctx, obj) + + return _concurrent + raise NotImplementedError( + "Concurrent decorator not supported for '%s' method." % fn.func_name) diff --git a/libmproxy/script/script.py b/libmproxy/script/script.py new file mode 100644 index 00000000..a58ba0af --- /dev/null +++ b/libmproxy/script/script.py @@ -0,0 +1,123 @@ +""" +The script object representing mitmproxy inline scripts. +Script objects know nothing about mitmproxy or mitmproxy's API - this knowledge is provided +by the mitmproxy-specific ScriptContext. +""" +from __future__ import absolute_import, print_function, division +import os +import shlex +import traceback +import sys +import blinker + +from watchdog.events import PatternMatchingEventHandler, FileModifiedEvent +from watchdog.observers import Observer + +from ..exceptions import ScriptException + +script_change = blinker.Signal() + + +class Script(object): + """ + Script object representing an inline script. + """ + + def __init__(self, command, context, use_reloader=True): + self.command = command + self.args = self.parse_command(command) + self.ctx = context + self.ns = None + self.load() + if use_reloader: + self.start_observe() + + @classmethod + def parse_command(cls, command): + if not command or not command.strip(): + raise ScriptException("Empty script command.") + if os.name == "nt": # Windows: escape all backslashes in the path. + backslashes = shlex.split(command, posix=False)[0].count("\\") + command = command.replace("\\", "\\\\", backslashes) + args = shlex.split(command) + args[0] = os.path.expanduser(args[0]) + if not os.path.exists(args[0]): + raise ScriptException( + ("Script file not found: %s.\r\n" + "If your script path contains spaces, " + "make sure to wrap it in additional quotes, e.g. -s \"'./foo bar/baz.py' --args\".") % + args[0]) + elif os.path.isdir(args[0]): + raise ScriptException("Not a file: %s" % args[0]) + return args + + def load(self): + """ + Loads an inline script. + + Returns: + The return value of self.run("start", ...) + + Raises: + ScriptException on failure + """ + if self.ns is not None: + self.unload() + script_dir = os.path.dirname(os.path.abspath(self.args[0])) + ns = {'__file__': os.path.abspath(self.args[0])} + sys.path.append(script_dir) + try: + execfile(self.args[0], ns, ns) + except Exception as e: + # Python 3: use exception chaining, https://www.python.org/dev/peps/pep-3134/ + raise ScriptException(traceback.format_exc(e)) + sys.path.pop() + self.ns = ns + return self.run("start", self.args) + + def unload(self): + ret = self.run("done") + self.ns = None + return ret + + def run(self, name, *args, **kwargs): + """ + Runs an inline script hook. + + Returns: + The return value of the method. + None, if the script does not provide the method. + + Raises: + ScriptException if there was an exception. + """ + f = self.ns.get(name) + if f: + try: + return f(self.ctx, *args, **kwargs) + except Exception as e: + raise ScriptException(traceback.format_exc(e)) + else: + return None + + def start_observe(self): + script_dir = os.path.dirname(self.args[0]) + event_handler = ScriptModified(self) + observer = Observer() + observer.schedule(event_handler, script_dir) + observer.start() + + def stop_observe(self): + raise NotImplementedError() # FIXME + + +class ScriptModified(PatternMatchingEventHandler): + def __init__(self, script): + # We could enumerate all relevant *.py files (as werkzeug does it), + # but our case looks like it isn't as simple as enumerating sys.modules. + # This should be good enough for now. + super(ScriptModified, self).__init__(ignore_directories=True, patterns=["*.py"]) + self.script = script + + def on_modified(self, event=FileModifiedEvent): + script_change.send(self.script) diff --git a/libmproxy/script/script_context.py b/libmproxy/script/script_context.py new file mode 100644 index 00000000..d8748cb2 --- /dev/null +++ b/libmproxy/script/script_context.py @@ -0,0 +1,59 @@ +""" +The mitmproxy script context provides an API to inline scripts. +""" +from __future__ import absolute_import, print_function, division +from .. import contentviews + + +class ScriptContext(object): + """ + The script context should be used to interact with the global mitmproxy state from within a + script. + """ + + def __init__(self, master): + self._master = master + + def log(self, message, level="info"): + """ + Logs an event. + + By default, only events with level "error" get displayed. This can be controlled with the "-v" switch. + How log messages are handled depends on the front-end. mitmdump will print them to stdout, + mitmproxy sends output to the eventlog for display ("e" keyboard shortcut). + """ + self._master.add_event(message, level) + + def kill_flow(self, f): + """ + Kills a flow immediately. No further data will be sent to the client or the server. + """ + f.kill(self._master) + + def duplicate_flow(self, f): + """ + Returns a duplicate of the specified flow. The flow is also + injected into the current state, and is ready for editing, replay, + etc. + """ + self._master.pause_scripts = True + f = self._master.duplicate_flow(f) + self._master.pause_scripts = False + return f + + def replay_request(self, f): + """ + Replay the request on the current flow. The response will be added + to the flow object. + """ + return self._master.replay_request(f, block=True, run_scripthooks=False) + + @property + def app_registry(self): + return self._master.apps + + def add_contentview(self, view_obj): + contentviews.add(view_obj) + + def remove_contentview(self, view_obj): + contentviews.remove(view_obj) @@ -15,25 +15,25 @@ with open(os.path.join(here, 'README.rst'), encoding='utf-8') as f: # Core dependencies deps = { "netlib>=%s, <%s" % (version.MINORVERSION, version.NEXT_MINORVERSION), - "pyasn1>0.1.2", - "tornado>=4.0.2", - "configargparse>=0.9.3", - "pyperclip>=1.5.8", - "blinker>=1.3", - "pyparsing>=1.5.2", - "html2text>=2015.4.14", - "construct>=2.5.2", - "six>=1.9.0", - "lxml>=3.3.6", - "Pillow>=2.3.0", + "pyasn1~=0.1.9", + "tornado~=4.3.0", + "configargparse~=0.10.0", + "pyperclip~=1.5.22", + "blinker~=1.4", + "pyparsing~=2.0.5", + "html2text~=2015.11.4", + "construct~=2.5.2", + "six~=1.10.0", + "lxml~=3.4.4", + "Pillow~=3.0.0", } # A script -> additional dependencies dict. scripts = { "mitmproxy": { - "urwid>=1.3", + "urwid~=1.3.1", }, "mitmdump": { - "click>=5.1", + "click~=5.1", }, "mitmweb": set() } @@ -50,9 +50,9 @@ dev_deps = { "sphinxcontrib-documentedlist>=0.2", } example_deps = { - "pytz", - "harparser", - "beautifulsoup4", + "pytz~=2015.7", + "harparser~=0.2", + "beautifulsoup4~=4.4.1", } # Add *all* script dependencies to developer dependencies. for script_deps in scripts.values(): @@ -61,14 +61,14 @@ for script_deps in scripts.values(): # Remove mitmproxy for Windows support. if os.name == "nt": del scripts["mitmproxy"] - deps.add("pydivert>=0.0.7") # Transparent proxying on Windows + deps.add("pydivert~=0.0.7") # Transparent proxying on Windows # Add dependencies for available scripts as core dependencies. for script_deps in scripts.values(): deps.update(script_deps) if sys.version_info < (3, 4): - example_deps.add("enum34") + example_deps.add("enum34~=1.0.4") console_scripts = ["%s = libmproxy.main:%s" % (s, s) for s in scripts.keys()] @@ -107,9 +107,9 @@ setup( extras_require={ 'dev': list(dev_deps), 'contentviews': [ - "pyamf>=0.6.1", - "protobuf>=2.5.0", - "cssutils>=1.0" + "pyamf~=0.7.2", + "protobuf~=2.6.1", + "cssutils~=1.0.1" ], 'examples': list(example_deps) } diff --git a/test/test_contentview.py b/test/test_contentview.py index 97608520..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(): @@ -233,7 +248,7 @@ if cv.ViewProtobuf.is_available(): p = tutils.test_data.path("data/protobuf01") content_type, output = v(file(p, "rb").read()) assert content_type == "Protobuf" - assert output[0].text == '1: "3bbc333c-e61c-433b-819a-0b9a8cc103b8"' + assert output.next()[0][1] == '1: "3bbc333c-e61c-433b-819a-0b9a8cc103b8"' def test_get_by_shortcut(): 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_examples.py b/test/test_examples.py index dce257cf..2a30f9d5 100644 --- a/test/test_examples.py +++ b/test/test_examples.py @@ -22,7 +22,7 @@ def test_load_scripts(): if "modify_response_body" in f: f += " foo bar" # two arguments required try: - s = script.Script(f, tmaster) # Loads the script file. + s = script.Script(f, script.ScriptContext(tmaster)) # Loads the script file. except Exception as v: if "ImportError" not in str(v): raise diff --git a/test/test_script.py b/test/test_script.py index 1b0e5a5b..fbe3e107 100644 --- a/test/test_script.py +++ b/test/test_script.py @@ -9,13 +9,13 @@ def test_simple(): s = flow.State() fm = flow.FlowMaster(None, s) sp = tutils.test_data.path("scripts/a.py") - p = script.Script("%s --var 40" % sp, fm) + p = script.Script("%s --var 40" % sp, script.ScriptContext(fm)) assert "here" in p.ns assert p.run("here") == 41 assert p.run("here") == 42 - tutils.raises(script.ScriptError, p.run, "errargs") + tutils.raises(script.ScriptException, p.run, "errargs") # Check reload p.load() @@ -36,29 +36,30 @@ def test_duplicate_flow(): def test_err(): s = flow.State() fm = flow.FlowMaster(None, s) + sc = script.ScriptContext(fm) tutils.raises( "not found", - script.Script, "nonexistent", fm + script.Script, "nonexistent", sc ) tutils.raises( "not a file", - script.Script, tutils.test_data.path("scripts"), fm + script.Script, tutils.test_data.path("scripts"), sc ) tutils.raises( - script.ScriptError, - script.Script, tutils.test_data.path("scripts/syntaxerr.py"), fm + script.ScriptException, + script.Script, tutils.test_data.path("scripts/syntaxerr.py"), sc ) tutils.raises( - script.ScriptError, - script.Script, tutils.test_data.path("scripts/loaderr.py"), fm + script.ScriptException, + script.Script, tutils.test_data.path("scripts/loaderr.py"), sc ) - scr = script.Script(tutils.test_data.path("scripts/unloaderr.py"), fm) - tutils.raises(script.ScriptError, scr.unload) + scr = script.Script(tutils.test_data.path("scripts/unloaderr.py"), sc) + tutils.raises(script.ScriptException, scr.unload) def test_concurrent(): @@ -84,7 +85,7 @@ def test_concurrent2(): fm = flow.FlowMaster(None, s) s = script.Script( tutils.test_data.path("scripts/concurrent_decorator.py"), - fm) + script.ScriptContext(fm)) s.load() m = mock.Mock() @@ -125,5 +126,6 @@ def test_command_parsing(): s = flow.State() fm = flow.FlowMaster(None, s) absfilepath = os.path.normcase(tutils.test_data.path("scripts/a.py")) - s = script.Script(absfilepath, fm) + s = script.Script(absfilepath, script.ScriptContext(fm)) assert os.path.isfile(s.args[0]) + |