Merge pull request #1372 from cortesi/setheaders

setheaders -> addon

mitmproxy/builtins/__init__.py

@@ -7,6 +7,7 @@ from mitmproxy.builtins import stickyauth
 from mitmproxy.builtins import stickycookie
 from mitmproxy.builtins import script
 from mitmproxy.builtins import replace
+from mitmproxy.builtins import setheaders
 
 
 def default_addons():
@@ -18,4 +19,5 @@ def default_addons():
         script.ScriptLoader(),
         filestreamer.FileStreamer(),
         replace.Replace(),
+        setheaders.SetHeaders(),
     ]

mitmproxy/builtins/setheaders.py

@@ -0,0 +1,39 @@
+from mitmproxy import exceptions
+from mitmproxy import filt
+
+
+class SetHeaders:
+    def __init__(self):
+        self.lst = []
+
+    def configure(self, options):
+        """
+            options.setheaders is a tuple of (fpatt, header, value)
+
+            fpatt: String specifying a filter pattern.
+            header: Header name.
+            value: Header value string
+        """
+        for fpatt, header, value in options.setheaders:
+            cpatt = filt.parse(fpatt)
+            if not cpatt:
+                raise exceptions.OptionsError(
+                    "Invalid setheader filter pattern %s" % fpatt
+                )
+            self.lst.append((fpatt, header, value, cpatt))
+
+    def run(self, f, hdrs):
+        for _, header, value, cpatt in self.lst:
+            if cpatt(f):
+                hdrs.pop(header, None)
+        for _, header, value, cpatt in self.lst:
+            if cpatt(f):
+                hdrs.add(header, value)
+
+    def request(self, flow):
+        if not flow.reply.acked:
+            self.run(flow, flow.request.headers)
+
+    def response(self, flow):
+        if not flow.reply.acked:
+            self.run(flow, flow.response.headers)

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.setheaders:
-            for i in options.setheaders:
-                self.setheaders.add(*i)
-
         r = self.set_intercept(options.intercept)
         if r:
             print("Intercept error: {}".format(r), file=sys.stderr)

mitmproxy/console/options.py

@@ -36,7 +36,7 @@ class Options(urwid.WidgetWrap):
                 select.Option(
                     "Header Set Patterns",
                     "H",
-                    lambda: master.setheaders.count(),
+                    lambda: len(master.options.setheaders),
                     self.setheaders
                 ),
                 select.Option(
@@ -156,7 +156,6 @@ class Options(urwid.WidgetWrap):
         self.master.showhost = False
         self.master.refresh_server_playback = True
         self.master.server.config.no_upstream_cert = False
-        self.master.setheaders.clear()
         self.master.set_ignore_filter([])
         self.master.set_tcp_filter([])
 
@@ -165,6 +164,7 @@ class Options(urwid.WidgetWrap):
             anticomp = False,
             replacements = [],
             scripts = [],
+            setheaders = [],
             stickyauth = None,
             stickycookie = None
         )
@@ -197,13 +197,12 @@ class Options(urwid.WidgetWrap):
         signals.update_settings.send(self)
 
     def setheaders(self):
-        def _set(*args, **kwargs):
+        def _set(shdrs):
-            self.master.setheaders.set(*args, **kwargs)
+            self.master.options.setheaders = shdrs
-            signals.update_settings.send(self)
         self.master.view_grideditor(
             grideditor.SetHeadersEditor(
                 self.master,
-                self.master.setheaders.get_specs(),
+                self.master.options.setheaders,
                 _set
             )
         )
@@ -211,7 +210,6 @@ class Options(urwid.WidgetWrap):
     def ignorepatterns(self):
         def _set(ignore):
             self.master.set_ignore_filter(ignore)
