diff options
67 files changed, 1059 insertions, 1056 deletions
@@ -22,3 +22,4 @@ sslkeylogfile.log .tox/ .python-version coverage.xml +web/coverage/ diff --git a/docs/install.rst b/docs/install.rst index b37d9c91..7753dc44 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 6ec0caaa..c333a98b 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 <api.html#mitmproxy.models.http.HTTPRequest>`_ and `HTTPResponse -<api.html#mitmproxy.models.http.HTTPResponse>`_ objects have a handy `replace -<api.html#mitmproxy.models.http.HTTPResponse.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 01e036b2..632783a7 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 69b64163..5d5c7902 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/custom_option.py b/examples/simple/custom_option.py index c8bc98d4..5b6070dd 100644 --- a/examples/simple/custom_option.py +++ b/examples/simple/custom_option.py @@ -6,6 +6,6 @@ def load(l): l.add_option("custom", bool, False, "A custom option") -def configure(options, updated): +def configure(updated): if "custom" in updated: - ctx.log.info("custom option value: %s" % options.custom) + ctx.log.info("custom option value: %s" % ctx.options.custom) diff --git a/examples/simple/filter_flows.py b/examples/simple/filter_flows.py index d2b735be..896fa54a 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 15e7693c..a0956e33 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 442a5118..dff72afa 100644 --- a/examples/simple/modify_body_inject_iframe.py +++ b/examples/simple/modify_body_inject_iframe.py @@ -1,29 +1,26 @@ -# Usage: mitmdump -s "iframe_injector.py url" # (this script works best with --anticache) -import sys from bs4 import BeautifulSoup +from mitmproxy import ctx class Injector: - def __init__(self, iframe_url): - self.iframe_url = iframe_url + def load(self, loader): + loader.add_option( + "iframe", str, "", "IFrame to inject" + ) 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 ctx.options.iframe: + html = BeautifulSoup(flow.response.content, "html.parser") + if html.body: + iframe = html.new_tag( + "iframe", + src=ctx.options.iframe, + 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 84292eb9..00000000 --- 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 241b3cde..13c90413 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", 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,56 @@ 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, call its load event, and then register all its + sub-addons. This should be used by addons that dynamically manage + addons. + + If the calling addon is already running, it should follow with + running and configure events. 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. + Remove an addon and all its sub-addons. + + 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. """ - self.chain = [i for i in self.chain if i is not addon] + 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(addon, "done") - - def done(self): - self.trigger("done") + self.invoke_addon(a, "done") def __len__(self): return len(self.chain) @@ -126,22 +210,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 +231,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/anticache.py b/mitmproxy/addons/anticache.py index 8d748a21..5b34d5a5 100644 --- a/mitmproxy/addons/anticache.py +++ b/mitmproxy/addons/anticache.py @@ -1,10 +1,7 @@ -class AntiCache: - def __init__(self): - self.enabled = False +from mitmproxy import ctx - def configure(self, options, updated): - self.enabled = options.anticache +class AntiCache: def request(self, flow): - if self.enabled: + if ctx.options.anticache: flow.request.anticache() diff --git a/mitmproxy/addons/anticomp.py b/mitmproxy/addons/anticomp.py index eaf01296..d7d1ca8d 100644 --- a/mitmproxy/addons/anticomp.py +++ b/mitmproxy/addons/anticomp.py @@ -1,10 +1,7 @@ -class AntiComp: - def __init__(self): - self.enabled = False +from mitmproxy import ctx - def configure(self, options, updated): - self.enabled = options.anticomp +class AntiComp: def request(self, flow): - if self.enabled: + if ctx.options.anticomp: flow.request.anticomp() diff --git a/mitmproxy/addons/check_alpn.py b/mitmproxy/addons/check_alpn.py index cb3c87e3..193159b2 100644 --- a/mitmproxy/addons/check_alpn.py +++ b/mitmproxy/addons/check_alpn.py @@ -7,7 +7,7 @@ class CheckALPN: def __init__(self): self.failed = False - def configure(self, options, updated): + def configure(self, updated): self.failed = mitmproxy.ctx.master.options.http2 and not tcp.HAS_ALPN if self.failed: ctx.log.warn( diff --git a/mitmproxy/addons/check_ca.py b/mitmproxy/addons/check_ca.py index a83ab8e1..f786af5a 100644 --- a/mitmproxy/addons/check_ca.py +++ b/mitmproxy/addons/check_ca.py @@ -5,7 +5,7 @@ class CheckCA: def __init__(self): self.failed = False - def configure(self, options, updated): + def configure(self, updated): has_ca = ( mitmproxy.ctx.master.server and mitmproxy.ctx.master.server.config and diff --git a/mitmproxy/addons/clientplayback.py b/mitmproxy/addons/clientplayback.py index 3345e65a..acb77bb2 100644 --- a/mitmproxy/addons/clientplayback.py +++ b/mitmproxy/addons/clientplayback.py @@ -20,12 +20,12 @@ class ClientPlayback: def load(self, flows: typing.Sequence[flow.Flow]): self.flows = flows - def configure(self, options, updated): + def configure(self, updated): if "client_replay" in updated: - if options.client_replay: - ctx.log.info("Client Replay: {}".format(options.client_replay)) + if ctx.options.client_replay: + ctx.log.info("Client Replay: {}".format(ctx.options.client_replay)) try: - flows = io.read_flows_from_paths(options.client_replay) + flows = io.read_flows_from_paths(ctx.options.client_replay) except exceptions.FlowReadException as e: raise exceptions.OptionsError(str(e)) self.load(flows) diff --git a/mitmproxy/addons/core_option_validation.py b/mitmproxy/addons/core_option_validation.py index fd5f2788..baeee764 100644 --- a/mitmproxy/addons/core_option_validation.py +++ b/mitmproxy/addons/core_option_validation.py @@ -4,12 +4,14 @@ """ from mitmproxy import exceptions from mitmproxy import platform +from mitmproxy import ctx from mitmproxy.net import server_spec from mitmproxy.utils import human class CoreOptionValidation: - def configure(self, opts, updated): + def configure(self, updated): + opts = ctx.options if opts.add_upstream_certs_to_client_chain and not opts.upstream_cert: raise exceptions.OptionsError( "The no-upstream-cert and add-upstream-certs-to-client-chain " diff --git a/mitmproxy/addons/disable_h2c.py b/mitmproxy/addons/disable_h2c.py index b43d5207..392a29a5 100644 --- a/mitmproxy/addons/disable_h2c.py +++ b/mitmproxy/addons/disable_h2c.py @@ -14,9 +14,6 @@ class DisableH2C: by sending the connection preface. We just kill those flows. """ - def configure(self, options, updated): - pass - def process_flow(self, f): if f.request.headers.get('upgrade', '') == 'h2c': mitmproxy.ctx.log.warn("HTTP/2 cleartext connections (h2c upgrade requests) are currently not supported.") diff --git a/mitmproxy/addons/dumper.py b/mitmproxy/addons/dumper.py index 5fd8408f..3c3e1c65 100644 --- a/mitmproxy/addons/dumper.py +++ b/mitmproxy/addons/dumper.py @@ -29,24 +29,18 @@ def colorful(line, styles): class Dumper: def __init__(self, outfile=sys.stdout): self.filter = None # type: flowfilter.TFilter - self.flow_detail = None # type: int self.outfp = outfile # type: typing.io.TextIO - self.showhost = None # type: bool - self.default_contentview = "auto" # type: str - def configure(self, options, updated): + def configure(self, updated): if "view_filter" in updated: - if options.view_filter: - self.filter = flowfilter.parse(options.view_filter) + if ctx.options.view_filter: + self.filter = flowfilter.parse(ctx.options.view_filter) if not self.filter: raise exceptions.OptionsError( - "Invalid filter expression: %s" % options.view_filter + "Invalid filter expression: %s" % ctx.options.view_filter ) else: self.filter = None - self.flow_detail = options.flow_detail - self.showhost = options.showhost - self.default_contentview = options.default_contentview def echo(self, text, ident=None, **style): if ident: @@ -67,13 +61,13 @@ class Dumper: def _echo_message(self, message): _, lines, error = contentviews.get_message_content_view( - self.default_contentview, + ctx.options.default_contentview, message ) if error: ctx.log.debug(error) - if self.flow_detail == 3: + if ctx.options.flow_detail == 3: lines_to_echo = itertools.islice(lines, 70) else: lines_to_echo = lines @@ -95,7 +89,7 @@ class Dumper: if next(lines, None): self.echo("(cut off)", ident=4, dim=True) - if self.flow_detail >= 2: + if ctx.options.flow_detail >= 2: self.echo("") def _echo_request_line(self, flow): @@ -121,12 +115,12 @@ class Dumper: fg=method_color, bold=True ) - if self.showhost: + if ctx.options.showhost: url = flow.request.pretty_url else: url = flow.request.url terminalWidthLimit = max(shutil.get_terminal_size()[0] - 25, 50) - if self.flow_detail < 1 and len(url) > terminalWidthLimit: + if ctx.options.flow_detail < 1 and len(url) > terminalWidthLimit: url = url[:terminalWidthLimit] + "…" url = click.style(strutils.escape_control_characters(url), bold=True) @@ -176,7 +170,7 @@ class Dumper: size = click.style(size, bold=True) arrows = click.style(" <<", bold=True) - if self.flow_detail == 1: + if ctx.options.flow_detail == 1: # This aligns the HTTP response code with the HTTP request method: # 127.0.0.1:59519: GET http://example.com/ # << 304 Not Modified 0b @@ -194,16 +188,16 @@ class Dumper: def echo_flow(self, f): if f.request: self._echo_request_line(f) - if self.flow_detail >= 2: + if ctx.options.flow_detail >= 2: self._echo_headers(f.request.headers) - if self.flow_detail >= 3: + if ctx.options.flow_detail >= 3: self._echo_message(f.request) if f.response: self._echo_response_line(f) - if self.flow_detail >= 2: + if ctx.options.flow_detail >= 2: self._echo_headers(f.response.headers) - if self.flow_detail >= 3: + if ctx.options.flow_detail >= 3: self._echo_message(f.response) if f.error: @@ -211,7 +205,7 @@ class Dumper: self.echo(" << {}".format(msg), bold=True, fg="red") def match(self, f): - if self.flow_detail == 0: + if ctx.options.flow_detail == 0: return False if not self.filter: return True @@ -239,7 +233,7 @@ class Dumper: if self.match(f): message = f.messages[-1] self.echo(f.message_info(message)) - if self.flow_detail >= 3: + if ctx.options.flow_detail >= 3: self._echo_message(message) def websocket_end(self, f): @@ -267,5 +261,5 @@ class Dumper: server=repr(f.server_conn.address), direction=direction, )) - if self.flow_detail >= 3: + if ctx.options.flow_detail >= 3: self._echo_message(message) diff --git a/mitmproxy/addons/intercept.py b/mitmproxy/addons/intercept.py index 4a3fe17e..ac8c4c88 100644 --- a/mitmproxy/addons/intercept.py +++ b/mitmproxy/addons/intercept.py @@ -1,20 +1,21 @@ from mitmproxy import flowfilter from mitmproxy import exceptions +from mitmproxy import ctx class Intercept: def __init__(self): self.filt = None - def configure(self, opts, updated): + def configure(self, updated): if "intercept" in updated: - if not opts.intercept: + if not ctx.options.intercept: self.filt = None return - self.filt = flowfilter.parse(opts.intercept) + self.filt = flowfilter.parse(ctx.options.intercept) if not self.filt: raise exceptions.OptionsError( - "Invalid interception filter: %s" % opts.intercept + "Invalid interception filter: %s" % ctx.options.intercept ) def process_flow(self, f): diff --git a/mitmproxy/addons/onboarding.py b/mitmproxy/addons/onboarding.py index cb57990f..07536c34 100644 --- a/mitmproxy/addons/onboarding.py +++ b/mitmproxy/addons/onboarding.py @@ -1,17 +1,18 @@ from mitmproxy.addons import wsgiapp from mitmproxy.addons.onboardingapp import app +from mitmproxy import ctx class Onboarding(wsgiapp.WSGIApp): + name = "onboarding" + def __init__(self): super().__init__(app.Adapter(app.application), None, None) - self.enabled = False - def configure(self, options, updated): - self.host = options.onboarding_host - self.port = options.onboarding_port - self.enabled = options.onboarding + def configure(self, updated): + self.host = ctx.options.onboarding_host + self.port = ctx.options.onboarding_port def request(self, f): - if self.enabled: + if ctx.options.onboarding: super().request(f) diff --git a/mitmproxy/addons/proxyauth.py b/mitmproxy/addons/proxyauth.py index 43677f61..fecdcb84 100644 --- a/mitmproxy/addons/proxyauth.py +++ b/mitmproxy/addons/proxyauth.py @@ -10,6 +10,7 @@ import mitmproxy.net.http from mitmproxy import connections # noqa from mitmproxy import exceptions from mitmproxy import http +from mitmproxy import ctx from mitmproxy.net.http import status_codes REALM = "mitmproxy" @@ -45,7 +46,6 @@ class ProxyAuth: self.nonanonymous = False self.htpasswd = None self.singleuser = None - self.mode = None self.authenticated = weakref.WeakKeyDictionary() # type: MutableMapping[connections.ClientConnection, Tuple[str, str]] """Contains all connections that are permanently authenticated after an HTTP CONNECT""" @@ -58,7 +58,7 @@ class ProxyAuth: - True, if authentication is done as if mitmproxy is a proxy - False, if authentication is done as if mitmproxy is a HTTP server """ - return self.mode in ("regular", "upstream") + return ctx.options.mode in ("regular", "upstream") def which_auth_header(self) -> str: if self.is_proxy_auth(): @@ -113,16 +113,16 @@ class ProxyAuth: return False # Handlers - def configure(self, options, updated): + def configure(self, updated): if "proxyauth" in updated: self.nonanonymous = False self.singleuser = None self.htpasswd = None - if options.proxyauth: - if options.proxyauth == "any": + if ctx.options.proxyauth: + if ctx.options.proxyauth == "any": self.nonanonymous = True - elif options.proxyauth.startswith("@"): - p = options.proxyauth[1:] + elif ctx.options.proxyauth.startswith("@"): + p = ctx.options.proxyauth[1:] try: self.htpasswd = passlib.apache.HtpasswdFile(p) except (ValueError, OSError) as v: @@ -130,20 +130,18 @@ class ProxyAuth: "Could not open htpasswd file: %s" % p ) else: - parts = options.proxyauth.split(':') + parts = ctx.options.proxyauth.split(':') if len(parts) != 2: raise exceptions.OptionsError( "Invalid single-user auth specification." ) self.singleuser = parts - if "mode" in updated: - self.mode = options.mode if self.enabled(): - if options.mode == "transparent": + if ctx.options.mode == "transparent": raise exceptions.OptionsError( "Proxy Authentication not supported in transparent mode." ) - if options.mode == "socks5": + if ctx.options.mode == "socks5": raise exceptions.OptionsError( "Proxy Authentication not supported in SOCKS mode. " "https://github.com/mitmproxy/mitmproxy/issues/738" diff --git a/mitmproxy/addons/readfile.py b/mitmproxy/addons/readfile.py index 03dcd084..949da15d 100644 --- a/mitmproxy/addons/readfile.py +++ b/mitmproxy/addons/readfile.py @@ -9,9 +9,6 @@ class ReadFile: """ An addon that handles reading from file on startup. """ - def __init__(self): - self.path = None - def load_flows_file(self, path: str) -> int: path = os.path.expanduser(path) cnt = 0 @@ -31,16 +28,11 @@ class ReadFile: ctx.log.error("Flow file corrupted.") raise exceptions.FlowReadException(v) - def configure(self, options, updated): - if "rfile" in updated and options.rfile: - self.path = options.rfile - def running(self): - if self.path: + if ctx.options.rfile: try: - self.load_flows_file(self.path) + self.load_flows_file(ctx.options.rfile) except exceptions.FlowReadException as v: raise exceptions.OptionsError(v) finally: - self.path = None ctx.master.addons.trigger("processing_complete") diff --git a/mitmproxy/addons/replace.py b/mitmproxy/addons/replace.py index d6c11ca4..054264fa 100644 --- a/mitmproxy/addons/replace.py +++ b/mitmproxy/addons/replace.py @@ -47,7 +47,7 @@ class Replace: def __init__(self): self.lst = [] - def configure(self, options, updated): + def configure(self, updated): """ .replacements is a list of tuples (fpat, rex, s): @@ -57,7 +57,7 @@ class Replace: """ if "replacements" in updated: lst = [] - for rep in options.replacements: + for rep in ctx.options.replacements: fpatt, rex, s = parse_hook(rep) flt = flowfilter.parse(fpatt) diff --git a/mitmproxy/addons/script.py b/mitmproxy/addons/script.py index eef21293..0497af85 100644 --- a/mitmproxy/addons/script.py +++ b/mitmproxy/addons/script.py @@ -1,223 +1,68 @@ -import contextlib import os -import shlex +import importlib +import time import sys -import threading -import traceback -import types 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) - - -class ReloadHandler(watchdog.events.FileSystemEventHandler): - def __init__(self, callback): - self.callback = callback - - def filter(self, event): - """ - Returns True only when .py file is changed - """ - if event.is_directory: - return False - if os.path.basename(event.src_path).startswith("."): - return False - if event.src_path.endswith(".py"): - return True - return False - - def on_modified(self, event): - if self.filter(event): - self.callback() - - def on_created(self, event): - if self.filter(event): - self.callback() + sys.path[:] = oldpath class Script: """ An addon that manages a single script. """ - def __init__(self, command): - self.name = command + ReloadInterval = 2 - 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 prox(*args, **kwargs): - self.run(evt, *args, **kwargs) - return prox - setattr(self, i, mkprox()) + self.last_load = 0 + self.last_mtime = 0 - 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) - - 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 + @property + def addons(self): + return [self.ns] if self.ns else [] 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() - - def configure(self, options, updated): - 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, updated) - - def done(self): - self.run("done") - self.dead = True + if time.time() - self.last_load > self.ReloadInterval: + mtime = os.stat(self.path).st_mtime + if mtime > self.last_mtime: + ctx.log.info("Loading script: %s" % self.name) + 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) + ctx.master.addons.invoke_addon(self.ns, "running") + ctx.master.addons.invoke_addon( + self.ns, + "configure", + ctx.options.keys() + ) + self.last_load = time.time() + self.last_mtime = mtime class ScriptLoader: @@ -226,73 +71,53 @@ 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): + def configure(self, updated): if "scripts" in updated: - for s in options.scripts: - if options.scripts.count(s) > 1: + for s in ctx.options.scripts: + if ctx.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 ctx.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 = [] - for s in options.scripts: + for s in ctx.options.scripts: if s in current: 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/serverplayback.py b/mitmproxy/addons/serverplayback.py index 7bb0c716..2255aaf2 100644 --- a/mitmproxy/addons/serverplayback.py +++ b/mitmproxy/addons/serverplayback.py @@ -10,8 +10,6 @@ from mitmproxy import io class ServerPlayback: def __init__(self): - self.options = None - self.flowmap = {} self.stop = False self.final_flow = None @@ -38,27 +36,27 @@ class ServerPlayback: queriesArray = urllib.parse.parse_qsl(query, keep_blank_values=True) key = [str(r.port), str(r.scheme), str(r.method), str(path)] # type: List[Any] - if not self.options.server_replay_ignore_content: - if self.options.server_replay_ignore_payload_params and r.multipart_form: + if not ctx.options.server_replay_ignore_content: + if ctx.options.server_replay_ignore_payload_params and r.multipart_form: key.extend( (k, v) for k, v in r.multipart_form.items(multi=True) - if k.decode(errors="replace") not in self.options.server_replay_ignore_payload_params + if k.decode(errors="replace") not in ctx.options.server_replay_ignore_payload_params ) - elif self.options.server_replay_ignore_payload_params and r.urlencoded_form: + elif ctx.options.server_replay_ignore_payload_params and r.urlencoded_form: key.extend( (k, v) for k, v in r.urlencoded_form.items(multi=True) - if k not in self.options.server_replay_ignore_payload_params + if k not in ctx.options.server_replay_ignore_payload_params ) else: key.append(str(r.raw_content)) - if not self.options.server_replay_ignore_host: + if not ctx.options.server_replay_ignore_host: key.append(r.host) filtered = [] - ignore_params = self.options.server_replay_ignore_params or [] + ignore_params = ctx.options.server_replay_ignore_params or [] for p in queriesArray: if p[0] not in ignore_params: filtered.append(p) @@ -66,9 +64,9 @@ class ServerPlayback: key.append(p[0]) key.append(p[1]) - if self.options.server_replay_use_headers: + if ctx.options.server_replay_use_headers: headers = [] - for i in self.options.server_replay_use_headers: + for i in ctx.options.server_replay_use_headers: v = r.headers.get(i) headers.append((i, v)) key.append(headers) @@ -83,7 +81,7 @@ class ServerPlayback: """ hsh = self._hash(request) if hsh in self.flowmap: - if self.options.server_replay_nopop: + if ctx.options.server_replay_nopop: return self.flowmap[hsh][0] else: ret = self.flowmap[hsh].pop(0) @@ -91,13 +89,12 @@ class ServerPlayback: del self.flowmap[hsh] return ret - def configure(self, options, updated): - self.options = options + def configure(self, updated): if "server_replay" in updated: self.clear() - if options.server_replay: + if ctx.options.server_replay: try: - flows = io.read_flows_from_paths(options.server_replay) + flows = io.read_flows_from_paths(ctx.options.server_replay) except exceptions.FlowReadException as e: raise exceptions.OptionsError(str(e)) self.load_flows(flows) @@ -112,13 +109,13 @@ class ServerPlayback: if rflow: response = rflow.response.copy() response.is_replay = True - if self.options.refresh_server_playback: + if ctx.options.refresh_server_playback: response.refresh() f.response = response if not self.flowmap: self.final_flow = f self.stop = True - elif self.options.replay_kill_extra: + elif ctx.options.replay_kill_extra: ctx.log.warn( "server_playback: killed non-replay request {}".format( f.request.url diff --git a/mitmproxy/addons/setheaders.py b/mitmproxy/addons/setheaders.py index 9e60eabd..d4d16e40 100644 --- a/mitmproxy/addons/setheaders.py +++ b/mitmproxy/addons/setheaders.py @@ -1,5 +1,6 @@ from mitmproxy import exceptions from mitmproxy import flowfilter +from mitmproxy import ctx def parse_setheader(s): @@ -43,17 +44,10 @@ class SetHeaders: def __init__(self): self.lst = [] - def configure(self, options, updated): - """ - options.setheaders is a tuple of (fpatt, header, value) - - fpatt: String specifying a filter pattern. - header: Header name. - value: Header value string - """ + def configure(self, updated): if "setheaders" in updated: self.lst = [] - for shead in options.setheaders: + for shead in ctx.options.setheaders: fpatt, header, value = parse_setheader(shead) flt = flowfilter.parse(fpatt) diff --git a/mitmproxy/addons/stickyauth.py b/mitmproxy/addons/stickyauth.py index 1a1d4fc4..24831d5b 100644 --- a/mitmproxy/addons/stickyauth.py +++ b/mitmproxy/addons/stickyauth.py @@ -1,5 +1,6 @@ from mitmproxy import exceptions from mitmproxy import flowfilter +from mitmproxy import ctx class StickyAuth: @@ -7,13 +8,13 @@ class StickyAuth: self.flt = None self.hosts = {} - def configure(self, options, updated): + def configure(self, updated): if "stickyauth" in updated: - if options.stickyauth: - flt = flowfilter.parse(options.stickyauth) + if ctx.options.stickyauth: + flt = flowfilter.parse(ctx.options.stickyauth) if not flt: raise exceptions.OptionsError( - "stickyauth: invalid filter expression: %s" % options.stickyauth + "stickyauth: invalid filter expression: %s" % ctx.options.stickyauth ) self.flt = flt else: diff --git a/mitmproxy/addons/stickycookie.py b/mitmproxy/addons/stickycookie.py index fb1c5072..04d99975 100644 --- a/mitmproxy/addons/stickycookie.py +++ b/mitmproxy/addons/stickycookie.py @@ -5,6 +5,7 @@ from mitmproxy.net.http import cookies from mitmproxy import exceptions from mitmproxy import flowfilter +from mitmproxy import ctx def ckey(attrs, f): @@ -33,13 +34,13 @@ class StickyCookie: self.jar = collections.defaultdict(dict) self.flt = None - def configure(self, options, updated): + def configure(self, updated): if "stickycookie" in updated: - if options.stickycookie: - flt = flowfilter.parse(options.stickycookie) + if ctx.options.stickycookie: + flt = flowfilter.parse(ctx.options.stickycookie) if not flt: raise exceptions.OptionsError( - "stickycookie: invalid filter expression: %s" % options.stickycookie + "stickycookie: invalid filter expression: %s" % ctx.options.stickycookie ) self.flt = flt else: diff --git a/mitmproxy/addons/streambodies.py b/mitmproxy/addons/streambodies.py index a10bdb93..181f0337 100644 --- a/mitmproxy/addons/streambodies.py +++ b/mitmproxy/addons/streambodies.py @@ -8,10 +8,10 @@ class StreamBodies: def __init__(self): self.max_size = None - def configure(self, options, updated): - if "stream_large_bodies" in updated and options.stream_large_bodies: + def configure(self, updated): + if "stream_large_bodies" in updated and ctx.options.stream_large_bodies: try: - self.max_size = human.parse_size(options.stream_large_bodies) + self.max_size = human.parse_size(ctx.options.stream_large_bodies) except ValueError as e: raise exceptions.OptionsError(e) diff --git a/mitmproxy/addons/streamfile.py b/mitmproxy/addons/streamfile.py index 183d2036..fde5a1c5 100644 --- a/mitmproxy/addons/streamfile.py +++ b/mitmproxy/addons/streamfile.py @@ -3,6 +3,7 @@ import os.path from mitmproxy import exceptions from mitmproxy import flowfilter from mitmproxy import io +from mitmproxy import ctx class StreamFile: @@ -20,26 +21,26 @@ class StreamFile: self.stream = io.FilteredFlowWriter(f, flt) self.active_flows = set() - def configure(self, options, updated): + def configure(self, updated): # We're already streaming - stop the previous stream and restart if "streamfile_filter" in updated: - if options.streamfile_filter: - self.filt = flowfilter.parse(options.streamfile_filter) + if ctx.options.streamfile_filter: + self.filt = flowfilter.parse(ctx.options.streamfile_filter) if not self.filt: raise exceptions.OptionsError( - "Invalid filter specification: %s" % options.streamfile_filter + "Invalid filter specification: %s" % ctx.options.streamfile_filter ) else: self.filt = None if "streamfile" in updated: if self.stream: self.done() - if options.streamfile: - if options.streamfile.startswith("+"): - path = options.streamfile[1:] + if ctx.options.streamfile: + if ctx.options.streamfile.startswith("+"): + path = ctx.options.streamfile[1:] mode = "ab" else: - path = options.streamfile + path = ctx.options.streamfile mode = "wb" self.start_stream_to_path(path, mode, self.filt) diff --git a/mitmproxy/addons/termlog.py b/mitmproxy/addons/termlog.py index 5fdb6245..4c37b005 100644 --- a/mitmproxy/addons/termlog.py +++ b/mitmproxy/addons/termlog.py @@ -2,23 +2,25 @@ import sys import click from mitmproxy import log +from mitmproxy import ctx + +# 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): - self.options = None self.outfile = outfile - def configure(self, options, updated): - self.options = options - 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): + if ctx.options.verbosity >= log.log_tier(e.level): click.secho( e.msg, file=outfile, diff --git a/mitmproxy/addons/termstatus.py b/mitmproxy/addons/termstatus.py index 951ddd3c..c3c91283 100644 --- a/mitmproxy/addons/termstatus.py +++ b/mitmproxy/addons/termstatus.py @@ -8,15 +8,8 @@ from mitmproxy.utils import human class TermStatus: - def __init__(self): - self.server = False - - def configure(self, options, updated): - if "server" in updated: - self.server = options.server - def running(self): - if self.server: + if ctx.options.server: ctx.log.info( "Proxy server listening at http://{}".format( human.format_address(ctx.master.server.address) diff --git a/mitmproxy/addons/upstream_auth.py b/mitmproxy/addons/upstream_auth.py index 9beecfc0..685494c2 100644 --- a/mitmproxy/addons/upstream_auth.py +++ b/mitmproxy/addons/upstream_auth.py @@ -2,6 +2,7 @@ import re import base64 from mitmproxy import exceptions +from mitmproxy import ctx from mitmproxy.utils import strutils @@ -26,20 +27,17 @@ class UpstreamAuth(): """ def __init__(self): self.auth = None - self.root_mode = None - def configure(self, options, updated): + def configure(self, updated): # FIXME: We're doing this because our proxy core is terminally confused # at the moment. Ideally, we should be able to check if we're in # reverse proxy mode at the HTTP layer, so that scripts can put the # proxy in reverse proxy mode for specific reuests. - if "mode" in updated: - self.root_mode = options.mode if "upstream_auth" in updated: - if options.upstream_auth is None: + if ctx.options.upstream_auth is None: self.auth = None else: - self.auth = parse_upstream_auth(options.upstream_auth) + self.auth = parse_upstream_auth(ctx.options.upstream_auth) def http_connect(self, f): if self.auth and f.mode == "upstream": @@ -49,5 +47,5 @@ class UpstreamAuth(): if self.auth: if f.mode == "upstream" and not f.server_conn.via: f.request.headers["Proxy-Authorization"] = self.auth - elif self.root_mode == "reverse": + elif ctx.options.mode == "reverse": f.request.headers["Proxy-Authorization"] = self.auth diff --git a/mitmproxy/addons/view.py b/mitmproxy/addons/view.py index 7e9d66a1..341958c2 100644 --- a/mitmproxy/addons/view.py +++ b/mitmproxy/addons/view.py @@ -18,6 +18,7 @@ import sortedcontainers import mitmproxy.flow from mitmproxy import flowfilter from mitmproxy import exceptions +from mitmproxy import ctx from mitmproxy import http # noqa # The underlying sorted list implementation expects the sort key to be stable @@ -302,26 +303,26 @@ class View(collections.Sequence): return self._store.get(flow_id) # Event handlers - def configure(self, opts, updated): + def configure(self, updated): if "view_filter" in updated: filt = None - if opts.view_filter: - filt = flowfilter.parse(opts.view_filter) + if ctx.options.view_filter: + filt = flowfilter.parse(ctx.options.view_filter) if not filt: raise exceptions.OptionsError( - "Invalid interception filter: %s" % opts.view_filter + "Invalid interception filter: %s" % ctx.options.view_filter ) self.set_filter(filt) if "console_order" in updated: - if opts.console_order not in self.orders: + if ctx.options.console_order not in self.orders: raise exceptions.OptionsError( - "Unknown flow order: %s" % opts.console_order + "Unknown flow order: %s" % ctx.options.console_order ) - self.set_order(self.orders[opts.console_order]) + self.set_order(self.orders[ctx.options.console_order]) if "console_order_reversed" in updated: - self.set_reversed(opts.console_order_reversed) + self.set_reversed(ctx.options.console_order_reversed) if "console_focus_follow" in updated: - self.focus_follow = opts.console_focus_follow + self.focus_follow = ctx.options.console_focus_follow def request(self, f): self.add(f) diff --git a/mitmproxy/addons/wsgiapp.py b/mitmproxy/addons/wsgiapp.py index c37fcb7b..155444fc 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/ctx.py b/mitmproxy/ctx.py index 7b5231e6..954edcb1 100644 --- a/mitmproxy/ctx.py +++ b/mitmproxy/ctx.py @@ -1,4 +1,7 @@ import mitmproxy.master # noqa import mitmproxy.log # noqa +import mitmproxy.options # noqa + master = None # type: "mitmproxy.master.Master" log = None # type: "mitmproxy.log.Log" +options = None # type: "mitmproxy.options.Options" diff --git a/mitmproxy/master.py b/mitmproxy/master.py index 946b25a4..94900915 100644 --- a/mitmproxy/master.py +++ b/mitmproxy/master.py @@ -50,11 +50,13 @@ class Master: return mitmproxy_ctx.master = self mitmproxy_ctx.log = log.Log(self) + mitmproxy_ctx.options = self.options try: yield finally: mitmproxy_ctx.master = None mitmproxy_ctx.log = None + mitmproxy_ctx.options = None def tell(self, mtype, m): m.reply = controller.DummyReply() @@ -103,7 +105,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 8bc174c7..39ebb2e6 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,19 @@ 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", + kwargs.keys() + ) + + def script(self, path): + """ + Loads a script from path, and returns the enclosed addon. + """ + sc = script.Script(path) + loader = addonmanager.Loader(self.master) + self.master.addons.invoke_addon(sc, "load", loader) + self.configure(sc) + self.master.addons.invoke_addon(sc, "tick") + return sc.addons[0] if sc.addons else None diff --git a/mitmproxy/tools/main.py b/mitmproxy/tools/main.py index 6db232fc..c0326739 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.keys()) remaining = opts.update_known(**unknown) if remaining and opts.verbosity > 1: print("Ignored options: %s" % remaining) @@ -80,7 +80,6 @@ setup( "ruamel.yaml>=0.13.2, <0.15", "tornado>=4.3, <4.6", "urwid>=1.3.1, <1.4", - "watchdog>=0.8.3, <0.9", "brotlipy>=0.5.1, <0.7", "sortedcontainers>=1.5.4, <1.6", # transitive from cryptography, we just blacklist here. diff --git a/test/examples/test_har_dump.py b/test/examples/_test_har_dump.py index e5cfd2e1..e5cfd2e1 100644 --- a/test/examples/test_har_dump.py +++ b/test/examples/_test_har_dump.py diff --git a/test/examples/test_examples.py b/test/examples/test_examples.py index 46fdcd36..4c1631ce 100644 --- a/test/examples/test_examples.py +++ b/test/examples/test_examples.py @@ -1,13 +1,7 @@ -import pytest - -from mitmproxy import options from mitmproxy import contentviews -from mitmproxy import proxy -from mitmproxy import master -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 @@ -15,106 +9,93 @@ from ..mitmproxy import tservers example_dir = tutils.test_data.push("../examples") -class ScriptError(Exception): - pass - - -class RaiseMaster(master.Master): - def add_log(self, e, level): - if level in ("warn", "error"): - raise ScriptError(e) - - -def tscript(cmd, args=""): - o = options.Options() - cmd = example_dir.path(cmd) + " " + args - m = RaiseMaster(o, proxy.DummyServer()) - sc = script.Script(cmd) - m.addons.add(sc) - return m, sc - - class TestScripts(tservers.MasterTest): def test_add_header(self): - m, _ = tscript("simple/add_header.py") - f = tflow.tflow(resp=tutils.tresp()) - m.addons.handle_lifecycle("response", f) - assert f.response.headers["newheader"] == "foo" + with taddons.context() as tctx: + a = tctx.script(example_dir.path("simple/add_header.py")) + f = tflow.tflow(resp=tutils.tresp()) + a.response(f) + assert f.response.headers["newheader"] == "foo" def test_custom_contentviews(self): - m, sc = tscript("simple/custom_contentview.py") - swapcase = contentviews.get("swapcase") - _, fmt = swapcase(b"<html>Test!</html>") - assert any(b'tEST!' in val[0][1] for val in fmt) + with taddons.context() as tctx: + tctx.script(example_dir.path("simple/custom_contentview.py")) + swapcase = contentviews.get("swapcase") + _, fmt = swapcase(b"<html>Test!</html>") + 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"<html><body>mitmproxy</body></html>")) - 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"<html><body>mitmproxy</body></html>") + ) + 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") + with taddons.context() as tctx: + sc = tctx.script(example_dir.path("simple/modify_form.py")) - form_header = Headers(content_type="application/x-www-form-urlencoded") - f = tflow.tflow(req=tutils.treq(headers=form_header)) - m.addons.handle_lifecycle("request", f) + form_header = Headers(content_type="application/x-www-form-urlencoded") + f = tflow.tflow(req=tutils.treq(headers=form_header)) + sc.request(f) - assert f.request.urlencoded_form["mitmproxy"] == "rocks" + assert f.request.urlencoded_form["mitmproxy"] == "rocks" - f.request.headers["content-type"] = "" - m.addons.handle_lifecycle("request", f) - assert list(f.request.urlencoded_form.items()) == [("foo", "bar")] + f.request.headers["content-type"] = "" + sc.request(f) + assert list(f.request.urlencoded_form.items()) == [("foo", "bar")] def test_modify_querystring(self): - m, sc = tscript("simple/modify_querystring.py") - f = tflow.tflow(req=tutils.treq(path="/search?q=term")) + with taddons.context() as tctx: + sc = tctx.script(example_dir.path("simple/modify_querystring.py")) + f = tflow.tflow(req=tutils.treq(path="/search?q=term")) - m.addons.handle_lifecycle("request", f) - assert f.request.query["mitmproxy"] == "rocks" + sc.request(f) + assert f.request.query["mitmproxy"] == "rocks" - f.request.path = "/" - 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" + f.request.path = "/" + sc.request(f) + assert f.request.query["mitmproxy"] == "rocks" def test_redirect_requests(self): - m, sc = tscript("simple/redirect_requests.py") - f = tflow.tflow(req=tutils.treq(host="example.org")) - m.addons.handle_lifecycle("request", f) - assert f.request.host == "mitmproxy.org" + with taddons.context() as tctx: + sc = tctx.script(example_dir.path("simple/redirect_requests.py")) + f = tflow.tflow(req=tutils.treq(host="example.org")) + sc.request(f) + assert f.request.host == "mitmproxy.org" def test_send_reply_from_proxy(self): - m, sc = tscript("simple/send_reply_from_proxy.py") - f = tflow.tflow(req=tutils.treq(host="example.com", port=80)) - m.addons.handle_lifecycle("request", f) - assert f.response.content == b"Hello World" + with taddons.context() as tctx: + sc = tctx.script(example_dir.path("simple/send_reply_from_proxy.py")) + f = tflow.tflow(req=tutils.treq(host="example.com", port=80)) + sc.request(f) + assert f.response.content == b"Hello World" def test_dns_spoofing(self): - m, sc = tscript("complex/dns_spoofing.py") - original_host = "example.com" + with taddons.context() as tctx: + sc = tctx.script(example_dir.path("complex/dns_spoofing.py")) + + original_host = "example.com" - host_header = Headers(host=original_host) - f = tflow.tflow(req=tutils.treq(headers=host_header, port=80)) + host_header = Headers(host=original_host) + f = tflow.tflow(req=tutils.treq(headers=host_header, port=80)) - m.addons.handle_lifecycle("requestheaders", f) + tctx.master.addons.invoke_addon(sc, "requestheaders", f) - # Rewrite by reverse proxy mode - f.request.scheme = "https" - f.request.port = 443 + # Rewrite by reverse proxy mode + f.request.scheme = "https" + f.request.port = 443 - m.addons.handle_lifecycle("request", f) + tctx.master.addons.invoke_addon(sc, "request", f) - assert f.request.scheme == "http" - assert f.request.port == 80 + assert f.request.scheme == "http" + assert f.request.port == 80 - assert f.request.headers["Host"] == original_host + assert f.request.headers["Host"] == original_host diff --git a/test/mitmproxy/addons/test_core_option_validation.py b/test/mitmproxy/addons/test_core_option_validation.py index 0bb2bb0d..6d6d5ba4 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_onboarding.py b/test/mitmproxy/addons/test_onboarding.py index 63125c23..42a3b574 100644 --- a/test/mitmproxy/addons/test_onboarding.py +++ b/test/mitmproxy/addons/test_onboarding.py @@ -1,4 +1,5 @@ from mitmproxy.addons import onboarding +from mitmproxy.test import taddons from .. import tservers @@ -7,10 +8,14 @@ class TestApp(tservers.HTTPProxyTest): return [onboarding.Onboarding()] def test_basic(self): - assert self.app("/").status_code == 200 + with taddons.context() as tctx: + tctx.configure(self.addons()[0]) + assert self.app("/").status_code == 200 def test_cert(self): - for ext in ["pem", "p12"]: - resp = self.app("/cert/%s" % ext) - assert resp.status_code == 200 - assert resp.content + with taddons.context() as tctx: + tctx.configure(self.addons()[0]) + for ext in ["pem", "p12"]: + resp = self.app("/cert/%s" % ext) + assert resp.status_code == 200 + assert resp.content diff --git a/test/mitmproxy/addons/test_proxyauth.py b/test/mitmproxy/addons/test_proxyauth.py index 513f3f08..86621709 100644 --- a/test/mitmproxy/addons/test_proxyauth.py +++ b/test/mitmproxy/addons/test_proxyauth.py @@ -66,9 +66,6 @@ def test_configure(): with pytest.raises(exceptions.OptionsError): ctx.configure(up, proxyauth="any", mode="socks5") - ctx.configure(up, mode="regular") - assert up.mode == "regular" - def test_check(): up = proxyauth.ProxyAuth() diff --git a/test/mitmproxy/addons/test_script.py b/test/mitmproxy/addons/test_script.py index c68981de..859d99f9 100644 --- a/test/mitmproxy/addons/test_script.py +++ b/test/mitmproxy/addons/test_script.py @@ -1,8 +1,6 @@ import traceback import sys import time -import re -import watchdog.events import pytest from unittest import mock @@ -14,159 +12,95 @@ 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): - self.called = False - - def __call__(self, *args, **kwargs): - self.called = True - - -def test_reloadhandler(): - rh = script.ReloadHandler(Called()) - assert not rh.filter(watchdog.events.DirCreatedEvent("path")) - assert not rh.filter(watchdog.events.FileModifiedEvent("/foo/.bar")) - assert not rh.filter(watchdog.events.FileModifiedEvent("/foo/bar")) - assert rh.filter(watchdog.events.FileModifiedEvent("/foo/bar.py")) - - assert not rh.callback.called - rh.on_modified(watchdog.events.FileModifiedEvent("/foo/bar")) - assert not rh.callback.called - rh.on_modified(watchdog.events.FileModifiedEvent("/foo/bar.py")) - assert rh.callback.called - rh.callback.called = False - - rh.on_created(watchdog.events.FileCreatedEvent("foo")) - assert not rh.callback.called - rh.on_created(watchdog.events.FileCreatedEvent("foo.py")) - 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) + tctx.configure(sc) + sc.tick() + + rec = tctx.master.addons.get("recorder") + + assert rec.call_log[0][0:2] == ("recorder", "load") - sc.ns.call_log = [] + 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(".") + sc.tick() + for _ in range(3): + sc.last_load, sc.last_mtime = 0, 0 sc.tick() time.sleep(0.1) - if tctx.master.logs: - return - raise AssertionError("Change event not detected.") + tctx.master.has_log("Loading") 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) + tctx.configure(sc) + sc.tick() + f = tflow.tflow(resp=True) - sc.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.addons.trigger("request", f) + + tctx.master.has_log("ValueError: Error!") + tctx.master.has_log("error.py") def test_addon(self): with taddons.context() as tctx: @@ -175,11 +109,11 @@ class TestScript: "mitmproxy/data/addonscripts/addon.py" ) ) - l = addonmanager.Loader(tctx.master) - sc.load(l) + tctx.master.addons.add(sc) tctx.configure(sc) + sc.tick() assert sc.ns.event_log == [ - 'scriptload', 'addonload', 'addonconfigure' + 'scriptload', 'addonload', 'scriptconfigure', 'addonconfigure' ] @@ -194,49 +128,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 +170,76 @@ 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, ] ) + tctx.master.addons.invoke_addon(sc, "tick") debug = [i.msg for i in tctx.master.logs if i.level == "debug"] assert debug == [ 'a load', - 'a configure', 'a running', + 'a configure', + 'a tick', 'b load', - 'b configure', 'b running', + 'b configure', + 'b tick', 'c load', - 'c configure', 'c running', + 'c configure', + 'c tick', ] + tctx.master.logs = [] tctx.configure( sc, scripts = [ - "%s %s" % (rec, "c"), - "%s %s" % (rec, "a"), - "%s %s" % (rec, "b"), + "%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 == [] + assert debug == [ + 'c configure', + 'a configure', + 'b configure', + ] tctx.master.logs = [] tctx.configure( sc, scripts = [ - "%s %s" % (rec, "x"), - "%s %s" % (rec, "a"), + "%s/e.py" % rec, + "%s/a.py" % rec, ] ) + tctx.master.addons.invoke_addon(sc, "tick") + 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', + 'a configure', + 'e load', + 'e running', + 'e configure', + 'e tick', + 'a tick', ] diff --git a/test/mitmproxy/addons/test_serverplayback.py b/test/mitmproxy/addons/test_serverplayback.py index 7078b66e..29de48a0 100644 --- a/test/mitmproxy/addons/test_serverplayback.py +++ b/test/mitmproxy/addons/test_serverplayback.py @@ -6,7 +6,6 @@ from mitmproxy.test import tflow import mitmproxy.test.tutils from mitmproxy.addons import serverplayback -from mitmproxy import options from mitmproxy import exceptions from mitmproxy import io @@ -39,86 +38,88 @@ def test_tick(): def test_server_playback(): sp = serverplayback.ServerPlayback() - sp.configure(options.Options(), []) - f = tflow.tflow(resp=True) + with taddons.context() as tctx: + tctx.configure(sp) + f = tflow.tflow(resp=True) - assert not sp.flowmap + assert not sp.flowmap - sp.load_flows([f]) - assert sp.flowmap - assert sp.next_flow(f) - assert not sp.flowmap + sp.load_flows([f]) + assert sp.flowmap + assert sp.next_flow(f) + assert not sp.flowmap - sp.load_flows([f]) - assert sp.flowmap - sp.clear() - assert not sp.flowmap + sp.load_flows([f]) + assert sp.flowmap + sp.clear() + assert not sp.flowmap def test_ignore_host(): sp = serverplayback.ServerPlayback() - sp.configure(options.Options(server_replay_ignore_host=True), []) + with taddons.context() as tctx: + tctx.configure(sp, server_replay_ignore_host=True) - r = tflow.tflow(resp=True) - r2 = tflow.tflow(resp=True) + r = tflow.tflow(resp=True) + r2 = tflow.tflow(resp=True) - r.request.host = "address" - r2.request.host = "address" - assert sp._hash(r) == sp._hash(r2) - r2.request.host = "wrong_address" - assert sp._hash(r) == sp._hash(r2) + r.request.host = "address" + r2.request.host = "address" + assert sp._hash(r) == sp._hash(r2) + r2.request.host = "wrong_address" + assert sp._hash(r) == sp._hash(r2) def test_ignore_content(): s = serverplayback.ServerPlayback() - s.configure(options.Options(server_replay_ignore_content=False), []) + with taddons.context() as tctx: + tctx.configure(s, server_replay_ignore_content=False) - r = tflow.tflow(resp=True) - r2 = tflow.tflow(resp=True) + r = tflow.tflow(resp=True) + r2 = tflow.tflow(resp=True) - r.request.content = b"foo" - r2.request.content = b"foo" - assert s._hash(r) == s._hash(r2) - r2.request.content = b"bar" - assert not s._hash(r) == s._hash(r2) + r.request.content = b"foo" + r2.request.content = b"foo" + assert s._hash(r) == s._hash(r2) + r2.request.content = b"bar" + assert not s._hash(r) == s._hash(r2) - s.configure(options.Options(server_replay_ignore_content=True), []) - r = tflow.tflow(resp=True) - r2 = tflow.tflow(resp=True) - r.request.content = b"foo" - r2.request.content = b"foo" - assert s._hash(r) == s._hash(r2) - r2.request.content = b"bar" - assert s._hash(r) == s._hash(r2) - r2.request.content = b"" - assert s._hash(r) == s._hash(r2) - r2.request.content = None - assert s._hash(r) == s._hash(r2) + tctx.configure(s, server_replay_ignore_content=True) + r = tflow.tflow(resp=True) + r2 = tflow.tflow(resp=True) + r.request.content = b"foo" + r2.request.content = b"foo" + assert s._hash(r) == s._hash(r2) + r2.request.content = b"bar" + assert s._hash(r) == s._hash(r2) + r2.request.content = b"" + assert s._hash(r) == s._hash(r2) + r2.request.content = None + assert s._hash(r) == s._hash(r2) def test_ignore_content_wins_over_params(): s = serverplayback.ServerPlayback() - s.configure( - options.Options( + with taddons.context() as tctx: + tctx.configure( + s, server_replay_ignore_content=True, server_replay_ignore_payload_params=[ "param1", "param2" ] - ), - [] - ) - # NOTE: parameters are mutually exclusive in options + ) - r = tflow.tflow(resp=True) - r.request.headers["Content-Type"] = "application/x-www-form-urlencoded" - r.request.content = b"paramx=y" + # NOTE: parameters are mutually exclusive in options + r = tflow.tflow(resp=True) + r.request.headers["Content-Type"] = "application/x-www-form-urlencoded" + r.request.content = b"paramx=y" - r2 = tflow.tflow(resp=True) - r2.request.headers["Content-Type"] = "application/x-www-form-urlencoded" - r2.request.content = b"paramx=x" + r2 = tflow.tflow(resp=True) + r2.request.headers["Content-Type"] = "application/x-www-form-urlencoded" + r2.request.content = b"paramx=x" - # same parameters - assert s._hash(r) == s._hash(r2) + # same parameters + assert s._hash(r) == s._hash(r2) def test_ignore_payload_params_other_content_type(): @@ -147,136 +148,139 @@ def test_ignore_payload_params_other_content_type(): def test_hash(): s = serverplayback.ServerPlayback() - s.configure(options.Options(), []) + with taddons.context() as tctx: + tctx.configure(s) - r = tflow.tflow() - r2 = tflow.tflow() + r = tflow.tflow() + r2 = tflow.tflow() - assert s._hash(r) - assert s._hash(r) == s._hash(r2) - r.request.headers["foo"] = "bar" - assert s._hash(r) == s._hash(r2) - r.request.path = "voing" - assert s._hash(r) != s._hash(r2) + assert s._hash(r) + assert s._hash(r) == s._hash(r2) + r.request.headers["foo"] = "bar" + assert s._hash(r) == s._hash(r2) + r.request.path = "voing" + assert s._hash(r) != s._hash(r2) - r.request.path = "path?blank_value" - r2.request.path = "path?" - assert s._hash(r) != s._hash(r2) + r.request.path = "path?blank_value" + r2.request.path = "path?" + assert s._hash(r) != s._hash(r2) def test_headers(): s = serverplayback.ServerPlayback() - s.configure(options.Options(server_replay_use_headers=["foo"]), []) + with taddons.context() as tctx: + tctx.configure(s, server_replay_use_headers=["foo"]) - r = tflow.tflow(resp=True) - r.request.headers["foo"] = "bar" - r2 = tflow.tflow(resp=True) - assert not s._hash(r) == s._hash(r2) - r2.request.headers["foo"] = "bar" - assert s._hash(r) == s._hash(r2) - r2.request.headers["oink"] = "bar" - assert s._hash(r) == s._hash(r2) + r = tflow.tflow(resp=True) + r.request.headers["foo"] = "bar" + r2 = tflow.tflow(resp=True) + assert not s._hash(r) == s._hash(r2) + r2.request.headers["foo"] = "bar" + assert s._hash(r) == s._hash(r2) + r2.request.headers["oink"] = "bar" + assert s._hash(r) == s._hash(r2) - r = tflow.tflow(resp=True) - r2 = tflow.tflow(resp=True) - assert s._hash(r) == s._hash(r2) + r = tflow.tflow(resp=True) + r2 = tflow.tflow(resp=True) + assert s._hash(r) == s._hash(r2) def test_load(): s = serverplayback.ServerPlayback() - s.configure(options.Options(), []) + with taddons.context() as tctx: + tctx.configure(s) - r = tflow.tflow(resp=True) - r.request.headers["key"] = "one" + r = tflow.tflow(resp=True) + r.request.headers["key"] = "one" - r2 = tflow.tflow(resp=True) - r2.request.headers["key"] = "two" + r2 = tflow.tflow(resp=True) + r2.request.headers["key"] = "two" - s.load_flows([r, r2]) + s.load_flows([r, r2]) - assert s.count() == 2 + assert s.count() == 2 - n = s.next_flow(r) - assert n.request.headers["key"] == "one" - assert s.count() == 1 + n = s.next_flow(r) + assert n.request.headers["key"] == "one" + assert s.count() == 1 - n = s.next_flow(r) - assert n.request.headers["key"] == "two" - assert not s.flowmap - assert s.count() == 0 + n = s.next_flow(r) + assert n.request.headers["key"] == "two" + assert not s.flowmap + assert s.count() == 0 - assert not s.next_flow(r) + assert not s.next_flow(r) def test_load_with_server_replay_nopop(): s = serverplayback.ServerPlayback() - s.configure(options.Options(server_replay_nopop=True), []) + with taddons.context() as tctx: + tctx.configure(s, server_replay_nopop=True) - r = tflow.tflow(resp=True) - r.request.headers["key"] = "one" + r = tflow.tflow(resp=True) + r.request.headers["key"] = "one" - r2 = tflow.tflow(resp=True) - r2.request.headers["key"] = "two" + r2 = tflow.tflow(resp=True) + r2.request.headers["key"] = "two" - s.load_flows([r, r2]) + s.load_flows([r, r2]) - assert s.count() == 2 - s.next_flow(r) - assert s.count() == 2 + assert s.count() == 2 + s.next_flow(r) + assert s.count() == 2 def test_ignore_params(): s = serverplayback.ServerPlayback() - s.configure( - options.Options( + with taddons.context() as tctx: + tctx.configure( + s, server_replay_ignore_params=["param1", "param2"] - ), - [] - ) + ) - r = tflow.tflow(resp=True) - r.request.path = "/test?param1=1" - r2 = tflow.tflow(resp=True) - r2.request.path = "/test" - assert s._hash(r) == s._hash(r2) - r2.request.path = "/test?param1=2" - assert s._hash(r) == s._hash(r2) - r2.request.path = "/test?param2=1" - assert s._hash(r) == s._hash(r2) - r2.request.path = "/test?param3=2" - assert not s._hash(r) == s._hash(r2) + r = tflow.tflow(resp=True) + r.request.path = "/test?param1=1" + r2 = tflow.tflow(resp=True) + r2.request.path = "/test" + assert s._hash(r) == s._hash(r2) + r2.request.path = "/test?param1=2" + assert s._hash(r) == s._hash(r2) + r2.request.path = "/test?param2=1" + assert s._hash(r) == s._hash(r2) + r2.request.path = "/test?param3=2" + assert not s._hash(r) == s._hash(r2) def thash(r, r2, setter): s = serverplayback.ServerPlayback() - s.configure( - options.Options( + with taddons.context() as tctx: + s = serverplayback.ServerPlayback() + tctx.configure( + s, server_replay_ignore_payload_params=["param1", "param2"] - ), - [] - ) - - setter(r, paramx="x", param1="1") - - setter(r2, paramx="x", param1="1") - # same parameters - assert s._hash(r) == s._hash(r2) - # ignored parameters != - setter(r2, paramx="x", param1="2") - assert s._hash(r) == s._hash(r2) - # missing parameter - setter(r2, paramx="x") - assert s._hash(r) == s._hash(r2) - # ignorable parameter added - setter(r2, paramx="x", param1="2") - assert s._hash(r) == s._hash(r2) - # not ignorable parameter changed - setter(r2, paramx="y", param1="1") - assert not s._hash(r) == s._hash(r2) - # not ignorable parameter missing - setter(r2, param1="1") - r2.request.content = b"param1=1" - assert not s._hash(r) == s._hash(r2) + ) + + setter(r, paramx="x", param1="1") + + setter(r2, paramx="x", param1="1") + # same parameters + assert s._hash(r) == s._hash(r2) + # ignored parameters != + setter(r2, paramx="x", param1="2") + assert s._hash(r) == s._hash(r2) + # missing parameter + setter(r2, paramx="x") + assert s._hash(r) == s._hash(r2) + # ignorable parameter added + setter(r2, paramx="x", param1="2") + assert s._hash(r) == s._hash(r2) + # not ignorable parameter changed + setter(r2, paramx="y", param1="1") + assert not s._hash(r) == s._hash(r2) + # not ignorable parameter missing + setter(r2, param1="1") + r2.request.content = b"param1=1" + assert not s._hash(r) == s._hash(r2) def test_ignore_payload_params(): diff --git a/test/mitmproxy/addons/test_termstatus.py b/test/mitmproxy/addons/test_termstatus.py index 7becc857..2debaff5 100644 --- a/test/mitmproxy/addons/test_termstatus.py +++ b/test/mitmproxy/addons/test_termstatus.py @@ -5,6 +5,7 @@ from mitmproxy.test import taddons def test_configure(): ts = termstatus.TermStatus() with taddons.context() as ctx: + ctx.configure(ts, server=False) ts.running() assert not ctx.master.logs ctx.configure(ts, server=True) diff --git a/test/mitmproxy/data/addonscripts/addon.py b/test/mitmproxy/data/addonscripts/addon.py index beef2ce7..8c834d82 100644 --- a/test/mitmproxy/data/addonscripts/addon.py +++ b/test/mitmproxy/data/addonscripts/addon.py @@ -9,14 +9,16 @@ class Addon: def load(self, opts): event_log.append("addonload") - def configure(self, options, updated): + def configure(self, updated): event_log.append("addonconfigure") -def configure(options, updated): - event_log.append("addonconfigure") +def configure(updated): + event_log.append("scriptconfigure") 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 162c00f4..d1ab6c6c 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 8e6988d4..2a7d300c 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/concurrent_decorator_err.py b/test/mitmproxy/data/addonscripts/concurrent_decorator_err.py index 7bc28182..4f80e98a 100644 --- a/test/mitmproxy/data/addonscripts/concurrent_decorator_err.py +++ b/test/mitmproxy/data/addonscripts/concurrent_decorator_err.py @@ -2,5 +2,5 @@ from mitmproxy.script import concurrent @concurrent -def start(opts): +def load(v): pass diff --git a/test/mitmproxy/data/addonscripts/recorder/a.py b/test/mitmproxy/data/addonscripts/recorder/a.py new file mode 100644 index 00000000..df81d86b --- /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 00000000..ccbae705 --- /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 00000000..b8b0915e --- /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 00000000..eb5eff5e --- /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 index fe497b05..a962d3df 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 447b15a5..b4bb46bb 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.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 86efdfc2..ceff9fb9 100644 --- a/test/mitmproxy/script/test_concurrent.py +++ b/test/mitmproxy/script/test_concurrent.py @@ -2,10 +2,7 @@ from mitmproxy.test import tflow from mitmproxy.test import tutils from mitmproxy.test import taddons -from mitmproxy import addonmanager from mitmproxy import controller -from mitmproxy.addons import script - import time from .. import tservers @@ -20,14 +17,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) @@ -39,25 +33,20 @@ class TestConcurrent(tservers.MasterTest): def test_concurrent_err(self): with taddons.context() as tctx: - sc = script.Script( + tctx.script( tutils.test_data.path( "mitmproxy/data/addonscripts/concurrent_decorator_err.py" ) ) - l = addonmanager.Loader(tctx.master) - sc.load(l) assert tctx.master.has_log("decorator not supported") def test_concurrent_class(self): with taddons.context() as tctx: - sc = script.Script( + sc = tctx.script( tutils.test_data.path( "mitmproxy/data/addonscripts/concurrent_decorator_class.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 e556bf7d..cba40412 100644 --- a/test/mitmproxy/test_addonmanager.py +++ b/test/mitmproxy/test_addonmanager.py @@ -6,14 +6,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 @@ -35,19 +38,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)) @@ -71,10 +61,15 @@ 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_defaults(): @@ -82,33 +77,30 @@ def test_defaults(): 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") - - 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): + 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 + + 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(): @@ -119,29 +111,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() + assert a.get("four") - a.add(AChain("one", AChain("two", AChain("three", AChain("four", None))))) - assert not a.get("one") - assert not a.get("two") + 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 a.get("four") - a.clear() + 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 42371cfe..5a4c99fc 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 44b9ff3f..c87c9e83 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.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 b737b82a..b8005529 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.keys()) self.addons.trigger("running") def reset(self, addons): diff --git a/web/src/js/__tests__/ducks/flowsSpec.js b/web/src/js/__tests__/ducks/flowsSpec.js index acfa3083..2dfe046f 100644 --- a/web/src/js/__tests__/ducks/flowsSpec.js +++ b/web/src/js/__tests__/ducks/flowsSpec.js @@ -1,31 +1,225 @@ jest.unmock('../../ducks/flows'); +jest.mock('../../utils') -import reduceFlows, * as flowActions from '../../ducks/flows' -import * as storeActions from '../../ducks/utils/store' +import reduceFlows from "../../ducks/flows" +import * as flowActions from "../../ducks/flows" +import reduceStore from "../../ducks/utils/store" +import {fetchApi} from "../../utils" +import {createStore} from "./tutils" - -describe('select flow', () => { - - let state = reduceFlows(undefined, {}) +describe('flow reducer', () => { + let state = undefined for (let i of [1, 2, 3, 4]) { - state = reduceFlows(state, storeActions.add({ id: i })) + state = reduceFlows(state, { type: flowActions.ADD, data: { id: i }, cmd: 'add' }) } - it('should be possible to select a single flow', () => { - expect(reduceFlows(state, flowActions.select(2))).toEqual( - { - ...state, - selected: [2], - } - ) - }) - - it('should be possible to deselect a flow', () => { - expect(reduceFlows({ ...state, selected: [1] }, flowActions.select())).toEqual( - { - ...state, - selected: [], - } - ) + it('should return initial state', () => { + expect(reduceFlows(undefined, {})).toEqual({ + highlight: null, + filter: null, + sort: { column: null, desc: false }, + selected: [], + ...reduceStore(undefined, {}) + }) + }) + + describe('selections', () => { + it('should be possible to select a single flow', () => { + expect(reduceFlows(state, flowActions.select(2))).toEqual( + { + ...state, + selected: [2], + } + ) + }) + + it('should be possible to deselect a flow', () => { + expect(reduceFlows({ ...state, selected: [1] }, flowActions.select())).toEqual( + { + ...state, + selected: [], + } + ) + }) + + it('should be possible to select relative', () => { + // haven't selected any flow + expect( + flowActions.selectRelative(state, 1) + ).toEqual( + flowActions.select(4) + ) + + // already selected some flows + expect( + flowActions.selectRelative({ ...state, selected: [2] }, 1) + ).toEqual( + flowActions.select(3) + ) + }) + + it('should update state.selected on remove', () => { + let next + next = reduceFlows({ ...state, selected: [2] }, { + type: flowActions.REMOVE, + data: 2, + cmd: 'remove' + }) + expect(next.selected).toEqual([3]) + + //last row + next = reduceFlows({ ...state, selected: [4] }, { + type: flowActions.REMOVE, + data: 4, + cmd: 'remove' + }) + expect(next.selected).toEqual([3]) + + //multiple selection + next = reduceFlows({ ...state, selected: [2, 3, 4] }, { + type: flowActions.REMOVE, + data: 3, + cmd: 'remove' + }) + expect(next.selected).toEqual([2, 4]) + }) + }) + + it('should be possible to set filter', () => { + let filt = "~u 123" + expect(reduceFlows(undefined, flowActions.setFilter(filt)).filter).toEqual(filt) + }) + + it('should be possible to set highlight', () => { + let key = "foo" + expect(reduceFlows(undefined, flowActions.setHighlight(key)).highlight).toEqual(key) + }) + + it('should be possible to set sort', () => { + let sort = { column: "TLSColumn", desc: 1 } + expect(reduceFlows(undefined, flowActions.setSort(sort.column, sort.desc)).sort).toEqual(sort) + }) + +}) + +describe('flows actions', () => { + + let store = createStore({reduceFlows}) + + let tflow = { id: 1 } + it('should handle resume action', () => { + store.dispatch(flowActions.resume(tflow)) + expect(fetchApi).toBeCalledWith('/flows/1/resume', { method: 'POST' }) + }) + + it('should handle resumeAll action', () => { + store.dispatch(flowActions.resumeAll()) + expect(fetchApi).toBeCalledWith('/flows/resume', { method: 'POST' }) + }) + + it('should handle kill action', () => { + store.dispatch(flowActions.kill(tflow)) + expect(fetchApi).toBeCalledWith('/flows/1/kill', { method: 'POST' }) + + }) + + it('should handle killAll action', () => { + store.dispatch(flowActions.killAll()) + expect(fetchApi).toBeCalledWith('/flows/kill', { method: 'POST' }) + }) + + it('should handle remove action', () => { + store.dispatch(flowActions.remove(tflow)) + expect(fetchApi).toBeCalledWith('/flows/1', { method: 'DELETE' }) + }) + + it('should handle duplicate action', () => { + store.dispatch(flowActions.duplicate(tflow)) + expect(fetchApi).toBeCalledWith('/flows/1/duplicate', { method: 'POST' }) + }) + + it('should handle replay action', () => { + store.dispatch(flowActions.replay(tflow)) + expect(fetchApi).toBeCalledWith('/flows/1/replay', { method: 'POST' }) + }) + + it('should handle revert action', () => { + store.dispatch(flowActions.revert(tflow)) + expect(fetchApi).toBeCalledWith('/flows/1/revert', { method: 'POST' }) + }) + + it('should handle update action', () => { + store.dispatch(flowActions.update(tflow, 'foo')) + expect(fetchApi.put).toBeCalledWith('/flows/1', 'foo') + }) + + it('should handle uploadContent action', () => { + let body = new FormData(), + file = new window.Blob(['foo'], { type: 'plain/text' }) + body.append('file', file) + store.dispatch(flowActions.uploadContent(tflow, 'foo', 'foo')) + expect(fetchApi).toBeCalledWith('/flows/1/foo/content', { method: 'POST', body}) + }) + + it('should handle clear action', () => { + store.dispatch(flowActions.clear()) + expect(fetchApi).toBeCalledWith('/clear', { method: 'POST'} ) + }) + + it('should handle download action', () => { + let state = reduceFlows(undefined, {}) + expect(reduceFlows(state, flowActions.download())).toEqual(state) + }) + + it('should handle upload action', () => { + let body = new FormData() + body.append('file', 'foo') + store.dispatch(flowActions.upload('foo')) + expect(fetchApi).toBeCalledWith('/flows/dump', { method: 'POST', body }) + }) +}) + +describe('makeSort', () => { + it('should be possible to sort by TLSColumn', () => { + let sort = flowActions.makeSort({ column: 'TLSColumn', desc: true }), + a = { request: { scheme: 'http' } }, + b = { request: { scheme: 'https' } } + expect(sort(a, b)).toEqual(1) + }) + + it('should be possible to sort by PathColumn', () => { + let sort = flowActions.makeSort({ column: 'PathColumn', desc: true }), + a = { request: {} }, + b = { request: {} } + expect(sort(a, b)).toEqual(0) + + }) + + it('should be possible to sort by MethodColumn', () => { + let sort = flowActions.makeSort({ column: 'MethodColumn', desc: true }), + a = { request: { method: 'GET' } }, + b = { request: { method: 'POST' } } + expect(sort(b, a)).toEqual(-1) + }) + + it('should be possible to sort by StatusColumn', () => { + let sort = flowActions.makeSort({ column: 'StatusColumn', desc: false }), + a = { response: { status_code: 200 } }, + b = { response: { status_code: 404 } } + expect(sort(a, b)).toEqual(-1) + }) + + it('should be possible to sort by TimeColumn', () => { + let sort = flowActions.makeSort({ column: 'TimeColumn', desc: false }), + a = { response: { timestamp_end: 9 }, request: { timestamp_start: 8 } }, + b = { response: { timestamp_end: 10 }, request: { timestamp_start: 8 } } + expect(sort(b, a)).toEqual(1) + }) + + it('should be possible to sort by SizeColumn', () => { + let sort = flowActions.makeSort({ column: 'SizeColumn', desc: true }), + a = { request: { contentLength: 1 }, response: { contentLength: 1 } }, + b = { request: { contentLength: 1 } } + expect(sort(a, b)).toEqual(-1) }) }) diff --git a/web/src/js/ducks/flows.js b/web/src/js/ducks/flows.js index 92408891..523ec396 100644 --- a/web/src/js/ducks/flows.js +++ b/web/src/js/ducks/flows.js @@ -1,5 +1,6 @@ import { fetchApi } from "../utils" -import reduceStore, * as storeActions from "./utils/store" +import reduceStore from "./utils/store" +import * as storeActions from "./utils/store" import Filt from "../filt/filt" import { RequestUtils } from "../flow/utils" @@ -29,8 +30,6 @@ export default function reduce(state = defaultState, action) { case UPDATE: case REMOVE: case RECEIVE: - // FIXME: Update state.selected on REMOVE: - // The selected flow may have been removed, we need to select the next one in the view. let storeAction = storeActions[action.cmd]( action.data, makeFilter(state.filter), @@ -152,22 +151,20 @@ export function setSort(column, desc) { return { type: SET_SORT, sort: { column, desc } } } -export function selectRelative(shift) { - return (dispatch, getState) => { - let currentSelectionIndex = getState().flows.viewIndex[getState().flows.selected[0]] - let minIndex = 0 - let maxIndex = getState().flows.view.length - 1 - let newIndex - if (currentSelectionIndex === undefined) { - newIndex = (shift < 0) ? minIndex : maxIndex - } else { - newIndex = currentSelectionIndex + shift - newIndex = window.Math.max(newIndex, minIndex) - newIndex = window.Math.min(newIndex, maxIndex) - } - let flow = getState().flows.view[newIndex] - dispatch(select(flow ? flow.id : undefined)) +export function selectRelative(flows, shift) { + let currentSelectionIndex = flows.viewIndex[flows.selected[0]] + let minIndex = 0 + let maxIndex = flows.view.length - 1 + let newIndex + if (currentSelectionIndex === undefined) { + newIndex = (shift < 0) ? minIndex : maxIndex + } else { + newIndex = currentSelectionIndex + shift + newIndex = window.Math.max(newIndex, minIndex) + newIndex = window.Math.min(newIndex, maxIndex) } + let flow = flows.view[newIndex] + return select(flow ? flow.id : undefined) } @@ -212,7 +209,7 @@ export function uploadContent(flow, file, type) { const body = new FormData() file = new window.Blob([file], { type: 'plain/text' }) body.append('file', file) - return dispatch => fetchApi(`/flows/${flow.id}/${type}/content`, { method: 'post', body }) + return dispatch => fetchApi(`/flows/${flow.id}/${type}/content`, { method: 'POST', body }) } @@ -228,7 +225,7 @@ export function download() { export function upload(file) { const body = new FormData() body.append('file', file) - return dispatch => fetchApi('/flows/dump', { method: 'post', body }) + return dispatch => fetchApi('/flows/dump', { method: 'POST', body }) } diff --git a/web/src/js/ducks/ui/keyboard.js b/web/src/js/ducks/ui/keyboard.js index 30fd76e1..0e3491fa 100644 --- a/web/src/js/ducks/ui/keyboard.js +++ b/web/src/js/ducks/ui/keyboard.js @@ -9,39 +9,40 @@ export function onKeyDown(e) { return () => { } } - var key = e.keyCode - var shiftKey = e.shiftKey + let key = e.keyCode, + shiftKey = e.shiftKey e.preventDefault() return (dispatch, getState) => { - const flow = getState().flows.byId[getState().flows.selected[0]] + const flows = getState().flows, + flow = flows.byId[getState().flows.selected[0]] switch (key) { case Key.K: case Key.UP: - dispatch(flowsActions.selectRelative(-1)) + dispatch(flowsActions.selectRelative(flows, -1)) break case Key.J: case Key.DOWN: - dispatch(flowsActions.selectRelative(+1)) + dispatch(flowsActions.selectRelative(flows, +1)) break case Key.SPACE: case Key.PAGE_DOWN: - dispatch(flowsActions.selectRelative(+10)) + dispatch(flowsActions.selectRelative(flows, +10)) break case Key.PAGE_UP: - dispatch(flowsActions.selectRelative(-10)) + dispatch(flowsActions.selectRelative(flows, -10)) break case Key.END: - dispatch(flowsActions.selectRelative(+1e10)) + dispatch(flowsActions.selectRelative(flows, +1e10)) break case Key.HOME: - dispatch(flowsActions.selectRelative(-1e10)) + dispatch(flowsActions.selectRelative(flows, -1e10)) break case Key.ESC: |