diff --git a/mitmproxy/builtins/__init__.py b/mitmproxy/builtins/__init__.py new file mode 100644 index 000000000..591731452 --- /dev/null +++ b/mitmproxy/builtins/__init__.py @@ -0,0 +1,9 @@ +from __future__ import absolute_import, print_function, division + +from mitmproxy.builtins import anticomp + + +def default_addons(): + return [ + anticomp.AntiComp(), + ] diff --git a/mitmproxy/builtins/anticomp.py b/mitmproxy/builtins/anticomp.py new file mode 100644 index 000000000..2820a85c6 --- /dev/null +++ b/mitmproxy/builtins/anticomp.py @@ -0,0 +1,12 @@ +from __future__ import absolute_import, print_function, division + +class AntiComp: + def __init__(self): + self.enabled = False + + def configure(self, options): + self.enabled = options.anticomp + + def request(self, flow): + if self.enabled: + flow.request.anticomp() diff --git a/mitmproxy/console/master.py b/mitmproxy/console/master.py index 00905f366..06a054647 100644 --- a/mitmproxy/console/master.py +++ b/mitmproxy/console/master.py @@ -15,6 +15,7 @@ import weakref import urwid +from mitmproxy import builtins from mitmproxy import contentviews from mitmproxy import controller from mitmproxy import exceptions @@ -217,6 +218,7 @@ class ConsoleMaster(flow.FlowMaster): def __init__(self, server, options): flow.FlowMaster.__init__(self, options, server, ConsoleState()) + self.addons.add(*builtins.default_addons()) self.stream_path = None self.options.errored.connect(self.options_error) @@ -251,7 +253,6 @@ class ConsoleMaster(flow.FlowMaster): self.refresh_server_playback = options.refresh_server_playback self.anticache = options.anticache - self.anticomp = options.anticomp self.killextra = options.kill self.rheaders = options.rheaders self.nopop = options.nopop diff --git a/mitmproxy/console/options.py b/mitmproxy/console/options.py index 6a4b8dd6d..db6d405a5 100644 --- a/mitmproxy/console/options.py +++ b/mitmproxy/console/options.py @@ -102,7 +102,7 @@ class Options(urwid.WidgetWrap): select.Option( "Anti-Compression", "o", - lambda: master.anticomp, + lambda: master.options.anticomp, self.toggle_anticomp ), select.Option( @@ -177,7 +177,7 @@ class Options(urwid.WidgetWrap): self.master.anticache = not self.master.anticache def toggle_anticomp(self): - self.master.anticomp = not self.master.anticomp + self.master.options.anticomp = not self.master.options.anticomp def toggle_killextra(self): self.master.killextra = not self.master.killextra diff --git a/mitmproxy/console/statusbar.py b/mitmproxy/console/statusbar.py index d1ab5906f..543c67713 100644 --- a/mitmproxy/console/statusbar.py +++ b/mitmproxy/console/statusbar.py @@ -189,7 +189,7 @@ class StatusBar(urwid.WidgetWrap): opts = [] if self.master.anticache: opts.append("anticache") - if self.master.anticomp: + if self.master.options.anticomp: opts.append("anticomp") if self.master.showhost: opts.append("showhost") diff --git a/mitmproxy/controller.py b/mitmproxy/controller.py index d09038f8f..8b968eb5c 100644 --- a/mitmproxy/controller.py +++ b/mitmproxy/controller.py @@ -199,7 +199,7 @@ def handler(f): if handling: # Python2/3 compatibility hack fn = getattr(f, "func_name", None) or getattr(f, "__name__") - master.addons(fn) + master.addons(fn, message) if handling and not message.reply.acked and not message.reply.taken: message.reply.ack() diff --git a/mitmproxy/dump.py b/mitmproxy/dump.py index cd5159450..cf7aa5c95 100644 --- a/mitmproxy/dump.py +++ b/mitmproxy/dump.py @@ -12,6 +12,7 @@ from mitmproxy import exceptions from mitmproxy import filt from mitmproxy import flow from mitmproxy import options +from mitmproxy import builtins from netlib import human from netlib import tcp from netlib import strutils @@ -59,6 +60,7 @@ class DumpMaster(flow.FlowMaster): def __init__(self, server, options, outfile=None): flow.FlowMaster.__init__(self, options, server, flow.State()) + self.addons.add(*builtins.default_addons()) self.outfile = outfile self.o = options self.anticache = options.anticache diff --git a/mitmproxy/flow/master.py b/mitmproxy/flow/master.py index b1951f944..ed2ee1384 100644 --- a/mitmproxy/flow/master.py +++ b/mitmproxy/flow/master.py @@ -46,7 +46,6 @@ class FlowMaster(controller.Master): self.stickyauth_txt = None self.anticache = False - self.anticomp = False self.stream_large_bodies = None # type: Optional[modules.StreamLargeBodies] self.refresh_server_playback = False self.replacehooks = modules.ReplaceHooks() @@ -332,8 +331,6 @@ class FlowMaster(controller.Master): if self.anticache: f.request.anticache() - if self.anticomp: - f.request.anticomp() if self.server_playback: pb = self.do_server_playback(f) diff --git a/mitmproxy/web/master.py b/mitmproxy/web/master.py index 22972c142..7c775c327 100644 --- a/mitmproxy/web/master.py +++ b/mitmproxy/web/master.py @@ -6,6 +6,7 @@ import collections import tornado.httpserver import tornado.ioloop +from mitmproxy import builtins from mitmproxy import controller from mitmproxy import exceptions from mitmproxy import flow @@ -148,6 +149,7 @@ class WebMaster(flow.FlowMaster): def __init__(self, server, options): super(WebMaster, self).__init__(options, server, WebState()) + self.addons.add(*builtins.default_addons()) self.app = app.Application( self, self.options.wdebug, self.options.wauthenticator ) diff --git a/test/mitmproxy/builtins/__init__.py b/test/mitmproxy/builtins/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/test/mitmproxy/builtins/test_anticomp.py b/test/mitmproxy/builtins/test_anticomp.py new file mode 100644 index 000000000..6bfd54bb5 --- /dev/null +++ b/test/mitmproxy/builtins/test_anticomp.py @@ -0,0 +1,22 @@ +from .. import tutils, mastertest +from mitmproxy.builtins import anticomp +from mitmproxy.flow import master +from mitmproxy.flow import state +from mitmproxy import options + + +class TestAntiComp(mastertest.MasterTest): + def test_simple(self): + s = state.State() + m = master.FlowMaster(options.Options(anticomp = True), None, s) + sa = anticomp.AntiComp() + m.addons.add(sa) + + f = tutils.tflow(resp=True) + self.invoke(m, "request", f) + + f = tutils.tflow(resp=True) + + f.request.headers["Accept-Encoding"] = "foobar" + self.invoke(m, "request", f) + assert f.request.headers["Accept-Encoding"] == "identity" diff --git a/test/mitmproxy/mastertest.py b/test/mitmproxy/mastertest.py index 9e726a322..06854e25f 100644 --- a/test/mitmproxy/mastertest.py +++ b/test/mitmproxy/mastertest.py @@ -3,10 +3,16 @@ import mock from . import tutils import netlib.tutils -from mitmproxy import flow, proxy, models +from mitmproxy import flow, proxy, models, controller class MasterTest: + def invoke(self, master, handler, message): + with master.handlecontext(): + func = getattr(master, handler) + func(message) + message.reply = controller.DummyReply() + def cycle(self, master, content): f = tutils.tflow(req=netlib.tutils.treq(content=content)) l = proxy.Log("connect") diff --git a/test/mitmproxy/test_flow.py b/test/mitmproxy/test_flow.py index eda01ad9d..6b2708721 100644 --- a/test/mitmproxy/test_flow.py +++ b/test/mitmproxy/test_flow.py @@ -870,7 +870,6 @@ class TestFlowMaster: s = flow.State() fm = flow.FlowMaster(None, None, s) fm.anticache = True - fm.anticomp = True f = tutils.tflow(req=None) fm.clientconnect(f.client_conn) f.request = HTTPRequest.wrap(netlib.tutils.treq())