diff --git a/mitmproxy/builtins/__init__.py b/mitmproxy/builtins/__init__.py index 591731452..867ebb22b 100644 --- a/mitmproxy/builtins/__init__.py +++ b/mitmproxy/builtins/__init__.py @@ -1,9 +1,11 @@ from __future__ import absolute_import, print_function, division from mitmproxy.builtins import anticomp +from mitmproxy.builtins import stickyauth def default_addons(): return [ anticomp.AntiComp(), + stickyauth.StickyAuth(), ] diff --git a/mitmproxy/builtins/stickyauth.py b/mitmproxy/builtins/stickyauth.py new file mode 100644 index 000000000..1309911c4 --- /dev/null +++ b/mitmproxy/builtins/stickyauth.py @@ -0,0 +1,28 @@ +from __future__ import absolute_import, print_function, division + +from mitmproxy import filt +from mitmproxy import exceptions + + +class StickyAuth: + def __init__(self): + # Compiled filter + self.flt = None + self.hosts = {} + + def configure(self, options): + if options.stickyauth: + flt = filt.parse(options.stickyauth) + if not flt: + raise exceptions.OptionsError( + "stickyauth: invalid filter expression: %s" % options.stickyauth + ) + self.flt = flt + + def request(self, flow): + host = flow.request.host + if "authorization" in flow.request.headers: + self.hosts[host] = flow.request.headers["authorization"] + elif flow.match(self.flt): + if host in self.hosts: + flow.request.headers["authorization"] = self.hosts[host] diff --git a/mitmproxy/console/master.py b/mitmproxy/console/master.py index 06a054647..8e5dae6b0 100644 --- a/mitmproxy/console/master.py +++ b/mitmproxy/console/master.py @@ -244,11 +244,6 @@ class ConsoleMaster(flow.FlowMaster): print("Sticky cookies error: {}".format(r), file=sys.stderr) sys.exit(1) - r = self.set_stickyauth(options.stickyauth) - if r: - print("Sticky auth error: {}".format(r), file=sys.stderr) - sys.exit(1) - self.set_stream_large_bodies(options.stream_large_bodies) self.refresh_server_playback = options.refresh_server_playback @@ -300,6 +295,9 @@ class ConsoleMaster(flow.FlowMaster): self.__dict__[name] = value signals.update_settings.send(self) + def set_stickyauth(self, txt): + self.options.stickyauth = txt + def options_error(self, opts, exc): signals.status_message.send( message=str(exc), diff --git a/mitmproxy/console/options.py b/mitmproxy/console/options.py index db6d405a5..331c28a58 100644 --- a/mitmproxy/console/options.py +++ b/mitmproxy/console/options.py @@ -120,7 +120,7 @@ class Options(urwid.WidgetWrap): select.Option( "Sticky Auth", "A", - lambda: master.stickyauth_txt, + lambda: master.options.stickyauth, self.sticky_auth ), select.Option( @@ -262,7 +262,7 @@ class Options(urwid.WidgetWrap): def sticky_auth(self): signals.status_prompt.send( prompt = "Sticky auth filter", - text = self.master.stickyauth_txt, + text = self.master.options.stickyauth, callback = self.master.set_stickyauth ) diff --git a/mitmproxy/console/statusbar.py b/mitmproxy/console/statusbar.py index 543c67713..d0a240184 100644 --- a/mitmproxy/console/statusbar.py +++ b/mitmproxy/console/statusbar.py @@ -177,10 +177,10 @@ class StatusBar(urwid.WidgetWrap): r.append("[") r.append(("heading_key", "t")) r.append(":%s]" % self.master.stickycookie_txt) - if self.master.stickyauth_txt: + if self.master.options.stickyauth: r.append("[") r.append(("heading_key", "u")) - r.append(":%s]" % self.master.stickyauth_txt) + r.append(":%s]" % self.master.options.stickyauth) if self.master.state.default_body_view.name != "Auto": r.append("[") r.append(("heading_key", "M")) diff --git a/mitmproxy/controller.py b/mitmproxy/controller.py index 8b968eb5c..2f0c8bf29 100644 --- a/mitmproxy/controller.py +++ b/mitmproxy/controller.py @@ -197,9 +197,7 @@ def handler(f): with master.handlecontext(): ret = f(master, message) if handling: - # Python2/3 compatibility hack - fn = getattr(f, "func_name", None) or getattr(f, "__name__") - master.addons(fn, message) + master.addons(f.__name__, 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 cf7aa5c95..f69e3777c 100644 --- a/mitmproxy/dump.py +++ b/mitmproxy/dump.py @@ -87,9 +87,6 @@ class DumpMaster(flow.FlowMaster): if options.stickycookie: self.set_stickycookie(options.stickycookie) - if options.stickyauth: - self.set_stickyauth(options.stickyauth) - if options.outfile: err = self.start_stream_to_path( options.outfile[0], diff --git a/mitmproxy/flow/__init__.py b/mitmproxy/flow/__init__.py index c14a0fecf..2a5d9b38e 100644 --- a/mitmproxy/flow/__init__.py +++ b/mitmproxy/flow/__init__.py @@ -5,7 +5,7 @@ from mitmproxy.flow.io import FlowWriter, FilteredFlowWriter, FlowReader, read_f from mitmproxy.flow.master import FlowMaster from mitmproxy.flow.modules import ( AppRegistry, ReplaceHooks, SetHeaders, StreamLargeBodies, ClientPlaybackState, - ServerPlaybackState, StickyCookieState, StickyAuthState + ServerPlaybackState, StickyCookieState ) from mitmproxy.flow.state import State, FlowView @@ -16,6 +16,5 @@ __all__ = [ "FlowWriter", "FilteredFlowWriter", "FlowReader", "read_flows_from_paths", "FlowMaster", "AppRegistry", "ReplaceHooks", "SetHeaders", "StreamLargeBodies", "ClientPlaybackState", - "ServerPlaybackState", "StickyCookieState", "StickyAuthState", - "State", "FlowView", + "ServerPlaybackState", "StickyCookieState", "State", "FlowView", ] diff --git a/mitmproxy/flow/master.py b/mitmproxy/flow/master.py index ed2ee1384..e469c4991 100644 --- a/mitmproxy/flow/master.py +++ b/mitmproxy/flow/master.py @@ -42,9 +42,6 @@ class FlowMaster(controller.Master): self.stickycookie_state = None # type: Optional[modules.StickyCookieState] self.stickycookie_txt = None - self.stickyauth_state = False # type: Optional[modules.StickyAuthState] - self.stickyauth_txt = None - self.anticache = False self.stream_large_bodies = None # type: Optional[modules.StreamLargeBodies] self.refresh_server_playback = False @@ -136,17 +133,6 @@ class FlowMaster(controller.Master): else: self.stream_large_bodies = False - def set_stickyauth(self, txt): - if txt: - flt = filt.parse(txt) - if not flt: - return "Invalid filter expression." - self.stickyauth_state = modules.StickyAuthState(flt) - self.stickyauth_txt = txt - else: - self.stickyauth_state = None - self.stickyauth_txt = None - def start_client_playback(self, flows, exit): """ flows: List of flows. @@ -326,8 +312,6 @@ class FlowMaster(controller.Master): def process_new_request(self, f): if self.stickycookie_state: self.stickycookie_state.handle_request(f) - if self.stickyauth_state: - self.stickyauth_state.handle_request(f) if self.anticache: f.request.anticache() diff --git a/mitmproxy/flow/modules.py b/mitmproxy/flow/modules.py index ab41da8d1..83228a042 100644 --- a/mitmproxy/flow/modules.py +++ b/mitmproxy/flow/modules.py @@ -350,20 +350,3 @@ class StickyCookieState: if l: f.request.stickycookie = True f.request.headers["cookie"] = "; ".join(l) - - -class StickyAuthState: - def __init__(self, flt): - """ - flt: Compiled filter. - """ - self.flt = flt - self.hosts = {} - - def handle_request(self, f): - host = f.request.host - if "authorization" in f.request.headers: - self.hosts[host] = f.request.headers["authorization"] - elif f.match(self.flt): - if host in self.hosts: - f.request.headers["authorization"] = self.hosts[host] diff --git a/test/mitmproxy/builtins/test_stickyauth.py b/test/mitmproxy/builtins/test_stickyauth.py new file mode 100644 index 000000000..9233f4353 --- /dev/null +++ b/test/mitmproxy/builtins/test_stickyauth.py @@ -0,0 +1,23 @@ +from .. import tutils, mastertest +from mitmproxy.builtins import stickyauth +from mitmproxy.flow import master +from mitmproxy.flow import state +from mitmproxy import options + + +class TestStickyAuth(mastertest.MasterTest): + def test_simple(self): + s = state.State() + m = master.FlowMaster(options.Options(stickyauth = ".*"), None, s) + sa = stickyauth.StickyAuth() + m.addons.add(sa) + + f = tutils.tflow(resp=True) + f.request.headers["authorization"] = "foo" + self.invoke(m, "request", f) + + assert "address" in sa.hosts + + f = tutils.tflow(resp=True) + self.invoke(m, "request", f) + assert f.request.headers["authorization"] == "foo" diff --git a/test/mitmproxy/test_flow.py b/test/mitmproxy/test_flow.py index 6b2708721..a6a3038c9 100644 --- a/test/mitmproxy/test_flow.py +++ b/test/mitmproxy/test_flow.py @@ -120,20 +120,6 @@ class TestStickyCookieState: assert "cookie" in f.request.headers -class TestStickyAuthState: - - def test_response(self): - s = flow.StickyAuthState(filt.parse(".*")) - f = tutils.tflow(resp=True) - f.request.headers["authorization"] = "foo" - s.handle_request(f) - assert "address" in s.hosts - - f = tutils.tflow(resp=True) - s.handle_request(f) - assert f.request.headers["authorization"] == "foo" - - class TestClientPlaybackState: def test_tick(self): @@ -1004,26 +990,6 @@ class TestFlowMaster: fm.request(f) assert f.request.headers["cookie"] == "foo=bar" - def test_stickyauth(self): - s = flow.State() - fm = flow.FlowMaster(None, None, s) - assert "Invalid" in fm.set_stickyauth("~h") - fm.set_stickyauth(".*") - assert fm.stickyauth_state - fm.set_stickyauth(None) - assert not fm.stickyauth_state - - fm.set_stickyauth(".*") - f = tutils.tflow(resp=True) - f.request.headers["authorization"] = "foo" - fm.request(f) - - f = tutils.tflow(resp=True) - assert fm.stickyauth_state.hosts - assert "authorization" not in f.request.headers - fm.request(f) - assert f.request.headers["authorization"] == "foo" - def test_stream(self): with tutils.tmpdir() as tdir: p = os.path.join(tdir, "foo")