From 1a1f7e6fd94b434c4440d3cbc9407dbe3e0db6f9 Mon Sep 17 00:00:00 2001 From: Aldo Cortesi Date: Sat, 16 Jul 2016 13:00:33 +1200 Subject: replacehooks -> addon Also fixes a bug in header replacements in netlib that resulted in a mutable multidict. --- mitmproxy/builtins/__init__.py | 2 ++ mitmproxy/builtins/replace.py | 49 +++++++++++++++++++++++++++++++ mitmproxy/console/master.py | 4 --- mitmproxy/console/options.py | 12 ++++---- mitmproxy/console/statusbar.py | 2 +- mitmproxy/dump.py | 4 --- mitmproxy/flow/master.py | 5 ---- netlib/http/headers.py | 4 +-- test/mitmproxy/builtins/test_replace.py | 52 +++++++++++++++++++++++++++++++++ test/mitmproxy/test_flow.py | 49 ------------------------------- test/mitmproxy/test_server.py | 17 ++++------- test/mitmproxy/tservers.py | 5 +++- 12 files changed, 122 insertions(+), 83 deletions(-) create mode 100644 mitmproxy/builtins/replace.py create mode 100644 test/mitmproxy/builtins/test_replace.py diff --git a/mitmproxy/builtins/__init__.py b/mitmproxy/builtins/__init__.py index 6b357902..b4d3c0ff 100644 --- a/mitmproxy/builtins/__init__.py +++ b/mitmproxy/builtins/__init__.py @@ -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(), ] diff --git a/mitmproxy/builtins/replace.py b/mitmproxy/builtins/replace.py new file mode 100644 index 00000000..83b96cee --- /dev/null +++ b/mitmproxy/builtins/replace.py @@ -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) diff --git a/mitmproxy/console/master.py b/mitmproxy/console/master.py index 7192c281..0ef12001 100644 --- a/mitmproxy/console/master.py +++ b/mitmproxy/console/master.py @@ -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) diff --git a/mitmproxy/console/options.py b/mitmproxy/console/options.py index d8824b05..f0cc4ef5 100644 --- a/mitmproxy/console/options.py +++ b/mitmproxy/console/options.py @@ -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 ) ) diff --git a/mitmproxy/console/statusbar.py b/mitmproxy/console/statusbar.py index 47cc99f8..1c3be19c 100644 --- a/mitmproxy/console/statusbar.py +++ b/mitmproxy/console/statusbar.py @@ -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]") diff --git a/mitmproxy/dump.py b/mitmproxy/dump.py index 18c24d61..de63ca10 100644 --- a/mitmproxy/dump.py +++ b/mitmproxy/dump.py @@ -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) diff --git a/mitmproxy/flow/master.py b/mitmproxy/flow/master.py index b52e8cb6..02ae7c74 100644 --- a/mitmproxy/flow/master.py +++ b/mitmproxy/flow/master.py @@ -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: diff --git a/netlib/http/headers.py b/netlib/http/headers.py index 413add87..c8cf3e43 100644 --- a/netlib/http/headers.py +++ b/netlib/http/headers.py @@ -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 diff --git a/test/mitmproxy/builtins/test_replace.py b/test/mitmproxy/builtins/test_replace.py new file mode 100644 index 00000000..f8010bec --- /dev/null +++ b/test/mitmproxy/builtins/test_replace.py @@ -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" diff --git a/test/mitmproxy/test_flow.py b/test/mitmproxy/test_flow.py index c58a9703..8197ba08 100644 --- a/test/mitmproxy/test_flow.py +++ b/test/mitmproxy/test_flow.py @@ -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") diff --git a/test/mitmproxy/test_server.py b/test/mitmproxy/test_server.py index a5196dae..2e580d47 100644 --- a/test/mitmproxy/test_server.py +++ b/test/mitmproxy/test_server.py @@ -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" diff --git a/test/mitmproxy/tservers.py b/test/mitmproxy/tservers.py index 9a66984b..9b830b2d 100644 --- a/test/mitmproxy/tservers.py +++ b/test/mitmproxy/tservers.py @@ -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() -- cgit v1.2.3