Merge pull request #1363 from cortesi/replace

replacehooks -> addon
This commit is contained in:
Aldo Cortesi 2016-07-16 18:57:41 +12:00 committed by GitHub
commit 903807292b
12 changed files with 122 additions and 83 deletions

View File

@ -6,6 +6,7 @@ from mitmproxy.builtins import stickyauth
from mitmproxy.builtins import stickycookie
from mitmproxy.builtins import script
from mitmproxy.builtins import stream
from mitmproxy.builtins import replace
def default_addons():
@ -16,4 +17,5 @@ def default_addons():
stickycookie.StickyCookie(),
script.ScriptLoader(),
stream.Stream(),
replace.Replace(),
]

View File

@ -0,0 +1,49 @@
import re
from mitmproxy import exceptions
from mitmproxy import filt
class Replace:
def __init__(self):
self.lst = []
def configure(self, options):
"""
.replacements is a list of tuples (fpat, rex, s):
fpatt: a string specifying a filter pattern.
rex: a regular expression.
s: the replacement string
"""
lst = []
for fpatt, rex, s in options.replacements:
cpatt = filt.parse(fpatt)
if not cpatt:
raise exceptions.OptionsError(
"Invalid filter pattern: %s" % fpatt
)
try:
re.compile(rex)
except re.error as e:
raise exceptions.OptionsError(
"Invalid regular expression: %s - %s" % (rex, str(e))
)
lst.append((rex, s, cpatt))
self.lst = lst
def execute(self, f):
for rex, s, cpatt in self.lst:
if cpatt(f):
if f.response:
f.response.replace(rex, s)
else:
f.request.replace(rex, s)
def request(self, flow):
if not flow.reply.acked:
self.execute(flow)
def response(self, flow):
if not flow.reply.acked:
self.execute(flow)

View File

@ -210,10 +210,6 @@ class ConsoleMaster(flow.FlowMaster):
self.options = self.options # type: Options
self.options.errored.connect(self.options_error)
if options.replacements:
for i in options.replacements:
self.replacehooks.add(*i)
if options.setheaders:
for i in options.setheaders:
self.setheaders.add(*i)

View File

@ -48,7 +48,7 @@ class Options(urwid.WidgetWrap):
select.Option(
"Replacement Patterns",
"R",
lambda: master.replacehooks.count(),
lambda: len(master.options.replacements),
self.replacepatterns
),
select.Option(
@ -157,14 +157,14 @@ class Options(urwid.WidgetWrap):
self.master.refresh_server_playback = True
self.master.server.config.no_upstream_cert = False
self.master.setheaders.clear()
self.master.replacehooks.clear()
self.master.set_ignore_filter([])
self.master.set_tcp_filter([])
self.master.options.update(
scripts = [],
anticache = False,
anticomp = False,
replacements = [],
scripts = [],
stickyauth = None,
stickycookie = None
)
@ -221,13 +221,13 @@ class Options(urwid.WidgetWrap):
)
def replacepatterns(self):
def _set(*args, **kwargs):
self.master.replacehooks.set(*args, **kwargs)
def _set(replacements):
self.master.options.replacements = replacements
signals.update_settings.send(self)
self.master.view_grideditor(
grideditor.ReplaceEditor(
self.master,
self.master.replacehooks.get_specs(),
self.master.options.replacements,
_set
)
)

View File

@ -141,7 +141,7 @@ class StatusBar(urwid.WidgetWrap):
r.append("[")
r.append(("heading_key", "H"))
r.append("eaders]")
if self.master.replacehooks.count():
if len(self.master.options.replacements):
r.append("[")
r.append(("heading_key", "R"))
r.append("eplacing]")

View File

@ -69,10 +69,6 @@ class DumpMaster(flow.FlowMaster):
else:
self.filt = None
if options.replacements:
for i in options.replacements:
self.replacehooks.add(*i)
if options.setheaders:
for i in options.setheaders:
self.setheaders.add(*i)

View File

@ -37,7 +37,6 @@ class FlowMaster(controller.Master):
self.stream_large_bodies = None # type: Optional[modules.StreamLargeBodies]
self.refresh_server_playback = False
self.replacehooks = modules.ReplaceHooks()
self.setheaders = modules.SetHeaders()
self.replay_ignore_params = False
self.replay_ignore_content = None
@ -328,8 +327,6 @@ class FlowMaster(controller.Master):
if f not in self.state.flows: # don't add again on replay
self.state.add_flow(f)
self.active_flows.add(f)
if not f.reply.acked:
self.replacehooks.run(f)
if not f.reply.acked:
self.setheaders.run(f)
if not f.reply.acked:
@ -350,8 +347,6 @@ class FlowMaster(controller.Master):
def response(self, f):
self.active_flows.discard(f)
self.state.update_flow(f)
if not f.reply.acked:
self.replacehooks.run(f)
if not f.reply.acked:
self.setheaders.run(f)
if not f.reply.acked:

