From e6eeab60946e61047ed858422badbda189a6f9e8 Mon Sep 17 00:00:00 2001 From: Aldo Cortesi Date: Tue, 25 Apr 2017 19:06:24 +1200 Subject: [PATCH] Revamp how addons work - Addons now nest, which means that addons can manage addons. This has a number of salutary effects - the scripts addon no longer has to poke into the global addons list, we no longer have to replace/remove/boot-outof parent addons when we load scripts, and this paves the way for making our top-level tools into addons themselves. - All addon calls are now wrapped in a safe execution environment where exceptions are caught, and output to stdout/stderr are intercepted and turned into logs. - We no longer support script arguments in sys.argv - creating an option properly is the only way to pass arguments. This means that all scripts are always directly controllable from interctive tooling, and that arguments are type-checked. For now, I've disabled testing of the har dump example - it needs to be moved to the new argument handling, and become a class addon. I'll address that in a separate patch. --- docs/install.rst | 6 +- docs/scripting/overview.rst | 18 +- examples/complex/dns_spoofing.py | 3 +- examples/simple/add_header_class.py | 3 +- examples/simple/filter_flows.py | 5 +- examples/simple/io_write_dumpfile.py | 5 +- examples/simple/modify_body_inject_iframe.py | 42 ++-- examples/simple/script_arguments.py | 17 -- mitmproxy/addonmanager.py | 162 ++++++++---- mitmproxy/addons/onboarding.py | 2 + mitmproxy/addons/script.py | 229 ++++------------- mitmproxy/addons/termlog.py | 9 +- mitmproxy/addons/wsgiapp.py | 4 + mitmproxy/master.py | 2 +- mitmproxy/test/taddons.py | 44 +++- mitmproxy/tools/main.py | 2 +- .../{test_har_dump.py => _test_har_dump.py} | 0 test/examples/test_examples.py | 31 ++- .../addons/test_core_option_validation.py | 2 +- test/mitmproxy/addons/test_script.py | 236 +++++++----------- test/mitmproxy/data/addonscripts/addon.py | 4 +- .../data/addonscripts/concurrent_decorator.py | 1 + .../concurrent_decorator_class.py | 3 +- .../mitmproxy/data/addonscripts/recorder/a.py | 3 + .../mitmproxy/data/addonscripts/recorder/b.py | 3 + .../mitmproxy/data/addonscripts/recorder/c.py | 3 + .../mitmproxy/data/addonscripts/recorder/e.py | 3 + .../addonscripts/{ => recorder}/recorder.py | 8 +- test/mitmproxy/proxy/test_server.py | 4 +- test/mitmproxy/script/test_concurrent.py | 5 +- test/mitmproxy/test_addonmanager.py | 120 +++++---- test/mitmproxy/test_taddons.py | 20 ++ test/mitmproxy/tools/console/test_master.py | 8 +- test/mitmproxy/tservers.py | 2 +- 34 files changed, 462 insertions(+), 547 deletions(-) delete mode 100644 examples/simple/script_arguments.py rename test/examples/{test_har_dump.py => _test_har_dump.py} (100%) create mode 100644 test/mitmproxy/data/addonscripts/recorder/a.py create mode 100644 test/mitmproxy/data/addonscripts/recorder/b.py create mode 100644 test/mitmproxy/data/addonscripts/recorder/c.py create mode 100644 test/mitmproxy/data/addonscripts/recorder/e.py rename test/mitmproxy/data/addonscripts/{ => recorder}/recorder.py (81%) diff --git a/docs/install.rst b/docs/install.rst index b37d9c915..7753dc448 100644 --- a/docs/install.rst +++ b/docs/install.rst @@ -123,12 +123,12 @@ You can check you Python version by running ``python3 --version``. sudo zypper install python3-pip python3-devel libffi-devel openssl-devel gcc-c++ sudo pip3 install mitmproxy - + .. _install-source-windows: -🐱💻 Installation from Source on Windows -^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +Installation from Source on Windows +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ .. note:: Mitmproxy's console interface is not supported on Windows, but you can use diff --git a/docs/scripting/overview.rst b/docs/scripting/overview.rst index 6ec0caaa7..c333a98b1 100644 --- a/docs/scripting/overview.rst +++ b/docs/scripting/overview.rst @@ -54,24 +54,8 @@ and is replaced by the class instance. Handling arguments ------------------ -Scripts can handle their own command-line arguments, just like any other Python -program. Let's build on the example above to do something slightly more -sophisticated - replace one value with another in all responses. Mitmproxy's -`HTTPRequest `_ and `HTTPResponse -`_ objects have a handy `replace -`_ method that takes care -of all the details for us. -.. literalinclude:: ../../examples/simple/script_arguments.py - :caption: :src:`examples/simple/script_arguments.py` - :language: python - -We can now call this script on the command-line like this: - ->>> mitmdump -dd -s "./script_arguments.py html faketml" - -Whenever a handler is called, mitpmroxy rewrites the script environment so that -it sees its own arguments as if it was invoked from the command-line. +FIXME Logging and the context diff --git a/examples/complex/dns_spoofing.py b/examples/complex/dns_spoofing.py index 01e036b2e..632783a77 100644 --- a/examples/complex/dns_spoofing.py +++ b/examples/complex/dns_spoofing.py @@ -54,5 +54,4 @@ class Rerouter: flow.request.port = port -def load(l): - l.boot_into(Rerouter()) +addons = [Rerouter()] diff --git a/examples/simple/add_header_class.py b/examples/simple/add_header_class.py index 69b64163f..5d5c7902a 100644 --- a/examples/simple/add_header_class.py +++ b/examples/simple/add_header_class.py @@ -3,5 +3,4 @@ class AddHeader: flow.response.headers["newheader"] = "foo" -def load(l): - return l.boot_into(AddHeader()) +addons = [AddHeader()] diff --git a/examples/simple/filter_flows.py b/examples/simple/filter_flows.py index d2b735be0..896fa54a3 100644 --- a/examples/simple/filter_flows.py +++ b/examples/simple/filter_flows.py @@ -17,7 +17,4 @@ class Filter: print(flow) -def load(l): - if len(sys.argv) != 2: - raise ValueError("Usage: -s 'filt.py FILTER'") - l.boot_into(Filter(sys.argv[1])) +addons = [Filter(sys.argv[1])] diff --git a/examples/simple/io_write_dumpfile.py b/examples/simple/io_write_dumpfile.py index 15e7693cf..a0956e335 100644 --- a/examples/simple/io_write_dumpfile.py +++ b/examples/simple/io_write_dumpfile.py @@ -23,7 +23,4 @@ class Writer: self.w.add(flow) -def load(l): - if len(sys.argv) != 2: - raise ValueError('Usage: -s "flowriter.py filename"') - l.boot_into(Writer(sys.argv[1])) +addons = [Writer(sys.argv[1])] diff --git a/examples/simple/modify_body_inject_iframe.py b/examples/simple/modify_body_inject_iframe.py index 442a51181..d54468d2a 100644 --- a/examples/simple/modify_body_inject_iframe.py +++ b/examples/simple/modify_body_inject_iframe.py @@ -1,29 +1,31 @@ -# Usage: mitmdump -s "iframe_injector.py url" # (this script works best with --anticache) -import sys from bs4 import BeautifulSoup class Injector: - def __init__(self, iframe_url): - self.iframe_url = iframe_url + def __init__(self): + self.iframe_url = None + + def load(self, loader): + loader.add_option( + "iframe", str, "", "IFrame to inject" + ) + + def configure(self, options, updated): + self.iframe_url = options.iframe def response(self, flow): - if flow.request.host in self.iframe_url: - return - html = BeautifulSoup(flow.response.content, "html.parser") - if html.body: - iframe = html.new_tag( - "iframe", - src=self.iframe_url, - frameborder=0, - height=0, - width=0) - html.body.insert(0, iframe) - flow.response.content = str(html).encode("utf8") + if self.iframe_url: + html = BeautifulSoup(flow.response.content, "html.parser") + if html.body: + iframe = html.new_tag( + "iframe", + src=self.iframe_url, + frameborder=0, + height=0, + width=0) + html.body.insert(0, iframe) + flow.response.content = str(html).encode("utf8") -def load(l): - if len(sys.argv) != 2: - raise ValueError('Usage: -s "iframe_injector.py url"') - return l.boot_into(Injector(sys.argv[1])) +addons = [Injector()] diff --git a/examples/simple/script_arguments.py b/examples/simple/script_arguments.py deleted file mode 100644 index 84292eb99..000000000 --- a/examples/simple/script_arguments.py +++ /dev/null @@ -1,17 +0,0 @@ -import argparse - - -class Replacer: - def __init__(self, src, dst): - self.src, self.dst = src, dst - - def response(self, flow): - flow.response.replace(self.src, self.dst) - - -def load(l): - parser = argparse.ArgumentParser() - parser.add_argument("src", type=str) - parser.add_argument("dst", type=str) - args = parser.parse_args() - l.boot_into(Replacer(args.src, args.dst)) diff --git a/mitmproxy/addonmanager.py b/mitmproxy/addonmanager.py index 241b3cde3..b6d7adb6d 100644 --- a/mitmproxy/addonmanager.py +++ b/mitmproxy/addonmanager.py @@ -1,4 +1,7 @@ import typing +import traceback +import contextlib +import sys from mitmproxy import exceptions from mitmproxy import eventsequence @@ -11,13 +14,66 @@ def _get_name(itm): return getattr(itm, "name", itm.__class__.__name__.lower()) +def cut_traceback(tb, func_name): + """ + Cut off a traceback at the function with the given name. + The func_name's frame is excluded. + + Args: + tb: traceback object, as returned by sys.exc_info()[2] + func_name: function name + + Returns: + Reduced traceback. + """ + tb_orig = tb + for _, _, fname, _ in traceback.extract_tb(tb): + tb = tb.tb_next + if fname == func_name: + break + return tb or tb_orig + + +class StreamLog: + """ + A class for redirecting output using contextlib. + """ + def __init__(self, log): + self.log = log + + def write(self, buf): + if buf.strip(): + self.log(buf) + + def flush(self): # pragma: no cover + # Click uses flush sometimes, so we dummy it up + pass + + +@contextlib.contextmanager +def safecall(): + stdout_replacement = StreamLog(ctx.log.warn) + try: + with contextlib.redirect_stdout(stdout_replacement): + yield + except exceptions.AddonHalt: + raise + except Exception as e: + etype, value, tb = sys.exc_info() + tb = cut_traceback(tb, "invoke_addon").tb_next + ctx.log.error( + "Addon error: %s" % "".join( + traceback.format_exception(etype, value, tb) + ) + ) + + class Loader: """ A loader object is passed to the load() event when addons start up. """ def __init__(self, master): self.master = master - self.boot_into_addon = None def add_option( self, @@ -35,25 +91,33 @@ class Loader: choices ) - def boot_into(self, addon): - self.boot_into_addon = addon - func = getattr(addon, "load", None) - if func: - func(self) + +def traverse(chain): + """ + Recursively traverse an addon chain. + """ + for a in chain: + yield a + if hasattr(a, "addons"): + yield from traverse(a.addons) class AddonManager: def __init__(self, master): + self.lookup = {} self.chain = [] self.master = master - master.options.changed.connect(self.configure_all) + master.options.changed.connect(self._configure_all) + + def _configure_all(self, options, updated): + self.trigger("configure", options, updated) def clear(self): """ Remove all addons. """ - self.done() - self.chain = [] + for i in self.chain: + self.remove(i) def get(self, name): """ @@ -61,36 +125,52 @@ class AddonManager: attribute on the instance, or the lower case class name if that does not exist. """ - for i in self.chain: - if name == _get_name(i): - return i + return self.lookup.get(name, None) - def configure_all(self, options, updated): - self.trigger("configure", options, updated) + def register(self, addon): + """ + Register an addon and all its sub-addons with the manager without + adding it to the chain. This should be used by addons that + dynamically manage addons. Must be called within a current context. + """ + for a in traverse([addon]): + name = _get_name(a) + if name in self.lookup: + raise exceptions.AddonError( + "An addon called '%s' already exists." % name + ) + l = Loader(self.master) + self.invoke_addon(addon, "load", l) + for a in traverse([addon]): + name = _get_name(a) + self.lookup[name] = a + return addon def add(self, *addons): """ - Add addons to the end of the chain, and run their startup events. + Add addons to the end of the chain, and run their load event. + If any addon has sub-addons, they are registered. """ with self.master.handlecontext(): for i in addons: - l = Loader(self.master) - self.invoke_addon(i, "load", l) - if l.boot_into_addon: - self.chain.append(l.boot_into_addon) - else: - self.chain.append(i) + self.chain.append(self.register(i)) def remove(self, addon): """ - Remove an addon from the chain, and run its done events. - """ - self.chain = [i for i in self.chain if i is not addon] - with self.master.handlecontext(): - self.invoke_addon(addon, "done") + Remove an addon and all its sub-addons. - def done(self): - self.trigger("done") + If the addon is not in the chain - that is, if it's managed by a + parent addon - it's the parent's responsibility to remove it from + its own addons attribute. + """ + for a in traverse([addon]): + n = _get_name(a) + if n not in self.lookup: + raise exceptions.AddonError("No such addon: %s" % n) + self.chain = [i for i in self.chain if i is not a] + del self.lookup[_get_name(a)] + with self.master.handlecontext(): + self.invoke_addon(a, "done") def __len__(self): return len(self.chain) @@ -126,22 +206,19 @@ class AddonManager: def invoke_addon(self, addon, name, *args, **kwargs): """ - Invoke an event on an addon. This method must run within an - established handler context. + Invoke an event on an addon and all its children. This method must + run within an established handler context. """ - if not ctx.master: - raise exceptions.AddonError( - "invoke_addon called without a handler context." - ) if name not in eventsequence.Events: name = "event_" + name - func = getattr(addon, name, None) - if func: - if not callable(func): - raise exceptions.AddonError( - "Addon handler %s not callable" % name - ) - func(*args, **kwargs) + for a in traverse([addon]): + func = getattr(a, name, None) + if func: + if not callable(func): + raise exceptions.AddonError( + "Addon handler %s not callable" % name + ) + func(*args, **kwargs) def trigger(self, name, *args, **kwargs): """ @@ -150,6 +227,7 @@ class AddonManager: with self.master.handlecontext(): for i in self.chain: try: - self.invoke_addon(i, name, *args, **kwargs) + with safecall(): + self.invoke_addon(i, name, *args, **kwargs) except exceptions.AddonHalt: return diff --git a/mitmproxy/addons/onboarding.py b/mitmproxy/addons/onboarding.py index cb57990fd..6552ec9e0 100644 --- a/mitmproxy/addons/onboarding.py +++ b/mitmproxy/addons/onboarding.py @@ -3,6 +3,8 @@ from mitmproxy.addons.onboardingapp import app class Onboarding(wsgiapp.WSGIApp): + name = "onboarding" + def __init__(self): super().__init__(app.Adapter(app.application), None, None) self.enabled = False diff --git a/mitmproxy/addons/script.py b/mitmproxy/addons/script.py index eef21293c..bda823b42 100644 --- a/mitmproxy/addons/script.py +++ b/mitmproxy/addons/script.py @@ -1,123 +1,31 @@ -import contextlib import os -import shlex -import sys +import importlib import threading -import traceback -import types +import sys from mitmproxy import addonmanager from mitmproxy import exceptions from mitmproxy import ctx -from mitmproxy import eventsequence - 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 ValueError("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 ValueError( - ("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 ValueError("Not a file: %s" % args[0]) - return args[0], args[1:] - - -def cut_traceback(tb, func_name): - """ - Cut off a traceback at the function with the given name. - The func_name's frame is excluded. - - Args: - tb: traceback object, as returned by sys.exc_info()[2] - func_name: function name - - Returns: - Reduced traceback. - """ - tb_orig = tb - - for _, _, fname, _ in traceback.extract_tb(tb): - tb = tb.tb_next - if fname == func_name: - break - - if tb is None: - # We could not find the method, take the full stack trace. - # This may happen on some Python interpreters/flavors (e.g. PyInstaller). - return tb_orig - else: - return tb - - -class StreamLog: - """ - A class for redirecting output using contextlib. - """ - def __init__(self, log): - self.log = log - - def write(self, buf): - if buf.strip(): - self.log(buf) - - -@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) - stdout_replacement = StreamLog(ctx.log.warn) +def load_script(actx, path): + if not os.path.exists(path): + ctx.log.info("No such file: %s" % path) + return + loader = importlib.machinery.SourceFileLoader(os.path.basename(path), path) try: - with contextlib.redirect_stdout(stdout_replacement): - yield - except SystemExit as v: - ctx.log.error("Script exited with code %s" % v.code) - except Exception: - etype, value, tb = sys.exc_info() - tb = cut_traceback(tb, "scriptenv").tb_next - ctx.log.error( - "Script error: %s" % "".join( - traceback.format_exception(etype, value, tb) - ) - ) + oldpath = sys.path + sys.path.insert(0, os.path.dirname(path)) + with addonmanager.safecall(): + m = loader.load_module() + if not getattr(m, "name", None): + m.name = path + return m 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) - return types.SimpleNamespace(**ns) + sys.path[:] = oldpath class ReloadHandler(watchdog.events.FileSystemEventHandler): @@ -149,59 +57,39 @@ 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) + def __init__(self, path): + self.name = "scriptmanager:" + path + self.path = path self.ns = None self.observer = None - self.dead = False self.last_options = None self.should_reload = threading.Event() - for i in eventsequence.Events: - if not hasattr(self, i): - def mkprox(): - evt = i + def load(self, l): + self.ns = load_script(ctx, self.path) - 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 = getattr(self.ns, name, None) - if func: - with scriptenv(self.path, self.args): - return func(*args, **kwargs) + @property + def addons(self): + if self.ns is not None: + return [self.ns] + return [] def reload(self): self.should_reload.set() - def load_script(self): - self.ns = load_script(self.path, self.args) - l = addonmanager.Loader(ctx.master) - self.run("load", l) - if l.boot_into_addon: - self.ns = l.boot_into_addon - 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, self.last_options.keys()) - else: - self.run("tick") - - def load(self, l): - self.last_options = ctx.master.options - self.load_script() + if self.ns: + ctx.master.addons.remove(self.ns) + self.ns = load_script(ctx, self.path) + if self.ns: + # We're already running, so we have to explicitly register and + # configure the addon + ctx.master.addons.register(self.ns) + self.configure(self.last_options, self.last_options.keys()) def configure(self, options, updated): self.last_options = options @@ -213,11 +101,6 @@ class Script: os.path.dirname(self.path) or "." ) self.observer.start() - self.run("configure", options, updated) - - def done(self): - self.run("done") - self.dead = True class ScriptLoader: @@ -226,21 +109,14 @@ class ScriptLoader: """ def __init__(self): self.is_running = False + self.addons = [] def running(self): self.is_running = True def run_once(self, command, flows): - try: - sc = Script(command) - except ValueError as e: - raise ValueError(str(e)) - sc.load_script() - for f in flows: - for evt, o in eventsequence.iterate(f): - sc.run(evt, o) - sc.done() - return sc + # Returning once we have proper commands + raise NotImplementedError def configure(self, options, updated): if "scripts" in updated: @@ -248,25 +124,21 @@ class ScriptLoader: 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: + for a in self.addons[:]: + if a.path not in options.scripts: ctx.log.info("Un-loading script: %s" % a.name) ctx.master.addons.remove(a) + self.addons.remove(a) # The machinations below are to ensure that: # - Scripts remain in the same order - # - Scripts are listed directly after the script addon. This is - # needed to ensure that interactions with, for instance, flow - # serialization remains correct. # - Scripts are not initialized un-necessarily. If only a - # script's order in the script list has changed, it should simply - # be moved. + # script's order in the script list has changed, it is just + # moved. current = {} - for a in ctx.master.addons.chain[:]: - if isinstance(a, Script): - current[a.name] = a - ctx.master.addons.chain.remove(a) + for a in self.addons: + current[a.path] = a ordered = [] newscripts = [] @@ -275,24 +147,15 @@ class ScriptLoader: ordered.append(current[s]) else: ctx.log.info("Loading script: %s" % s) - try: - sc = Script(s) - except ValueError as e: - raise exceptions.OptionsError(str(e)) + sc = Script(s) ordered.append(sc) newscripts.append(sc) - ochain = ctx.master.addons.chain - pos = ochain.index(self) - ctx.master.addons.chain = ochain[:pos + 1] + ordered + ochain[pos + 1:] + self.addons = ordered for s in newscripts: - l = addonmanager.Loader(ctx.master) - ctx.master.addons.invoke_addon(s, "load", l) + ctx.master.addons.register(s) if self.is_running: # If we're already running, we configure and tell the addon # we're up and running. - ctx.master.addons.invoke_addon( - s, "configure", options, options.keys() - ) ctx.master.addons.invoke_addon(s, "running") diff --git a/mitmproxy/addons/termlog.py b/mitmproxy/addons/termlog.py index 5fdb6245f..0d42b18e7 100644 --- a/mitmproxy/addons/termlog.py +++ b/mitmproxy/addons/termlog.py @@ -3,6 +3,11 @@ import click from mitmproxy import log +# These get over-ridden by the save execution context. Keep them around so we +# can log directly. +realstdout = sys.stdout +realstderr = sys.stderr + class TermLog: def __init__(self, outfile=None): @@ -14,9 +19,9 @@ class TermLog: def log(self, e): if log.log_tier(e.level) == log.log_tier("error"): - outfile = self.outfile or sys.stderr + outfile = self.outfile or realstderr else: - outfile = self.outfile or sys.stdout + outfile = self.outfile or realstdout if self.options.verbosity >= log.log_tier(e.level): click.secho( diff --git a/mitmproxy/addons/wsgiapp.py b/mitmproxy/addons/wsgiapp.py index c37fcb7ba..155444fc2 100644 --- a/mitmproxy/addons/wsgiapp.py +++ b/mitmproxy/addons/wsgiapp.py @@ -13,6 +13,10 @@ class WSGIApp: def __init__(self, app, host, port): self.app, self.host, self.port = app, host, port + @property + def name(self): + return "wsgiapp:%s:%s" % (self.host, self.port) + def serve(self, app, flow): """ Serves app on flow, and prevents further handling of the flow. diff --git a/mitmproxy/master.py b/mitmproxy/master.py index 946b25a45..46fdb585b 100644 --- a/mitmproxy/master.py +++ b/mitmproxy/master.py @@ -103,7 +103,7 @@ class Master: def shutdown(self): self.server.shutdown() self.should_exit.set() - self.addons.done() + self.addons.trigger("done") def create_request(self, method, url): """ diff --git a/mitmproxy/test/taddons.py b/mitmproxy/test/taddons.py index 8bc174c7c..3dbccba2d 100644 --- a/mitmproxy/test/taddons.py +++ b/mitmproxy/test/taddons.py @@ -1,3 +1,4 @@ +import sys import contextlib import mitmproxy.master @@ -5,6 +6,7 @@ import mitmproxy.options from mitmproxy import proxy from mitmproxy import addonmanager from mitmproxy import eventsequence +from mitmproxy.addons import script class TestAddons(addonmanager.AddonManager): @@ -26,6 +28,10 @@ class RecordingMaster(mitmproxy.master.Master): self.events = [] self.logs = [] + def dump_log(self, outf=sys.stdout): + for i in self.logs: + print("%s: %s" % (i.level, i.msg), file=outf) + def has_log(self, txt, level=None): for i in self.logs: if level and i.level != level: @@ -51,14 +57,21 @@ class context: provides a number of helper methods for common testing scenarios. """ def __init__(self, master = None, options = None): - self.options = options or mitmproxy.options.Options() + options = options or mitmproxy.options.Options() self.master = master or RecordingMaster( options, proxy.DummyServer(options) ) + self.options = self.master.options self.wrapped = None + def ctx(self): + """ + Returns a new handler context. + """ + return self.master.handlecontext() + def __enter__(self): - self.wrapped = self.master.handlecontext() + self.wrapped = self.ctx() self.wrapped.__enter__() return self @@ -75,11 +88,13 @@ class context: """ f.reply._state = "start" for evt, arg in eventsequence.iterate(f): - h = getattr(addon, evt, None) - if h: - h(arg) - if f.reply.state == "taken": - return + self.master.addons.invoke_addon( + addon, + evt, + arg + ) + if f.reply.state == "taken": + return def configure(self, addon, **kwargs): """ @@ -89,4 +104,17 @@ class context: """ with self.options.rollback(kwargs.keys(), reraise=True): self.options.update(**kwargs) - addon.configure(self.options, kwargs.keys()) + self.master.addons.invoke_addon( + addon, + "configure", + self.options, + kwargs.keys() + ) + + def script(self, path): + sc = script.Script(path) + loader = addonmanager.Loader(self.master) + sc.load(loader) + for a in addonmanager.traverse(sc.addons): + getattr(a, "load", lambda x: None)(loader) + return sc diff --git a/mitmproxy/tools/main.py b/mitmproxy/tools/main.py index 6db232fc9..de9032c89 100644 --- a/mitmproxy/tools/main.py +++ b/mitmproxy/tools/main.py @@ -76,7 +76,7 @@ def run(MasterKlass, args, extra=None): # pragma: no cover unknown = optmanager.load_paths(opts, args.conf) server = process_options(parser, opts, args) master = MasterKlass(opts, server) - master.addons.configure_all(opts, opts.keys()) + master.addons.trigger("configure", opts, opts.keys()) remaining = opts.update_known(**unknown) if remaining and opts.verbosity > 1: print("Ignored options: %s" % remaining) diff --git a/test/examples/test_har_dump.py b/test/examples/_test_har_dump.py similarity index 100% rename from test/examples/test_har_dump.py rename to test/examples/_test_har_dump.py diff --git a/test/examples/test_examples.py b/test/examples/test_examples.py index 46fdcd362..4b691df25 100644 --- a/test/examples/test_examples.py +++ b/test/examples/test_examples.py @@ -1,5 +1,3 @@ -import pytest - from mitmproxy import options from mitmproxy import contentviews from mitmproxy import proxy @@ -8,6 +6,7 @@ from mitmproxy.addons import script from mitmproxy.test import tflow from mitmproxy.test import tutils +from mitmproxy.test import taddons from mitmproxy.net.http import Headers from ..mitmproxy import tservers @@ -27,7 +26,7 @@ class RaiseMaster(master.Master): def tscript(cmd, args=""): o = options.Options() - cmd = example_dir.path(cmd) + " " + args + cmd = example_dir.path(cmd) m = RaiseMaster(o, proxy.DummyServer()) sc = script.Script(cmd) m.addons.add(sc) @@ -48,14 +47,18 @@ class TestScripts(tservers.MasterTest): assert any(b'tEST!' in val[0][1] for val in fmt) def test_iframe_injector(self): - with pytest.raises(ScriptError): - tscript("simple/modify_body_inject_iframe.py") - - m, sc = tscript("simple/modify_body_inject_iframe.py", "http://example.org/evil_iframe") - f = tflow.tflow(resp=tutils.tresp(content=b"mitmproxy")) - m.addons.handle_lifecycle("response", f) - content = f.response.content - assert b'iframe' in content and b'evil_iframe' in content + with taddons.context() as tctx: + sc = tctx.script(example_dir.path("simple/modify_body_inject_iframe.py")) + tctx.configure( + sc, + iframe = "http://example.org/evil_iframe" + ) + f = tflow.tflow( + resp=tutils.tresp(content=b"mitmproxy") + ) + tctx.master.addons.invoke_addon(sc, "response", f) + content = f.response.content + assert b'iframe' in content and b'evil_iframe' in content def test_modify_form(self): m, sc = tscript("simple/modify_form.py") @@ -81,12 +84,6 @@ class TestScripts(tservers.MasterTest): m.addons.handle_lifecycle("request", f) assert f.request.query["mitmproxy"] == "rocks" - def test_arguments(self): - m, sc = tscript("simple/script_arguments.py", "mitmproxy rocks") - f = tflow.tflow(resp=tutils.tresp(content=b"I <3 mitmproxy")) - m.addons.handle_lifecycle("response", f) - assert f.response.content == b"I <3 rocks" - def test_redirect_requests(self): m, sc = tscript("simple/redirect_requests.py") f = tflow.tflow(req=tutils.treq(host="example.org")) diff --git a/test/mitmproxy/addons/test_core_option_validation.py b/test/mitmproxy/addons/test_core_option_validation.py index 0bb2bb0d7..6d6d5ba41 100644 --- a/test/mitmproxy/addons/test_core_option_validation.py +++ b/test/mitmproxy/addons/test_core_option_validation.py @@ -11,7 +11,7 @@ def test_simple(): with pytest.raises(exceptions.OptionsError): tctx.configure(sa, body_size_limit = "invalid") tctx.configure(sa, body_size_limit = "1m") - assert tctx.options._processed["body_size_limit"] + assert tctx.master.options._processed["body_size_limit"] with pytest.raises(exceptions.OptionsError, match="mutually exclusive"): tctx.configure( diff --git a/test/mitmproxy/addons/test_script.py b/test/mitmproxy/addons/test_script.py index c68981de8..4a86fad2c 100644 --- a/test/mitmproxy/addons/test_script.py +++ b/test/mitmproxy/addons/test_script.py @@ -1,7 +1,6 @@ import traceback import sys import time -import re import watchdog.events import pytest @@ -14,23 +13,8 @@ from mitmproxy import exceptions from mitmproxy import options from mitmproxy import proxy from mitmproxy import master -from mitmproxy import utils from mitmproxy.addons import script -from ...conftest import skip_not_windows - - -def test_scriptenv(): - with taddons.context() as tctx: - with script.scriptenv("path", []): - raise SystemExit - assert tctx.master.has_log("exited", "error") - - tctx.master.clear() - with script.scriptenv("path", []): - raise ValueError("fooo") - assert tctx.master.has_log("fooo", "error") - class Called: def __init__(self): @@ -60,113 +44,86 @@ def test_reloadhandler(): assert rh.callback.called -class TestParseCommand: - def test_empty_command(self): - with pytest.raises(ValueError): - script.parse_command("") - - with pytest.raises(ValueError): - script.parse_command(" ") - - def test_no_script_file(self, tmpdir): - with pytest.raises(Exception, match="not found"): - script.parse_command("notfound") - - with pytest.raises(Exception, match="Not a file"): - script.parse_command(str(tmpdir)) - - def test_parse_args(self): - with utils.chdir(tutils.test_data.dirname): - assert script.parse_command( - "mitmproxy/data/addonscripts/recorder.py" - ) == ("mitmproxy/data/addonscripts/recorder.py", []) - assert script.parse_command( - "mitmproxy/data/addonscripts/recorder.py foo bar" - ) == ("mitmproxy/data/addonscripts/recorder.py", ["foo", "bar"]) - assert script.parse_command( - "mitmproxy/data/addonscripts/recorder.py 'foo bar'" - ) == ("mitmproxy/data/addonscripts/recorder.py", ["foo bar"]) - - @skip_not_windows - def test_parse_windows(self): - with utils.chdir(tutils.test_data.dirname): - assert script.parse_command( - "mitmproxy/data\\addonscripts\\recorder.py" - ) == ("mitmproxy/data\\addonscripts\\recorder.py", []) - assert script.parse_command( - "mitmproxy/data\\addonscripts\\recorder.py 'foo \\ bar'" - ) == ("mitmproxy/data\\addonscripts\\recorder.py", ['foo \\ bar']) - - def test_load_script(): - with taddons.context(): + with taddons.context() as tctx: ns = script.load_script( + tctx.ctx(), tutils.test_data.path( - "mitmproxy/data/addonscripts/recorder.py" - ), [] + "mitmproxy/data/addonscripts/recorder/recorder.py" + ) ) - assert ns.load + assert ns.addons + + ns = script.load_script( + tctx.ctx(), + "nonexistent" + ) + assert not ns def test_script_print_stdout(): with taddons.context() as tctx: with mock.patch('mitmproxy.ctx.log.warn') as mock_warn: - with script.scriptenv("path", []): + with addonmanager.safecall(): ns = script.load_script( + tctx.ctx(), tutils.test_data.path( "mitmproxy/data/addonscripts/print.py" - ), [] + ) ) ns.load(addonmanager.Loader(tctx.master)) mock_warn.assert_called_once_with("stdoutprint") class TestScript: + def test_notfound(self): + with taddons.context() as tctx: + sc = script.Script("nonexistent") + tctx.master.addons.add(sc) + def test_simple(self): - with taddons.context(): + with taddons.context() as tctx: sc = script.Script( tutils.test_data.path( - "mitmproxy/data/addonscripts/recorder.py" + "mitmproxy/data/addonscripts/recorder/recorder.py" ) ) - sc.load_script() - assert sc.ns.call_log[0][0:2] == ("solo", "load") + tctx.master.addons.add(sc) - sc.ns.call_log = [] + rec = tctx.master.addons.get("recorder") + + assert rec.call_log[0][0:2] == ("recorder", "load") + + rec.call_log = [] f = tflow.tflow(resp=True) - sc.request(f) + tctx.master.addons.trigger("request", f) - recf = sc.ns.call_log[0] - assert recf[1] == "request" + assert rec.call_log[0][1] == "request" def test_reload(self, tmpdir): with taddons.context() as tctx: f = tmpdir.join("foo.py") f.ensure(file=True) + f.write("\n") sc = script.Script(str(f)) tctx.configure(sc) - for _ in range(100): - f.write(".") + for _ in range(5): + sc.reload() sc.tick() time.sleep(0.1) - if tctx.master.logs: - return - raise AssertionError("Change event not detected.") def test_exception(self): with taddons.context() as tctx: sc = script.Script( tutils.test_data.path("mitmproxy/data/addonscripts/error.py") ) - l = addonmanager.Loader(tctx.master) - sc.load(l) + tctx.master.addons.add(sc) f = tflow.tflow(resp=True) - sc.request(f) + tctx.master.addons.trigger("request", f) + assert tctx.master.logs[0].level == "error" - assert len(tctx.master.logs[0].msg.splitlines()) == 6 - assert re.search(r'addonscripts[\\/]error.py", line \d+, in request', tctx.master.logs[0].msg) - assert re.search(r'addonscripts[\\/]error.py", line \d+, in mkerr', tctx.master.logs[0].msg) - assert tctx.master.logs[0].msg.endswith("ValueError: Error!\n") + tctx.master.has_log("ValueError: Error!") + tctx.master.has_log("error.py") def test_addon(self): with taddons.context() as tctx: @@ -175,11 +132,9 @@ class TestScript: "mitmproxy/data/addonscripts/addon.py" ) ) - l = addonmanager.Loader(tctx.master) - sc.load(l) - tctx.configure(sc) + tctx.master.addons.add(sc) assert sc.ns.event_log == [ - 'scriptload', 'addonload', 'addonconfigure' + 'scriptload', 'addonload' ] @@ -194,49 +149,33 @@ class TestCutTraceback: self.raise_(4) except RuntimeError: tb = sys.exc_info()[2] - tb_cut = script.cut_traceback(tb, "test_simple") + tb_cut = addonmanager.cut_traceback(tb, "test_simple") assert len(traceback.extract_tb(tb_cut)) == 5 - tb_cut2 = script.cut_traceback(tb, "nonexistent") + tb_cut2 = addonmanager.cut_traceback(tb, "nonexistent") assert len(traceback.extract_tb(tb_cut2)) == len(traceback.extract_tb(tb)) class TestScriptLoader: - def test_run_once(self): - o = options.Options(scripts=[]) - m = master.Master(o, proxy.DummyServer()) - sl = script.ScriptLoader() - m.addons.add(sl) - - f = tflow.tflow(resp=True) - with m.handlecontext(): - sc = sl.run_once( - tutils.test_data.path( - "mitmproxy/data/addonscripts/recorder.py" - ), [f] - ) - evts = [i[1] for i in sc.ns.call_log] - assert evts == ['load', 'requestheaders', 'request', 'responseheaders', 'response', 'done'] - - f = tflow.tflow(resp=True) - with m.handlecontext(): - with pytest.raises(Exception, match="file not found"): - sl.run_once("nonexistent", [f]) - def test_simple(self): o = options.Options(scripts=[]) m = master.Master(o, proxy.DummyServer()) sc = script.ScriptLoader() + sc.running() m.addons.add(sc) assert len(m.addons) == 1 o.update( scripts = [ - tutils.test_data.path("mitmproxy/data/addonscripts/recorder.py") + tutils.test_data.path( + "mitmproxy/data/addonscripts/recorder/recorder.py" + ) ] ) - assert len(m.addons) == 2 + assert len(m.addons) == 1 + assert len(sc.addons) == 1 o.update(scripts = []) assert len(m.addons) == 1 + assert len(sc.addons) == 0 def test_dupes(self): sc = script.ScriptLoader() @@ -252,65 +191,70 @@ class TestScriptLoader: sc = script.ScriptLoader() with taddons.context() as tctx: tctx.master.addons.add(sc) - with pytest.raises(exceptions.OptionsError): - tctx.configure( - sc, - scripts = ["nonexistent"] - ) + tctx.configure(sc, scripts = ["nonexistent"]) + tctx.master.has_log("nonexistent: file not found") def test_order(self): - rec = tutils.test_data.path("mitmproxy/data/addonscripts/recorder.py") + rec = tutils.test_data.path("mitmproxy/data/addonscripts/recorder") sc = script.ScriptLoader() + sc.is_running = True with taddons.context() as tctx: - tctx.master.addons.add(sc) - sc.running() tctx.configure( sc, scripts = [ - "%s %s" % (rec, "a"), - "%s %s" % (rec, "b"), - "%s %s" % (rec, "c"), + "%s/a.py" % rec, + "%s/b.py" % rec, + "%s/c.py" % rec, + ] + ) + + debug = [i.msg for i in tctx.master.logs if i.level == "debug"] + assert debug == [ + 'a load', + 'a running', + + 'b load', + 'b running', + + 'c load', + 'c running', + + 'a configure', + 'b configure', + 'c configure', + ] + + tctx.master.logs = [] + tctx.configure( + sc, + scripts = [ + "%s/c.py" % rec, + "%s/a.py" % rec, + "%s/b.py" % rec, ] ) debug = [i.msg for i in tctx.master.logs if i.level == "debug"] assert debug == [ - 'a load', - 'a configure', - 'a running', - - 'b load', - 'b configure', - 'b running', - - 'c load', 'c configure', - 'c running', + 'a configure', + 'b configure', ] - tctx.master.logs = [] - tctx.configure( - sc, - scripts = [ - "%s %s" % (rec, "c"), - "%s %s" % (rec, "a"), - "%s %s" % (rec, "b"), - ] - ) - debug = [i.msg for i in tctx.master.logs if i.level == "debug"] - assert debug == [] tctx.master.logs = [] tctx.configure( sc, scripts = [ - "%s %s" % (rec, "x"), - "%s %s" % (rec, "a"), + "%s/e.py" % rec, + "%s/a.py" % rec, ] ) + debug = [i.msg for i in tctx.master.logs if i.level == "debug"] assert debug == [ 'c done', 'b done', - 'x load', - 'x configure', - 'x running', + 'e load', + 'e running', + 'e configure', + 'a configure', ] diff --git a/test/mitmproxy/data/addonscripts/addon.py b/test/mitmproxy/data/addonscripts/addon.py index beef2ce7a..42e28a933 100644 --- a/test/mitmproxy/data/addonscripts/addon.py +++ b/test/mitmproxy/data/addonscripts/addon.py @@ -19,4 +19,6 @@ def configure(options, updated): def load(l): event_log.append("scriptload") - l.boot_into(Addon()) + + +addons = [Addon()] diff --git a/test/mitmproxy/data/addonscripts/concurrent_decorator.py b/test/mitmproxy/data/addonscripts/concurrent_decorator.py index 162c00f4b..d1ab6c6c1 100644 --- a/test/mitmproxy/data/addonscripts/concurrent_decorator.py +++ b/test/mitmproxy/data/addonscripts/concurrent_decorator.py @@ -1,4 +1,5 @@ import time +import sys from mitmproxy.script import concurrent diff --git a/test/mitmproxy/data/addonscripts/concurrent_decorator_class.py b/test/mitmproxy/data/addonscripts/concurrent_decorator_class.py index 8e6988d46..2a7d300ce 100644 --- a/test/mitmproxy/data/addonscripts/concurrent_decorator_class.py +++ b/test/mitmproxy/data/addonscripts/concurrent_decorator_class.py @@ -9,5 +9,4 @@ class ConcurrentClass: time.sleep(0.1) -def load(l): - l.boot_into(ConcurrentClass()) +addons = [ConcurrentClass()] diff --git a/test/mitmproxy/data/addonscripts/recorder/a.py b/test/mitmproxy/data/addonscripts/recorder/a.py new file mode 100644 index 000000000..df81d86b0 --- /dev/null +++ b/test/mitmproxy/data/addonscripts/recorder/a.py @@ -0,0 +1,3 @@ +import recorder + +addons = [recorder.Recorder("a")] diff --git a/test/mitmproxy/data/addonscripts/recorder/b.py b/test/mitmproxy/data/addonscripts/recorder/b.py new file mode 100644 index 000000000..ccbae7059 --- /dev/null +++ b/test/mitmproxy/data/addonscripts/recorder/b.py @@ -0,0 +1,3 @@ +import recorder + +addons = [recorder.Recorder("b")] diff --git a/test/mitmproxy/data/addonscripts/recorder/c.py b/test/mitmproxy/data/addonscripts/recorder/c.py new file mode 100644 index 000000000..b8b0915e4 --- /dev/null +++ b/test/mitmproxy/data/addonscripts/recorder/c.py @@ -0,0 +1,3 @@ +import recorder + +addons = [recorder.Recorder("c")] diff --git a/test/mitmproxy/data/addonscripts/recorder/e.py b/test/mitmproxy/data/addonscripts/recorder/e.py new file mode 100644 index 000000000..eb5eff5e8 --- /dev/null +++ b/test/mitmproxy/data/addonscripts/recorder/e.py @@ -0,0 +1,3 @@ +import recorder + +addons = [recorder.Recorder("e")] diff --git a/test/mitmproxy/data/addonscripts/recorder.py b/test/mitmproxy/data/addonscripts/recorder/recorder.py similarity index 81% rename from test/mitmproxy/data/addonscripts/recorder.py rename to test/mitmproxy/data/addonscripts/recorder/recorder.py index fe497b05a..a962d3df6 100644 --- a/test/mitmproxy/data/addonscripts/recorder.py +++ b/test/mitmproxy/data/addonscripts/recorder/recorder.py @@ -1,13 +1,12 @@ from mitmproxy import controller from mitmproxy import eventsequence from mitmproxy import ctx -import sys -class CallLogger: +class Recorder: call_log = [] - def __init__(self, name = "solo"): + def __init__(self, name = "recorder"): self.name = name def __getattr__(self, attr): @@ -22,5 +21,4 @@ class CallLogger: raise AttributeError -def load(l): - l.boot_into(CallLogger(*sys.argv[1:])) +addons = [Recorder()] diff --git a/test/mitmproxy/proxy/test_server.py b/test/mitmproxy/proxy/test_server.py index 447b15a5d..b54a764f7 100644 --- a/test/mitmproxy/proxy/test_server.py +++ b/test/mitmproxy/proxy/test_server.py @@ -296,8 +296,8 @@ class TestHTTP(tservers.HTTPProxyTest, CommonMixin): class TestHTTPAuth(tservers.HTTPProxyTest): def test_auth(self): self.master.addons.add(proxyauth.ProxyAuth()) - self.master.addons.configure_all( - self.master.options, self.master.options.keys() + self.master.addons.trigger( + "configure", self.master.options, self.master.options.keys() ) self.master.options.proxyauth = "test:test" assert self.pathod("202").status_code == 407 diff --git a/test/mitmproxy/script/test_concurrent.py b/test/mitmproxy/script/test_concurrent.py index 86efdfc22..d24f96a24 100644 --- a/test/mitmproxy/script/test_concurrent.py +++ b/test/mitmproxy/script/test_concurrent.py @@ -20,14 +20,11 @@ class Thing: class TestConcurrent(tservers.MasterTest): def test_concurrent(self): with taddons.context() as tctx: - sc = script.Script( + sc = tctx.script( tutils.test_data.path( "mitmproxy/data/addonscripts/concurrent_decorator.py" ) ) - l = addonmanager.Loader(tctx.master) - sc.load(l) - f1, f2 = tflow.tflow(), tflow.tflow() tctx.cycle(sc, f1) tctx.cycle(sc, f2) diff --git a/test/mitmproxy/test_addonmanager.py b/test/mitmproxy/test_addonmanager.py index 5bb88eb6f..8391e7218 100644 --- a/test/mitmproxy/test_addonmanager.py +++ b/test/mitmproxy/test_addonmanager.py @@ -5,14 +5,17 @@ from mitmproxy import exceptions from mitmproxy import options from mitmproxy import master from mitmproxy import proxy +from mitmproxy.test import taddons from mitmproxy.test import tflow class TAddon: - def __init__(self, name): + def __init__(self, name, addons=None): self.name = name self.tick = True self.custom_called = False + if addons: + self.addons = addons def __repr__(self): return "Addon(%s)" % self.name @@ -34,19 +37,6 @@ class AOption: l.add_option("custom_option", bool, False, "help") -class AChain: - def __init__(self, name, next): - self.name = name - self.next = next - - def load(self, l): - if self.next: - l.boot_into(self.next) - - def __repr__(self): - return "<%s>" % self.name - - def test_halt(): o = options.Options() m = master.Master(o, proxy.DummyServer(o)) @@ -70,40 +60,42 @@ def test_lifecycle(): a = addonmanager.AddonManager(m) a.add(TAddon("one")) + with pytest.raises(exceptions.AddonError): + a.add(TAddon("one")) + with pytest.raises(exceptions.AddonError): + a.remove(TAddon("nonexistent")) + f = tflow.tflow() a.handle_lifecycle("request", f) - a.configure_all(o, o.keys()) + a._configure_all(o, o.keys()) def test_simple(): - o = options.Options() - m = master.Master(o, proxy.DummyServer(o)) - a = addonmanager.AddonManager(m) - with pytest.raises(exceptions.AddonError): - a.invoke_addon(TAddon("one"), "done") + with taddons.context() as tctx: + a = tctx.master.addons - assert len(a) == 0 - a.add(TAddon("one")) - assert a.get("one") - assert not a.get("two") - assert len(a) == 1 - a.clear() - assert len(a) == 0 - assert not a.chain + assert len(a) == 0 + a.add(TAddon("one")) + assert a.get("one") + assert not a.get("two") + assert len(a) == 1 + a.clear() + assert len(a) == 0 + assert not a.chain - a.add(TAddon("one")) - a.trigger("done") - with pytest.raises(exceptions.AddonError): + a.add(TAddon("one")) + a.trigger("done") a.trigger("tick") + tctx.master.has_log("not callable") - a.remove(a.get("one")) - assert not a.get("one") + a.remove(a.get("one")) + assert not a.get("one") - ta = TAddon("one") - a.add(ta) - a.trigger("custom") - assert ta.custom_called + ta = TAddon("one") + a.add(ta) + a.trigger("custom") + assert ta.custom_called def test_load_option(): @@ -114,29 +106,47 @@ def test_load_option(): assert "custom_option" in m.options._options -def test_loadchain(): +def test_nesting(): o = options.Options() m = master.Master(o, proxy.DummyServer(o)) a = addonmanager.AddonManager(m) - a.add(AChain("one", None)) + a.add( + TAddon( + "one", + addons=[ + TAddon("two"), + TAddon("three", addons=[TAddon("four")]) + ] + ) + ) + assert len(a.chain) == 1 assert a.get("one") - a.clear() - - a.add(AChain("one", AChain("two", None))) - assert not a.get("one") assert a.get("two") - a.clear() - - a.add(AChain("one", AChain("two", AChain("three", None)))) - assert not a.get("one") - assert not a.get("two") assert a.get("three") - a.clear() - - a.add(AChain("one", AChain("two", AChain("three", AChain("four", None))))) - assert not a.get("one") - assert not a.get("two") - assert not a.get("three") assert a.get("four") - a.clear() + + a.trigger("custom") + assert a.get("one").custom_called + assert a.get("two").custom_called + assert a.get("three").custom_called + assert a.get("four").custom_called + + a.remove(a.get("three")) + assert not a.get("three") + assert not a.get("four") + + +class D: + def __init__(self): + self.w = None + + def log(self, x): + self.w = x + + +def test_streamlog(): + dummy = D() + s = addonmanager.StreamLog(dummy.log) + s.write("foo") + assert dummy.w == "foo" diff --git a/test/mitmproxy/test_taddons.py b/test/mitmproxy/test_taddons.py index 42371cfef..5a4c99fc9 100644 --- a/test/mitmproxy/test_taddons.py +++ b/test/mitmproxy/test_taddons.py @@ -1,4 +1,6 @@ +import io from mitmproxy.test import taddons +from mitmproxy.test import tutils from mitmproxy import ctx @@ -9,3 +11,21 @@ def test_recordingmaster(): ctx.log.error("foo") assert not tctx.master.has_log("foo", level="debug") assert tctx.master.has_log("foo", level="error") + + +def test_dumplog(): + with taddons.context() as tctx: + ctx.log.info("testing") + s = io.StringIO() + tctx.master.dump_log(s) + assert s.getvalue() + + +def test_load_script(): + with taddons.context() as tctx: + s = tctx.script( + tutils.test_data.path( + "mitmproxy/data/addonscripts/recorder/recorder.py" + ) + ) + assert s diff --git a/test/mitmproxy/tools/console/test_master.py b/test/mitmproxy/tools/console/test_master.py index 44b9ff3f7..add8c4d3b 100644 --- a/test/mitmproxy/tools/console/test_master.py +++ b/test/mitmproxy/tools/console/test_master.py @@ -30,7 +30,7 @@ class TestMaster(tservers.MasterTest): opts["verbosity"] = 1 o = options.Options(**opts) m = console.master.ConsoleMaster(o, proxy.DummyServer()) - m.addons.configure_all(o, o.keys()) + m.addons.trigger("configure", o, o.keys()) return m def test_basic(self): @@ -42,12 +42,6 @@ class TestMaster(tservers.MasterTest): pass assert len(m.view) == i - def test_run_script_once(self): - m = self.mkmaster() - f = tflow.tflow(resp=True) - m.run_script_once("nonexistent", [f]) - assert any("Input error" in str(l) for l in m.logbuffer) - def test_intercept(self): """regression test for https://github.com/mitmproxy/mitmproxy/issues/1605""" m = self.mkmaster(intercept="~b bar") diff --git a/test/mitmproxy/tservers.py b/test/mitmproxy/tservers.py index b737b82a8..d8dda5eb8 100644 --- a/test/mitmproxy/tservers.py +++ b/test/mitmproxy/tservers.py @@ -74,7 +74,7 @@ class TestMaster(taddons.RecordingMaster): self.state = TestState() self.addons.add(self.state) self.addons.add(*addons) - self.addons.configure_all(self.options, self.options.keys()) + self.addons.trigger("configure", self.options, self.options.keys()) self.addons.trigger("running") def reset(self, addons):