Merge pull request #1349 from cortesi/addons

Addons: StickyAuth

12 files changed, 63 insertions, 85 deletions

diff --git a/mitmproxy/builtins/__init__.py b/mitmproxy/builtins/__init__.py
index 59173145..867ebb22 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 00000000..1309911c
--- /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 06a05464..8e5dae6b 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 db6d405a..331c28a5 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 543c6771..d0a24018 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 8b968eb5..2f0c8bf2 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 cf7aa5c9..f69e3777 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 c14a0fec..2a5d9b38 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 ed2ee138..e469c499 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 ab41da8d..83228a04 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 00000000..9233f435
--- /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 6b270872..a6a3038c 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")