From 914659e888792441cafb6373da53d7e15456f32c Mon Sep 17 00:00:00 2001 From: Aldo Cortesi Date: Wed, 1 Feb 2017 12:19:51 +1300 Subject: [PATCH] Regularise setheaders options As per replacements: - Make the option type a string/tuple union - Localise parsing strictly within the addon - Adapt the console editor (we'll find a more elegant solution later) --- mitmproxy/addons/setheaders.py | 44 ++++++++++++++- mitmproxy/options.py | 2 +- mitmproxy/tools/cmdline.py | 68 ++---------------------- mitmproxy/tools/console/options.py | 9 +++- test/mitmproxy/addons/test_setheaders.py | 25 ++++++--- test/mitmproxy/test_cmdline.py | 18 ------- 6 files changed, 73 insertions(+), 93 deletions(-) diff --git a/mitmproxy/addons/setheaders.py b/mitmproxy/addons/setheaders.py index 601e75216..95cf9a09d 100644 --- a/mitmproxy/addons/setheaders.py +++ b/mitmproxy/addons/setheaders.py @@ -2,6 +2,43 @@ from mitmproxy import exceptions from mitmproxy import flowfilter +def parse_setheader(s): + """ + Returns a (pattern, regex, replacement) tuple. + + The general form for a replacement hook is as follows: + + /patt/regex/replacement + + The first character specifies the separator. Example: + + :~q:foo:bar + + If only two clauses are specified, the pattern is set to match + universally (i.e. ".*"). Example: + + /foo/bar/ + + Clauses are parsed from left to right. Extra separators are taken to be + part of the final clause. For instance, the replacement clause below is + "foo/bar/": + + /one/two/foo/bar/ + """ + sep, rem = s[0], s[1:] + parts = rem.split(sep, 2) + if len(parts) == 2: + patt = ".*" + a, b = parts + elif len(parts) == 3: + patt, a, b = parts + else: + raise exceptions.OptionsError( + "Invalid replacement specifier: %s" % s + ) + return patt, a, b + + class SetHeaders: def __init__(self): self.lst = [] @@ -16,7 +53,12 @@ class SetHeaders: """ if "setheaders" in updated: self.lst = [] - for fpatt, header, value in options.setheaders: + for shead in options.setheaders: + if isinstance(shead, str): + fpatt, header, value = parse_setheader(shead) + else: + fpatt, header, value = shead + flt = flowfilter.parse(fpatt) if not flt: raise exceptions.OptionsError( diff --git a/mitmproxy/options.py b/mitmproxy/options.py index 3b64cc6a4..630d29645 100644 --- a/mitmproxy/options.py +++ b/mitmproxy/options.py @@ -41,7 +41,7 @@ class Options(optmanager.OptManager): replacements: Sequence[Union[Tuple[str, str, str], str]] = [], replacement_files: Sequence[Union[Tuple[str, str, str], str]] = [], server_replay_use_headers: Sequence[str] = [], - setheaders: Sequence[Tuple[str, str, str]] = [], + setheaders: Sequence[Union[Tuple[str, str, str], str]] = [], server_replay: Sequence[str] = [], stickycookie: Optional[str] = None, stickyauth: Optional[str] = None, diff --git a/mitmproxy/tools/cmdline.py b/mitmproxy/tools/cmdline.py index bb11b9c23..1c620fd60 100644 --- a/mitmproxy/tools/cmdline.py +++ b/mitmproxy/tools/cmdline.py @@ -2,7 +2,6 @@ import argparse import os from mitmproxy import exceptions -from mitmproxy import flowfilter from mitmproxy import options from mitmproxy import platform from mitmproxy.utils import human @@ -18,57 +17,6 @@ class ParseException(Exception): pass -def _parse_hook(s): - sep, rem = s[0], s[1:] - parts = rem.split(sep, 2) - if len(parts) == 2: - patt = ".*" - a, b = parts - elif len(parts) == 3: - patt, a, b = parts - else: - raise ParseException( - "Malformed hook specifier - too few clauses: %s" % s - ) - - if not a: - raise ParseException("Empty clause: %s" % str(patt)) - - if not flowfilter.parse(patt): - raise ParseException("Malformed filter pattern: %s" % patt) - - return patt, a, b - - -def parse_setheader(s): - """ - Returns a (pattern, header, value) tuple. - - The general form for a replacement hook is as follows: - - /patt/header/value - - The first character specifies the separator. Example: - - :~q:foo:bar - - If only two clauses are specified, the pattern is set to match - universally (i.e. ".*"). Example: - - /foo/bar/ - - Clauses are parsed from left to right. Extra separators are taken to be - part of the final clause. For instance, the value clause below is - "foo/bar/": - - /one/two/foo/bar/ - - Checks that pattern and regex are both well-formed. Raises - ParseException on error. - """ - return _parse_hook(s) - - def get_common_options(args): stickycookie, stickyauth = None, None if args.stickycookie_filt: @@ -81,14 +29,6 @@ def get_common_options(args): if stream_large_bodies: stream_large_bodies = human.parse_size(stream_large_bodies) - setheaders = [] - for i in args.setheader or []: - try: - p = parse_setheader(i) - except ParseException as e: - raise exceptions.OptionsError(e) - setheaders.append(p) - if args.streamfile and args.streamfile[0] == args.rfile: if args.streamfile[1] == "wb": raise exceptions.OptionsError( @@ -171,7 +111,7 @@ def get_common_options(args): rfile=args.rfile, replacements=args.replacements, replacement_files=args.replacement_files, - setheaders=setheaders, + setheaders=args.setheaders, server_replay=args.server_replay, scripts=args.scripts, stickycookie=stickycookie, @@ -648,7 +588,7 @@ def set_headers(parser): ) group.add_argument( "--setheader", - action="append", type=str, dest="setheader", + action="append", type=str, dest="setheaders", metavar="PATTERN", help="Header set pattern." ) @@ -708,8 +648,8 @@ def common_options(parser): def mitmproxy(): - # Don't import mitmproxy.tools.console for mitmdump, urwid is not available on all - # platforms. + # Don't import mitmproxy.tools.console for mitmdump, urwid is not available + # on all platforms. from .console import palettes parser = argparse.ArgumentParser(usage="%(prog)s [options]") diff --git a/mitmproxy/tools/console/options.py b/mitmproxy/tools/console/options.py index e88006fef..4115bd180 100644 --- a/mitmproxy/tools/console/options.py +++ b/mitmproxy/tools/console/options.py @@ -7,6 +7,7 @@ from mitmproxy.tools.console import select from mitmproxy.tools.console import signals from mitmproxy.addons import replace +from mitmproxy.addons import setheaders footer = [ ('heading_key', "enter/space"), ":toggle ", @@ -190,10 +191,16 @@ class Options(urwid.WidgetWrap): ) def setheaders(self): + data = [] + for d in self.master.options.setheaders: + if isinstance(d, str): + data.append(setheaders.parse_setheader(d)) + else: + data.append(d) self.master.view_grideditor( grideditor.SetHeadersEditor( self.master, - self.master.options.setheaders, + data, self.master.options.setter("setheaders") ) ) diff --git a/test/mitmproxy/addons/test_setheaders.py b/test/mitmproxy/addons/test_setheaders.py index 34395ddfc..a721c180a 100644 --- a/test/mitmproxy/addons/test_setheaders.py +++ b/test/mitmproxy/addons/test_setheaders.py @@ -3,19 +3,28 @@ from mitmproxy.test import tutils from mitmproxy.test import taddons from mitmproxy.addons import setheaders -from mitmproxy import options class TestSetHeaders: + def test_parse_setheaders(self): + x = setheaders.parse_setheader("/foo/bar/voing") + assert x == ("foo", "bar", "voing") + x = setheaders.parse_setheader("/foo/bar/vo/ing/") + assert x == ("foo", "bar", "vo/ing/") + x = setheaders.parse_setheader("/bar/voing") + assert x == (".*", "bar", "voing") + tutils.raises("invalid replacement", setheaders.parse_setheader, "/") + def test_configure(self): sh = setheaders.SetHeaders() - o = options.Options( - setheaders = [("~b", "one", "two")] - ) - tutils.raises( - "invalid setheader filter pattern", - sh.configure, o, o.keys() - ) + with taddons.context() as tctx: + tutils.raises( + "invalid setheader filter pattern", + tctx.configure, + sh, + setheaders = [("~b", "one", "two")] + ) + tctx.configure(sh, setheaders = ["/foo/bar/voing"]) def test_setheaders(self): sh = setheaders.SetHeaders() diff --git a/test/mitmproxy/test_cmdline.py b/test/mitmproxy/test_cmdline.py index fe0373d12..96d5ae312 100644 --- a/test/mitmproxy/test_cmdline.py +++ b/test/mitmproxy/test_cmdline.py @@ -1,11 +1,5 @@ import argparse from mitmproxy.tools import cmdline -from mitmproxy.test import tutils - - -def test_parse_setheaders(): - x = cmdline.parse_setheader("/foo/bar/voing") - assert x == ("foo", "bar", "voing") def test_common(): @@ -21,18 +15,6 @@ def test_common(): assert v["stickycookie"] == "foo" assert v["stickyauth"] == "foo" - opts.setheader = ["/foo/bar/voing"] - v = cmdline.get_common_options(opts) - assert v["setheaders"] == [("foo", "bar", "voing")] - - opts.setheader = ["//"] - tutils.raises( - "empty clause", - cmdline.get_common_options, - opts - ) - opts.setheader = [] - def test_mitmproxy(): ap = cmdline.mitmproxy()