aboutsummaryrefslogtreecommitdiffstats
diff options
context:
space:
mode:
authorAldo Cortesi <aldo@corte.si>2016-07-15 16:48:01 +1200
committerGitHub <noreply@github.com>2016-07-15 16:48:01 +1200
commit64e16f5112669a31f86183be640eeee3a47c2961 (patch)
tree2e8ecc38f022c2b302d5ce67fd282dab45941fd9
parente051928f26afbc58c01d7854b62f8bfb0d4ed5ee (diff)
parentc7d0850d8f697915b183f4fafd5ede7df2245569 (diff)
downloadmitmproxy-64e16f5112669a31f86183be640eeee3a47c2961.tar.gz
mitmproxy-64e16f5112669a31f86183be640eeee3a47c2961.tar.bz2
mitmproxy-64e16f5112669a31f86183be640eeee3a47c2961.zip
Merge pull request #1356 from cortesi/script
Scripts to addon
-rw-r--r--examples/har_extractor.py8
-rw-r--r--examples/stub.py9
-rw-r--r--mitmproxy/addons.py1
-rw-r--r--mitmproxy/builtins/__init__.py2
-rw-r--r--mitmproxy/builtins/script.py185
-rw-r--r--mitmproxy/console/grideditor.py11
-rw-r--r--mitmproxy/console/master.py32
-rw-r--r--mitmproxy/console/options.py16
-rw-r--r--mitmproxy/console/statusbar.py9
-rw-r--r--mitmproxy/controller.py50
-rw-r--r--mitmproxy/dump.py8
-rw-r--r--mitmproxy/exceptions.py4
-rw-r--r--mitmproxy/flow/master.py101
-rw-r--r--mitmproxy/script/__init__.py6
-rw-r--r--mitmproxy/script/concurrent.py2
-rw-r--r--mitmproxy/script/reloader.py47
-rw-r--r--mitmproxy/script/script.py136
-rw-r--r--netlib/utils.py7
-rw-r--r--test/mitmproxy/builtins/test_script.py187
-rw-r--r--test/mitmproxy/data/addonscripts/concurrent_decorator.py (renamed from test/mitmproxy/data/scripts/concurrent_decorator.py)1
-rw-r--r--test/mitmproxy/data/addonscripts/concurrent_decorator_err.py (renamed from test/mitmproxy/data/scripts/concurrent_decorator_err.py)0
-rw-r--r--test/mitmproxy/data/addonscripts/duplicate_flow.py6
-rw-r--r--test/mitmproxy/data/addonscripts/error.py7
-rw-r--r--test/mitmproxy/data/addonscripts/recorder.py25
-rw-r--r--test/mitmproxy/data/addonscripts/stream_modify.py8
-rw-r--r--test/mitmproxy/data/addonscripts/tcp_stream_modify.py5
-rw-r--r--test/mitmproxy/mastertest.py10
-rw-r--r--test/mitmproxy/script/test_concurrent.py44
-rw-r--r--test/mitmproxy/script/test_reloader.py34
-rw-r--r--test/mitmproxy/script/test_script.py83
-rw-r--r--test/mitmproxy/test_dump.py4
-rw-r--r--test/mitmproxy/test_examples.py241
-rw-r--r--test/mitmproxy/test_flow.py64
-rw-r--r--test/mitmproxy/test_script.py13
-rw-r--r--test/mitmproxy/test_server.py18
35 files changed, 676 insertions, 708 deletions
diff --git a/examples/har_extractor.py b/examples/har_extractor.py
index 2a69b9af..90412ec0 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 10b34283..e5b4a39a 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 7ac65a09..c779aaf8 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 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..14f7dd5f
--- /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 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..503cdcd3 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 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/concurrent.py b/mitmproxy/script/concurrent.py
index 010a5fa0..0cc0514e 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 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..da60b64c
--- /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
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..b6ac8d89
--- /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 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..080746e8 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 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..f8646336 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
-
-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:
- test_data = json.load(fp)
- assert json.loads(ex.ns["context"].HARLog.json()) == test_data["test_response"]
+ 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", 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)
+
+ 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 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):