diff options
40 files changed, 578 insertions, 713 deletions
diff --git a/examples/custom_contentviews.py b/examples/custom_contentviews.py index 5a63e2a0..b10d936f 100644 --- a/examples/custom_contentviews.py +++ b/examples/custom_contentviews.py @@ -62,7 +62,7 @@ class ViewPigLatin(contentviews.View): pig_view = ViewPigLatin() -def start(): +def configure(options): contentviews.add(pig_view) diff --git a/examples/filt.py b/examples/filt.py index 21744edd..102d1274 100644 --- a/examples/filt.py +++ b/examples/filt.py @@ -6,7 +6,7 @@ from mitmproxy import filt state = {} -def start(): +def configure(options): if len(sys.argv) != 2: raise ValueError("Usage: -s 'filt.py FILTER'") state["filter"] = filt.parse(sys.argv[1]) diff --git a/examples/flowwriter.py b/examples/flowwriter.py index 07c7ca20..d8fbc1f4 100644 --- a/examples/flowwriter.py +++ b/examples/flowwriter.py @@ -6,7 +6,7 @@ from mitmproxy.flow import FlowWriter state = {} -def start(): +def configure(options): if len(sys.argv) != 2: raise ValueError('Usage: -s "flowriter.py filename"') diff --git a/examples/har_extractor.py b/examples/har_extractor.py index 2a69b9af..23deb43a 100644 --- a/examples/har_extractor.py +++ b/examples/har_extractor.py @@ -61,7 +61,7 @@ class Context(object): context = Context() -def start(): +def configure(options): """ On start we create a HARLog instance. You will have to adapt this to suit your actual needs of HAR generation. As it will probably be diff --git a/examples/iframe_injector.py b/examples/iframe_injector.py index 70247d31..40934dd3 100644 --- a/examples/iframe_injector.py +++ b/examples/iframe_injector.py @@ -7,7 +7,7 @@ from mitmproxy.models import decoded iframe_url = None -def start(): +def configure(options): if len(sys.argv) != 2: raise ValueError('Usage: -s "iframe_injector.py url"') global iframe_url diff --git a/examples/modify_response_body.py b/examples/modify_response_body.py index 23ad0151..8b6908a4 100644 --- a/examples/modify_response_body.py +++ b/examples/modify_response_body.py @@ -8,7 +8,7 @@ from mitmproxy.models import decoded state = {} -def start(): +def configure(options): if len(sys.argv) != 3: raise ValueError('Usage: -s "modify_response_body.py old new"') # You may want to use Python's argparse for more sophisticated argument diff --git a/examples/proxapp.py b/examples/proxapp.py index 2935b587..095f412a 100644 --- a/examples/proxapp.py +++ b/examples/proxapp.py @@ -16,7 +16,7 @@ def hello_world(): # Register the app using the magic domain "proxapp" on port 80. Requests to # this domain and port combination will now be routed to the WSGI app instance. -def start(): +def configure(options): mitmproxy.ctx.master.apps.add(app, "proxapp", 80) # SSL works too, but the magic domain needs to be resolvable from the mitmproxy machine due to mitmproxy's design. diff --git a/examples/stub.py b/examples/stub.py index 10b34283..614acee2 100644 --- a/examples/stub.py +++ b/examples/stub.py @@ -4,11 +4,11 @@ import mitmproxy """ -def start(): +def configure(options): """ - Called once on script startup, before any other events. + Called once on script startup before any other events, and whenever options changes. """ - mitmproxy.ctx.log("start") + mitmproxy.ctx.log("configure") def clientconnect(root_layer): diff --git a/examples/tls_passthrough.py b/examples/tls_passthrough.py index 20e8f9be..306f55f6 100644 --- a/examples/tls_passthrough.py +++ b/examples/tls_passthrough.py @@ -113,7 +113,7 @@ class TlsFeedback(TlsLayer): tls_strategy = None -def start(): +def configure(options): global tls_strategy if len(sys.argv) == 2: tls_strategy = ProbabilisticStrategy(float(sys.argv[1])) diff --git a/mitmproxy/builtins/__init__.py b/mitmproxy/builtins/__init__.py index 8021c20f..6b357902 100644 --- a/mitmproxy/builtins/__init__.py +++ b/mitmproxy/builtins/__init__.py @@ -4,6 +4,7 @@ from mitmproxy.builtins import anticache from mitmproxy.builtins import anticomp from mitmproxy.builtins import stickyauth from mitmproxy.builtins import stickycookie +from mitmproxy.builtins import script from mitmproxy.builtins import stream @@ -13,5 +14,6 @@ def default_addons(): anticomp.AntiComp(), stickyauth.StickyAuth(), stickycookie.StickyCookie(), + script.ScriptLoader(), stream.Stream(), ] diff --git a/mitmproxy/builtins/script.py b/mitmproxy/builtins/script.py new file mode 100644 index 00000000..34801ff7 --- /dev/null +++ b/mitmproxy/builtins/script.py @@ -0,0 +1,162 @@ +from __future__ import absolute_import, print_function, division + +import contextlib +import os +import shlex +import sys +import threading +import traceback + +from mitmproxy import exceptions +from mitmproxy import controller +from mitmproxy import ctx + + +import watchdog.events +# The OSX reloader in watchdog 0.8.3 breaks when unobserving paths. +# We use the PollingObserver instead. +if sys.platform == 'darwin': # pragma: no cover + from watchdog.observers.polling import PollingObserver as Observer +else: + from watchdog.observers import Observer + + +def parse_command(command): + """ + Returns a (path, args) tuple. + """ + if not command or not command.strip(): + raise exceptions.AddonError("Empty script command.") + # Windows: escape all backslashes in the path. + if os.name == "nt": # pragma: no cover + backslashes = shlex.split(command, posix=False)[0].count("\\") + command = command.replace("\\", "\\\\", backslashes) + args = shlex.split(command) # pragma: no cover + args[0] = os.path.expanduser(args[0]) + if not os.path.exists(args[0]): + raise exceptions.AddonError( + ("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 exceptions.AddonError("Not a file: %s" % args[0]) + return args[0], args[1:] + + +@contextlib.contextmanager +def scriptenv(path, args): + oldargs = sys.argv + sys.argv = [path] + args + script_dir = os.path.dirname(os.path.abspath(path)) + sys.path.append(script_dir) + try: + yield + except Exception: + _, _, tb = sys.exc_info() + scriptdir = os.path.dirname(os.path.abspath(path)) + for i, s in enumerate(reversed(traceback.extract_tb(tb))): + if not os.path.abspath(s[0]).startswith(scriptdir): + break + else: + tb = tb.tb_next + ctx.log.warn("".join(traceback.format_tb(tb))) + finally: + sys.argv = oldargs + sys.path.pop() + + +def load_script(path, args): + ns = {'__file__': os.path.abspath(path)} + with scriptenv(path, args): + with open(path, "rb") as f: + code = compile(f.read(), path, 'exec') + exec(code, ns, ns) + return ns + + +class ReloadHandler(watchdog.events.FileSystemEventHandler): + def __init__(self, callback): + self.callback = callback + + def on_modified(self, event): + self.callback() + + def on_created(self, event): + self.callback() + + +class Script: + """ + An addon that manages a single script. + """ + def __init__(self, command): + self.name = command + + self.command = command + self.path, self.args = parse_command(command) + self.ns = None + self.observer = None + + self.last_options = None + self.should_reload = threading.Event() + + for i in controller.Events - set(["tick"]): + def mkprox(): + evt = i + + def prox(*args, **kwargs): + self.run(evt, *args, **kwargs) + return prox + setattr(self, i, mkprox()) + + def run(self, name, *args, **kwargs): + # It's possible for ns to be un-initialised if we failed during + # configure + if self.ns is not None: + func = self.ns.get(name) + if func: + with scriptenv(self.path, self.args): + func(*args, **kwargs) + + def reload(self): + self.should_reload.set() + + def tick(self): + if self.should_reload.is_set(): + self.should_reload.clear() + self.ns = None + ctx.log.info("Reloading script: %s" % self.name) + self.configure(self.last_options) + else: + self.run("tick") + + def configure(self, options): + self.last_options = options + if not self.observer: + self.observer = Observer() + # Bind the handler to the real underlying master object + self.observer.schedule( + ReloadHandler(self.reload), + os.path.dirname(self.path) or "." + ) + self.observer.start() + if not self.ns: + self.ns = load_script(self.path, self.args) + self.run("configure", options) + + +class ScriptLoader(): + """ + An addon that manages loading scripts from options. + """ + def configure(self, options): + for s in options.scripts or []: + if not ctx.master.addons.has_addon(s): + ctx.log.info("Loading script: %s" % s) + sc = Script(s) + ctx.master.addons.add(sc) + for a in ctx.master.addons.chain: + if isinstance(a, Script): + if a.name not in options.scripts or []: + ctx.master.addons.remove(a) diff --git a/mitmproxy/console/grideditor.py b/mitmproxy/console/grideditor.py index 9fa51ccb..f304de57 100644 --- a/mitmproxy/console/grideditor.py +++ b/mitmproxy/console/grideditor.py @@ -6,11 +6,12 @@ import re import urwid +from mitmproxy import exceptions from mitmproxy import filt -from mitmproxy import script -from mitmproxy import utils +from mitmproxy.builtins import script from mitmproxy.console import common from mitmproxy.console import signals +from netlib import strutils from netlib.http import cookies from netlib.http import user_agents @@ -55,7 +56,7 @@ class TextColumn: o = editor.walker.get_current_value() if o is not None: n = editor.master.spawn_editor(o.encode("string-escape")) - n = utils.clean_hanging_newline(n) + n = strutils.clean_hanging_newline(n) editor.walker.set_current_value(n, False) editor.walker._modified() elif key in ["enter"]: @@ -643,8 +644,8 @@ class ScriptEditor(GridEditor): def is_error(self, col, val): try: - script.Script.parse_command(val) - except script.ScriptException as e: + script.parse_command(val) + except exceptions.AddonError as e: return str(e) diff --git a/mitmproxy/console/master.py b/mitmproxy/console/master.py index bc373a2b..64bd9f0a 100644 --- a/mitmproxy/console/master.py +++ b/mitmproxy/console/master.py @@ -248,23 +248,6 @@ class ConsoleMaster(flow.FlowMaster): if options.server_replay: self.server_playback_path(options.server_replay) - if options.scripts: - for i in options.scripts: - try: - self.load_script(i) - except exceptions.ScriptException as e: - print("Script load error: {}".format(e), file=sys.stderr) - sys.exit(1) - - if options.outfile: - err = self.start_stream_to_path( - options.outfile[0], - options.outfile[1] - ) - if err: - print("Stream file error: {}".format(err), file=sys.stderr) - sys.exit(1) - self.view_stack = [] if options.app: @@ -685,20 +668,7 @@ class ConsoleMaster(flow.FlowMaster): self.refresh_focus() def edit_scripts(self, scripts): - commands = [x[0] for x in scripts] # remove outer array - if commands == [s.command for s in self.scripts]: - return - - self.unload_scripts() - for command in commands: - try: - self.load_script(command) - except exceptions.ScriptException as e: - signals.status_message.send( - message='Error loading "{}".'.format(command) - ) - signals.add_event('Error loading "{}":\n{}'.format(command, e), "error") - signals.update_settings.send(self) + self.options.scripts = [x[0] for x in scripts] def stop_client_playback_prompt(self, a): if a != "n": diff --git a/mitmproxy/console/options.py b/mitmproxy/console/options.py index d363ba74..d8824b05 100644 --- a/mitmproxy/console/options.py +++ b/mitmproxy/console/options.py @@ -54,7 +54,7 @@ class Options(urwid.WidgetWrap): select.Option( "Scripts", "S", - lambda: master.scripts, + lambda: master.options.scripts, self.scripts ), @@ -160,12 +160,14 @@ class Options(urwid.WidgetWrap): self.master.replacehooks.clear() self.master.set_ignore_filter([]) self.master.set_tcp_filter([]) - self.master.scripts = [] - self.master.options.anticache = False - self.master.options.anticomp = False - self.master.options.stickyauth = None - self.master.options.stickycookie = None + self.master.options.update( + scripts = [], + anticache = False, + anticomp = False, + stickyauth = None, + stickycookie = None + ) self.master.state.default_body_view = contentviews.get("Auto") @@ -234,7 +236,7 @@ class Options(urwid.WidgetWrap): self.master.view_grideditor( grideditor.ScriptEditor( self.master, - [[i.command] for i in self.master.scripts], + [[i] for i in self.master.options.scripts], self.master.edit_scripts ) ) diff --git a/mitmproxy/console/statusbar.py b/mitmproxy/console/statusbar.py index fc41869c..e7a700a6 100644 --- a/mitmproxy/console/statusbar.py +++ b/mitmproxy/console/statusbar.py @@ -218,14 +218,13 @@ class StatusBar(urwid.WidgetWrap): dst.address.host, dst.address.port )) - if self.master.scripts: + if self.master.options.scripts: r.append("[") r.append(("heading_key", "s")) - r.append("cripts:%s]" % len(self.master.scripts)) - # r.append("[lt:%0.3f]"%self.master.looptime) + r.append("cripts:%s]" % len(self.master.options.scripts)) - if self.master.stream: - r.append("[W:%s]" % self.master.stream_path) + if self.master.options.outfile: + r.append("[W:%s]" % self.master.outfile[0]) return r diff --git a/mitmproxy/controller.py b/mitmproxy/controller.py index 2f0c8bf2..d3ae1baa 100644 --- a/mitmproxy/controller.py +++ b/mitmproxy/controller.py @@ -32,6 +32,8 @@ Events = frozenset([ "error", "log", + "done", + "tick", "script_change", ]) @@ -44,7 +46,17 @@ class Log(object): def __call__(self, text, level="info"): self.master.add_event(text, level) - # We may want to add .log(), .warn() etc. here at a later point in time + def debug(self, txt): + self(txt, "debug") + + def info(self, txt): + self(txt, "info") + + def warn(self, txt): + self(txt, "warn") + + def error(self, txt): + self(txt, "error") class Master(object): @@ -97,26 +109,25 @@ class Master(object): self.shutdown() def tick(self, timeout): + with self.handlecontext(): + self.addons("tick") changed = False try: - # This endless loop runs until the 'Queue.Empty' - # exception is thrown. - while True: - mtype, obj = self.event_queue.get(timeout=timeout) - if mtype not in Events: - raise exceptions.ControlException("Unknown event %s" % repr(mtype)) - handle_func = getattr(self, mtype) - if not callable(handle_func): - raise exceptions.ControlException("Handler %s not callable" % mtype) - if not handle_func.__dict__.get("__handler"): - raise exceptions.ControlException( - "Handler function %s is not decorated with controller.handler" % ( - handle_func - ) + mtype, obj = self.event_queue.get(timeout=timeout) + if mtype not in Events: + raise exceptions.ControlException("Unknown event %s" % repr(mtype)) + handle_func = getattr(self, mtype) + if not callable(handle_func): + raise exceptions.ControlException("Handler %s not callable" % mtype) + if not handle_func.__dict__.get("__handler"): + raise exceptions.ControlException( + "Handler function %s is not decorated with controller.handler" % ( + handle_func ) - handle_func(obj) - self.event_queue.task_done() - changed = True + ) + handle_func(obj) + self.event_queue.task_done() + changed = True except queue.Empty: pass return changed diff --git a/mitmproxy/dump.py b/mitmproxy/dump.py index 274e01f3..999a709a 100644 --- a/mitmproxy/dump.py +++ b/mitmproxy/dump.py @@ -93,13 +93,6 @@ class DumpMaster(flow.FlowMaster): not options.keepserving ) - scripts = options.scripts or [] - for command in scripts: - try: - self.load_script(command, use_reloader=True) - except exceptions.ScriptException as e: - raise DumpError(str(e)) - if options.rfile: try: self.load_flows_file(options.rfile) @@ -335,6 +328,5 @@ class DumpMaster(flow.FlowMaster): def run(self): # pragma: no cover if self.options.rfile and not self.options.keepserving: - self.unload_scripts() # make sure to trigger script unload events. return super(DumpMaster, self).run() diff --git a/mitmproxy/exceptions.py b/mitmproxy/exceptions.py index 282784b6..3b41fe1c 100644 --- a/mitmproxy/exceptions.py +++ b/mitmproxy/exceptions.py @@ -99,3 +99,7 @@ class ControlException(ProxyException): class OptionsError(Exception): pass + + +class AddonError(Exception): + pass diff --git a/mitmproxy/flow/master.py b/mitmproxy/flow/master.py index 27ceee87..aa09e109 100644 --- a/mitmproxy/flow/master.py +++ b/mitmproxy/flow/master.py @@ -9,7 +9,6 @@ import netlib.exceptions from mitmproxy import controller from mitmproxy import exceptions from mitmproxy import models -from mitmproxy import script from mitmproxy.flow import io from mitmproxy.flow import modules from mitmproxy.onboarding import app @@ -35,8 +34,6 @@ class FlowMaster(controller.Master): self.server_playback = None # type: Optional[modules.ServerPlaybackState] self.client_playback = None # type: Optional[modules.ClientPlaybackState] self.kill_nonreplay = False - self.scripts = [] # type: List[script.Script] - self.pause_scripts = False self.stream_large_bodies = None # type: Optional[modules.StreamLargeBodies] self.refresh_server_playback = False @@ -60,44 +57,6 @@ class FlowMaster(controller.Master): level: debug, info, error """ - def unload_scripts(self): - for s in self.scripts[:]: - self.unload_script(s) - - def unload_script(self, script_obj): - try: - script_obj.unload() - except script.ScriptException as e: - self.add_event("Script error:\n" + str(e), "error") - script.reloader.unwatch(script_obj) - self.scripts.remove(script_obj) - - def load_script(self, command, use_reloader=False): - """ - Loads a script. - - Raises: - ScriptException - """ - s = script.Script(command) - s.load() - if use_reloader: - s.reply = controller.DummyReply() - script.reloader.watch(s, lambda: self.event_queue.put(("script_change", s))) - self.scripts.append(s) - - def _run_single_script_hook(self, script_obj, name, *args, **kwargs): - if script_obj and not self.pause_scripts: - try: - script_obj.run(name, *args, **kwargs) - except script.ScriptException as e: - self.add_event("Script error:\n{}".format(e), "error") - - def run_scripts(self, name, msg): - for script_obj in self.scripts: - if not msg.reply.acked: - self._run_single_script_hook(script_obj, name, msg) - def get_ignore_filter(self): return self.server.config.check_ignore.patterns @@ -298,11 +257,11 @@ class FlowMaster(controller.Master): if not pb and self.kill_nonreplay: f.kill(self) - def replay_request(self, f, block=False, run_scripthooks=True): + def replay_request(self, f, block=False): """ Returns None if successful, or error message if not. """ - if f.live and run_scripthooks: + if f.live: return "Can't replay live request." if f.intercepted: return "Can't replay while intercepting..." @@ -319,7 +278,7 @@ class FlowMaster(controller.Master): rt = http_replay.RequestReplayThread( self.server.config, f, - self.event_queue if run_scripthooks else False, + self.event_queue, self.should_exit ) rt.start() # pragma: no cover @@ -332,28 +291,27 @@ class FlowMaster(controller.Master): @controller.handler def clientconnect(self, root_layer): - self.run_scripts("clientconnect", root_layer) + pass @controller.handler def clientdisconnect(self, root_layer): - self.run_scripts("clientdisconnect", root_layer) + pass @controller.handler def serverconnect(self, server_conn): - self.run_scripts("serverconnect", server_conn) + pass @controller.handler def serverdisconnect(self, server_conn): - self.run_scripts("serverdisconnect", server_conn) + pass @controller.handler def next_layer(self, top_layer): - self.run_scripts("next_layer", top_layer) + pass @controller.handler def error(self, f): self.state.update_flow(f) - self.run_scripts("error", f) if self.client_playback: self.client_playback.clear(f) return f @@ -381,8 +339,6 @@ class FlowMaster(controller.Master): self.setheaders.run(f) if not f.reply.acked: self.process_new_request(f) - if not f.reply.acked: - self.run_scripts("request", f) return f @controller.handler @@ -393,7 +349,6 @@ class FlowMaster(controller.Master): except netlib.exceptions.HttpException: f.reply.kill() return - self.run_scripts("responseheaders", f) return f @controller.handler @@ -404,7 +359,6 @@ class FlowMaster(controller.Master): self.replacehooks.run(f) if not f.reply.acked: self.setheaders.run(f) - self.run_scripts("response", f) if not f.reply.acked: if self.client_playback: self.client_playback.clear(f) @@ -417,45 +371,14 @@ class FlowMaster(controller.Master): self.state.update_flow(f) @controller.handler - def script_change(self, s): - """ - Handle a script whose contents have been changed on the file system. - - Args: - s (script.Script): the changed script - - Returns: - True, if reloading was successful. - False, otherwise. - """ - ok = True - # We deliberately do not want to fail here. - # In the worst case, we have an "empty" script object. - try: - s.unload() - except script.ScriptException as e: - ok = False - self.add_event('Error reloading "{}":\n{}'.format(s.path, e), 'error') - try: - s.load() - except script.ScriptException as e: - ok = False - self.add_event('Error reloading "{}":\n{}'.format(s.path, e), 'error') - else: - self.add_event('"{}" reloaded.'.format(s.path), 'info') - return ok - - @controller.handler def tcp_open(self, flow): # TODO: This would break mitmproxy currently. # self.state.add_flow(flow) self.active_flows.add(flow) - self.run_scripts("tcp_open", flow) @controller.handler def tcp_message(self, flow): - # type: (TCPFlow) -> None - self.run_scripts("tcp_message", flow) + pass @controller.handler def tcp_error(self, flow): @@ -463,13 +386,7 @@ class FlowMaster(controller.Master): repr(flow.server_conn.address), flow.error ), "info") - self.run_scripts("tcp_error", flow) @controller.handler def tcp_close(self, flow): self.active_flows.discard(flow) - self.run_scripts("tcp_close", flow) - - def shutdown(self): - super(FlowMaster, self).shutdown() - self.unload_scripts() diff --git a/mitmproxy/script/__init__.py b/mitmproxy/script/__init__.py index 9a3985ab..e75f282a 100644 --- a/mitmproxy/script/__init__.py +++ b/mitmproxy/script/__init__.py @@ -1,11 +1,5 @@ -from . import reloader from .concurrent import concurrent -from .script import Script -from ..exceptions import ScriptException __all__ = [ - "Script", "concurrent", - "ScriptException", - "reloader" ] diff --git a/mitmproxy/script/reloader.py b/mitmproxy/script/reloader.py deleted file mode 100644 index 857d76cd..00000000 --- a/mitmproxy/script/reloader.py +++ /dev/null @@ -1,47 +0,0 @@ -from __future__ import absolute_import, print_function, division - -import os - -from watchdog.events import RegexMatchingEventHandler - -from watchdog.observers.polling import PollingObserver as Observer -# We occasionally have watchdog errors on Windows, Linux and Mac when using the native observers. -# After reading through the watchdog source code and issue tracker, -# we may want to replace this with a very simple implementation of our own. - -_observers = {} - - -def watch(script, callback): - if script in _observers: - raise RuntimeError("Script already observed") - script_dir = os.path.dirname(os.path.abspath(script.path)) - script_name = os.path.basename(script.path) - event_handler = _ScriptModificationHandler(callback, filename=script_name) - observer = Observer() - observer.schedule(event_handler, script_dir) - observer.start() - _observers[script] = observer - - -def unwatch(script): - observer = _observers.pop(script, None) - if observer: - observer.stop() - observer.join() - - -class _ScriptModificationHandler(RegexMatchingEventHandler): - - def __init__(self, callback, filename='.*'): - - super(_ScriptModificationHandler, self).__init__( - ignore_directories=True, - regexes=['.*' + filename] - ) - self.callback = callback - - def on_modified(self, event): - self.callback() - -__all__ = ["watch", "unwatch"] diff --git a/mitmproxy/script/script.py b/mitmproxy/script/script.py deleted file mode 100644 index db4909ca..00000000 --- a/mitmproxy/script/script.py +++ /dev/null @@ -1,136 +0,0 @@ -""" -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. -""" -# Do not import __future__ here, this would apply transitively to the inline scripts. -from __future__ import absolute_import, print_function, division - -import os -import shlex -import sys -import contextlib - -import six -from typing import List # noqa - -from mitmproxy import exceptions - - -@contextlib.contextmanager -def scriptenv(path, args): - # type: (str, List[str]) -> None - oldargs = sys.argv - script_dir = os.path.dirname(os.path.abspath(path)) - - sys.argv = [path] + args - sys.path.append(script_dir) - try: - yield - finally: - sys.argv = oldargs - sys.path.pop() - - -class Script(object): - """ - Script object representing an inline script. - """ - - def __init__(self, command): - self.command = command - self.path, self.args = self.parse_command(command) - self.ns = None - - def __enter__(self): - self.load() - return self - - def __exit__(self, exc_type, exc_val, exc_tb): - if exc_val: - return False # re-raise the exception - self.unload() - - @staticmethod - def parse_command(command): - # type: (str) -> Tuple[str,List[str]] - """ - Returns a (path, args) tuple. - """ - if not command or not command.strip(): - raise exceptions.ScriptException("Empty script command.") - # Windows: escape all backslashes in the path. - if os.name == "nt": # pragma: no cover - backslashes = shlex.split(command, posix=False)[0].count("\\") - command = command.replace("\\", "\\\\", backslashes) - args = shlex.split(command) # pragma: no cover - args[0] = os.path.expanduser(args[0]) - if not os.path.exists(args[0]): - raise exceptions.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 exceptions.ScriptException("Not a file: %s" % args[0]) - return args[0], args[1:] - - 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: - raise exceptions.ScriptException("Script is already loaded") - self.ns = {'__file__': os.path.abspath(self.path)} - - with scriptenv(self.path, self.args): - try: - with open(self.path) as f: - code = compile(f.read(), self.path, 'exec') - exec(code, self.ns, self.ns) - except Exception: - six.reraise( - exceptions.ScriptException, - exceptions.ScriptException.from_exception_context(), - sys.exc_info()[2] - ) - return self.run("start") - - def unload(self): - try: - return self.run("done") - finally: - self.ns = None - - 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. - """ - if self.ns is None: - raise exceptions.ScriptException("Script not loaded.") - f = self.ns.get(name) - if f: - try: - with scriptenv(self.path, self.args): - return f(*args, **kwargs) - except Exception: - six.reraise( - exceptions.ScriptException, - exceptions.ScriptException.from_exception_context(), - sys.exc_info()[2] - ) - else: - return None diff --git a/netlib/utils.py b/netlib/utils.py index 23c16dc3..9eebf22c 100644 --- a/netlib/utils.py +++ b/netlib/utils.py @@ -56,6 +56,13 @@ class Data(object): dirname = os.path.dirname(inspect.getsourcefile(m)) self.dirname = os.path.abspath(dirname) + def push(self, subpath): + """ + Change the data object to a path relative to the module. + """ + self.dirname = os.path.join(self.dirname, subpath) + return self + def path(self, path): """ Returns a path to the package data housed at 'path' under this diff --git a/test/mitmproxy/builtins/test_script.py b/test/mitmproxy/builtins/test_script.py new file mode 100644 index 00000000..7c9787f8 --- /dev/null +++ b/test/mitmproxy/builtins/test_script.py @@ -0,0 +1,128 @@ +import time + +from mitmproxy.builtins import script +from mitmproxy import exceptions +from mitmproxy.flow import master +from mitmproxy.flow import state +from mitmproxy.flow import options + +from .. import tutils, mastertest + + +class TestParseCommand: + def test_empty_command(self): + with tutils.raises(exceptions.AddonError): + script.parse_command("") + + with tutils.raises(exceptions.AddonError): + script.parse_command(" ") + + def test_no_script_file(self): + with tutils.raises("not found"): + script.parse_command("notfound") + + with tutils.tmpdir() as dir: + with tutils.raises("not a file"): + script.parse_command(dir) + + def test_parse_args(self): + with tutils.chdir(tutils.test_data.dirname): + assert script.parse_command("data/scripts/a.py") == ("data/scripts/a.py", []) + assert script.parse_command("data/scripts/a.py foo bar") == ("data/scripts/a.py", ["foo", "bar"]) + assert script.parse_command("data/scripts/a.py 'foo bar'") == ("data/scripts/a.py", ["foo bar"]) + + @tutils.skip_not_windows + def test_parse_windows(self): + with tutils.chdir(tutils.test_data.dirname): + assert script.parse_command("data\\scripts\\a.py") == ("data\\scripts\\a.py", []) + assert script.parse_command("data\\scripts\\a.py 'foo \\ bar'") == ("data\\scripts\\a.py", 'foo \\ bar', []) + + +def test_load_script(): + ns = script.load_script( + tutils.test_data.path( + "data/addonscripts/recorder.py" + ), [] + ) + assert ns["configure"] + + +class TestScript(mastertest.MasterTest): + def test_simple(self): + s = state.State() + m = master.FlowMaster(options.Options(), None, s) + sc = script.Script( + tutils.test_data.path( + "data/addonscripts/recorder.py" + ) + ) + m.addons.add(sc) + assert sc.ns["call_log"] == [("configure", (options.Options(),), {})] + + sc.ns["call_log"] = [] + f = tutils.tflow(resp=True) + self.invoke(m, "request", f) + + recf = sc.ns["call_log"][0] + assert recf[0] == "request" + + def test_reload(self): + s = state.State() + m = mastertest.RecordingMaster(options.Options(), None, s) + with tutils.tmpdir(): + with open("foo.py", "w"): + pass + sc = script.Script("foo.py") + m.addons.add(sc) + + for _ in range(100): + with open("foo.py", "a") as f: + f.write(".") + m.addons.invoke_with_context(sc, "tick") + time.sleep(0.1) + if m.event_log: + return + raise AssertionError("Change event not detected.") + + def test_exception(self): + s = state.State() + m = mastertest.RecordingMaster(options.Options(), None, s) + sc = script.Script( + tutils.test_data.path("data/addonscripts/error.py") + ) + m.addons.add(sc) + f = tutils.tflow(resp=True) + self.invoke(m, "request", f) + assert m.event_log[0][0] == "warn" + + def test_duplicate_flow(self): + s = state.State() + fm = master.FlowMaster(None, None, s) + fm.addons.add( + script.Script( + tutils.test_data.path("data/addonscripts/duplicate_flow.py") + ) + ) + f = tutils.tflow() + fm.request(f) + assert fm.state.flow_count() == 2 + assert not fm.state.view[0].request.is_replay + assert fm.state.view[1].request.is_replay + + +class TestScriptLoader(mastertest.MasterTest): + def test_simple(self): + s = state.State() + o = options.Options(scripts=[]) + m = master.FlowMaster(o, None, s) + sc = script.ScriptLoader() + m.addons.add(sc) + assert len(m.addons) == 1 + o.update( + scripts = [ + tutils.test_data.path("data/addonscripts/recorder.py") + ] + ) + assert len(m.addons) == 2 + o.update(scripts = []) + assert len(m.addons) == 1 diff --git a/test/mitmproxy/data/scripts/concurrent_decorator.py b/test/mitmproxy/data/addonscripts/concurrent_decorator.py index 162c00f4..a56c2af1 100644 --- a/test/mitmproxy/data/scripts/concurrent_decorator.py +++ b/test/mitmproxy/data/addonscripts/concurrent_decorator.py @@ -1,7 +1,6 @@ import time from mitmproxy.script import concurrent - @concurrent def request(flow): time.sleep(0.1) diff --git a/test/mitmproxy/data/scripts/concurrent_decorator_err.py b/test/mitmproxy/data/addonscripts/concurrent_decorator_err.py index 756869c8..756869c8 100644 --- a/test/mitmproxy/data/scripts/concurrent_decorator_err.py +++ b/test/mitmproxy/data/addonscripts/concurrent_decorator_err.py diff --git a/test/mitmproxy/data/addonscripts/duplicate_flow.py b/test/mitmproxy/data/addonscripts/duplicate_flow.py new file mode 100644 index 00000000..b466423c --- /dev/null +++ b/test/mitmproxy/data/addonscripts/duplicate_flow.py @@ -0,0 +1,6 @@ +from mitmproxy import ctx + + +def request(flow): + f = ctx.master.duplicate_flow(flow) + ctx.master.replay_request(f, block=True) diff --git a/test/mitmproxy/data/addonscripts/error.py b/test/mitmproxy/data/addonscripts/error.py new file mode 100644 index 00000000..8ece9fce --- /dev/null +++ b/test/mitmproxy/data/addonscripts/error.py @@ -0,0 +1,7 @@ + +def mkerr(): + raise ValueError("Error!") + + +def request(flow): + mkerr() diff --git a/test/mitmproxy/data/addonscripts/recorder.py b/test/mitmproxy/data/addonscripts/recorder.py new file mode 100644 index 00000000..728203e3 --- /dev/null +++ b/test/mitmproxy/data/addonscripts/recorder.py @@ -0,0 +1,18 @@ +from mitmproxy import controller +from mitmproxy import ctx + +call_log = [] + +# Keep a log of all possible event calls +evts = list(controller.Events) + ["configure"] +for i in evts: + def mkprox(): + evt = i + + def prox(*args, **kwargs): + lg = (evt, args, kwargs) + if evt != "log": + ctx.log.info(str(lg)) + call_log.append(lg) + return prox + globals()[i] = mkprox() diff --git a/test/mitmproxy/data/addonscripts/stream_modify.py b/test/mitmproxy/data/addonscripts/stream_modify.py new file mode 100644 index 00000000..bc616342 --- /dev/null +++ b/test/mitmproxy/data/addonscripts/stream_modify.py @@ -0,0 +1,8 @@ + +def modify(chunks): + for chunk in chunks: + yield chunk.replace(b"foo", b"bar") + + +def responseheaders(flow): + flow.response.stream = modify diff --git a/test/mitmproxy/data/addonscripts/tcp_stream_modify.py b/test/mitmproxy/data/addonscripts/tcp_stream_modify.py new file mode 100644 index 00000000..af4ccf7e --- /dev/null +++ b/test/mitmproxy/data/addonscripts/tcp_stream_modify.py @@ -0,0 +1,5 @@ + +def tcp_message(flow): + message = flow.messages[-1] + if not message.from_client: + message.content = message.content.replace(b"foo", b"bar") diff --git a/test/mitmproxy/mastertest.py b/test/mitmproxy/mastertest.py index 9754d3a9..240f6a73 100644 --- a/test/mitmproxy/mastertest.py +++ b/test/mitmproxy/mastertest.py @@ -3,6 +3,7 @@ import mock from . import tutils import netlib.tutils +from mitmproxy.flow import master from mitmproxy import flow, proxy, models, controller @@ -39,3 +40,12 @@ class MasterTest: t = tutils.tflow(resp=True) fw.add(t) f.close() + + +class RecordingMaster(master.FlowMaster): + def __init__(self, *args, **kwargs): + master.FlowMaster.__init__(self, *args, **kwargs) + self.event_log = [] + + def add_event(self, e, level): + self.event_log.append((level, e)) diff --git a/test/mitmproxy/script/test_concurrent.py b/test/mitmproxy/script/test_concurrent.py index 57eeca19..92d1153b 100644 --- a/test/mitmproxy/script/test_concurrent.py +++ b/test/mitmproxy/script/test_concurrent.py @@ -1,28 +1,46 @@ -from mitmproxy.script import Script -from test.mitmproxy import tutils +from test.mitmproxy import tutils, mastertest from mitmproxy import controller +from mitmproxy.builtins import script +from mitmproxy import options +from mitmproxy.flow import master +from mitmproxy.flow import state import time class Thing: def __init__(self): self.reply = controller.DummyReply() + self.live = True -@tutils.skip_appveyor -def test_concurrent(): - with Script(tutils.test_data.path("data/scripts/concurrent_decorator.py")) as s: - f1, f2 = Thing(), Thing() - s.run("request", f1) - s.run("request", f2) +class TestConcurrent(mastertest.MasterTest): + @tutils.skip_appveyor + def test_concurrent(self): + s = state.State() + m = master.FlowMaster(options.Options(), None, s) + sc = script.Script( + tutils.test_data.path( + "data/addonscripts/concurrent_decorator.py" + ) + ) + m.addons.add(sc) + f1, f2 = tutils.tflow(), tutils.tflow() + self.invoke(m, "request", f1) + self.invoke(m, "request", f2) start = time.time() while time.time() - start < 5: if f1.reply.acked and f2.reply.acked: return raise ValueError("Script never acked") - -def test_concurrent_err(): - s = Script(tutils.test_data.path("data/scripts/concurrent_decorator_err.py")) - with tutils.raises("Concurrent decorator not supported for 'start' method"): - s.load() + def test_concurrent_err(self): + s = state.State() + m = mastertest.RecordingMaster(options.Options(), None, s) + sc = script.Script( + tutils.test_data.path( + "data/addonscripts/concurrent_decorator_err.py" + ) + ) + with m.handlecontext(): + sc.configure(options.Options()) + assert "decorator not supported" in m.event_log[0][1] diff --git a/test/mitmproxy/script/test_reloader.py b/test/mitmproxy/script/test_reloader.py deleted file mode 100644 index e33903b9..00000000 --- a/test/mitmproxy/script/test_reloader.py +++ /dev/null @@ -1,34 +0,0 @@ -import mock -from mitmproxy.script.reloader import watch, unwatch -from test.mitmproxy import tutils -from threading import Event - - -def test_simple(): - with tutils.tmpdir(): - with open("foo.py", "w"): - pass - - script = mock.Mock() - script.path = "foo.py" - - e = Event() - - def _onchange(): - e.set() - - watch(script, _onchange) - with tutils.raises("already observed"): - watch(script, _onchange) - - # Some reloaders don't register a change directly after watching, because they first need to initialize. - # To test if watching works at all, we do repeated writes every 100ms. - for _ in range(100): - with open("foo.py", "a") as f: - f.write(".") - if e.wait(0.1): - break - else: - raise AssertionError("No change detected.") - - unwatch(script) diff --git a/test/mitmproxy/script/test_script.py b/test/mitmproxy/script/test_script.py deleted file mode 100644 index 48fe65c9..00000000 --- a/test/mitmproxy/script/test_script.py +++ /dev/null @@ -1,83 +0,0 @@ -from mitmproxy.script import Script -from mitmproxy.exceptions import ScriptException -from test.mitmproxy import tutils - - -class TestParseCommand: - def test_empty_command(self): - with tutils.raises(ScriptException): - Script.parse_command("") - - with tutils.raises(ScriptException): - Script.parse_command(" ") - - def test_no_script_file(self): - with tutils.raises("not found"): - Script.parse_command("notfound") - - with tutils.tmpdir() as dir: - with tutils.raises("not a file"): - Script.parse_command(dir) - - def test_parse_args(self): - with tutils.chdir(tutils.test_data.dirname): - assert Script.parse_command("data/scripts/a.py") == ("data/scripts/a.py", []) - assert Script.parse_command("data/scripts/a.py foo bar") == ("data/scripts/a.py", ["foo", "bar"]) - assert Script.parse_command("data/scripts/a.py 'foo bar'") == ("data/scripts/a.py", ["foo bar"]) - - @tutils.skip_not_windows - def test_parse_windows(self): - with tutils.chdir(tutils.test_data.dirname): - assert Script.parse_command("data\\scripts\\a.py") == ("data\\scripts\\a.py", []) - assert Script.parse_command("data\\scripts\\a.py 'foo \\ bar'") == ("data\\scripts\\a.py", ['foo \\ bar']) - - -def test_simple(): - with tutils.chdir(tutils.test_data.path("data/scripts")): - s = Script("a.py --var 42") - assert s.path == "a.py" - assert s.ns is None - - s.load() - assert s.ns["var"] == 42 - - s.run("here") - assert s.ns["var"] == 43 - - s.unload() - assert s.ns is None - - with tutils.raises(ScriptException): - s.run("here") - - with Script("a.py --var 42") as s: - s.run("here") - - -def test_script_exception(): - with tutils.chdir(tutils.test_data.path("data/scripts")): - s = Script("syntaxerr.py") - with tutils.raises(ScriptException): - s.load() - - s = Script("starterr.py") - with tutils.raises(ScriptException): - s.load() - - s = Script("a.py") - s.load() - with tutils.raises(ScriptException): - s.load() - - s = Script("a.py") - with tutils.raises(ScriptException): - s.run("here") - - with tutils.raises(ScriptException): - with Script("reqerr.py") as s: - s.run("request", None) - - s = Script("unloaderr.py") - s.load() - with tutils.raises(ScriptException): - s.unload() diff --git a/test/mitmproxy/test_dump.py b/test/mitmproxy/test_dump.py index 9686be84..201386e3 100644 --- a/test/mitmproxy/test_dump.py +++ b/test/mitmproxy/test_dump.py @@ -245,12 +245,12 @@ class TestDumpMaster(mastertest.MasterTest): assert "XRESPONSE" in ret assert "XCLIENTDISCONNECT" in ret tutils.raises( - dump.DumpError, + exceptions.AddonError, self.mkmaster, None, scripts=["nonexistent"] ) tutils.raises( - dump.DumpError, + exceptions.AddonError, self.mkmaster, None, scripts=["starterr.py"] ) diff --git a/test/mitmproxy/test_examples.py b/test/mitmproxy/test_examples.py index bdadcd11..ef97219c 100644 --- a/test/mitmproxy/test_examples.py +++ b/test/mitmproxy/test_examples.py @@ -1,151 +1,125 @@ -import glob import json -import mock -import os -import sys -from contextlib import contextmanager -from mitmproxy import script +import sys +import os.path +from mitmproxy.flow import master +from mitmproxy.flow import state +from mitmproxy import options +from mitmproxy import contentviews +from mitmproxy.builtins import script import netlib.utils from netlib import tutils as netutils from netlib.http import Headers -from . import tutils - -example_dir = netlib.utils.Data(__name__).path("../../examples") - - -@contextmanager -def example(command): - command = os.path.join(example_dir, command) - with script.Script(command) as s: - yield s - - -@mock.patch("mitmproxy.ctx.master") -@mock.patch("mitmproxy.ctx.log") -def test_load_scripts(log, master): - scripts = glob.glob("%s/*.py" % example_dir) - - for f in scripts: - if "har_extractor" in f: - continue - if "flowwriter" in f: - f += " -" - if "iframe_injector" in f: - f += " foo" # one argument required - if "filt" in f: - f += " ~a" - if "modify_response_body" in f: - f += " foo bar" # two arguments required - - s = script.Script(f) - try: - s.load() - except Exception as v: - if "ImportError" not in str(v): - raise - else: - s.unload() - - -def test_add_header(): - flow = tutils.tflow(resp=netutils.tresp()) - with example("add_header.py") as ex: - ex.run("response", flow) - assert flow.response.headers["newheader"] == "foo" - - -@mock.patch("mitmproxy.contentviews.remove") -@mock.patch("mitmproxy.contentviews.add") -def test_custom_contentviews(add, remove): - with example("custom_contentviews.py"): - assert add.called - pig = add.call_args[0][0] - _, fmt = pig(b"<html>test!</html>") - assert any(b'esttay!' in val[0][1] for val in fmt) - assert not pig(b"gobbledygook") - assert remove.called - - -def test_iframe_injector(): - with tutils.raises(script.ScriptException): - with example("iframe_injector.py"): - pass - - flow = tutils.tflow(resp=netutils.tresp(content=b"<html>mitmproxy</html>")) - with example("iframe_injector.py http://example.org/evil_iframe") as ex: - ex.run("response", flow) - content = flow.response.content - assert b'iframe' in content and b'evil_iframe' in content +from . import tutils, mastertest +example_dir = netlib.utils.Data(__name__).push("../../examples") -def test_modify_form(): - form_header = Headers(content_type="application/x-www-form-urlencoded") - flow = tutils.tflow(req=netutils.treq(headers=form_header)) - with example("modify_form.py") as ex: - ex.run("request", flow) - assert flow.request.urlencoded_form[b"mitmproxy"] == b"rocks" - flow.request.headers["content-type"] = "" - ex.run("request", flow) - assert list(flow.request.urlencoded_form.items()) == [(b"foo", b"bar")] +class ScriptError(Exception): + pass -def test_modify_querystring(): - flow = tutils.tflow(req=netutils.treq(path=b"/search?q=term")) - with example("modify_querystring.py") as ex: - ex.run("request", flow) - assert flow.request.query["mitmproxy"] == "rocks" +class RaiseMaster(master.FlowMaster): + def add_event(self, e, level): + if level in ("warn", "error"): + raise ScriptError(e) - flow.request.path = "/" - ex.run("request", flow) - assert flow.request.query["mitmproxy"] == "rocks" +def tscript(cmd, args=""): + cmd = example_dir.path(cmd) + " " + args + m = RaiseMaster(options.Options(), None, state.State()) + sc = script.Script(cmd) + m.addons.add(sc) + return m, sc -def test_modify_response_body(): - with tutils.raises(script.ScriptException): - with example("modify_response_body.py"): - assert True - flow = tutils.tflow(resp=netutils.tresp(content=b"I <3 mitmproxy")) - with example("modify_response_body.py mitmproxy rocks") as ex: - assert ex.ns["state"]["old"] == b"mitmproxy" and ex.ns["state"]["new"] == b"rocks" - ex.run("response", flow) - assert flow.response.content == b"I <3 rocks" +class TestScripts(mastertest.MasterTest): + def test_add_header(self): + m, _ = tscript("add_header.py") + f = tutils.tflow(resp=netutils.tresp()) + self.invoke(m, "response", f) + assert f.response.headers["newheader"] == "foo" + def test_custom_contentviews(self): + m, sc = tscript("custom_contentviews.py") + pig = contentviews.get("pig_latin_HTML") + _, fmt = pig(b"<html>test!</html>") + assert any(b'esttay!' in val[0][1] for val in fmt) + assert not pig(b"gobbledygook") -def test_redirect_requests(): - flow = tutils.tflow(req=netutils.treq(host=b"example.org")) - with example("redirect_requests.py") as ex: - ex.run("request", flow) - assert flow.request.host == "mitmproxy.org" - - -@mock.patch("mitmproxy.ctx.log") -def test_har_extractor(log): - if sys.version_info >= (3, 0): - with tutils.raises("does not work on Python 3"): - with example("har_extractor.py -"): - pass - return - - with tutils.raises(script.ScriptException): - with example("har_extractor.py"): - pass - - times = dict( - timestamp_start=746203272, - timestamp_end=746203272, - ) - - flow = tutils.tflow( - req=netutils.treq(**times), - resp=netutils.tresp(**times) - ) + def test_iframe_injector(self): + with tutils.raises(ScriptError): + tscript("iframe_injector.py") - with example("har_extractor.py -") as ex: - ex.run("response", flow) + m, sc = tscript("iframe_injector.py", "http://example.org/evil_iframe") + flow = tutils.tflow(resp=netutils.tresp(content=b"<html>mitmproxy</html>")) + self.invoke(m, "response", flow) + content = flow.response.content + assert b'iframe' in content and b'evil_iframe' in content - with open(tutils.test_data.path("data/har_extractor.har")) as fp: + def test_modify_form(self): + m, sc = tscript("modify_form.py") + + form_header = Headers(content_type="application/x-www-form-urlencoded") + f = tutils.tflow(req=netutils.treq(headers=form_header)) + self.invoke(m, "request", f) + + assert f.request.urlencoded_form[b"mitmproxy"] == b"rocks" + + f.request.headers["content-type"] = "" + self.invoke(m, "request", f) + assert list(f.request.urlencoded_form.items()) == [(b"foo", b"bar")] + + def test_modify_querystring(self): + m, sc = tscript("modify_querystring.py") + f = tutils.tflow(req=netutils.treq(path="/search?q=term")) + + self.invoke(m, "request", f) + assert f.request.query["mitmproxy"] == "rocks" + + f.request.path = "/" + self.invoke(m, "request", f) + assert f.request.query["mitmproxy"] == "rocks" + + def test_modify_response_body(self): + with tutils.raises(ScriptError): + tscript("modify_response_body.py") + + m, sc = tscript("modify_response_body.py", "mitmproxy rocks") + f = tutils.tflow(resp=netutils.tresp(content=b"I <3 mitmproxy")) + self.invoke(m, "response", f) + assert f.response.content == b"I <3 rocks" + + def test_redirect_requests(self): + m, sc = tscript("redirect_requests.py") + f = tutils.tflow(req=netutils.treq(host="example.org")) + self.invoke(m, "request", f) + assert f.request.host == "mitmproxy.org" + + def test_har_extractor(self): + if sys.version_info >= (3, 0): + with tutils.raises("does not work on Python 3"): + tscript("har_extractor.py") + return + + with tutils.raises(ScriptError): + tscript("har_extractor.py") + + with tutils.tmpdir() as tdir: + times = dict( + timestamp_start=746203272, + timestamp_end=746203272, + ) + + path = os.path.join(tdir, "file") + m, sc = tscript("har_extractor.py", path) + f = tutils.tflow( + req=netutils.treq(**times), + resp=netutils.tresp(**times) + ) + self.invoke(m, "response", f) + m.addons.remove(sc) + + fp = open(path, "rb") test_data = json.load(fp) - assert json.loads(ex.ns["context"].HARLog.json()) == test_data["test_response"] + assert len(test_data["log"]["pages"]) == 1 diff --git a/test/mitmproxy/test_flow.py b/test/mitmproxy/test_flow.py index 1a07f74d..c58a9703 100644 --- a/test/mitmproxy/test_flow.py +++ b/test/mitmproxy/test_flow.py @@ -5,7 +5,7 @@ import netlib.utils from netlib.http import Headers from mitmproxy import filt, controller, flow from mitmproxy.contrib import tnetstring -from mitmproxy.exceptions import FlowReadException, ScriptException +from mitmproxy.exceptions import FlowReadException from mitmproxy.models import Error from mitmproxy.models import Flow from mitmproxy.models import HTTPFlow @@ -674,21 +674,6 @@ class TestSerialize: class TestFlowMaster: - def test_load_script(self): - s = flow.State() - fm = flow.FlowMaster(None, None, s) - - fm.load_script(tutils.test_data.path("data/scripts/a.py")) - fm.load_script(tutils.test_data.path("data/scripts/a.py")) - fm.unload_scripts() - with tutils.raises(ScriptException): - fm.load_script("nonexistent") - try: - fm.load_script(tutils.test_data.path("data/scripts/starterr.py")) - except ScriptException as e: - assert "ValueError" in str(e) - assert len(fm.scripts) == 0 - def test_getset_ignore(self): p = mock.Mock() p.config.check_ignore = HostMatcher() @@ -708,51 +693,7 @@ class TestFlowMaster: assert "intercepting" in fm.replay_request(f) f.live = True - assert "live" in fm.replay_request(f, run_scripthooks=True) - - def test_script_reqerr(self): - s = flow.State() - fm = flow.FlowMaster(None, None, s) - fm.load_script(tutils.test_data.path("data/scripts/reqerr.py")) - f = tutils.tflow() - fm.clientconnect(f.client_conn) - assert fm.request(f) - - def test_script(self): - s = flow.State() - fm = flow.FlowMaster(None, None, s) - fm.load_script(tutils.test_data.path("data/scripts/all.py")) - f = tutils.tflow(resp=True) - - f.client_conn.acked = False - fm.clientconnect(f.client_conn) - assert fm.scripts[0].ns["log"][-1] == "clientconnect" - f.server_conn.acked = False - fm.serverconnect(f.server_conn) - assert fm.scripts[0].ns["log"][-1] == "serverconnect" - f.reply.acked = False - fm.request(f) - assert fm.scripts[0].ns["log"][-1] == "request" - f.reply.acked = False - fm.response(f) - assert fm.scripts[0].ns["log"][-1] == "response" - # load second script - fm.load_script(tutils.test_data.path("data/scripts/all.py")) - assert len(fm.scripts) == 2 - f.server_conn.reply.acked = False - fm.clientdisconnect(f.server_conn) - assert fm.scripts[0].ns["log"][-1] == "clientdisconnect" - assert fm.scripts[1].ns["log"][-1] == "clientdisconnect" - - # unload first script - fm.unload_scripts() - assert len(fm.scripts) == 0 - fm.load_script(tutils.test_data.path("data/scripts/all.py")) - - f.error = tutils.terr() - f.reply.acked = False - fm.error(f) - assert fm.scripts[0].ns["log"][-1] == "error" + assert "live" in fm.replay_request(f) def test_duplicate_flow(self): s = flow.State() @@ -789,7 +730,6 @@ class TestFlowMaster: f.error.reply = controller.DummyReply() fm.error(f) - fm.load_script(tutils.test_data.path("data/scripts/a.py")) fm.shutdown() def test_client_playback(self): diff --git a/test/mitmproxy/test_script.py b/test/mitmproxy/test_script.py deleted file mode 100644 index 1e8220f1..00000000 --- a/test/mitmproxy/test_script.py +++ /dev/null @@ -1,13 +0,0 @@ -from mitmproxy import flow -from . import tutils - - -def test_duplicate_flow(): - s = flow.State() - fm = flow.FlowMaster(None, None, s) - fm.load_script(tutils.test_data.path("data/scripts/duplicate_flow.py")) - f = tutils.tflow() - fm.request(f) - assert fm.state.flow_count() == 2 - assert not fm.state.view[0].request.is_replay - assert fm.state.view[1].request.is_replay diff --git a/test/mitmproxy/test_server.py b/test/mitmproxy/test_server.py index 9dd8b79c..a5196dae 100644 --- a/test/mitmproxy/test_server.py +++ b/test/mitmproxy/test_server.py @@ -13,6 +13,7 @@ from netlib.http import authentication, http1 from netlib.tutils import raises from pathod import pathoc, pathod +from mitmproxy.builtins import script from mitmproxy import controller from mitmproxy.proxy.config import HostMatcher from mitmproxy.models import Error, HTTPResponse, HTTPFlow @@ -287,10 +288,13 @@ class TestHTTP(tservers.HTTPProxyTest, CommonMixin, AppMixin): self.master.set_stream_large_bodies(None) def test_stream_modify(self): - self.master.load_script(tutils.test_data.path("data/scripts/stream_modify.py")) + s = script.Script( + tutils.test_data.path("data/addonscripts/stream_modify.py") + ) + self.master.addons.add(s) d = self.pathod('200:b"foo"') assert d.content == b"bar" - self.master.unload_scripts() + self.master.addons.remove(s) class TestHTTPAuth(tservers.HTTPProxyTest): @@ -512,15 +516,15 @@ class TestTransparent(tservers.TransparentProxyTest, CommonMixin, TcpMixin): ssl = False def test_tcp_stream_modify(self): - self.master.load_script(tutils.test_data.path("data/scripts/tcp_stream_modify.py")) - + s = script.Script( + tutils.test_data.path("data/addonscripts/tcp_stream_modify.py") + ) + self.master.addons.add(s) self._tcpproxy_on() d = self.pathod('200:b"foo"') self._tcpproxy_off() - assert d.content == b"bar" - - self.master.unload_scripts() + self.master.addons.remove(s) class TestTransparentSSL(tservers.TransparentProxyTest, CommonMixin, TcpMixin): |