Add the addons mechanism

Demonstrate how it works and interacts with Options by adding our first new
builtin: anticomp.
This commit is contained in:
Aldo Cortesi 2016-07-13 23:45:49 +12:00
parent 1f72532cae
commit 255e1eb00b
13 changed files with 60 additions and 10 deletions

View File

@ -0,0 +1,9 @@
from __future__ import absolute_import, print_function, division
from mitmproxy.builtins import anticomp
def default_addons():
return [
anticomp.AntiComp(),
]

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

View File

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

View File

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

View File

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