diff --git a/examples/har_extractor.py b/examples/har_extractor.py index 2a69b9afd..90412ec0f 100644 --- a/examples/har_extractor.py +++ b/examples/har_extractor.py @@ -2,7 +2,7 @@ This inline script utilizes harparser.HAR from https://github.com/JustusW/harparser to generate a HAR log object. """ -import mitmproxy +import mitmproxy.ctx import six import sys import pytz @@ -221,9 +221,11 @@ def done(): if context.dump_file == '-': mitmproxy.ctx.log(pprint.pformat(json.loads(json_dump))) elif context.dump_file.endswith('.zhar'): - file(context.dump_file, "w").write(compressed_json_dump) + with open(context.dump_file, "wb") as f: + f.write(compressed_json_dump) else: - file(context.dump_file, "w").write(json_dump) + with open(context.dump_file, "wb") as f: + f.write(json_dump) mitmproxy.ctx.log( "HAR log finished with %s bytes (%s bytes compressed)" % ( len(json_dump), len(compressed_json_dump) diff --git a/examples/stub.py b/examples/stub.py index 10b342831..e5b4a39ab 100644 --- a/examples/stub.py +++ b/examples/stub.py @@ -6,11 +6,18 @@ import mitmproxy def start(): """ - Called once on script startup, before any other events. + Called once on script startup before any other events """ mitmproxy.ctx.log("start") +def configure(options): + """ + Called once on script startup before any other events, and whenever options changes. + """ + mitmproxy.ctx.log("configure") + + def clientconnect(root_layer): """ Called when a client initiates a connection to the proxy. Note that a diff --git a/mitmproxy/addons.py b/mitmproxy/addons.py index 7ac65a09d..c779aaf8c 100644 --- a/mitmproxy/addons.py +++ b/mitmproxy/addons.py @@ -21,6 +21,7 @@ class Addons(object): def add(self, *addons): self.chain.extend(addons) for i in addons: + self.invoke_with_context(i, "start") self.invoke_with_context(i, "configure", self.master.options) def remove(self, addon): diff --git a/mitmproxy/builtins/__init__.py b/mitmproxy/builtins/__init__.py index 8021c20ff..6b357902a 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 000000000..14f7dd5f2 --- /dev/null +++ b/mitmproxy/builtins/script.py @@ -0,0 +1,185 @@ +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 +from watchdog.observers import polling + + +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))): + tb = tb.tb_next + if not os.path.abspath(s[0]).startswith(scriptdir): + break + ctx.log.error("Script error: %s" % "".join(traceback.format_tb(tb))) + finally: + sys.argv = oldargs + sys.path.pop() + + +def load_script(path, args): + with open(path, "rb") as f: + try: + code = compile(f.read(), path, 'exec') + except SyntaxError as e: + ctx.log.error( + "Script error: %s line %s: %s" % ( + e.filename, e.lineno, e.msg + ) + ) + return + ns = {'__file__': os.path.abspath(path)} + with scriptenv(path, args): + 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.dead = False + + self.last_options = None + self.should_reload = threading.Event() + + for i in controller.Events: + if not hasattr(self, i): + 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 and not self.dead: + 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() + ctx.log.info("Reloading script: %s" % self.name) + self.ns = load_script(self.path, self.args) + self.configure(self.last_options) + else: + self.run("tick") + + def start(self): + self.ns = load_script(self.path, self.args) + self.run("start") + + def configure(self, options): + self.last_options = options + if not self.observer: + self.observer = polling.PollingObserver() + # Bind the handler to the real underlying master object + self.observer.schedule( + ReloadHandler(self.reload), + os.path.dirname(self.path) or "." + ) + self.observer.start() + self.run("configure", options) + + def done(self): + self.run("done") + self.dead = True + + +class ScriptLoader(): + """ + An addon that manages loading scripts from options. + """ + def configure(self, options): + for s in options.scripts: + if options.scripts.count(s) > 1: + raise exceptions.OptionsError("Duplicate script: %s" % s) + + for a in ctx.master.addons.chain[:]: + if isinstance(a, Script) and a.name not in options.scripts: + ctx.log.info("Un-loading script: %s" % a.name) + ctx.master.addons.remove(a) + + current = {} + for a in ctx.master.addons.chain[:]: + if isinstance(a, Script): + current[a.name] = a + ctx.master.addons.chain.remove(a) + + for s in options.scripts: + if s in current: + ctx.master.addons.chain.append(current[s]) + else: + ctx.log.info("Loading script: %s" % s) + sc = Script(s) + ctx.master.addons.add(sc) diff --git a/mitmproxy/console/grideditor.py b/mitmproxy/console/grideditor.py index 9fa51ccbe..f304de57a 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 bc373a2b0..64bd9f0af 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 d363ba747..d8824b056 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 fc41869c3..e7a700a6e 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 2f0c8bf29..503cdcd3c 100644 --- a/mitmproxy/controller.py +++ b/mitmproxy/controller.py @@ -33,6 +33,11 @@ Events = frozenset([ "error", "log", + "start", + "configure", + "done", + "tick", + "script_change", ]) @@ -44,7 +49,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 +112,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 274e01f35..999a709ae 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 282784b6d..3b41fe1ca 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 27ceee87b..aa09e1097 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) @@ -416,46 +370,15 @@ class FlowMaster(controller.Master): def handle_accept_intercept(self, f): 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 9a3985ab8..e75f282ae 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/concurrent.py b/mitmproxy/script/concurrent.py index 010a5fa09..0cc0514e7 100644 --- a/mitmproxy/script/concurrent.py +++ b/mitmproxy/script/concurrent.py @@ -13,7 +13,7 @@ class ScriptThread(basethread.BaseThread): def concurrent(fn): - if fn.__name__ not in controller.Events: + if fn.__name__ not in controller.Events - set(["start", "configure", "tick"]): raise NotImplementedError( "Concurrent decorator not supported for '%s' method." % fn.__name__ ) diff --git a/mitmproxy/script/reloader.py b/mitmproxy/script/reloader.py deleted file mode 100644 index 857d76cd5..000000000 --- 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 db4909cab..000000000 --- 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 23c16dc36..9eebf22cb 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 000000000..da60b64cd --- /dev/null +++ b/test/mitmproxy/builtins/test_script.py @@ -0,0 +1,187 @@ +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"] == [ + ("solo", "start", (), {}), + ("solo", "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[1] == "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] == "error" + + 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 + + def test_dupes(self): + s = state.State() + o = options.Options(scripts=["one", "one"]) + m = master.FlowMaster(o, None, s) + sc = script.ScriptLoader() + tutils.raises(exceptions.OptionsError, m.addons.add, sc) + + def test_order(self): + rec = tutils.test_data.path("data/addonscripts/recorder.py") + + s = state.State() + o = options.Options( + scripts = [ + "%s %s" % (rec, "a"), + "%s %s" % (rec, "b"), + "%s %s" % (rec, "c"), + ] + ) + m = mastertest.RecordingMaster(o, None, s) + sc = script.ScriptLoader() + m.addons.add(sc) + + debug = [(i[0], i[1]) for i in m.event_log if i[0] == "debug"] + assert debug == [ + ('debug', 'a start'), ('debug', 'a configure'), + ('debug', 'b start'), ('debug', 'b configure'), + ('debug', 'c start'), ('debug', 'c configure') + ] + m.event_log[:] = [] + + o.scripts = [ + "%s %s" % (rec, "c"), + "%s %s" % (rec, "a"), + "%s %s" % (rec, "b"), + ] + debug = [(i[0], i[1]) for i in m.event_log if i[0] == "debug"] + assert debug == [ + ('debug', 'c configure'), + ('debug', 'a configure'), + ('debug', 'b configure'), + ] + m.event_log[:] = [] + + o.scripts = [ + "%s %s" % (rec, "x"), + "%s %s" % (rec, "a"), + ] + debug = [(i[0], i[1]) for i in m.event_log if i[0] == "debug"] + assert debug == [ + ('debug', 'c done'), + ('debug', 'b done'), + ('debug', 'x start'), + ('debug', 'x configure'), + ('debug', 'a configure'), + ] diff --git a/test/mitmproxy/data/scripts/concurrent_decorator.py b/test/mitmproxy/data/addonscripts/concurrent_decorator.py similarity index 99% rename from test/mitmproxy/data/scripts/concurrent_decorator.py rename to test/mitmproxy/data/addonscripts/concurrent_decorator.py index 162c00f4b..a56c2af1c 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 similarity index 100% rename from test/mitmproxy/data/scripts/concurrent_decorator_err.py rename to 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 000000000..b466423cd --- /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 000000000..8ece9fcea --- /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 000000000..b6ac8d89d --- /dev/null +++ b/test/mitmproxy/data/addonscripts/recorder.py @@ -0,0 +1,25 @@ +from mitmproxy import controller +from mitmproxy import ctx +import sys + +call_log = [] + +if len(sys.argv) > 1: + name = sys.argv[1] +else: + name = "solo" + +# 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 = (name, evt, args, kwargs) + if evt != "log": + ctx.log.info(str(lg)) + call_log.append(lg) + ctx.log.debug("%s %s" % (name, evt)) + 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 000000000..bc6163421 --- /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 000000000..af4ccf7e1 --- /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 9754d3a99..240f6a730 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 57eeca191..080746e89 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.start() + 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 e33903b95..000000000 --- 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 48fe65c95..000000000 --- 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 9686be84d..201386e3e 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 bdadcd110..f86463361 100644 --- a/test/mitmproxy/test_examples.py +++ b/test/mitmproxy/test_examples.py @@ -1,151 +1,126 @@ -import glob import json -import mock -import os -import sys -from contextlib import contextmanager -from mitmproxy import script +import six +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 +from . import tutils, mastertest -example_dir = netlib.utils.Data(__name__).path("../../examples") +example_dir = netlib.utils.Data(__name__).push("../../examples") -@contextmanager -def example(command): - command = os.path.join(example_dir, command) - with script.Script(command) as s: - yield s +class ScriptError(Exception): + pass -@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() +class RaiseMaster(master.FlowMaster): + def add_event(self, e, level): + if level in ("warn", "error"): + raise ScriptError(e) -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" +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 -@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] +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"test!") assert any(b'esttay!' in val[0][1] for val in fmt) assert not pig(b"gobbledygook") - assert remove.called + def test_iframe_injector(self): + with tutils.raises(ScriptError): + tscript("iframe_injector.py") -def test_iframe_injector(): - with tutils.raises(script.ScriptException): - with example("iframe_injector.py"): - pass - - flow = tutils.tflow(resp=netutils.tresp(content=b"mitmproxy")) - with example("iframe_injector.py http://example.org/evil_iframe") 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"mitmproxy")) + self.invoke(m, "response", flow) content = flow.response.content assert b'iframe' in content and b'evil_iframe' in content + def test_modify_form(self): + m, sc = tscript("modify_form.py") -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" + form_header = Headers(content_type="application/x-www-form-urlencoded") + f = tutils.tflow(req=netutils.treq(headers=form_header)) + self.invoke(m, "request", f) - flow.request.headers["content-type"] = "" - ex.run("request", flow) - assert list(flow.request.urlencoded_form.items()) == [(b"foo", b"bar")] + 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(): - 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" + def test_modify_querystring(self): + m, sc = tscript("modify_querystring.py") + f = tutils.tflow(req=netutils.treq(path="/search?q=term")) - flow.request.path = "/" - ex.run("request", flow) - assert flow.request.query["mitmproxy"] == "rocks" + 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(): - with tutils.raises(script.ScriptException): - with example("modify_response_body.py"): - assert True + def test_modify_response_body(self): + with tutils.raises(ScriptError): + tscript("modify_response_body.py") - 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" + 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_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" + 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") -@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.tmpdir() as tdir: + times = dict( + timestamp_start=746203272, + timestamp_end=746203272, + ) - with tutils.raises(script.ScriptException): - with example("har_extractor.py"): - pass + path = os.path.join(tdir, "file") + m, sc = tscript("har_extractor.py", six.moves.shlex_quote(path)) + f = tutils.tflow( + req=netutils.treq(**times), + resp=netutils.tresp(**times) + ) + self.invoke(m, "response", f) + m.addons.remove(sc) - times = dict( - timestamp_start=746203272, - timestamp_end=746203272, - ) - - flow = tutils.tflow( - req=netutils.treq(**times), - resp=netutils.tresp(**times) - ) - - with example("har_extractor.py -") as ex: - ex.run("response", flow) - - with open(tutils.test_data.path("data/har_extractor.har")) as fp: - test_data = json.load(fp) - assert json.loads(ex.ns["context"].HARLog.json()) == test_data["test_response"] + with open(path, "rb") as f: + test_data = json.load(f) + assert len(test_data["log"]["pages"]) == 1 diff --git a/test/mitmproxy/test_flow.py b/test/mitmproxy/test_flow.py index 1a07f74d3..c58a97038 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 1e8220f10..000000000 --- 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 9dd8b79c2..a5196daec 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):