Revamp replacement hooks

- Replacement specifiers can be either strings or tuples. This lets us cope
gracefully with command-line parsing (and posible quick interactive
specification) without having to special-case replacement hooks, or have
knowledge of hook specification leak outside the addon. We can also now use the
same command-line spec format in config files.
- Split replacement and replacement from file into separate addons and options.
Files are now read on each replacement, so you can edit replacement files in
place without restart.
- Modernise the test suite to use addon test helpers.

TODO: editing and displaying replace-from-file in console app
This commit is contained in:
Aldo Cortesi 2017-02-01 11:01:47 +13:00
parent 02f51d043d
commit 4cc75a9560
7 changed files with 159 additions and 169 deletions

View File

@ -28,6 +28,7 @@ def default_addons():
onboarding.Onboarding(),
proxyauth.ProxyAuth(),
replace.Replace(),
replace.ReplaceFile(),
script.ScriptLoader(),
serverplayback.ServerPlayback(),
setheaders.SetHeaders(),

View File

@ -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)

View File

@ -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

View File

@ -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

View File

@ -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")
)
)

View File

@ -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)
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()
f.request.content = b"foo"
m.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

View File

@ -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()