View File

@ -183,8 +183,8 @@ class Headers(multidict.MultiDict):
pass
else:
replacements += n
fields.append([name, value])
self.fields = fields
fields.append((name, value))
self.fields = tuple(fields)
return replacements

View File

@ -0,0 +1,52 @@
from .. import tutils, mastertest
from mitmproxy.builtins import replace
from mitmproxy.flow import master
from mitmproxy.flow import state
from mitmproxy.flow import options
class TestReplace(mastertest.MasterTest):
def test_configure(self):
r = replace.Replace()
r.configure(options.Options(
replacements=[("one", "two", "three")]
))
tutils.raises(
"invalid filter pattern",
r.configure,
options.Options(
replacements=[("~b", "two", "three")]
)
)
tutils.raises(
"invalid regular expression",
r.configure,
options.Options(
replacements=[("foo", "+", "three")]
)
)
def test_simple(self):
s = state.State()
m = master.FlowMaster(
options.Options(
replacements = [
("~q", "foo", "bar"),
("~s", "foo", "bar"),
]
),
None,
s
)
sa = replace.Replace()
m.addons.add(sa)
f = tutils.tflow()
f.request.content = b"foo"
self.invoke(m, "request", f)
assert f.request.content == b"bar"
f = tutils.tflow(resp=True)
f.response.content = b"foo"
self.invoke(m, "response", f)
assert f.response.content == b"bar"

View File

@ -959,55 +959,6 @@ class TestClientConnection:
assert str(c)
def test_replacehooks():
h = flow.ReplaceHooks()
h.add("~q", "foo", "bar")
assert h.lst
h.set(
[
(".*", "one", "two"),
(".*", "three", "four"),
]
)
assert h.count() == 2
h.clear()
assert not h.lst
h.add("~q", "foo", "bar")
h.add("~s", "foo", "bar")
v = h.get_specs()
assert v == [('~q', 'foo', 'bar'), ('~s', 'foo', 'bar')]
assert h.count() == 2
h.clear()
assert h.count() == 0
f = tutils.tflow()
f.request.content = b"foo"
h.add("~s", "foo", "bar")
h.run(f)
assert f.request.content == b"foo"
f = tutils.tflow(resp=True)
f.request.content = b"foo"
f.response.content = b"foo"
h.run(f)
assert f.response.content == b"bar"
assert f.request.content == b"foo"
f = tutils.tflow()
h.clear()
h.add("~q", "foo", "bar")
f.request.content = b"foo"
h.run(f)
assert f.request.content == b"bar"
assert not h.add("~", "foo", "bar")
assert not h.add("foo", "*", "bar")
def test_setheaders():
h = flow.SetHeaders()
h.add("~q", "foo", "bar")

View File

@ -839,17 +839,12 @@ class TestUpstreamProxy(tservers.HTTPUpstreamProxyTest, CommonMixin, AppMixin):
ssl = False
def test_order(self):
self.proxy.tmaster.replacehooks.add(
"~q",
"foo",
"bar") # replace in request
self.chain[0].tmaster.replacehooks.add("~q", "bar", "baz")
self.chain[1].tmaster.replacehooks.add("~q", "foo", "oh noes!")
self.chain[0].tmaster.replacehooks.add(
"~s",
"baz",
"ORLY") # replace in response
self.proxy.tmaster.options.replacements = [
("~q", "foo", "bar"),
("~q", "bar", "baz"),
("~q", "foo", "oh noes!"),
("~s", "baz", "ORLY")
]
p = self.pathoc()
req = p.request("get:'%s/p/418:b\"foo\"'" % self.server.urlbase)
assert req.content == b"ORLY"

View File

@ -9,7 +9,9 @@ from mitmproxy.proxy.server import ProxyServer
import pathod.test
import pathod.pathoc
from mitmproxy import flow, controller
from mitmproxy.flow import options
from mitmproxy.cmdline import APP_HOST, APP_PORT
from mitmproxy import builtins
testapp = flask.Flask(__name__)
@ -34,7 +36,8 @@ class TestMaster(flow.FlowMaster):
config.port = 0
s = ProxyServer(config)
state = flow.State()
flow.FlowMaster.__init__(self, None, s, state)
flow.FlowMaster.__init__(self, options.Options(), s, state)
self.addons.add(*builtins.default_addons())
self.apps.add(testapp, "testapp", 80)
self.apps.add(errapp, "errapp", 80)
self.clear_log()