diff options
| -rw-r--r-- | mitmproxy/addons/__init__.py | 1 | ||||
| -rw-r--r-- | mitmproxy/addons/replace.py | 72 | ||||
| -rw-r--r-- | mitmproxy/options.py | 6 | ||||
| -rw-r--r-- | mitmproxy/tools/cmdline.py | 62 | ||||
| -rw-r--r-- | mitmproxy/tools/console/options.py | 10 | ||||
| -rw-r--r-- | test/mitmproxy/addons/test_replace.py | 113 | ||||
| -rw-r--r-- | test/mitmproxy/test_cmdline.py | 64 |
7 files changed, 159 insertions, 169 deletions
diff --git a/mitmproxy/addons/__init__.py b/mitmproxy/addons/__init__.py index 2e1d1c67..97fa2dcd 100644 --- a/mitmproxy/addons/__init__.py +++ b/mitmproxy/addons/__init__.py @@ -28,6 +28,7 @@ def default_addons(): onboarding.Onboarding(), proxyauth.ProxyAuth(), replace.Replace(), + replace.ReplaceFile(), script.ScriptLoader(), serverplayback.ServerPlayback(), setheaders.SetHeaders(), diff --git a/mitmproxy/addons/replace.py b/mitmproxy/addons/replace.py index 09200d5d..34bb40c2 100644 --- a/mitmproxy/addons/replace.py +++ b/mitmproxy/addons/replace.py @@ -2,9 +2,47 @@ import re from mitmproxy import exceptions from mitmproxy import flowfilter +from mitmproxy import ctx -class Replace: +def parse_hook(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 _ReplaceBase: def __init__(self): self.lst = [] @@ -16,9 +54,14 @@ class Replace: rex: a regular expression, as bytes. s: the replacement string, as bytes """ - if "replacements" in updated: + if self.optionName in updated: lst = [] - for fpatt, rex, s in options.replacements: + for rep in getattr(options, self.optionName): + if isinstance(rep, str): + fpatt, rex, s = parse_hook(rep) + else: + fpatt, rex, s = rep + flt = flowfilter.parse(fpatt) if not flt: raise exceptions.OptionsError( @@ -37,9 +80,9 @@ class Replace: for rex, s, flt in self.lst: if flt(f): if f.response: - f.response.replace(rex, s, flags=re.DOTALL) + self.replace(f.response, rex, s) else: - f.request.replace(rex, s, flags=re.DOTALL) + self.replace(f.request, rex, s) def request(self, flow): if not flow.reply.has_message: @@ -48,3 +91,22 @@ class Replace: def response(self, flow): if not flow.reply.has_message: self.execute(flow) + + +class Replace(_ReplaceBase): + optionName = "replacements" + + def replace(self, obj, rex, s): + obj.replace(rex, s, flags=re.DOTALL) + + +class ReplaceFile(_ReplaceBase): + optionName = "replacement_files" + + def replace(self, obj, rex, s): + try: + v = open(s, "rb").read() + except IOError as e: + ctx.log.warn("Could not read replacement file: %s" % s) + return + obj.replace(rex, v, flags=re.DOTALL) diff --git a/mitmproxy/options.py b/mitmproxy/options.py index db7bd437..3b64cc6a 100644 --- a/mitmproxy/options.py +++ b/mitmproxy/options.py @@ -1,4 +1,4 @@ -from typing import Tuple, Optional, Sequence +from typing import Tuple, Optional, Sequence, Union from mitmproxy import optmanager @@ -38,7 +38,8 @@ class Options(optmanager.OptManager): rfile: Optional[str] = None, scripts: Sequence[str] = [], showhost: bool = False, - replacements: Sequence[Tuple[str, str, str]] = [], + 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]] = [], server_replay: Sequence[str] = [], @@ -124,6 +125,7 @@ class Options(optmanager.OptManager): self.scripts = scripts self.showhost = showhost self.replacements = replacements + self.replacement_files = replacement_files self.server_replay_use_headers = server_replay_use_headers self.setheaders = setheaders self.server_replay = server_replay diff --git a/mitmproxy/tools/cmdline.py b/mitmproxy/tools/cmdline.py index 1d5bd017..bb11b9c2 100644 --- a/mitmproxy/tools/cmdline.py +++ b/mitmproxy/tools/cmdline.py @@ -1,5 +1,4 @@ import argparse -import re import os from mitmproxy import exceptions @@ -41,40 +40,6 @@ def _parse_hook(s): return patt, a, b -def parse_replace_hook(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/ - - Checks that pattern and regex are both well-formed. Raises - ParseException on error. - """ - patt, regex, replacement = _parse_hook(s) - try: - re.compile(regex) - except re.error as e: - raise ParseException("Malformed replacement regex: %s" % str(e)) - return patt, regex, replacement - - def parse_setheader(s): """ Returns a (pattern, header, value) tuple. @@ -116,26 +81,6 @@ def get_common_options(args): if stream_large_bodies: stream_large_bodies = human.parse_size(stream_large_bodies) - reps = [] - for i in args.replace or []: - try: - p = parse_replace_hook(i) - except ParseException as e: - raise exceptions.OptionsError(e) - reps.append(p) - for i in args.replace_file or []: - try: - patt, rex, path = parse_replace_hook(i) - except ParseException as e: - raise exceptions.OptionsError(e) - try: - v = open(path, "rb").read() - except IOError as e: - raise exceptions.OptionsError( - "Could not read replace file: %s" % path - ) - reps.append((patt, rex, v)) - setheaders = [] for i in args.setheader or []: try: @@ -224,7 +169,8 @@ def get_common_options(args): refresh_server_playback=not args.norefresh, server_replay_use_headers=args.server_replay_use_headers, rfile=args.rfile, - replacements=reps, + replacements=args.replacements, + replacement_files=args.replacement_files, setheaders=setheaders, server_replay=args.server_replay, scripts=args.scripts, @@ -676,13 +622,13 @@ def replacements(parser): ) group.add_argument( "--replace", - action="append", type=str, dest="replace", + action="append", type=str, dest="replacements", metavar="PATTERN", help="Replacement pattern." ) group.add_argument( "--replace-from-file", - action="append", type=str, dest="replace_file", + action="append", type=str, dest="replacement_files", metavar="PATH", help=""" Replacement pattern, where the replacement clause is a path to a diff --git a/mitmproxy/tools/console/options.py b/mitmproxy/tools/console/options.py index 8c953e8e..e88006fe 100644 --- a/mitmproxy/tools/console/options.py +++ b/mitmproxy/tools/console/options.py @@ -6,6 +6,8 @@ from mitmproxy.tools.console import grideditor from mitmproxy.tools.console import select from mitmproxy.tools.console import signals +from mitmproxy.addons import replace + footer = [ ('heading_key', "enter/space"), ":toggle ", ('heading_key', "C"), ":clear all ", @@ -215,10 +217,16 @@ class Options(urwid.WidgetWrap): ) def replacepatterns(self): + data = [] + for d in self.master.options.replacements: + if isinstance(d, str): + data.append(replace.parse_hook(d)) + else: + data.append(d) self.master.view_grideditor( grideditor.ReplaceEditor( self.master, - self.master.options.replacements, + data, self.master.options.setter("replacements") ) ) diff --git a/test/mitmproxy/addons/test_replace.py b/test/mitmproxy/addons/test_replace.py index 34fa5017..ec38b77d 100644 --- a/test/mitmproxy/addons/test_replace.py +++ b/test/mitmproxy/addons/test_replace.py @@ -1,57 +1,59 @@ +import os.path from mitmproxy.test import tflow from mitmproxy.test import tutils from .. import tservers from mitmproxy.addons import replace -from mitmproxy import master -from mitmproxy import options -from mitmproxy import proxy +from mitmproxy.test import taddons class TestReplace: + def test_parse_hook(self): + x = replace.parse_hook("/foo/bar/voing") + assert x == ("foo", "bar", "voing") + x = replace.parse_hook("/foo/bar/vo/ing/") + assert x == ("foo", "bar", "vo/ing/") + x = replace.parse_hook("/bar/voing") + assert x == (".*", "bar", "voing") + tutils.raises("invalid replacement", replace.parse_hook, "/") + def test_configure(self): r = replace.Replace() - updated = set(["replacements"]) - r.configure(options.Options( - replacements=[("one", "two", "three")] - ), updated) - tutils.raises( - "invalid filter pattern", - r.configure, - options.Options( + with taddons.context() as tctx: + tctx.configure(r, replacements=[("one", "two", "three")]) + tutils.raises( + "invalid filter pattern", + tctx.configure, + r, replacements=[("~b", "two", "three")] - ), - updated - ) - tutils.raises( - "invalid regular expression", - r.configure, - options.Options( + ) + tutils.raises( + "invalid regular expression", + tctx.configure, + r, replacements=[("foo", "+", "three")] - ), - updated - ) + ) + tctx.configure(r, replacements=["/a/b/c/"]) def test_simple(self): - o = options.Options( - replacements = [ - ("~q", "foo", "bar"), - ("~s", "foo", "bar"), - ] - ) - m = master.Master(o, proxy.DummyServer()) - sa = replace.Replace() - m.addons.add(sa) - - f = tflow.tflow() - f.request.content = b"foo" - m.request(f) - assert f.request.content == b"bar" + r = replace.Replace() + with taddons.context() as tctx: + tctx.configure( + r, + replacements = [ + ("~q", "foo", "bar"), + ("~s", "foo", "bar"), + ] + ) + f = tflow.tflow() + f.request.content = b"foo" + r.request(f) + assert f.request.content == b"bar" - f = tflow.tflow(resp=True) - f.response.content = b"foo" - m.response(f) - assert f.response.content == b"bar" + f = tflow.tflow(resp=True) + f.response.content = b"foo" + r.response(f) + assert f.response.content == b"bar" class TestUpstreamProxy(tservers.HTTPUpstreamProxyTest): @@ -72,3 +74,36 @@ class TestUpstreamProxy(tservers.HTTPUpstreamProxyTest): req = p.request("get:'%s/p/418:b\"foo\"'" % self.server.urlbase) assert req.content == b"ORLY" assert req.status_code == 418 + + +class TestReplaceFile: + def test_simple(self): + r = replace.ReplaceFile() + with tutils.tmpdir() as td: + rp = os.path.join(td, "replacement") + with open(rp, "w") as f: + f.write("bar") + with taddons.context() as tctx: + tctx.configure( + r, + replacement_files = [ + ("~q", "foo", rp), + ("~s", "foo", rp), + ("~b nonexistent", "nonexistent", "nonexistent"), + ] + ) + f = tflow.tflow() + f.request.content = b"foo" + r.request(f) + assert f.request.content == b"bar" + + f = tflow.tflow(resp=True) + f.response.content = b"foo" + r.response(f) + assert f.response.content == b"bar" + + f = tflow.tflow() + f.request.content = b"nonexistent" + assert not tctx.master.event_log + r.request(f) + assert tctx.master.event_log diff --git a/test/mitmproxy/test_cmdline.py b/test/mitmproxy/test_cmdline.py index d2e0c8a5..fe0373d1 100644 --- a/test/mitmproxy/test_cmdline.py +++ b/test/mitmproxy/test_cmdline.py @@ -3,38 +3,6 @@ from mitmproxy.tools import cmdline from mitmproxy.test import tutils -def test_parse_replace_hook(): - x = cmdline.parse_replace_hook("/foo/bar/voing") - assert x == ("foo", "bar", "voing") - - x = cmdline.parse_replace_hook("/foo/bar/vo/ing/") - assert x == ("foo", "bar", "vo/ing/") - - x = cmdline.parse_replace_hook("/bar/voing") - assert x == (".*", "bar", "voing") - - tutils.raises( - cmdline.ParseException, - cmdline.parse_replace_hook, - "/foo" - ) - tutils.raises( - "replacement regex", - cmdline.parse_replace_hook, - "patt/[/rep" - ) - tutils.raises( - "filter pattern", - cmdline.parse_replace_hook, - "/~/foo/rep" - ) - tutils.raises( - "empty clause", - cmdline.parse_replace_hook, - "//" - ) - - def test_parse_setheaders(): x = cmdline.parse_setheader("/foo/bar/voing") assert x == ("foo", "bar", "voing") @@ -65,38 +33,6 @@ def test_common(): ) opts.setheader = [] - opts.replace = ["/foo/bar/voing"] - v = cmdline.get_common_options(opts) - assert v["replacements"] == [("foo", "bar", "voing")] - - opts.replace = ["//"] - tutils.raises( - "empty clause", - cmdline.get_common_options, - opts - ) - - opts.replace = [] - opts.replace_file = [("/foo/bar/nonexistent")] - tutils.raises( - "could not read replace file", - cmdline.get_common_options, - opts - ) - - opts.replace_file = [("/~/bar/nonexistent")] - tutils.raises( - "filter pattern", - cmdline.get_common_options, - opts - ) - - p = tutils.test_data.path("mitmproxy/data/replace") - opts.replace_file = [("/foo/bar/%s" % p)] - v = cmdline.get_common_options(opts)["replacements"] - assert len(v) == 1 - assert v[0][2].strip() == b"replacecontents" - def test_mitmproxy(): ap = cmdline.mitmproxy() |