-            signals.update_settings.send(self)
         self.master.view_grideditor(
             grideditor.HostPatternEditor(
                 self.master,

mitmproxy/console/statusbar.py

@@ -137,7 +137,7 @@ class StatusBar(urwid.WidgetWrap):
     def get_status(self):
         r = []
 
-        if self.master.setheaders.count():
+        if len(self.master.options.setheaders):
             r.append("[")
             r.append(("heading_key", "H"))
             r.append("eaders]")

mitmproxy/dump.py

@@ -58,10 +58,6 @@ class DumpMaster(flow.FlowMaster):
                   "HTTP/2 is disabled. Use --no-http2 to silence this warning.",
                   file=sys.stderr)
 
-        if options.setheaders:
-            for i in options.setheaders:
-                self.setheaders.add(*i)
-
         if options.server_replay:
             self.start_server_playback(
                 self._readflow(options.server_replay),

mitmproxy/flow/__init__.py

@@ -4,8 +4,7 @@ from mitmproxy.flow import export, modules
 from mitmproxy.flow.io import FlowWriter, FilteredFlowWriter, FlowReader, read_flows_from_paths
 from mitmproxy.flow.master import FlowMaster
 from mitmproxy.flow.modules import (
-    AppRegistry, SetHeaders, StreamLargeBodies, ClientPlaybackState,
+    AppRegistry, StreamLargeBodies, ClientPlaybackState, ServerPlaybackState
-    ServerPlaybackState
 )
 from mitmproxy.flow.state import State, FlowView
 from mitmproxy.flow import options
@@ -16,7 +15,6 @@ __all__ = [
     "export", "modules",
     "FlowWriter", "FilteredFlowWriter", "FlowReader", "read_flows_from_paths",
     "FlowMaster",
-    "AppRegistry", "SetHeaders", "StreamLargeBodies", "ClientPlaybackState",
+    "AppRegistry", "StreamLargeBodies", "ClientPlaybackState",
-    "ServerPlaybackState", "State", "FlowView",
+    "ServerPlaybackState", "State", "FlowView", "options",
-    "options",
 ]

mitmproxy/flow/master.py

@@ -36,7 +36,6 @@ class FlowMaster(controller.Master):
 
         self.stream_large_bodies = None  # type: Optional[modules.StreamLargeBodies]
         self.refresh_server_playback = False
-        self.setheaders = modules.SetHeaders()
         self.replay_ignore_params = False
         self.replay_ignore_content = None
         self.replay_ignore_host = False
@@ -328,8 +327,6 @@ class FlowMaster(controller.Master):
                 return
         if f not in self.state.flows:  # don't add again on replay
             self.state.add_flow(f)
-        if not f.reply.acked:
-            self.setheaders.run(f)
         if not f.reply.acked:
             self.process_new_request(f)
         return f
@@ -347,8 +344,6 @@ class FlowMaster(controller.Master):
     @controller.handler
     def response(self, f):
         self.state.update_flow(f)
-        if not f.reply.acked:
-            self.setheaders.run(f)
         if not f.reply.acked:
             if self.client_playback:
                 self.client_playback.clear(f)

mitmproxy/flow/modules.py

@@ -5,7 +5,6 @@ import hashlib
 from six.moves import urllib
 
 from mitmproxy import controller
-from mitmproxy import filt
 from netlib import wsgi
 from netlib import version
 from netlib import strutils
@@ -39,60 +38,6 @@ class AppRegistry:
             return self.apps.get((host, request.port), None)
 
 
-class SetHeaders:
-    def __init__(self):
-        self.lst = []
-
-    def set(self, r):
-        self.clear()
-        for i in r:
-            self.add(*i)
-
-    def add(self, fpatt, header, value):
-        """
-            Add a set header hook.
-
-            fpatt: String specifying a filter pattern.
-            header: Header name.
-            value: Header value string
-
-            Returns True if hook was added, False if the pattern could not be
-            parsed.
-        """
-        cpatt = filt.parse(fpatt)
-        if not cpatt:
-            return False
-        self.lst.append((fpatt, header, value, cpatt))
-        return True
-
-    def get_specs(self):
-        """
-            Retrieve the hook specifcations. Returns a list of (fpatt, rex, s)
-            tuples.
-        """
-        return [i[:3] for i in self.lst]
-
-    def count(self):
-        return len(self.lst)
-
-    def clear(self):
-        self.lst = []
-
-    def run(self, f):
-        for _, header, value, cpatt in self.lst:
-            if cpatt(f):
-                if f.response:
-                    f.response.headers.pop(header, None)
-                else:
-                    f.request.headers.pop(header, None)
-        for _, header, value, cpatt in self.lst:
-            if cpatt(f):
-                if f.response:
-                    f.response.headers.add(header, value)
-                else:
-                    f.request.headers.add(header, value)
-
-
 class StreamLargeBodies(object):
     def __init__(self, max_size):
         self.max_size = max_size

test/mitmproxy/builtins/test_setheaders.py

@@ -0,0 +1,64 @@
+from .. import tutils, mastertest
+
+from mitmproxy.builtins import setheaders
+from mitmproxy.flow import state
+from mitmproxy.flow import options
+
+
+class TestSetHeaders(mastertest.MasterTest):
+    def mkmaster(self, **opts):
+        s = state.State()
+        m = mastertest.RecordingMaster(options.Options(**opts), None, s)
+        sh = setheaders.SetHeaders()
+        m.addons.add(sh)
+        return m, sh
+
+    def test_configure(self):
+        sh = setheaders.SetHeaders()
+        tutils.raises(
+            "invalid setheader filter pattern",
+            sh.configure,
+            options.Options(
+                setheaders = [("~b", "one", "two")]
+            )
+        )
+
+    def test_setheaders(self):
+        m, sh = self.mkmaster(
+            setheaders = [
+                ("~q", "one", "two"),
+                ("~s", "one", "three")
+            ]
+        )
+        f = tutils.tflow()
+        f.request.headers["one"] = "xxx"
+        self.invoke(m, "request", f)
+        assert f.request.headers["one"] == "two"
+
+        f = tutils.tflow(resp=True)
+        f.response.headers["one"] = "xxx"
+        self.invoke(m, "response", f)
+        assert f.response.headers["one"] == "three"
+
+        m, sh = self.mkmaster(
+            setheaders = [
+                ("~s", "one", "two"),
+                ("~s", "one", "three")
+            ]
+        )
+        f = tutils.tflow(resp=True)
+        f.request.headers["one"] = "xxx"
+        f.response.headers["one"] = "xxx"
+        self.invoke(m, "response", f)
+        assert f.response.headers.get_all("one") == ["two", "three"]
+
+        m, sh = self.mkmaster(
+            setheaders = [
+                ("~q", "one", "two"),
+                ("~q", "one", "three")
+            ]
+        )
+        f = tutils.tflow()
+        f.request.headers["one"] = "xxx"
+        self.invoke(m, "request", f)
+        assert f.request.headers.get_all("one") == ["two", "three"]

test/mitmproxy/test_flow.py

@@ -961,55 +961,3 @@ class TestClientConnection:
         assert c3.get_state() == c.get_state()
 
         assert str(c)
-
-
-def test_setheaders():
-    h = flow.SetHeaders()
-    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"
-
-    h.clear()
-    h.add("~s", "one", "two")
-    h.add("~s", "one", "three")
-    f = tutils.tflow(resp=True)
-    f.request.headers["one"] = "xxx"
-    f.response.headers["one"] = "xxx"
-    h.run(f)
-    assert f.request.headers["one"] == "xxx"
-    assert f.response.headers.get_all("one") == ["two", "three"]
-
-    h.clear()
-    h.add("~q", "one", "two")
-    h.add("~q", "one", "three")
-    f = tutils.tflow()
-    f.request.headers["one"] = "xxx"
-    h.run(f)
-    assert f.request.headers.get_all("one") == ["two", "three"]
-
-    assert not h.add("~", "foo", "bar")